Emulating ALibi with RoPE

LLMs usually use some sort of positional encoding for helping the model understand where different tokens — or more precisely, KV cache entries — are. Two of the main strategies for doing this are RoPE and ALiBi. ALiBi employs a linear bias to the scores based on relative positions, whereas RoPE rotates the key and query vector dimensions pairwise at different rates, in a manner such that their relative rotations depend on relative positions. It came up for me that I wanted to emulate ALiBi via RoPE, so this is how.

The construction

Take your already-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, at rate θ. For query, key positions i, j, and distance d, 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, so as long as you make $\theta$ small enough for long contexts, which 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}$ and $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 1/\mathrm{base}$) and hence its phase over the window ($\approx 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.