英伟达韩松团队新作：具有后神经架构搜索的高效语言模型

时令 发自 凹非寺
量子位 | 公众号 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

[2]https://github.com/NVlabs/Jet-Nemotron

[3]https://x.com/iScienceLuvr/status/1959832287073403137