Jet-Long: 動的二焦点RoPEによる効率的な長文脈拡張

要旨

近年、大規模言語モデル（LLM）は、検索拡張生成（Retrieval-Augmented Generation）、リポジトリレベルのコーディング、エージェントワークフローといった長文コンテキストアプリケーションにますます導入されている。これらのアプリケーションでは、累積された推論やツールのトレースにより、入力が事前学習ウィンドウを一桁超えることが日常的であり、その結果、ゼロショットコンテキスト拡張がオープンウェイトチェックポイントの主要な展開経路となっている。既存のゼロショット手法のほとんどは、事前に単一の再スケーリング係数を固定するため、積極的な係数は短いコンテキストの忠実性を犠牲にし、保守的な係数は長いコンテキストでは機能しなくなる。本論文では、チューニング不要のゼロショット手法であるJet-Longを提案する。これは、局所的なRoPE忠実ウィンドウと、現在の系列長に動的に適応する再スケーリング係数を持つ長距離ウィンドウを組み合わせることで、短い入力ではベースモデルを完全に復元し、長い入力ではクリーンに外挿する。包除注意マージと、オンザフライで行われるRoPE補正回転により、この二焦点構成は推論時に実質的にオーバーヘッドを生じない。単一のCuTeカーネルに統合され、長コンテキストプレフィルはH100上でFA2スループットの最大1.39倍に達し（Hopper専用のFA4に迫る）、シングルバッチ生成はあらゆる長さにおいて4％以下のオーバーヘッドしか発生させない。最大128KコンテキストのQwen3-1.7B/4B/8Bにおいて、Jet-LongはRULERで最強のベースラインを1.7B/4B/8Bそれぞれ+4.79/+2.18/+2.03pp上回り、HELMET-RAG（HELMETが下流の長コンテキスト性能の最も効率的な予測指標として特定したベンチマーク）で最高の全体精度を達成し、PG-19パープレキシティで最低値を記録した。Jet-Longはまた、Jet-Nemotronのようなハイブリッド注意アーキテクチャにも一般化し、再学習なしで長コンテキスト改善をさらに進め、ハイパーパラメータに対して頑健であり、展開の容易さを維持する。

English

Modern LLMs are increasingly deployed in long-context applications such as retrieval-augmented generation, repository-level coding, and agentic workflows whose accumulated reasoning and tool traces routinely push the input an order of magnitude past the pretraining window, making zero-shot context extension the dominant deployment path for open-weight checkpoints. Most existing zero-shot methods fix a single rescaling factor up front, so an aggressive factor sacrifices short-context fidelity while a conservative one breaks down at long contexts. We propose Jet-Long, a tuning-free zero-shot method that pairs a local RoPE-faithful window with a long-range window whose rescaling factor adapts dynamically to the current sequence length, recovering the base model exactly at short inputs while extrapolating cleanly at long ones. An inclusion-exclusion attention merge and an on-the-fly RoPE correction rotation make the bifocal construction essentially free at inference; fused into a single CuTe kernel, long-context prefill reaches up to 1.39times FA2 throughput on H100 (approaching the Hopper-only FA4), and single-batch generation incurs le 4% overhead at every length. On Qwen3-1.7B/4B/8B up to 128K context, Jet-Long leads RULER by +4.79/+2.18/+2.03~pp over the strongest baseline at 1.7B/4B/8B, achieves the best overall accuracy on HELMET-RAG (a benchmark identified by HELMET as the most efficient predictor of downstream long-context performance) and attains the lowest PG-19 perplexity. Jet-Long also generalizes to hybrid attention architectures such as Jet-Nemotron for further long-context improvement without retraining, and remains hyperparameter-resilient for ease of deployment.