🧪Эксперименты

Подход проверяют на Qwen-2.5-1.5B, Qwen-2.5-3B получая модели Jet-Nemotron-2B, Jet-Nemotron-4B, соотвественно.

На этапе поиска архитектуры обучают на 50B токенах, а финальное дообучение на 400B.

По метрикам все типа хорошо, заметный прирост качества на бенчах по сравнению с безлайном (MMLU-Pro, GSM8k, GPQA, CommonSense, Retreival задачах).

Ускорение растет с длиной контекста. На 4к/8к ускорение порядка 15% на префилл, но может достигать 6 раз на длинах контекста 256k. Ускорение на декодировании еще больше, до 53 с лишним раз 😱 на длинных контекстах, где оно ограничено в любом случае наличием небольшого количества full-attention.

💡 Выводы

Результат смотрится довольно впечатляюще. Однако есть опасение, что хорошие метрики обусловлены удачным протоколом дообучения 😺. Кроме того, мне кажется, что процедура слишком трудоемкая, чтобы получить широкое распространение. Да и паттерны важности attention, справедливые для рассмотренных задач, могут не переноситься на продвинутый ризонинг или ворочание репозиториев с кодом.

Тем не менее - сильная инженерная работа.