Jet-Nemotron：基于后神经架构搜索的高效语言模型

我们提出 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 倍的预填充（prefilling）加速。此外，尽管参数规模较小（总参数量 15B，激活参数量 2.2B），Jet-Nemotron-2B 在 MMLU 和 MMLU-Pro 上的精度仍高于近期先进的混合专家（MoE）全注意力模型，如 DeepSeek-V3-Small 和 Moonlight。