今回の論文
今回取り上げるのは、Bowen Peng、Jeffrey Quesnelle、Honglu Fan、Enrico Shippole による論文「YaRN: Efficient Context Window Extension of Large Language Models」です。2023年8月31日に arXiv で公開されました。研究分野としては、長文コンテキストLLM、位置埋め込み、コンテキスト長拡張に当たります。URL は https://arxiv.org/abs/2309.00071 です。
この論文を選んだ理由は、長文対応を「モデルを作り直す話」ではなく、既存のRoPEベースLLMを比較的少ない追加学習で長文化する話 として扱っているからです。RAG や圧縮で逃がす前に、まずモデル自体の文脈長をどこまで現実的に伸ばせるかを考えるうえで、かなり実務的なヒントがあります。
どんな技術か
YaRN は、RoPE を使うLLMの位置表現を調整し、元の学習時よりずっと長いコンテキストを扱えるようにする技術です。
長文対応というと、単純には「位置番号をもっと先まで振ればよい」と思いがちです。しかし実際には、学習時に見ていない長さへそのまま外挿すると、位置表現の振る舞いが崩れて性能が急落しやすいです。そこで過去には Position Interpolation（PI）が提案され、位置を圧縮して既存の範囲に収める方法が使われてきました。
YaRN はこの流れをさらに進め、RoPE の全次元を一律に引き伸ばすのではなく、周波数ごとに扱いを変える うえで、さらに attention 側の温度も調整します。結果として、少ない追加学習でも長文性能を伸ばしやすく、しかも元の短文性能を比較的保ちやすい設計になっています。
一言でいうと、YaRN は 「RoPEを雑に拡大せず、壊れやすい成分と残したい成分を分けて長文化する方法」 です。
課題
この技術が解決しようとしているのは、RoPEベースLLMが学習時より長い入力をうまく扱えない問題です。
何が難しいのかというと、Transformer の位置表現は単に「何番目のトークンか」を示すだけではなく、近いトークン同士の相対関係と、遠いトークンを見分けるための絶対寄りの情報を同時に持っているからです。長文対応のために位置を単純圧縮すると、近距離で効いていた細かい位置差まで鈍ってしまいます。
既存の Position Interpolation では、全てのRoPE次元を同じ比率でスケーリングします。これは実装しやすい反面、周波数の高い成分まで一律に圧縮してしまい、局所的な順序や近傍関係を読む能力が落ちやすい という限界があります。論文でも、PI 系の方法はスケール係数が大きくなると出力品質が崩れやすいと指摘されています。
なぜこの課題を解く必要があるのかというと、実際のAIシステムでは長い入力をそのまま読みたい場面が増えているからです。長いコードベース、複数文書の横断要約、長い会話履歴、レポートや議事録の読解などでは、検索だけで必要箇所を完全に切り出せるとは限りません。モデル自身が長い文脈を扱えると、前処理の難しさを減らせます。
また、既存モデルを生かしたまま長文化できるなら、ゼロから長文モデルを事前学習し直すより現実的です。特にオープンモデルを自前で運用するチームにとっては、追加ファインチューニングと位置表現の差し替えだけで文脈長を伸ばせるか は重要な論点です。
用語解説
- RoPE（Rotary Position Embedding）
- クエリとキーに回転をかける形で位置情報を埋め込む方式です。LLaMA 系を含む多くのLLMで使われており、YaRN はこの RoPE の拡張方法そのものを改良します。
- Position Interpolation（PI）
- 長い位置番号を、元の学習済みコンテキスト範囲へ圧縮して入れる方法です。長文化の出発点として重要ですが、全次元を一律に縮めるため、高周波成分まで弱めてしまう点が YaRN の比較対象になります。
- 周波数成分
- RoPE の各次元が持つ回転の細かさです。高周波ほど近いトークン間の差を細かく表現しやすく、低周波ほど長距離の位置変化をゆるく表します。YaRN はこの違いを使って、どの次元をどれだけ補間するかを変えます。
- コンテキスト長拡張
- 学習時の上限より長い入力を扱えるようモデルを調整することです。単にメモリを増やす話ではなく、位置表現が長さの外挿に耐えるよう再設計する必要があります。
- Sliding Window Perplexity
- 長文全体を少しずつ窓でずらしながら計算する perplexity 指標です。コンテキスト上限より長い文書に対しても、長文でどれだけ自然に予測できるかを比較しやすく、YaRN の長文評価でも中心的に使われています。
技術の仕組み
YaRN の仕組みは、RoPE のスケーリングをより細かく設計し、そのうえで attention の温度も合わせて調整する、という二段構えです。
基本アイデア
ベースラインの PI では、元のコンテキスト長を L、拡張後を L' とすると、スケール係数 s = L' / L を使って位置を一律に圧縮します。これは「全部まとめて元の座標系へ押し込む」やり方です。
YaRN 系の考え方では、それでは粗すぎると見ます。RoPE の次元ごとに波長が異なるため、もともと短距離の相対関係を細かく持っていた次元まで強く圧縮すると、近場の情報がにじみます。逆に、長距離用のゆっくりした次元は、長文化のためにしっかり補間したほうがよいです。
モデル構造
YaRN は新しい層や外部モジュールを追加する手法ではありません。対象は RoPE の計算方法と attention のスケーリング です。そのため、LLaMA や Llama 2 のような既存のRoPEベースモデルを大きく改造せずに使えます。
論文では、Llama 2 の 7B と 13B を 4k から 64k、さらに 128k へ拡張して評価しています。モデル本体のアーキテクチャは変えず、埋め込み周波数の計算と attention の扱いだけを変更しています。
周波数ごとに補間の強さを変える
YaRN の土台になっているのは、論文中で先に説明される NTK-by-parts という考え方です。ここでは各RoPE次元について、元の文脈長 L に対してその次元が何回回転するかを比率 r = L / λ で見ます。λ はその次元の波長です。
この比率を使うと、ざっくり次の方針が立てられます。
高周波の次元はなるべく触らない
波長が短い、つまり高周波の次元は、近い位置同士の違いを細かく表す役割が強いです。ここを大きく補間すると、局所的な順序情報が崩れやすくなります。そのため YaRN 系では、高周波側は補間を弱めるか、ほぼ触らない方向に寄せます。
低周波の次元はしっかり補間する
一方で低周波の次元は、もともと長い距離をゆるく表しているので、長文化のために補間しても壊れにくいです。そこで低周波側では PI に近い形でスケーリングをかけ、長い位置まで届くようにします。
中間帯は線形に混ぜる
論文では、比率 r に対してしきい値 α と β を置き、その間では線形ランプ関数で補間率をなめらかに切り替えます。これにより、「全部補間する」「全部そのまま」の二択ではなく、周波数帯ごとに滑らかに最適化する 形になります。
Attention 温度の調整
YaRN の本体は、上の NTK-by-parts に加えて attention 側へ温度 t を入れる点です。論文では attention の softmax を
softmax(q^T k / (t sqrt(D)))
の形に調整します。
これをそのまま読むと attention 実装を直接いじるように見えますが、論文の工夫は、RoPE の複素埋め込み側を一定倍率でスケーリングすることで、attention カーネル自体を変更せずに同等の効果を出す ところにあります。つまり FlashAttention 2 のような高速実装と相性を崩しにくいです。
LLaMA / Llama 2 系に対しては、論文は経験則として
sqrt(1 / t) = 0.1 ln(s) + 1
を推奨しています。拡張倍率 s が大きいほど、attention のエントロピーを調整する必要がある、という見方です。
学習方法
論文では、Llama 2 7B / 13B を PG19 データセットで追加学習しています。学習率は 2e-5、ウォームアップ 20 ステップ、AdamW を使用し、s=16 のモデルは 400 ステップ、s=32 のモデルは s=16 のチェックポイントから追加 200 ステップで学習しています。
ここで重要なのは、128k モデルでも 128k 文脈で直接長く学習していない点です。s=32 モデルも 64k 文脈データで学習しつつ、128k へ外挿できることを示しています。これは、長文学習データや学習コストが限られる現場ではかなり重要です。
推論方法
推論時は、拡張後のコンテキスト長に対応した RoPE と attention スケーリングをそのまま使います。さらに論文は、シーケンス長に応じてスケール係数を動的に変える Dynamic Scaling にも触れています。これは YaRN そのものではありませんが、長さが毎回変わる自動回帰生成で、短文側の性能低下を抑えつつ長文側で滑らかに劣化させる考え方です。
実装上の注意として、KV キャッシュを使う場合は、RoPE 適用後ではなく、適用前のキー・バリューをどう扱うかに注意が必要だと論文は述べています。スケール係数が変わると、過去トークンの RoPE も変わるためです。
実験と結果
論文では、長文で本当に読めるようになるのか、短文性能を壊しすぎないか、計算コストはどうか、という3点を中心に検証しています。
何を検証したのか
主な検証は次の4つです。長文データ上で perplexity がどう変わるか、passkey retrieval のような長文中の情報検索課題で機能するか、一般ベンチマークで元の性能をどれだけ保てるか、そして既存法よりどれだけ少ない計算で学習できるかです。
どんなデータセットや評価指標を使ったのか
長文言語モデリングでは GovReport と Proof-pile を使い、sliding window perplexity で評価しています。Passkey Retrieval では、長い無意味文中に埋め込まれた5桁の数字を取り出せるかを見ています。短文側の汎化には Hugging Face Open LLM Leaderboard の ARC-c、HellaSwag、MMLU、TruthfulQA を使っています。
長文 perplexity の結果
Llama 2 7B を 4k から 64k に拡張した s=16 モデルは、Proof-pile 128k 文書評価で 8k から 64k までは perplexity が 3.51 -> 2.42 と安定し、64k 学習だけで長文に耐えています。一方、同じ 7B の s=32 モデルは 400+200 ステップで 128k 評価時 perplexity 2.37 を出しており、64k データから 128k へ外挿できています。
13B でも傾向は同じで、s=32 モデルは 128k 評価時に perplexity 2.24 を記録しています。論文の主張は、より長い学習文脈を大量に用意しなくても、段階的に長文化できる という点です。
既存法との比較
LLaMA 7B を 2k から 32k に拡張した条件で、PI、NTK-aware、NTK-by-parts、YaRN を比較すると、YaRN は同じ 400 ステップ学習で他手法より低い perplexity を維持したと報告されています。図では、非ファインチューニング時もファインチューニング時も YaRN が一貫してよい傾向を示しています。
passkey retrieval でも、同じ学習予算なら YaRN が高い精度を示しました。つまり YaRN は単に次トークン予測が自然なだけでなく、長文の中から必要な位置情報を引き当てる能力 も改善しています。
短文ベンチマークへの影響
文脈長を伸ばすと、元の短文性能が落ちるのではないかという懸念があります。Llama 2 7B では、元モデルの HellaSwag 77.8 に対して YaRN s=16 は 78.8、MMLU は 43.8 から 42.5、TruthfulQA は 39.0 から 38.2 でした。13B でも大きな崩れはなく、短文性能の劣化は比較的小さいです。
もちろん完全に無傷ではありませんが、PI のように広い指標で大きく落とすよりはかなり穏やかです。短文ベンチマークを見ても、YaRN は 長文化と元性能維持のバランスがよい と言えます。
計算効率
論文の要約では、YaRN は以前の方法より 必要トークン数が約10分の1、学習ステップが約2.5分の1 で済むとされています。計算時間の表でも、LLaMA 7B を 32k 相当に伸ばす YaRN は 128 A100 GPU-hours、Llama 2 7B を 64k にする YaRN は 256 A100 GPU-hours とされ、PI 系の 640 A100 GPU-hours よりかなり軽いです。
ここで効いているのは、新規アーキテクチャを入れていないことと、RoPE と attention スケーリングの変更が推論・学習の追加オーバーヘッドをほぼ生まないことです。
何に使える？
YaRN が役立つのは、「長い入力をそのまま読ませたいが、長文専用モデルをゼロから作るほどではない」場面です。
長い社内文書や議事録の読解
RAG では切り出しに失敗しやすいケースでも、モデル自体が 64k や 128k を読めれば、複数資料をまとめて解釈しやすくなります。社内規程、議事録、仕様変更履歴のように文脈が散っている情報で効きやすいです。
コードベース全体の解析
大きなリポジトリでは、1ファイルだけ見ても判断できないことが多いです。YaRN 的な長文化を施したコード向けLLMなら、依存関係や関連ファイルをまとめて入れやすくなります。特にアーキテクチャ理解、影響範囲分析、長いPRレビュー支援で使い道があります。
長い会話履歴を持つAIエージェント
会話履歴、ツール実行結果、途中メモを長く保持したいエージェントでは、短い文脈ではすぐ取りこぼしが起きます。YaRN は推論方式ではなくモデル側の文脈長を伸ばす発想なので、メモリ管理を複雑にしすぎずに履歴保持量を増やせます。
RAG の前提条件を緩める
これは論文そのものの主張ではなく実務上の考察ですが、モデルの文脈長が伸びると、RAG の検索粒度をそこまで攻めなくてよくなる場面があります。多少広めに候補を入れても処理しやすくなるため、検索漏れへの耐性を持たせやすいです。
開発や事業へのヒント
この論文から得られるヒントは、長文対応を「新しい巨大モデルを待つ話」ではなく、既存モデルの再利用戦略として考えられることです。
小さな改造でも体験は大きく変わる
YaRN はモデルの本体構造を変えません。つまり、既存のオープンモデル資産や推論基盤をできるだけ維持しながら、長文対応のプロトタイプを作れます。小規模チームでも、ゼロから事前学習するのではなく、位置表現の再設定と軽い追加学習で勝負する という選択肢を持てます。
長文対応は検索だけでなくモデル側でも稼げる
長文アプリを作ると、どうしても RAG や要約圧縮ばかり考えがちです。しかし YaRN は、モデル自体の受け皿を広げる価値を示しています。検索と要約の設計コストが高い領域では、先にモデル側の文脈長を伸ばすほうが全体設計が単純になるかもしれません。
段階的拡張はプロダクト運営と相性がよい
論文では 64k 学習から 128k へ外挿しています。これは事業的にも重要で、最初から最大長を狙わず、32k、64k、128k と順に広げる運用がしやすいことを意味します。利用状況を見ながら段階投資できるので、GPU コストと需要のバランスを取りやすいです。
FlashAttention など既存最適化と併用しやすい
YaRN は attention 実装そのものを特殊化する手法ではないため、既存の高速カーネルと組み合わせやすいです。プロダクト開発では、精度向上策が推論基盤を壊すと導入が止まりがちですが、YaRN はそこが比較的扱いやすいです。
限界
まず、計算量そのものは減りません。YaRN は「より長く読めるようにする」技術であって、長文推論の二乗コストを解消するものではありません。64k や 128k を読ませれば、その分のメモリとレイテンシは重くなります。
次に、RoPE ベースモデルが前提です。ALiBi や別系統の位置表現にはそのまま適用できません。また、論文での中心検証は LLaMA / Llama 2 系なので、他モデルで同じハイパーパラメータがそのまま最適とは限りません。
さらに、長文化の効果は追加学習データの質にも依存します。論文では PG19 を使っていますが、実務でコードモデルや業務文書モデルへ適用するなら、長文側の分布がどれだけ近いかは気にする必要があります。
実装面では、KV キャッシュや dynamic scaling を含む推論コードの扱いに注意が必要です。特に長さに応じてスケールを変える実装では、キャッシュの再利用方法を誤ると理論どおり動きません。
最後に、短文性能は「ほぼ維持」であって「完全維持」ではありません。短文タスクの比重が高いプロダクトでは、長文化で得る価値と、短文ベンチマークの微小な低下を比較して判断する必要があります。
よくある質問
Q. YaRN は Position Interpolation と何が違うのですか？
A. PI は RoPE の全次元を同じ比率で圧縮しますが、YaRN は高周波と低周波で扱いを分けます。さらに attention 温度も調整するので、長文性能と短文性能のバランスがよくなりやすいです。
Q. YaRN を入れると推論が速くなるのですか？
A. 直接速くなるわけではありません。YaRN は文脈長を伸ばす技術です。長文を扱えるようにはなりますが、長文を読む計算コスト自体は依然として重いです。
Q. RAG があれば YaRN は不要ですか？
A. 必ずしも不要ではありません。RAG は必要箇所を絞るのに強いですが、検索漏れや文脈分断が起きることがあります。YaRN のような長文化は、モデル自身がより広い文脈を保持できるため、RAG を補完する役割を持てます。
Q. 小規模チームでも活かせますか？
A. はい。論文の完全再現は計算資源が要りますが、「既存のRoPEモデルを段階的に長文化する」「最初は 32k や 64k から試す」という考え方自体は、小規模チームでも十分参考になります。
Q. どんなケースでは効果が薄いですか？
A. そもそも数千トークン程度しか扱わないアプリでは恩恵が小さいです。また、検索で十分に問題が解けるタスクでは、モデル側を長文化するより RAG 改善のほうが費用対効果が高い場合があります。
今日の学び
この論文は、RoPEベースLLMが学習時より長い文脈へうまく一般化しにくい課題を扱っています。そこに対して YaRN は、周波数ごとに補間の強さを変える RoPE 拡張と attention 温度調整を組み合わせ、少ない追加学習で長文化しようとしました。
そこから得られるヒントは、長文対応は新しいモデルを待つだけでなく、既存モデルの位置表現設計を見直すだけでも前進できるということです。長文RAG、コード解析、エージェント開発では、検索だけでなくモデル側の文脈長拡張も設計の選択肢として持っておく価値があります。