用后置神经架构搜索打造高效语言模型，生成速度提升53倍
Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search
- 基于冻结MLP的预训练模型，高效探索注意力结构设计
- 20亿参数模型在多项测试中超越大模型，生成速度提升53.6倍
- 适合追求高吞吐、低延迟部署的AI工程师
我们提出Jet-Nemotron，一种新型混合架构语言模型，其精度与领先全注意力模型相当或更优，同时显著提升生成吞吐量。该模型通过后置神经架构搜索（PostNAS）构建，该方法从预训练的全注意力模型出发，冻结MLP权重，从而高效探索注意力模块设计。PostNAS包含四个核心组件：(1) 学习最优全注意力层的放置与移除，(2) 线性注意力模块选择，(3) 新型注意力模块设计，(4) 硬件感知超参数搜索。Jet-Nemotron-2B模型在综合基准测试中表现媲美甚至超越Qwen3、Qwen2.5、Gemma3和Llama3.2，生成吞吐量最高提升53.6倍，预填充速度提升6.1倍。其在MMLU和MMLU-Pro上的准确率也高于近期先进MoE全注意力模型DeepSeek-V3-Small（15B总量，2.2B激活）和Moonlight，尽管后者规模更大。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。