技術要約: Jet-Long

問題提起

現代の大規模言語モデル（LLM）は、検索拡張生成（RAG）、リポジトリレベルのコーディング、エージェント的ワークフローといった長文脈アプリケーションへの導入が進んでいます。これらのシナリオでは、蓄積された推論トレースやツールの出力により、入力長が事前学習ウィンドウ（例：32Kの学習ウィンドウに対し100K以上）を1桁上回ることが頻繁にあります。

長文脈に対して直接モデルを学習させることは、二次的な計算コストと長文脈データの不足により、極めて高価です。また、ファインチューニングを行うと短文脈での性能が低下することがよくあります。その結果、ゼロショット・コンテキスト拡張（ファインチューニングなしでのコンテキスト拡張）が、オープンウェイトのチェックポイントにおける主要なデプロイメントパスとなっています。

しかし、回転位置埋め込み（RoPE）を使用する標準的なTransformerベースのLLMは、以下の2つの主要な失敗モードにより、学習ウィンドウを超えて汎化することに失敗します：

1. 位置の分布外（OOD）: 学習時に見られなかったシーケンス位置において、低周波のRoPEコンポーネントが学習分布の外側の回転角を生成し、注意スコアを不安定にします。
2. アテンションの拡散: キーセットが増大するにつれ、ソフトマックス分布が平坦化し、無関係なトークンに確率質量が分散します。これに伴い、中間コンテキストへの注意が不足する「U字型」の positional bias が発生することがよくあります。

既存のゼロショット手法（NTK、YaRN、Self-Extend、DCAなど）は、通常、単一のスケーリング係数またはグルーピングサイズを事前に固定します。これはトレードオフを強いることになります。つまり、アグレッシブな係数は短文脈の忠実度を犠牲にし、保守的な係数は長文脈でのクリーンな外挿に失敗します。

手法: Jet-Long

Jet-Longは、ローカルなRoPEに忠実なウィンドウと、現在のシーケンス長に応じてスケーリング係数が動的に適応するロングレンジ・ウィンドウを組み合わせた、チューニングフリーかつゼロショットのコンテキスト拡張手法を提案します。

1. 動的バイフォーカルRoPE (Dynamic Bifocal RoPE)

核心となる革新技術は、固定係数のトレードオフを回避する動的バイフォーカルマッピングです。

• ローカル・ウィンドウ: サイズ $w_0$ のウィンドウは従来のRoPEを保持し、短文脈入力に対してベースモデルの事前学習挙動を正確に維持します。
• リモート・ウィンドウ: 位置 $w_0$ を超える位置に対して、動的な位置マッピング $f(\cdot)$ が適用されます。連続的に周波数を補間する手法とは異なり、Jet-Longは**離散的な位置エイリアシング（discrete position aliasing）**を使用します。
• 動的グルーピング: グループサイズ $G$ は、現在のシーケンス長 $L$ と事前学習ウィンドウ $w_{pretrained}$ に基づいて動的に計算されます：
  $$G = \max\left(1, \left\lceil \frac{L}{w_{pretrained}} \right\rceil\right)$$
  リモートマッピングは床関数（floor division）を用います：$f(x) = \lfloor x/G \rfloor$。これにより、すべてのリモート回転角が分布内（モデルの事前学習済みグリッド上）に収まりつつ、最小限必要な圧縮を適用することで位置解像度を最大化します。

2. 推論効率化メカニズム

計算効率を確保し、再学習や重いKVキャッシュの変更を必要としないよう、Jet-Longは2つの手法を採用しています。

• オンザフライ修正回転 (On-the-Fly Correction Rotation): KVキャッシュは、絶対位置における未圧縮のベースキーを保持します。生成中、Jet-Longはキャッシュを書き換える代わりに、アクティブなクエリとキャッシュされたキーに対して、オンザフライで修正回転 $\Delta$ を適用します。RoPEの回転は加法的に合成されるため、物理的なキャッシュに触れることなく、レジスタ内で外挿されたベクトルを再構成できます。これにより、ストリーミング生成中に停止することなく、長さの境界を越えた生成が可能になります。
• 包含–除外アテンション・マージ (Inclusion–Exclusion Attention Merge): プリフィル中に完全な二次アテンション行列を実体化することを避けるため、Jet-Longは包含–除外原理を用いて3つの標準的なFlashAttentionパスを結合します：
  1. ローカル・ウィンドウ・アテンション（ベースのクエリ/キー）。
  2. フル・コーザル・アテンション（リモートのクエリ/キー）。
  3. ローカル・ウィンドウ・アテンション（リモートのクエリ/キー）。
     LogSumExp統計量を用いてマージを安定させることで、ブール値のマスク行列を使用せずに、正確な距離ベースのルーティング出力を計算し、FlashAttentionのメモリ効率を維持します。

これらのコンポーネントは単一の CuTe カーネルに融合され、カーネル起動のオーバーヘッドを最小限に抑えます。

主な貢献

1. Jet-Long アルゴリズム: リモートウィンドウのスケーリング係数がシーケンス長に応じて動的に適応する、チューニングフリーのバイフォーカル手法。これにより、リモート回転を分布内に保ちつつ、ネイティブなコンテキスト内ではベースモデルを正確に再現します。
2. 効率的な実装: KVキャッシュに手を加えない包含–除外アテンション・マージとオンザフライ修正回転を採用。これらを単一のCuTeカーネルに融合することで、H100における長文脈時のFlashAttention-2 (FA2) プリフィル・スループットの1.28–1.39倍を達成し、シングルバッチ生成時のオーバーヘッドを4%以下に抑えます。
3. 実証的検証: Qwen3-1.7B/4B/8Bを用いた最大128Kコンテキストの広範な評価により、RULER、HELMET-RAG、PG-19ベンチマークにおいて、強力なゼロショット・ベースライン（DNTK、YaRN、DCA、Self-Extend）に対する優位性を実証しました。
4. 汎用性: 本手法は変更なしでハイブリッド・アテンション・アーキテクチャ（例：Jet-Nemotron）に転移可能であり、単一のハイパーパラメータ $w_0$ の選択に対して堅牢です。

実験結果

Qwen3ベースモデル（1.7B, 4B, 8B）を用い、32Kのネイティブウィンドウを128Kまで拡張して評価しました：

• RULER ベンチマーク: Jet-Longは、1.7B, 4B, 8Bモデルにおいて、最も強力なベースラインに対し、それぞれ**+4.79, +2.18, +2.03 パーセントポイント (pp)** 上回りました。ベースモデルの性能を32Kの学習ウィンドウ内で維持しつつ、拡張された長さにおいてすべてのベースラインを凌駕しています。
• HELMET-RAG: このベンチマークにおいて最高の総合精度を達成しました。HELMETの研究によれば、これはダウンストリームの長文脈性能を予測する上で最も効率的な指標です。
• PG-19 Perplexity: すべての外挿手法の中で最も低いパープレキシティを達成しました。ベースモデルが崩壊する一方で（例：1.7Bモデルで128K時にパープレキシティ >70）、Jet-Longは低いパープレキシティを維持しています（例：1.7Bで11.41）。
• ハイブリッド・アーキテクチャ: Jet-Nemotron（ソフトマックス・アテンションと線形アテンションの交互配置）に適用したところ、ベースのハイブリッドモデルが崩壊する128Kにおいて、Jet-LongはRULERの精度を**+10.01 pp (2B)** および +11.31 pp (4B) 向上させました。
• アブレーション研究:
  • ハイパーパラメータの耐性: ローカル・ウィンドウサイズ $w_0$ は頑健であり、512から4096の範囲の値であれば、最高性能の±1–2 pp以内に収まります。
  • エイリアシング vs 補間: 位置エイリアシング（Jet-Longのアプローチ）は、周波数補間（YaRNスタイル）を大幅に上回る性能を示しました（例：64Kで+6.99 pp）。これは、回転を離散的な事前学習済みグリッド上に保持することの利点を確認するものです。

意義

本論文は、Jet-Longが、長文脈学習の計算コストや固定係数外挿による忠実度の喪失を伴うことなく、RoPEレベルでの位置OOD失敗モードに対処することを主張しています。スケーリング係数を動的にし、融合カーネルを通じて推論構築を「実質的に無料」にすることで、Jet-Longは、オープンウェイトLLMを長文脈シナリオで展開するための実用的かつ高性能なパスを提供します。これは、アテンションの拡散に対処するアーキテクチャ上の代替案（スパースアテンションや状態空間モデルなど）を補完するものであり、再学習を必要とせず、モデルサイズやハイブリッドアーキテクチャを横断して汎用するゼロショット・ソリューションを提供します。