Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search
2508.15884v1
cs.CL, cs.AI, cs.LG
2025-08-25
Авторы:
Yuxian Gu, Qinghao Hu, Shang Yang, Haocheng Xi, Junyu Chen, Song Han, Han Cai
Резюме на русском
#### Контекст
Современные языковые модели широко применяются в различных областях, включая обработку естественного языка, генерацию текста и анализ данных. Однако существуют значимые проблемы, связанные с оптимальным дизайном этих моделей. Они часто требуют больших вычислительных ресурсов, что усложняет их развертывание в реальных системах. Также существуют ограничения в скорости подготовки данных и вывода, что сказывается на общем эффективности этих моделей в практических задачах.
**Jet-Nemotron** является решением этих проблем, предлагая новую гибридную архитектуру языковых моделей, которая сочетает в себе высокую точность и высокую скорость генерации текста. Модель разработана с помощью **Post Neural Architecture Search** (PostNAS), новый подход к эффективному дизайну архитектур нервных сетей, который позволяет эффективно оптимизировать модели для различных задач.
#### Метод
**Jet-Nemotron** разработана с использованием Postnas, представляющего собой алгоритм поиска лучшей архитектуры с помощью пост-оптимизации. Этот подход характеризуется четырьмя ключевыми компонентами:
1. **Оптимальное размещение и удаление полносвязных слоёв (full-attention layers)**: позволяет оптимизировать структуру модели, удаляя ненужные слои и сохраняя только самые важные.
2. **Выбор линейных блоков аттенции (linear attention blocks)**: упрощает модель, заменяя сложные блоки аттенции на более простые, но эффективные.
3. **Дизайн новых блоков аттенции**: разработка новых блоков, которые позволяют улучшить производительность и точность модели.
4. **Поиск гиперпараметров, ориентированный на железо (hardware-aware hyperparameter search)**: оптимизация модели с учетом ограничений ресурсов, таких как вычислительная мощность и память.
Таким образом, **Jet-Nemotron** использует предварительно обученную модель с полносвязной архитектурой, замороженную во время обучения, что позволяет эффективно экспериментировать с архитектурой без потерь в точности.
#### Результаты
**Jet-Nemotron** была проверена на различных наборах данных и бенчмарках, включая **MMLU** и **MMLU-Pro**. Модель показала существенное улучшение производительности по сравнению с другими языковыми моделями, такими как **Qwen3**, **Qwen2.5**, **Gemma3** и **Llama3.2**. Она достигла точности, сопоставимой или лучшей, чем у этих моделей, при этом показала значительные повышения в скорости генерации текста: до **53.6x** в сравнении с другими моделями. Также **Jet-Nemotron** показала выигрыш в скорости подготовки данных (prefilling) до **6.1x**. Эти результаты достигнуты благодаря оптимальному дизайну архитектуры, который позволяет использовать ресурсы
Abstract
We present Jet-Nemotron, a new family of hybrid-architecture language models,
which matches or exceeds the accuracy of leading full-attention models while
significantly improving generation throughput. Jet-Nemotron is developed using
Post Neural Architecture Search (PostNAS), a novel neural architecture
exploration pipeline that enables efficient model design. Unlike prior
approaches, PostNAS begins with a pre-trained full-attention model and freezes
its MLP weights, allowing efficient exploration of attention block designs. The
pipeline includes four key components: (1) learning optimal full-attention
layer placement and elimination, (2) linear attention block selection, (3)
designing new attention blocks, and (4) performing hardware-aware
hyperparameter search. Our Jet-Nemotron-2B model achieves comparable or
superior accuracy to Qwen3, Qwen2.5, Gemma3, and Llama3.2 across a
comprehensive suite of benchmarks while delivering up to 53.6x generation
throughput speedup and 6.1x prefilling speedup. It also achieves higher
accuracy on MMLU and MMLU-Pro than recent advanced MoE full-attention models,
such as DeepSeek-V3-Small and Moonlight, despite their larger scale with 15B
total and 2.2B activated parameters.
Ссылки и действия
Дополнительные ресурсы: