Jet-Long: новый метод расширения контекста LLM до 128K токенов без дообучения

Расширение контекстного окна больших языковых моделей (LLM) — одна из ключевых задач современного NLP. Большинство существующих подходов требуют дообучения или фиксируют коэффициент масштабирования, что приводит к потере точности на коротких контекстах. Новый метод Jet-Long предлагает элегантное решение: динамическую бифокальную RoPE, которая адаптируется к длине последовательности без дополнительного обучения. Это открывает возможности для эффективной работы с длинными документами, RAG и агентными сценариями.

Что такое Jet-Long и как он работает

Jet-Long — это метод расширения контекста для LLM, основанный на динамической бифокальной Rotary Position Embedding (RoPE). В отличие от традиционных подходов, которые используют фиксированный коэффициент масштабирования, Jet-Long комбинирует два окна внимания: локальное с верным RoPE и длинное с адаптивным коэффициентом. Это позволяет модели точно обрабатывать как короткие, так и длинные последовательности без потери качества.

Технически метод использует включительно-исключающее слияние внимания и корректирующее вращение RoPE. Такая конструкция практически не увеличивает вычислительные затраты на этапе вывода. Одно CuTe-ядро обеспечивает ускорение предзаполнения до 1.39× относительно FlashAttention 2 на H100, приближаясь к производительности FlashAttention 4. Это делает Jet-Long не только точным, но и быстрым решением.

Почему это важно для разработчиков и исследователей

Современные LLM всё чаще применяются в задачах с длинным контекстом: RAG, работа с репозиториями кода, агентные сценарии. Большинство существующих методов фиксируют коэффициент масштабирования, что приводит к потере точности на коротких контекстах или сбоям на длинных. Jet-Long решает эту проблему, динамически подстраиваясь под текущую длину последовательности.

Например, в сценариях RAG, где контекст может варьироваться от нескольких предложений до десятков тысяч токенов, Jet-Long обеспечивает стабильную точность. Агентные системы, которые требуют обработки длинных диалогов или кода, также выигрывают от адаптивного масштабирования. Разработчикам не нужно дообучать модель — достаточно применить Jet-Long к уже существующей LLM.

Какие результаты показывает Jet-Long на бенчмарках

Jet-Long протестирован на моделях Qwen3 с размерами 1.7B, 4B и 8B параметров на контекстах до 128K токенов. На бенчмарке RULER метод превосходит лучшие бейзлайны на 4.79, 2.18 и 2.03 процентных пункта соответственно. Также Jet-Long показывает лучшую общую точность на HELMET-RAG, который является наиболее эффективным предиктором downstream-производительности.

Кроме того, метод демонстрирует наименьшую перплексию на PG-19 — бенчмарке, измеряющем способность модели предсказывать текст. Это говорит о том, что Jet-Long не только расширяет контекст, но и сохраняет языковые способности модели. Сравнение с FlashAttention показывает, что скорость вывода практически не страдает, а в некоторых сценариях даже увеличивается.

Как Jet-Long масштабируется на контексты более 128K токенов

На данный момент Jet-Long протестирован только до 128K токенов. Вопрос масштабирования на более длинные контексты остаётся открытым. Однако архитектура бифокальной RoPE теоретически позволяет работать с последовательностями произвольной длины, если адаптивный коэффициент правильно настроен. Возможно, потребуются дополнительные оптимизации для контекстов в несколько миллионов токенов.

Исследователи также не проверяли совместимость Jet-Long с другими архитектурами внимания, помимо гибридной Jet-Nemotron. Для широкого внедрения потребуется адаптация под популярные модели, такие как LLaMA или Mistral. Пока что метод доступен только для определённых архитектур, но его принципы могут быть обобщены.

Кому будет полезен Jet-Long

Разработчикам и исследователям, работающим с открытыми LLM и нуждающимся в эффективном расширении контекста без дообучения. Метод особенно актуален для сценариев RAG, агентных систем и анализа больших документов. Если вы используете Qwen3 или другие совместимые модели, Jet-Long может стать простым способом увеличить контекстное окно без потери качества.

Для коммерческих проектов, где важна скорость вывода, Jet-Long предлагает ускорение предзаполнения до 1.39×. Это может снизить задержки в реальном времени. Однако перед внедрением стоит проверить совместимость с вашей архитектурой и протестировать на ваших данных.

Что остаётся неизвестным

Несмотря на впечатляющие результаты, остаются вопросы. Насколько метод масштабируется на контексты длиной более 128K токенов? Как он поведёт себя на других архитектурах внимания, помимо гибридной Jet-Nemotron? Также неясно, как Jet-Long работает с моделями, обученными с другими позиционными эмбеддингами. Будущие исследования должны ответить на эти вопросы.

Тем не менее, Jet-Long представляет собой значительный шаг вперёд. Он предлагает практичное решение для расширения контекста без дообучения, что экономит время и ресурсы. Если вас интересует работа с длинными контекстами, стоит обратить внимание на этот метод.