コンテキストウィンドウ・ウォーズ:SinusoidalからYaRNへの進化

The Context Window Wars From Sinusoidal to YaRN

ミステリー小説を読んでいると想像してみてください。しかし、100ページ目に到達した時点で、5ページ目で誰が殺人を犯したのかを完全に忘れてしまっていたらどうでしょう?これが大規模言語モデル(LLM)における「金魚の記憶力(Goldfish Memory)」問題です。

長年、LLMは固定されたコンテキストウィンドウ(2kまたは4kトークン)という足枷に縛られていました。巨大なPDFを貼り付けても、モデルは末尾を切り捨てるか、あるいは幻覚(ハルシネーション)を見るだけでした。この限界の原因は、単なるメモリ容量の問題だけではありません。それはモデルに単語の順序を教える「数学的なGPS」、すなわち Positional Embeddings(位置埋め込み) にありました。

私たちは今、モデルが「空間」と「時間」を扱う方法において、急速な進化を目の当たりにしています。硬直的な絶対位置から、柔軟で無限に近いコンテキストスケーリングへの移行です。ここでは、基本的なサイン波(Sinusoidal)から、最先端の YaRN に至るまでの道のりを解説します。

核となる概念:「解像度」の問題

コンテキストの拡張を理解するには、「解像度」のジレンマを理解する必要があります。

古い手法では、位置情報を特定の「番地」(1、2、3…)のように扱っていました。100軒の家を理解するように訓練されたモデルは、101番目の家を見た瞬間にパニックに陥ります。

新しい手法(RoPEやYaRNなど)は、位置を柔軟なゴムバンドのように扱います。100軒用に設計された通りに200軒の家を収めたい場合、通りを新しく作るのではなく、家を縮小(補間:Interpolation)して詰め込むのです。しかし、縮小しすぎると家はぼやけてしまい、モデルは隣同士の区別がつかなくなります。現代の技術の目標は、そのきめ細かな解像度を失うことなく、コンテキストウィンドウを引き伸ばすことにあります。


位置情報の進化

LLMがより長いドキュメントを「読む」方法を学習していった過程のヒエラルキーがこちらです。

1. SIN (Sinusoidal Embeddings:正弦波埋め込み)

  • 起源: Attention Is All You Need (Transformers) で導入。
  • 仕組み: 単語の埋め込み(Embeddings)に、固定されたサイン/コサイン波のパターンを加算します。
  • 欠点: これは絶対的(Absolute)です。モデルは「位置1から5単語後」ではなく、「位置5」そのものを学習します。学習時の長さを超えると壊滅的に機能しなくなります。

2. ALiBi (Attention with Linear Biases)

  • 転換点: 入力に埋め込みを加算するのではなく、ALiBiはアテンションスコアを直接修正します。
  • 仕組み: 2つの単語が離れているほど、アテンションスコアからペナルティを差し引きます。距離が遠いほど、注目度(Attention)が下がります。
  • 利点: 外挿(Extrapolation)が完璧に機能します。2kトークンで学習し、8kトークンでテストすることが可能です。
  • 欠点: 遠くのトークンに重いペナルティを科すため、「干し草の中の針(Needle in a haystack)」のような、長距離の依存関係を見つけるタスクには苦戦します。

3. RoPE (Rotary Positional Embeddings)

  • 標準: Llama、Mistral、PaLMなどで採用されています。
  • 仕組み: QueryベクトルとKeyベクトルを2次元空間で回転させます。回転角度が位置を表します。
  • 数式:
  • 利点: 相対的(Relative)な位置関係を見事に捉えます(2つのベクトル間の角度は、絶対位置ではなく、それらの距離にのみ依存します)。

4. PI (Position Interpolation:位置補間)

  • ハック: 研究者たちはLlama 1(2kコンテキスト)を32kまで拡張したいと考えました。
  • 仕組み: モデルが見たことのない角度へ外挿するのではなく、PIは32k分の位置情報を0〜2kの範囲に「押し込み(線形スケーリング)」ます。
  • 欠点: 高周波情報の損失を引き起こします。これは4k画像を480pにリサイズするようなもので、鋭いエッジ(詳細な情報)が失われてしまいます。

5. NTK-Aware Scaling (NTK認識スケーリング)

  • 改良: Neural Tangent Kernel (NTK) 理論は、ディープニューラルネットワークが高周波成分の学習に苦労することを示唆しています。
  • 仕組み: すべてを一律に押し込む(線形PI)のではなく、NTKスケーリングは低周波(長距離)を圧縮しつつ、高周波(局所範囲)はそのまま維持します。
  • 結果: ファインチューニングなしでコンテキストウィンドウを拡張でき、かつ局所的なタスクでのモデルの鋭さを保つことができます。

6. YaRN (Yet another RoPE Nontrivial extension)

  • 最先端(SOTA): Nous-Hermesや最新のLlamaファインチューンモデルで使用されています。
  • 問題点: RoPEを補間すると、アテンション分布の「エントロピー」が変化してしまいます。つまり、どれくらい注意を「集中」させるべきかについて、モデルが混乱してしまうのです。
  • 解決策: YaRNはNTKスケーリングに加えて、アテンションのソフトマックス(Softmax)に温度補正(Temperature correction)を組み合わせます。これにより、コンテキストを引き伸ばすことで生じる混乱を効果的に「再調整(Re-calibrate)」します。

コード:YaRNスケーリングの実装

ゼロから数式を書く必要はありません。transformersのような最近のライブラリの多くは、この設定を標準でサポートしています。

以下は、Llamaモデルを設定し、YaRNスケーリングを使用して4kモデルを128kコンテキストまで拡張する方法です。

from transformers import AutoConfig, AutoModelForCausalLM

# ベースとなる設定をロード
config = AutoConfig.from_pretrained("meta-llama/Llama-2-7b-hf")

# YaRN用にRoPEスケーリング設定を修正
# original_max_position_embeddings は元々 4096
config.rope_scaling = {
    "type": "yarn",              # YaRN方式を指定
    "factor": 32.0,              # スケール係数 (128k / 4k = 32)
    "original_max_position_embeddings": 4096,
    "finetuned": False           # 既に長文脈でファインチューニング済みのモデルならTrue
}

# 新しい設定でモデルをロード
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    config=config,
    device_map="auto"
)

print(f"Model context window extended to: {config.max_position_embeddings * config.rope_scaling['factor']}")

ステップ・バイ・ステップ:戦略の選び方

LLMを構築またはファインチューニングする場合、以下の決定マトリックスに従ってください。

  1. 標準コンテキスト (< 4k): 通常のRoPE (Vanilla RoPE) を使いましょう。堅牢であり、あらゆる場所でサポートされています。
  2. ゼロショット拡張: 事前学習済みのLlama 2/3を使って、追加学習なしですぐに長いPDFを読ませたい場合は、Dynamic NTK を使用してください。「箱から出してすぐ」驚くほど良く機能します。
  3. 長文脈のためのファインチューニング: 計算リソースに余裕があるなら、YaRN を使いましょう。線形補間と比較して収束が早く、元の短いコンテキストのタスクにおいても高い精度を維持します。
  4. ハードウェアの制約: メモリに制約がある場合は、ALiBi(またはRing Attentionのような新しい変種)を検討してください。これらは推論時のメモリ制限をより優雅に処理します。

結論

「Sinusoidal」から「YaRN」への飛躍は、位置の単なる暗記から、情報の相対性の理解へのシフトを意味します。数段落で混乱してしまうモデルから、小説一冊を一度に消化できるモデルへと移行したのです。

  • RoPE が「回転」を与え、
  • PI がそれを「引き伸ばせる」ことを示し、
  • NTK が「賢く引き伸ばす」方法を教え、
  • YaRN がその「焦点」を完成させました。