RoPE を用いた ALiBi のエミュレーション

LLM は通常、異なるトークン、より正確には KV キャッシュエントリの位置をモデルが理解するのを助けるために、何らかの位置エンコーディングを使用します。これを行うための主要な戦略は RoPE と ALiBi の 2 つです。ALiBi は相対位置に基づいたスコアへの線形バイアスを採用する一方、RoPE はキーとクエリのベクトル次元をペアごとに異なるレートで回転させ、その相対的な回転が相対位置に依存するようにします。 ALiBi を RoPE でエミュレートしたいという考えが浮かんだので、その方法を説明します。 構築 既にトレーニング済みの ALiBi アテンションヘッドを用意します。各ヘッドについて、Q と K に 2 つの次元を追加し、ゼロの重みとバイアスを持つプロジェクションを適用します。 $q_{fixed} = (-N, 0), k_{fixed} = (0, 1)$ すべての元の次元に対する RoPE の周波数はゼロ（恒等回転）に設定します。新しいペアのみが、レート θ で回転します。 クエリ、キーの位置 i, j、および距離 d に対して、RoPE の寄与は $(R_i q_{fixed})ullet (R_{j} k_{fixed}) = −N ext{sin}( heta d)$ スケーリングを適用すると $−rac{N}{ ext{head_dim}} ext{sin}( heta d)$ ALiBi は $−m ullet d$ を必要とするため、局所的な一致は $N = rac{m ext{head_dim}}}{ heta}$ となります。 $ heta ullet d$ が小さい場合、$ ext{sin}( heta d) ext{approx} heta d$ となり、$-rac{N}{ ext{head_dim}} ext{sin}( heta d) ext{approx} rac{N heta}{ ext{head_dim}} ullet d = -md$ したがって、ALiBi は固定された RoPE ペアの小角極限に相当します。ただし、長いコンテキストのために $ heta$ を十分に小さく設定する必要があり、その結果 N が大きくなります（BLOOM の最も急峻なヘッドの場合、$m ext{approx} 0.707$、$ ext{head_dim} = 64$、65K ウィンドウでの最大位相が $0.1$ の場合、$ heta ext{approx} 1.5 imes 10^{-6}$、$N ext{approx} 3.7 imes 10^6$ となります）。 もちろん、実際には特定のモデルの RoPE に収めるために、正確な $ heta$ を選択することはできず、RoPE によって与えられる利用可能な周波数の中から選択する必要があります。 $$ heta_r = ext{base}^{(rac{-2r}{ ext{rope_dim}})} $$ ここでもし既存のセットを使用している場合、すべての次元が回転してしまい、これは状況を大きく悪化させます。しかし、現実的に考えるために、後でベースのサイズによって課されるエミュレート可能性の制約を考慮します。ベースのサイズは、最も遅い RoPE ペアの周波数（$ ext{approx} 1/ ext{base}$）を決定し、それによってウィンドウ全体での位相（コンテキスト長 L に対して $ ext{approx} L/ ext{base}$）を決定します。 実際に機能するか？ BLOOM-560M でテストしました。BLOOM はネイティブで ALiBi を使用しているためです。予想通り、結果を多かれ少なかれ正確に再現することができました。 | Source | Contexts | Phase span at 65K (rad) | Output TV/token | Max output TV/token | KL/token | Top-1 agreement | Bias-only attention TV | Max attention TV | | ← Back to Blog | Home