長文脈LLMの技術 — RoPE補間からリング注意まで

「コンテキスト長128K」の中身は、性格の違う二つの壁を別々に越えた結果です。位置の壁を越える位置補間・NTK・YaRN、計算の壁を越える窓とリング注意、そして越えられたかを測るneedle試験の読み方を、前提知識ゼロから順に追います。

Extending Context Window of Large Language Models via Positional Interpolation

「128Kトークン対応」は、何を対応させたのか

モデルカードに「コンテキスト長 128K」とある。ではそのモデルに 128,000 トークンを流し込むと何が起きるのか。答えは「性格の違う二つの壁を、別々の方法で越えている」です。この二つを混ぜて考えると、長文脈まわりの話は必ずこんがらがります。

一つ目は位置の壁。モデルは 4K や 8K の長さで学習しており、見たことのない「位置 90000」を入れると、位置の表し方そのものが未知の領域に入って出力が崩れます。計算量とは無関係で、メモリが無限にあっても起きます。

二つ目は計算の壁。注意機構は全トークンが全トークンを見るので、計算もメモリも長さの2乗で膨らむ。位置の壁を完璧に越えたとしても、行列が1台のGPUに載りません。

前者を解くのが位置補間とYaRNの系統、後者が窓を切る系とリング注意の系統です。そして「本当に越えられたのか」を測るのが needle 試験。この3つを順に追います。

比喩: 30cmのものさしと、机の広さ

位置の壁は、ものさしの問題です。30cmのものさしだけで訓練を受けた測定係が、ある日「1mの板を測れ」と言われる。取れる手は二つです。

外挿は、30の先も同じ間隔で目盛りが続くと信じて、板の上に想像で目盛りを延ばすこと。見たことのない領域なので、たいてい狂います。

内挿は、板を1/3に縮小コピーして手元のものさしで測ること。目盛りは全部見慣れた範囲に収まりますが、代わりに細かい違いが潰れます。1mm差だったものが0.33mm差になり、隣り合う2点の区別がつきにくくなる。

位置まわりの長文脈手法は、ほぼすべて「外挿か内挿か、どこをどちらにするか」の話です。一方、計算の壁は机の広さの問題——全員分の答案を1枚の机に広げられないという物理の話。系統が二つに分かれるのは、性格がこれだけ違うからです。

おさらい: RoPEは「位置に比例して回す」

現代のオープンLLMはほぼRoPE（回転位置埋め込み）を使っています。詳細は位置エンコーディングを1から理解するに譲り、ここでは長文脈の議論に必要な一点だけ確認します。

RoPEはQueryとKeyのベクトルを2次元ずつのペアに分け、位置 $m$ に比例した角度だけ回します。$i$ 番目のペアが回る速さは、次の式で決まります。

$b$ は「底」と呼ばれる定数（多くのモデルで 10000）、$d$ は1ヘッドの次元数、$\theta_i$ は1トークン進むごとの回転角、$\lambda_i$ はそのペアが1周するのに要するトークン数、つまり波長です。式(1)が言っているのは「ペアごとに回る速さが違い、$i$ が大きいペアほどゆっくり回る」ということだけ。

そして位置 $m$ のQueryと位置 $n$ のKeyの内積を取ると、$m$ と $n$ が別々に消えて差 $m-n$ の関数だけが残ります。だからこそ「位置番号を詰める」という乱暴に見える操作が意味を持つのです。

位置補間 — 目盛りを詰めて、内側へ押し込む

いちばん単純な手が Position Interpolation（PI） です。学習時の長さを $L$、伸ばしたい長さを $L'$ として $s = L'/L$ を決め、位置番号そのものを縮めます。

式(2)は「位置90000を、$s=32$ なら位置2812.5として扱う」と言っています。整数でなくてよいのがポイントで、RoPEの回転角は連続なので小数の位置がそのまま使えます。これですべての回転角が学習済みの範囲に収まり、未知の角度は一つも現れません。

比喩でいう縮小コピーです。目盛りは足りるようになったが、その意味は薄まった。実際、PIをかけた直後のモデルはそのままでは性能が落ちるため、原論文も少量のファインチューニングとセットで提案しています。

波長で見ると、何が壊れたのかが分かる

PIの副作用の正体は、式(1)の波長 $\lambda_i$ を並べると見えてきます。$i$ が小さいペアの波長は数トークンぶんしかなく、隣り合う単語の区別に効いています。$i$ が大きいペアの波長は数万トークンにもなり、学習長 $L$ の中では1周もしません。

PIはこれを一律に $1/s$ へ縮めます。しかし困っているのは後者だけです。波長の長いペアは学習中に角度の一部しか見ておらず、外挿すれば未知の領域に出る。逆に波長が数トークンのペアは何万周もしており、角度の全域を見尽くしている——外挿しても未知は存在しない。それを一律に縮めれば、局所の解像度だけが理由もなく落ちます。

改良の方向はこれで決まります。速く回るペアはそのまま、遅いペアだけ縮める。