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.
Abstract (translated)
我们介绍了Jet-Nemotron，这是一种新的混合架构语言模型系列。它在准确率上与领先的全注意力模型持平或超越它们的同时，显著提高了生成吞吐量。Jet-Nemotron 使用后神经网络结构搜索（Post Neural Architecture Search，简称 PostNAS），这是一种新颖的神经架构探索流水线，能够高效地进行模型设计。不同于以往的方法，PostNAS 从一个预训练好的全注意力模型开始，并冻结其 MLP 权重，从而可以高效地探索注意力模块的设计。该流水线包含四个关键组件：(1) 学习最优全注意力层的放置和消除；(2) 线性注意力块的选择；(3) 设计新的注意力模块；以及 (4) 进行硬件感知的超参数搜索。 我们的 Jet-Nemotron-2B 模型在全面的基准测试套件中，实现了与 Qwen3、Qwen2.5、Gemma3 和 Llama3.2 相当或更好的准确率，同时提供了高达 53.6 倍的生成吞吐量加速和 6.1 倍的预填充速度提升。此外，在 MMLU 和 MMLU-Pro 测试中，Jet-Nemotron-2B 的表现优于最近先进的 MoE 全注意力模型（如 DeepSeek-V3-Small 和 Moonlight），尽管这些模型拥有更大规模的参数总数和激活参数数（分别为 150 亿和 22 亿）。
URL
https://arxiv.org/abs/2508.15884
