时令发自凹非寺
量子位 | 公众号 QbitAI
英伟达开源又放大招了！
韩松团队推出了一款全新的基于后神经架构搜索的高效语言模型——Jet-Nemotron。
该模型在一系列基准测试中，不仅表现出与 Qwen3、Qwen2.5、Gemma 3 和 Llama 3.2 相当甚至更优的准确率，还在生成吞吐量上实现最高 53.6 倍加速，在预填充阶段达到 6.1 倍加速。
值得一提的是，在 MMLU、MMLU-Pro 和 BBH 基准上，Jet-Nemotron-2B 相比 Qwen3-1.7B-Base 吞吐量提高了 47 倍，缓存大小缩小至1/47。
同时，它还实现了比 DeepSeek-V3-Small 和 Moonlight（共 150 亿参数，22 亿激活参数）更高的准确率。
代码和预训练模型都将开源，我们先来看看 Jet-Nemotron 是如何构建的。
Jet-Nemotron：基于后神经架构搜索构建
首先，Jet-Nemotron 是在后神经架构搜索（Post Neural Architecture Search，PostNAS）的基础上构建的。
其中，后神经架构搜索（PostNAS）模型是一种“站在大模型肩膀上做改造”的架构搜索方法。
它从一个预训练的全注意力模型出发，并直接继承其多层感知机权重，且在整个过程中保持这些权重被冻结（不再更新）。
Jet-Nemotron 就是将 PostNAS 通过以下 4 个步骤优化得到的：
全注意力层的放置和消除
在模型中保留少数几个全注意力层，对于在检索等高难度任务上保持高准确率至关重要。
然而，这些层的最佳放置位置一直不明确。
因此，研究团队引入了一种新方法，通过训练一个“一次性”超级网络 (once-for-all super network)，自动学习应该在哪些位置使用全注意力层。
实验结果表明，与常用的均匀放置策略相比，这种学习到的放置方式在 MMLU 基准上的准确率有显著提升。
选择线性注意力模块
在确定了全注意力层的放置方案后，研究团队进行注意力模块搜索，以确定最优的线性注意力模块。
在实验中，他们评估了 6 个最先进的线性注意力模块（RWKV7 由于训练吞吐量过低排除在外），结果如下。
由上表观察到，Gated DeltaNet 实现了最优的整体准确率。因此，研究团队在后续实验中都采用 Gated DeltaNet。
设计新型注意力模块
添加卷积是增强线性注意力能力的一种常用策略。然而，以往的方法仅仅依赖于静态卷积核，缺乏动态适应卷积核特征提取模式的能力。
于是，研究团队引入一种名为JetBlock的新型线性注意力模块。
此模块使用一个卷积核生成器 (kernel generator)，能够根据输入内容动态地生成因果卷积核 (dynamic causal convolution kernels)，然后将这些卷积核应用于 V (值) 词元上。此外，它还移除了在 Q (查询) 和 K (键) 上的冗余静态卷积，从而简化了计算流程。
执行硬件感知架构搜索
传统上，参数量被用作语言模型效率的代理指标。然而，参数数量与硬件效率并不直接相关。
基于“KV 缓存大小是影响长上下文和长生成吞吐量的最关键因素”的发现。
研究团队将 KV 缓存大小固定为原始设计的规格，并对 key 维度、value 维度以及注意力头数进行了小规模的网格搜索。
这种硬件感知搜索能够在保持相似生成吞吐量的同时，利用更多参数以实现更高准确率。
好消息是，研究团队计划在 GitHub 上公开代码和模型，目前正等待法律合规审核。
显著的效率提升
Jet-Nemotron-2B 和 Jet-Nemotron-4B 分别基于 Qwen2.5-1.5B 和 Qwen2.5-3B 模型构建。
为了全面评估模型性能，研究团队在数学、常识、检索、编码以及长上下文中都进行了测试。
数学任务上，Jet-Nemotron-2B 取得了 49.6 的平均准确率，比 Qwen3-1.7B-Base 高 6.3，同时速度快 47 倍。
相比之下，之前的线性注意力和混合模型在数学任务上远远落后于 Qwen3-1.7B-Base。
常识推理任务上，Jet-Nemotron-2B 平均准确率达到 62.0，超越所有基线模型。
检索任务上，Jet-Nemotron-2B 的表现优于除 Qwen3-1.7B-Base 之外的所有基线模型。
当扩展到 4B 时，Jet-Nemotron-4B 达到了 76.2 的最佳平均准确率，同时与 Qwen3 相比仍保持 21 倍的速度提升。
编码任务上，Jet-Nemotron-2B 的平均准确率高于所有基线模型。
同时，Jet-Nemotron-4B 在所有编码任务中都实现了更高的准确率。
长下文任务上，可以看出 Jet-Nemotron-2B 虽然只有两个全注意力层，但性能堪比拥有更多全注意力层的 Qwen2.5-1.5B 和 Gemma3n-E2B 等领先模型。
综合来看，Jet-Nemotron-2B 和 Jet-Nemotron-4B 在这些领域的表现均与 Qwen3-1.7B-Base 相当，甚至更胜一筹。
而由于全注意力层显著减少且 KV 缓存规模更小，Jet-Nemotron 与 Qwen3 相比有明显优势。
团队介绍
值得一提的是，此研究团队全为华人。
Yuxian Gu，本科与博士均就读于清华大学计算机科学与技术系，导师为黄民烈教授。
此前，他还在微软亚洲研究院实习，导师为研究员董力。
他的研究兴趣主要集中在语言模型的全生命周期，包括预训练、下游任务适配以及推理阶段的高效方法。
最近，他的研究重点是面向预训练大语言模型的数据构建理论与算法（如 PDS、指令预训练、Learning Law），以及利用知识蒸馏进行语言模型压缩（如 MiniLLM、MiniPLM）。
胡擎昊，本科毕业于浙江大学，硕士毕业于新加坡国立大学，现为麻省理工学院韩松教授的博士后研究员。
Shang Yang，现为麻省理工学院三年级博士生，导师为韩松教授。在此之前，他本科毕业于清华大学电子工程系。
Haochen Xi，本科毕业于清华大学姚班，导师为姚期智院士，目前博士就读于美国加州大学伯克利分校计算机科学专业，现为加州大学伯克利分校 MLsys 研究员。
Junyu Chen，现为清华大学姚班的一名本科生。曾在麻省理工学院 HAN 实验室担任研究实习生，导师为韩松教授。此前，还曾在清华大学与李毅教授合作研究 3D 视觉感知和人机交互。
韩松，本科毕业于清华大学电子工程系，在斯坦福大学获得博士学位，目前是麻省理工学院电子工程学院副教授。
他提出了被广泛用于高效人工智能计算的“深度压缩”技术，并且首次给现代人工智能芯片带来权重稀疏性的“高效推理机”，这些技术影响了 NVIDIA 的安培 GPU 架构等。
韩松还是 TinyML 研究的先驱，这项研究将深度学习带到物联网设备上，使边缘端机器学习成为可能。
2023 年，韩松创办的专注边缘设备机器学习优化的 OmniML 被英伟达收购，他也因此加入英伟达成为杰出科学家，其公司的 CEO 吴迪和 CTO 毛慧子同样也入职英伟达。
蔡涵，NVIDIA 研究院研究科学家。在上海交通大学获得硕士和学士学位，在麻省理工学院电子工程与计算机科学系获得博士学位。
参考链接：
[1]https://arxiv.org/abs/2508.15884
