YaRN(ヤーン)
Yet another RoPE extensioN(YaRN)の略。RoPEベースのLLMのコンテキスト長を、わずか400ステップ程度の短時間ファインチューニングで最大64〜128倍に拡張する位置エンコーディング補間手法。NTK-Aware ScalingとNTK-by-partsの長所を統合した手法。
YaRNとは
YaRN(Yet another RoPE extensioN、ヤーン)は、2023年8月にEleutherAIの研究者らによって提案されたRoPE(Rotary Position Embedding)の拡張手法である。LLMのコンテキスト長を学習コストを最小限に抑えつつ大幅に伸ばすことを目的とし、既存のPosition Interpolation、NTK-Aware Scaling、NTK-by-partsの知見を統合して改良した。
背景と動機
RoPEは現代のオープンソースLLM(LLaMA、Mistral、Qwen等)で広く採用されている位置エンコーディングだが、学習時のコンテキスト長を超えた外挿(Extrapolation)では性能が急落する。既存のスケーリング手法にはそれぞれ弱点があった。
- Position Interpolation(PI): 短文性能が劣化する
- NTK-Aware Scaling: 追加学習なしで使えるが拡張倍率に限界がある
- NTK-by-parts: 高周波・低周波の分離処理が不完全
YaRNの3つの要素
1. NTK-by-parts補間
RoPEの各次元の周波数を3つの領域に分類し、それぞれ異なるスケーリングを適用する。
| 周波数領域 | 波長 vs コンテキスト長 | スケーリング方針 |
|---|---|---|
| 高周波(近距離) | 波長 < 元のコンテキスト長 | スケーリングなし(そのまま保持) |
| 中周波 | 波長 ≈ 元のコンテキスト長 | 線形補間(PI的) |
| 低周波(遠距離) | 波長 > 元のコンテキスト長 | NTKスケーリング |
2. Attention分布補正(√t スケーリング)
コンテキスト長拡張に伴いAttentionのエントロピーが増大する問題を、温度パラメータtでスケーリングして補正する。具体的にはAttention logitsを√tで除算し、元の分布の鋭さを維持する。
3. 動的スケーリング対応
推論時のコンテキスト長に応じてスケーリング係数を動的に調整するDynamic YaRNモードも提供される。短い入力では元の精度を維持しつつ、長い入力でのみ拡張パラメータを適用する。
実験結果
Mistral-7B(4K→128K)
| ベンチマーク | 元モデル(4K) | PI(128K) | YaRN(128K) |
|---|---|---|---|
| Perplexity(PG-19) | 8.15 | 8.98 | 8.29 |
| Passkey Retrieval | 100%@4K | 89%@128K | 99.8%@128K |
| MMLU(5-shot) | 62.1 | 58.3 | 61.7 |
LLaMA 2-7B(4K→64K)
400ステップのファインチューニング(A100×4で約1時間)で4Kから64Kへの拡張を達成。Perplexityの劣化はPosition Interpolation比で約50%改善。
実装と利用
YaRNはHugging Face Transformersに統合されており、rope_scalingパラメータで指定できる。
主要な設定パラメータは以下の通り。
| パラメータ | 説明 | 推奨値 |
|---|---|---|
| factor | 拡張倍率 | 目標長 / 学習長 |
| original_max_position_embeddings | 元の学習コンテキスト長 | モデル依存 |
| beta_fast | 高周波の閾値 | 32(デフォルト) |
| beta_slow | 低周波の閾値 | 1(デフォルト) |
| attention_factor | √tスケーリング係数 | 自動計算 |
FAQ
Q1: YaRNのファインチューニングにはどのくらいのコストがかかりますか?
A1: 代表的なケースでは400ステップ程度で十分であり、A100×4で約1時間、A100×1でも4時間程度で完了する。学習データは拡張後のコンテキスト長に対応した長文テキストを含む数千サンプルが推奨される。
Q2: どのモデルでYaRNを使えますか?
A2: RoPEベースの位置エンコーディングを使用するモデル全般に適用できる。具体的にはLLaMA系列、Mistral、Qwen、Phi、CodeLlama、Yi等が対応している。ALiBiやSinusoidal Positional Encodingを使用するモデルには適用できない。
Q3: YaRNとContext Window Scalingの他の手法の違いは?
A3: YaRNはNTK-by-parts + Attention分布補正 + 動的スケーリングの3要素を統合しており、単独手法(PI、NTK-Aware)より幅広い拡張倍率で安定した性能を示す。特に16倍以上の大幅拡張ではYaRNの優位性が顕著になる。