Skip to main content

Long Context and Position Encoding

Transformer の self-attention は、それ自体には 順序の概念がありません。そのため、token の位置情報を別途与える必要があります。Position encoding の設計は、long context への外挿性能と直結します。

なぜ position が必要か​

Attention は集合演算的で、token を並び替えても出力が同じになります。これを避けるために、各 token に位置情報を加えます。古典的には、

  • Absolute positional embedding (GPT-2): 位置ごとに学習 vector を加える
  • Sinusoidal encoding (原論文): 正弦 / 余弦の組

がありました。これらは training context 長を超えると性能が大きく落ちる弱点があります。

Relative position と RoPE​

近年の LLM では、RoPE (Rotary Position Embedding) が事実上の標準です。RoPE は query と key を 2D ペアに分け、位置に応じた回転を適用します。

q~m=Rmqm,k~n=Rnkn\tilde{q}_m = R_m q_m, \quad \tilde{k}_n = R_n k_n ⟨q~m,k~n⟩=⟨qm,Rn−m kn⟩\langle \tilde{q}_m, \tilde{k}_n \rangle = \langle q_m, R_{n - m}\, k_n \rangle

このため、attention score が 位置の相対差 n−mn - m にだけ依存するようになります。Relative position の自然な実装として、Llama、Qwen、Mistral などで採用されています。

ALiBi​

ALiBi (Attention with Linear Biases) は、position embedding を加える代わりに、attention score に位置差に比例した 線形 bias を加えます。学習なしで extrapolation に強いことが特徴です。

attention_score(i, j) -= m * |i - j|

Long context への拡張​

Training 時の context 長 (例: 4k) を超えて、推論時に 32k、128k、1M token を扱いたいことがあります。代表的な手法:

手法内容
Position InterpolationRoPE の角度を縮めて長い context に対応
NTK-aware scaling高周波と低周波を別 scale で扱う
YaRNNTK + dynamic + temperature の組合せ
Continued pretraining長い context で追加学習
Long-context SFT長い文書での fine-tuning

注意点​

Long context は context 長を伸ばすだけでは性能が出ません。

  • Lost in the middle: 中盤の情報が無視されやすい
  • Needle-in-a-haystack だけでは実用性能は測れない
  • KV cache サイズが線形に増える (メモリ問題)
  • Attention 計算が二乗 (FlashAttention 等で緩和)

実用では、long context と RAG を組み合わせるのが現実的です。

数式で見る RoPE と相対位置​

RoPE は、query と key の各 2 次元ペアに位置 mm に応じた回転行列 Rm\mathbf{R}_m をかけます。

q~m=Rmqm,k~n=Rnkn\tilde{\mathbf{q}}_m=\mathbf{R}_m\mathbf{q}_m, \qquad \tilde{\mathbf{k}}_n=\mathbf{R}_n\mathbf{k}_n

Attention score は次のように相対位置 n−mn-m の形で現れます。

q~m⊤k~n=qm⊤Rn−mkn\tilde{\mathbf{q}}_m^\top\tilde{\mathbf{k}}_n =\mathbf{q}_m^\top\mathbf{R}_{n-m}\mathbf{k}_n

この式の気持ちは、「絶対位置そのものを token に足すのではなく、query と key の角度を位置に応じて回転させることで、内積の中に相対距離を自然に入れる」というものです。

Long context extension では、推論時の位置 index を training 時よりも圧縮して使う position interpolation がよく使われます。Training context 長を LtrainL_{train}、推論 context 長を LtestL_{test} とすると、位置 mm を次のように縮めます。

m′=m⋅LtrainLtestm' = m\cdot\frac{L_{train}}{L_{test}}

この式の気持ちは、「長い sequence の位置を、training で見た範囲の角度に押し込む」というものです。ただし、単純な補間は高周波成分の扱いを崩すことがあるため、NTK-aware scaling や YaRN のような周波数ごとの補正が使われます。

関連ページ​

主なソース​