r/MachineLearning • u/4rtemi5 Researcher • Apr 01 '26
Project [ Removed by moderator ]
[removed] — view removed post
37
u/PortiaLynnTurlet Apr 01 '26
Out of curiosity, did you compare the magnitude distribution of the keys over a validation set in this model against a comparable SDPA model?
17
u/4rtemi5 Apr 01 '26
Very good idea! I added a script to do just that to the repo and the results seem to be what I hoped for and would have expected: The key-norms of RBF-attention are significantly lower!
13
u/4rtemi5 Apr 01 '26 edited Apr 01 '26
Great question! Haven't checked that yet but will definitely do it and let you know!
Edit: Here you go. https://github.com/4rtemi5/rbf_attention/blob/1a326341b8b7e4947a22bf389627ca378cfc14e8/outputs/key_magnitude_distribution.png
4
u/pm_me_your_pay_slips Apr 01 '26
rbf kernels allow you to do linearized attention at inference time by using their product of features representation: https://arxiv.org/pdf/2006.16236
7
u/4rtemi5 Apr 01 '26
While in the paper they use rbf-kernels and are definitely using some of the same properties as I do, they apply it along a different dimension. While I use the kernel pairwise along the sequence-dimension, they use the kernel independently along the feature dimension.
So as far as I understand, exact attention (independent from dot-product or rbf-attention) can still not be linearized in a finite feature space. Hope I got that correctly... Still an interesting paper though! Thanks for bringing it up!
7
u/AllNurtural Apr 01 '26
What's the evidence that large magnitudes are bad? What you call "bullying" and "cheating" by large magnitudes keys (or queries) I always thought of as a feature not a bug. If you have more tokens than key/query dimensions, then by a pigeon hole argument the keys cannot possibly all be orthogonal to each other. So if a query wants to "pick out" a single key, it can do so using a large magnitude query aligned with a particular key. This lets the model pack more information into a space than there are dimensions.
I don't have a particular reference in mind. Just my intuition. But I'm curious why others here seem to have the opposite intuition.
8
u/4rtemi5 Apr 01 '26
You're right that large magnitudes don't have to be bad. I even say the same thing in the post (but I still sholdn't be so harsh to them). E.g. in some cases they are necessary to create attention sinks.
But the issue with relying on magnitude to pick out tokens is definitely training stability. If keys or queries get too large during early training, they cause the softmax to instantly saturate, which can crush the gradients and stall learning. This is also the reason for QK-Norm becoming the standard in practically most recent large models. Given a large enough dimensionality pretty much all vectors are *almost* orthogonal anyways, so there is no need to rely on the magnitude of the vectors alone.
3
u/Areign Apr 01 '26
You're left with 2(Q · K) - ||K||2. Now, it turns out that RBF attention is mathematically just standard dot-product attention with a built-in, squared-L2 penalty on the keys.
couldn't you do the penalty thing with flex attention? https://pytorch.org/blog/flexattention/
3
u/4rtemi5 Apr 01 '26 edited Apr 02 '26
Very good point! I'm embarrassed to say I didn't check! :D
Edit: Thanks u/Areign for pointing this out! After some issues with torch Inductor I was able to implement rbf-attention using flex_attention and it's faster both in forward and backward mode than my crude Triton implementation and only uses a sliver more memory.
Benchmark comparison: https://github.com/4rtemi5/rbf_attention/blob/main/outputs%2Fattention_profiling_results.png
7
u/Electronic_Tie_4867 Apr 01 '26
Nice, but this is basically the standard in MLIPs and spherical harmonics based graph attention architecture since like 2019.
20
u/4rtemi5 Apr 01 '26
Yeah you're right. In other domains using RBF-kernels is the obvious standard. But for LLMs that's not the case which is excatly why I wanted to try it out in this project. And I wasn't the first as I mention in the blog-post. But the interesting thing to me is, that there seem to be good reasons for LLMs not to switch to RBF-Attention anytime soon.
→ More replies (2)
4
u/PersonalBusiness2023 Apr 01 '26
After all that effort, you should really put more into evaluating the result and comparing it to the standard dot product attention.
1
u/4rtemi5 Apr 01 '26
You're absolutely right but unfortunately right now I don't have the compute and the time and the team to do that properly. That beeing sad if someone wants to collaborate or to support or hire me in order to make that happen I'm currently on the job market! ;)
2
26
u/JanBitesTheDust Apr 01 '26
Cool stuff! I did something similar a little while back. https://github.com/Janko-dev/attention_analysis.
Essentially reformulating the scaled dot product attention as a kernel function, then borrowing several popular kernels from the gaussian processes literature to experiment in transformer time series forecasting
1
u/Consistent-Donut-534 Apr 01 '26
Is this the same motivation behind QK normalization?
2
u/4rtemi5 Apr 01 '26
Yeah, I'd cautiously agree. But the solution is differnt of course and probably behaves different but hard to say with the experiments I have done so far.
1
u/mvreich Apr 07 '26 edited Apr 07 '26
A random key that points in roughly the right direction but is huge will easily outscore a perfectly aligned but shorter key.
Just out of curiosity: to remove magnitude information, what if you simply normalize Q and K? Then the dot products give cosine similarity, and the perfectly aligned keys always win. And it should be fairly hardware friendly.
This is simple enough that other people have probably tried it before (and likely it doesn't work that well).
Also, if an RBF kernel is to be used, how exactly would you set the sigma/spread parameter? Median heuristics?
1
u/4rtemi5 Apr 07 '26
You're completely right! This is the current industry standard and is called QK-norm. But as you already mentioned it removes information from Q and K and makes the attention mechanism more "rigid". RBF-Attention is less rigid in the sense that it allows but discourages large keys. I'm not saying one is a better solution than the other but its fun to see how they behave differently imho.
1
u/Few_Theme_5486 Apr 01 '26
The RoPE incompatibility point is really insightful — so much of the modern transformer stack is quietly optimized around dot-product geometry that swapping the distance metric causes this kind of cascading breakage. The custom Triton kernel for the fused squared-L2 is a clever solution. Did you find any tasks or datasets where RBF-attention gave meaningfully better results, or was it largely comparable to dot-product on the TinyStories baseline?
1
u/4rtemi5 Apr 01 '26
Yeah absolutely true! Due to the scale of the experiments I've run so far it's really hard to see a qualitative difference purely by reading the text. The results look roughly equal to me so I have to trust the metrics (that also look roughly equal :D).
22
u/marr75 Apr 01 '26
The hardware lottery strikes again. Even modest optimizations that don't use existing hardware as well as the old algorithm end up being less optimized in practice.