机器之心报道
机器之心编辑部
又一个真正轻量、快速、强悍的大语言模型闪亮登场!
Transformer 架构对计算和内存的巨大需求使得大模型效率的提升成为一大难题。为应对这一挑战,研究者们投入了大量精力来设计更高效的 LM 架构。
与此同时,大量工作致力于构建混合模型,将全注意力和线性注意力相结合,以在准确性和效率之间取得平衡。虽然这些模型比全注意力架构具有更高的效率,但其准确性仍明显落后于 SOTA 全注意力模型。
近日,来自英伟达的研究者提出了一种新的混合架构语言模型新系列 ——Jet-Nemotron。其在达到 SOTA 全注意力模型精度的同时,还具备卓越的效率。
具体来说,2B 版本的 Jet-Nemotron 性能就能赶超 Qwen3、Qwen2.5、Gemma3 和 Llama3.2 等最 SOTA 开源全注意力语言模型,同时实现了显著的效率提升。在 H100 GPU 上,其生成吞吐量实现了高达 53.6 倍的加速(上下文长度为 256K,最大 batch size)。
此外,在 MMLU 和 MMLU-Pro 基准上,Jet-Nemotron 的准确率也超过了近期一些先进的 MoE 全注意力模型(如 DeepSeek-V3-Small 和 Moonlight),尽管这些模型的参数规模更大。
论文标题:Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search论文地址:https://www.arxiv.org/pdf/2508.15884
下图将 Jet-Nemotron 与之前的高效大语言模型进行了对比。
值得注意的是,Jet-Nemotron-2B 在 MMLU-Pro 上的准确率高于 Qwen3-1.7B-Base,并且在 64K 上下文长度下,在英伟达 H100 GPU 上的生成吞吐量是后者的 47 倍。
Jet-Nemotron 建立在两项核心创新之上:
后神经架构搜索 (Post Neural Architecture Search,PostNAS):一种高效的后训练架构探索与自适应 pipeline,可适用于任意预训练的 Transformer 模型。JetBlock:一种新型的线性注意力模块,其性能显著优于 Mamba2 等先前的设计。
英伟达研究科学家 Han Cai 以及 MIT 副教授韩松都各自在推特上「安利」了这项研究,其中韩松表示「一个轻量级且可以快速运行的大语言模型来了。」
PostNAS —— 后训练架构探索与自适应
与以往从零开始训练模型、以探索新架构的方法不同,PostNAS 的思路是:在已有的预训练 Transformer 模型上,灵活尝试不同的注意力(attention)模块设计。这样不仅大大降低了开发新型大语言模型架构的成本和风险,还提高了研究效率。
当然,在这一框架下设计出的新架构,如果直接从零训练,可能并不能达到最优结果。但研究者认为,它们依然非常有价值:
立即带来收益—— 如图 1 所示,这些架构能在现有全注意力模型的基础上,实现效率和精度的立刻提升,从而带来实际好处,例如服务质量改善和运维成本下降。快速的创新试验场—— 如果一个新设计在该框架下表现不佳,那么它在完整的预训练过程中成功的可能性也极低。这个「过滤机制」帮助研究人员避免在无望的架构上浪费大量算力和资源。
PostNAS 首先确定全注意力层的最佳位置,然后再搜索更优的注意力模块设计。
同时,研究者提出了一种自动化方法,用来高效确定全注意力层的放置位置。整体方法如下图 4 所示。通过在预训练的全注意力模型中加入可选的线性注意力路径,研究者构建了一个 once-for-all 超网络。训练练过程中的每一步都随机采样一条激活路径,从而形成一个子网络,并使用特征蒸馏损失进行训练。
训练完成后,研究者采用束搜索来确定给定约束条件下(例如仅允许 2 层全注意力层)的最优放置方式。
搜索目标与任务相关:对于 MMLU,研究者选择在正确答案上损失最低的配置(即最大化−loss);而对于数学与检索类任务,研究者则选择准确率最高的配置。如下图 5 (b) 所示,PostNAS 在精度上显著优于均匀放置策略。
在预训练的 Transformer 模型中,并非所有注意力层都具有同等贡献。PostNAS 揭示了其中最关键的注意力层。
PostNAS 精度提升分解。通过将 PostNAS 应用于基线模型,论文在所有基准测试上都取得了显著的精度提升。
此外,KV 缓存大小是影响长上下文和长文本生成吞吐量的最关键因素。PostNAS 的硬件感知搜索能够发掘这样的架构:在保持相似生成吞吐量的同时,拥有更多参数并取得更高精度。下表 2 为硬件感知架构搜索的详细结果。
JetBlock —— 具备SOTA 精度的全新线性注意力模块
借助 PostNAS,研究者提出了 JetBlock。这是一种新颖的线性注意力模块,可以将动态卷积与硬件感知的架构搜索相结合,从而增强线性注意力。
结果显示,在保持与现有设计相近训练与推理吞吐量的同时,JetBlock 在精度上实现了显著提升。在相同训练数据与训练方案情况下,下图对 Mamba2 Block 与 JetBlock 的各性能指标(包括通用知识、数学、常识和检索)进行了比较。
主要结果如下图所示:在全面的基准测试套件中,Jet-Nemotron-2B 和 Jet-Nemotron-4B 的精度能够媲美甚至超越领先的高效语言模型(例如 Qwen3),同时运行速度显著更快,它们分别比 Qwen3-1.7B-Base 快了 21 倍和 47 倍。
更多技术细节与实验结果请参阅原论文。
参考链接:https://hanlab.mit.edu/projects/jet-nemotron