Jet-Long: Efficiënte lange-contextuitbreiding met dynamische bifocale RoPE

Samenvatting

Moderne LLMs worden steeds vaker ingezet in toepassingen met lange context, zoals retrieval-augmented generation, repository-niveau coderen en agentische workflows, waarbij de opgebouwde redeneer- en gereedschapsporen de input routinematig een orde van grootte voorbij het pretrainingsvenster duwen, waardoor nul-shot contextuitbreiding het dominante implementatiepad wordt voor open-gewichts checkpoints. De meeste bestaande nul-shot methoden leggen vooraf één enkele herschalingsfactor vast, waardoor een agressieve factor de getrouwheid voor korte context opoffert, terwijl een conservatieve factor bij lange context faalt. Wij stellen Jet-Long voor, een afstemmingsvrije nul-shot methode die een lokaal RoPE-getrouw venster combineert met een venster voor lange afstand waarvan de herschalingsfactor zich dynamisch aanpast aan de huidige sequentielengte, waarbij het basismodel exact wordt hersteld bij korte invoer en soepel wordt geëxtrapoleerd bij lange invoer. Een inclusie-exclusie aandachtsmerge en een on-the-fly RoPE-correctierotatie maken de bifocale constructie in wezen gratis tijdens inferentie; samengevoegd in een enkele CuTe-kernel bereikt lange-context prefill tot 1,39 keer de FA2-doorvoer op H100 (benadert de Hopper-only FA4), en enkel-batch generatie introduceert ≤ 4% overhead bij elke lengte. Op Qwen3-1.7B/4B/8B tot 128K context leidt Jet-Long tot +4,79/+2,18/+2,03 pp op RULER ten opzichte van de sterkste basislijn bij 1.7B/4B/8B, behaalt het de beste algehele nauwkeurigheid op HELMET-RAG (een benchmark die door HELMET is geïdentificeerd als de meest efficiënte voorspeller van downstream lange-context prestaties) en bereikt het de laagste PG-19 perplexiteit. Jet-Long generaliseert ook naar hybride aandachtsarchitecturen zoals Jet-Nemotron voor verdere verbetering van lange context zonder hertraining, en blijft hyperparameter-resistent voor eenvoudige implementatie.

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.