用旋转位置编码（RoPE）模拟 ALiBi

用旋转位置编码（RoPE）模拟 ALiBi
Emulating ALiBi with Rope

原始链接: https://alexlitzenberger.com/blog/emulating_alibi_with_rope

本摘要探讨了如何利用旋转位置编码（RoPE）来模拟 ALiBi 位置编码。ALiBi 通过基于相对位置的线性偏置作用于注意力分数，而 RoPE 则通过成对旋转来处理查询（query）和键（key）向量。作者证明了 ALiBi 在数学上等同于固定 RoPE 对的小角度极限。通过为查询和键向量增加两个具有特定固定权重的维度，并将旋转频率（$\theta$）设置得足够低，RoPE 的贡献即可近似于 ALiBi 所需的线性惩罚（$−m \cdot d$）。关键技术要求包括： * 增加零权重投影的维度，以确保仅有新的向量对发生旋转。 * 校准 $N$ 和 $\theta$ 以匹配 ALiBi 的斜率（$m$）。 * 处理 RoPE 预定义频率基准所带来的局限性，这些基准决定了可用旋转速率的范围。在原生使用 ALiBi 的 BLOOM-560M 模型上的测试证实，该方法能够以高精度成功复现 ALiBi 的注意力行为。这种方法为在不同位置编码方案之间转换或结合的模型提供了一种可行的桥梁。

抱歉。

原文

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.