Jet-Long: Эффективное расширение длинного контекста с динамическим бифокальным RoPE

Аннотация

Современные LLM всё чаще применяются в приложениях с длинным контекстом, таких как генерация с дополнением через поиск, программирование на уровне репозиториев и агентные рабочие процессы, чьи накопленные цепочки рассуждений и следы инструментов регулярно увеличивают входные данные на порядок за пределы окна предварительного обучения, что делает расширение контекста без дополнительного обучения доминирующим путём развёртывания для моделей с открытыми весами. Большинство существующих методов без обучения фиксируют единый коэффициент перемасштабирования заранее, поэтому агрессивный коэффициент жертвует точностью на коротких контекстах, а консервативный — выходит из строя на длинных. Мы предлагаем Jet-Long — метод без дообучения, который сочетает локальное окно, верное RoPE, с дальним окном, коэффициент перемасштабирования которого динамически адаптируется к текущей длине последовательности, точно восстанавливая базовую модель на коротких входных данных и чисто экстраполируя на длинных. Объединение внимания по принципу включения-исключения и поворот коррекции RoPE на лету делают бифокальную конструкцию практически бесплатной при инференсе; при слиянии в единое ядро CuTe префиллинг длинного контекста достигает до 1.39× пропускной способности FA2 на H100 (приближаясь к FA4, доступному только на Hopper), а однопакетная генерация вносит накладные расходы не более 4% на любой длине. На Qwen3-1.7B/4B/8B с контекстом до 128K Jet-Long превосходит RULER на +4.79/+2.18/+2.03~процентных пункта по сравнению с сильнейшим базовым методом при 1.7B/4B/8B, достигает наилучшей общей точности на 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.