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

Группа исследователей опубликовала в arXiv препринт, описывающий новый метод расширения длинного контекста для больших языковых моделей (LLM) под названием Jet-Long. Решение позволяет использовать модели с контекстом до 128 тысяч токенов без дополнительного обучения, что актуально для таких задач, как RAG, анализ кода репозиториев и агентные системы.

Jet-Long относится к классу zero-shot методов, то есть не требует дообучения модели. Его ключевая особенность — динамический бифокальный RoPE: метод сочетает локальное окно с сохранением исходного RoPE и длинное окно, коэффициент рескалинга которого адаптируется к текущей длине последовательности. На коротких входах модель работает как исходная, на длинных — корректно экстраполирует.

Для эффективной реализации Jet-Long использует специальное объединение внимания с включением-исключением и вращение коррекции RoPE на лету. Благодаря интеграции в одно CuTe-ядро prefilling достигает пропускной способности до 1,39? от FA2 на H100, приближаясь к показателям Hopper-специфичного FA4. Генерация одного батча даёт не более 4% накладных расходов при любой длине.

Эксперименты проводились на моделях Qwen3 с 1,7, 4 и 8 миллиардами параметров при контексте до 128K. Jet-Long превзошёл сильнейшие baseline на тесте RULER на 4,79%, 2,18% и 2,03% соответственно. Кроме того, на бенчмарке HELMET-RAG (признанном наиболее эффективным предиктором downstream-производительности) метод показал лучшую общую точность, а на тесте перплексии PG-19 — минимальные значения.

Авторы отмечают, что Jet-Long обобщается на гибридные архитектуры внимания, такие как Jet-Nemotron, и позволяет улучшить работу с длинным контекстом без переобучения. Метод также устойчив к выбору гиперпараметров, что упрощает его развёртывание.

Разработка может быть интересна специалистам по внедрению LLM в продакшн, где требуется эффективная обработка длинных последовательностей без дополнительных затрат на обучение моделей.