🤖 machine learning

Jet-Long: Efficient Long-Context Extension with Dynamic Bifocal RoPE

Jet-Long is a tuning-free, zero-shot method that dynamically adapts RoPE rescaling factors to balance short-context fidelity and long-context extrapolation, achieving state-of-the-art performance on long-context benchmarks with minimal inference overhead through an efficient bifocal attention mechanism.

Original authors: Haozhan Tang, Zerui Wang, Yuxian Gu, Song Han, Han Cai

Published 2026-07-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haozhan Tang, Zerui Wang, Yuxian Gu, Song Han, Han Cai

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-smart robot brain (a Large Language Model) that was trained to read stories up to a certain length, say 32,000 words. Now, you want to ask it to read a whole novel or a massive code repository that is four times longer. If you just hand it the long text, the robot gets confused. It's like trying to read a book where the page numbers suddenly jump from 100 to 1,000,000; the robot's internal compass (called RoPE) spins wildly, and it starts hallucinating or forgetting the middle of the story.

For a while, scientists tried to fix this by picking one "magic number" to stretch the robot's vision. But this was a lose-lose game: if you stretched it too much to see the end of the book, the robot forgot the beginning. If you kept it tight to remember the start, it couldn't see the end.

Enter Jet-Long, a new method that acts like a dynamic bifocal lens for the robot's eyes.

The Bifocal Glasses Trick

Instead of forcing the robot to use one stretched view for the whole book, Jet-Long gives it two pairs of glasses at the same time:

The Close-Up Lens: For the first chunk of text (the "local window"), the robot sees everything exactly as it was trained to see it. No stretching, no distortion. This keeps the robot's memory sharp for the recent past.
The Zoom-Out Lens: For the rest of the text (the "remote window"), the robot uses a special trick. It doesn't just stretch the view; it dynamically groups the pages together. As the book gets longer, the robot automatically adjusts how many pages it groups into one "super-page."

The magic is that this grouping factor isn't fixed. It changes on the fly depending on how long the text is right now. If the text is short, the robot sees every word individually. If the text is huge, it groups words just enough so the robot's internal compass stays calm and doesn't spin out of control. This means the robot stays perfectly accurate for short tasks but can still read massive documents without getting lost.

The "Free Lunch" of Speed

Usually, doing two different calculations at once (looking close and looking far) would make the robot twice as slow. But the authors of this paper found a clever way to merge these calculations. They built a special engine (a "fused kernel") that does the math for both lenses in a single pass.

Think of it like a chef who usually has to chop vegetables, then boil them, then fry them in three separate pots. Jet-Long is like a magical pot that does all three steps at once without losing any flavor.

The Result: When the robot reads a massive 128,000-token context, Jet-Long is actually 1.39 times faster at the start (the "prefill" phase) than the standard method (FlashAttention-2) on an H100 chip.
The Catch: When the robot is generating new text one word at a time, it's only about 4% slower than the standard method. That's a tiny price to pay for being able to read a whole library instead of just a pamphlet.

What They Tested (And What They Didn't)

The team tested this on three different sizes of robot brains (1.7 billion, 4 billion, and 8 billion parameters) and found that Jet-Long consistently beat the previous best methods.

On a test called RULER (which checks if the robot can find hidden needles in a haystack of text), Jet-Long scored 4.79 points higher on the smallest model and 2.03 points higher on the largest model compared to the best competitor.
On a test called HELMET-RAG (which simulates real-world search tasks), it achieved the best overall accuracy.
On a test called PG-19 (reading old books), it kept the robot's confusion (perplexity) at the lowest levels, whereas other methods let the robot's confusion skyrocket as the text got longer.

Important Note on Limits: The paper explicitly rules out the idea that you need to retrain the robot from scratch to make it work. Jet-Long works on "zero-shot" models, meaning you can take a pre-trained robot and just plug in this lens. However, the paper also notes that while Jet-Long fixes the "spinning compass" problem, it doesn't fix a different problem called "attention diffusion" (where the robot gets overwhelmed by too many options). That problem is usually solved by changing the robot's architecture entirely, not just by adding a lens.

The Verdict

Jet-Long isn't a magic wand that makes robots infinitely smart, but it is a highly efficient, tuning-free tool that lets existing robots handle text lengths up to 128,000 tokens without losing their mind. It keeps the robot's short-term memory perfect while dynamically adjusting its long-term vision, all while running almost as fast as the standard methods. The authors measured this on real hardware (H100 GPUs) and found it works across different model sizes and even on hybrid robot brains that mix different types of attention. It's a solid, proven upgrade for anyone trying to make AI read longer.

Technical Summary: Jet-Long

Problem Statement

Modern Large Language Models (LLMs) are increasingly deployed in long-context applications such as retrieval-augmented generation (RAG), repository-level coding, and agentic workflows. In these scenarios, accumulated reasoning traces and tool outputs often push input lengths an order of magnitude beyond the pretraining window (e.g., 100K+ tokens vs. a 32K training window).

Directly training models on long contexts remains prohibitively expensive due to quadratic compute costs and scarce long-context data, while fine-tuning often degrades short-context performance. Consequently, zero-shot context extension (extending context without fine-tuning) has become the dominant deployment path for open-weight checkpoints.

However, vanilla Transformer-based LLMs using Rotary Position Embedding (RoPE) fail to generalize beyond their training window due to two primary failure modes:

Position Out-of-Distribution (OOD): At sequence positions unseen during training, low-frequency RoPE components produce rotation angles outside the training distribution, causing erratic attention scores.
Attention Diffusion: As the key set grows, the softmax distribution flattens, dispersing probability mass over irrelevant tokens, often accompanied by a "U-shaped" positional bias that under-attends to middle-context information.

Existing zero-shot methods (e.g., NTK, YaRN, Self-Extend, DCA) typically fix a single rescaling factor or grouping size upfront. This forces a trade-off: an aggressive factor sacrifices short-context fidelity, while a conservative one fails to extrapolate cleanly at long contexts.

Methodology: Jet-Long

Jet-Long proposes a tuning-free, zero-shot context extension method that pairs a local RoPE-faithful window with a long-range window whose rescaling factor adapts dynamically to the current sequence length.

1. Dynamic Bifocal RoPE

The core innovation is a dynamic bifocal mapping that avoids the fixed-factor trade-off:

Local Window: A window of size $w_0$ retains classic RoPE, preserving the base model's pretraining behavior exactly for short inputs.
Remote Window: For positions beyond $w_0$ , a dynamic position mapping $f(\cdot)$ is applied. Unlike methods that interpolate frequencies continuously, Jet-Long uses discrete position aliasing.
Dynamic Grouping: The group size $G$ is calculated dynamically based on the current sequence length $L$ and the pretrained window $w_{pretrained}$ :
$G = \max\left(1, \left\lceil \frac{L}{w_{pretrained}} \right\rceil\right)$
The remote mapping is a floor division: $f(x) = \lfloor x/G \rfloor$ . This ensures that every remote rotation angle remains in-distribution (on the model's pretrained grid) while maximizing positional resolution by applying the minimum necessary compression.

2. Inference Efficiency Mechanisms

To ensure the method is computationally efficient and does not require retraining or heavy KV cache modifications, Jet-Long employs two key techniques:

On-the-Fly Correction Rotation: The KV cache stores uncompressed base keys at their exact absolute positions. During generation, instead of rewriting the cache, Jet-Long applies a correction rotation $\Delta$ to the active query and cached keys on the fly. Since RoPE rotations compose additively, this reconstructs the extrapolated vectors in registers without touching the physical cache, allowing streaming generation across length boundaries without stalling.
Inclusion–Exclusion Attention Merge: To avoid materializing the full quadratic attention matrix during prefill, Jet-Long combines three standard FlashAttention passes using an inclusion–exclusion principle:
1. Local window attention (base queries/keys).
2. Full causal attention (remote queries/keys).
3. Local window attention (remote queries/keys).
  By stabilizing the merge with LogSumExp statistics, the method computes the exact distance-based routing output without a boolean mask matrix, retaining FlashAttention's memory efficiency.

These components are fused into a single CuTe kernel, minimizing kernel launch overhead.

Key Contributions

Jet-Long Algorithm: A tuning-free bifocal method where the remote-window rescaling factor adapts dynamically to sequence length, keeping remote rotations in-distribution while reproducing the base model exactly within its native context.
Efficient Implementation: An inclusion–exclusion attention merge and on-the-fly correction rotation that leave the KV cache untouched. Fused into a single CuTe kernel, this achieves 1.28–1.39× FlashAttention-2 (FA2) prefill throughput on H100 at long contexts and incurs ≤4% overhead during single-batch generation.
Empirical Validation: Extensive evaluation on Qwen3-1.7B/4B/8B up to 128K context, demonstrating superior performance over strong zero-shot baselines (DNTK, YaRN, DCA, Self-Extend) on RULER, HELMET-RAG, and PG-19 benchmarks.
Generalization: The method transfers unchanged to hybrid attention architectures (e.g., Jet-Nemotron) and remains robust to the choice of the single hyperparameter $w_0$ .

Experimental Results

Evaluated on Qwen3 base models (1.7B, 4B, 8B) with a 32K native window extended to 128K:

RULER Benchmark: Jet-Long leads the strongest baselines by +4.79, +2.18, and +2.03 percentage points (pp) on the 1.7B, 4B, and 8B models, respectively. It preserves base model performance within the 32K training window and outperforms all baselines at extended lengths.
HELMET-RAG: Achieves the best overall accuracy on this benchmark, identified by the HELMET study as the most efficient predictor of downstream long-context performance.
PG-19 Perplexity: Attains the lowest perplexity among all extrapolation methods. While the base model collapses (e.g., perplexity >70 at 128K for 1.7B), Jet-Long maintains low perplexity (e.g., 11.41 for 1.7B).
Hybrid Architectures: Applied to Jet-Nemotron (interleaving softmax and linear attention), Jet-Long lifts RULER accuracy by +10.01 pp (2B) and +11.31 pp (4B) at 128K, where the base hybrid model otherwise collapses.
Ablation Studies:
- Hyperparameter Resilience: The local window size $w_0$ is robust; values between 512 and 4096 stay within 1–2 pp of the best performance.
- Aliasing vs. Interpolation: Position aliasing (Jet-Long's approach) outperforms frequency interpolation (YaRN-style) by significant margins (e.g., +6.99 pp at 64K), confirming the benefit of keeping rotations on the discrete pretrained grid.

Significance

The paper claims that Jet-Long addresses the position-OOD failure mode at the RoPE level without the computational cost of long-context training or the fidelity loss of fixed-factor extrapolation. By making the rescaling factor dynamic and the inference construction "essentially free" via fused kernels, Jet-Long provides a practical, high-performance deployment path for open-weight LLMs in long-context scenarios. It complements architectural alternatives (like sparse attention or state-space models) that address attention diffusion, offering a zero-shot solution that requires no retraining and generalizes across model sizes and hybrid architectures.