Emulating ALiBi with RoPE

LLMs typically use some sort of positional encoding for helping the model understand where different tokens (or more precisely, KV cache entries) are. Two of the most common strategies for doing this are RoPE and ALiBi. ALiBi employs a linear bias to the scores based on relative positions, whereas RoPE rotates key and query vector dimensions pairwise at different rates according to their position in the sequence.

It came up for me that I wanted to emulate ALiBi via RoPE, so this is how.

How

Take your trained ALiBi attention heads, for each head, add two dimensions to Q and K with projections that have zero weights and biases such that $q_{fixed} = (-N, 0)$, $k_{fixed} = (0, 1)$.

The RoPE frequencies for all the original dimensions are set to zero (identity rotation). Only the new pair rotates, which it does at rate θ. For query, key positions i, j, and distance d(=j-i), the RoPE contributions $(R_i q_{fixed})\cdot (R_{j} k_{fixed}) = −N \sin(\theta d)$ With scaling: $−\frac{N}{\sqrt{head\_dim}} · sin(\theta d)$ ALiBi needs $−m \cdot d$, so the local match is $N = \frac{m\sqrt{\mathrm{head\_dim}}}{\theta}.$ When $\theta \cdot d$ is small, $\sin(\theta d) \approx \theta d$, and $-\frac{N}{\sqrt{head\_dim}} \sin{\theta d} \approx \frac{N \theta}{\sqrt{head\_dim}} \cdot d = -md$

So ALiBi is equivalent to the small-angle limit of a fixed RoPE pair. As long as you make $\theta$ small enough for long contexts, doing this in turn causes $N$ to be large (for BLOOM's steepest head, $m \approx 0.707$ with $\mathrm{head\_dim} = 64$, a max phase of 0.1 over a 65K window means $\theta \approx 1.5 \times 10^{-6}$, $N \approx 3.7 \times 10^6$). Of course practically if you want to fit it into a specific model's rope you can't choose your exact $\theta$ you have to choose among the available frequencies given by rope where

$$\theta_r = base^{(\frac{-2r}{rope\_dim})}$$ Then again if you are using an existing set all your dimensions are turning which really screws things up. However we want to be realistic here so later we will consider the limitations on emulatability imposed by the size of the base, which sets the frequency of the slowest RoPE pair ($\approx \frac{1}{\mathrm{base}}$) and hence its phase over the window ($\approx \frac{L}{\mathrm{base}}$, for context length $L$).

Does it Actually Work?

I tested it with BLOOM-560M since Bloom natively uses ALiBi. Sure enough it was able to reproduce the results more or less exactly.

TV is total variation the amount of probability mass that would need to be moved from one bin to another to match the distributions. For the KV TV that is total variation over the attention distributions. For the output TV that is total variation over the output distribution. The rad span is the maximum rotation in radians the RoPE pair used to emulate ALiBi is rotated in the sequence measured in Radians.

The TV measurement is using prefill rather than autoregression i.e. the token sequence is fixed. Both native BLOOM and the RoPE patched version are evaluated on the same sequence. Errors do accumulate across layers but not due to differences in autoregressive results since then the sequences would not remain comparable beyond that point. The results are helped by distant pairs generally having low attention weight/influence so I also tested on a retrieval focused.

Is This Really Realistic?

Q: Isn't this is just glueing on one very slow rotating RoPE pair and keeping the rest of the dimension no-rope. Furthermore is a rotation that slow realistic to what is actually found in real rope models?

A: It depends what "realistic" means. Historically RoPE uses a base like 10k these days long-context models often use much larger rope bases (e.g. 500k-1m). This construction cares about the angular frequency of the slowest RoPE pair $\approx \frac{1}{base}$. Appendix 1 shows that extreme degradation occurs around a rotation of $\pi$ (you can see why by considering that it makes the adjustment for very distant contexts the same as the immediate context). Staying approximately linear requires $\theta \cdot d << 1$ but we can see performance remains strong far beyond that regime. Therefore to practically be able to use a RoPE pair to approximate ALiBi over a context window, its slowest angular frequency must satisfy $\theta \cdot \mathrm{max\_distance} = \frac{max\_distance}{base} < \pi$, (since we want recent items to be weighted higher than far distant ones).

P

Some concrete cached model configs (TV here is BLOOM-560M patched with the sine bias at that config's slowest frequency, versus native BLOOM follow

Of course though it also helps that this approximation is most off on far distant items that don't get much attention most of the time anyways.

Reflections

RoPE can simulate ALiBi rather simply, (perhaps even pretty obviously) so it raises the question: what are they actually doing? How are the models learning to use them? Why can such different things be used to do the same job — are they really all that different? I have not yet thought of a good clean and efficient way to simulate RoPE with ALiBi. While this certainly speaks to its' relative complexity I don't know that I would characterize it as expressiveness. Since being locked into having dims that rotate at different rates feels much more like a handicap under many circumstances.

Code can be found here.

Appendix 1

Degradation around $\pi$