北京
时令 发自 凹非寺
量子位 | 公众号 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
特别声明：以上内容(如有图片或视频亦包括在内)为自媒体平台“网易号”用户上传并发布，本平台仅提供信息存储服务。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.
