💬 NLP

Nemotron-TwoTower: Diffusion Language Modeling with Pretrained Autoregressive Context

Nemotron-TwoTower 引入了一种新颖的分块自回归扩散架构，将上下文处理与去噪解耦为两个专门的塔，在保留其自回归基准 98.7% 质量的同时，实现了 2.42 倍的生成吞吐量提升。

原作者： Fitsum Reda, John Kamalu, Roger Waleffe, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro

发布于 2026-06-26

📖 1 分钟阅读☕ 轻松阅读

原作者： Fitsum Reda, John Kamalu, Roger Waleffe, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图写一个长篇故事，但你有一个严格的规则：你一次只能写一个词，并且必须在前一个词完全写完之前，不能开始写下一个词。这就是大多数当前 AI 语言模型（称为“自回归”模型）的工作方式。这就像一个非常细心、动作缓慢的打字员，从不出错，但完成一个句子需要很长时间。

现在，想象一种新的方法，你可以一次性输入整个句子，观察它，然后同时修正所有单词中的错误。这种方法速度更快，但通常会导致故事的质量下降，因为 AI 在试图同时修正所有内容时会感到困惑。

Nemotron-TwoTower 是一项新发明，它试图实现两全其美：既拥有同时输入所有内容的效率，又拥有那种细心的、一次写一个词的打字员的高质量。

以下是它的工作原理，使用一个简单的类比：

双塔结构：图书管理员与编辑

研究人员构建了一个由两个截然不同的“塔”（AI 的组成部分）组成的系统，它们像一个两人团队一样协作：

冻结上下文塔（图书管理员）：
- 角色： 这是一个“专家”，它了解目前为止的故事内容。它是一个预训练的 AI，并且是冻结（锁定）状态的。它就像一位严谨的图书管理员，从头到尾逐词阅读故事。
- 职责： 它不会改变。它只是完整、有序地记录下目前为止写下的所有内容。它会告诉另一个塔：“这是直到目前为止的故事背景信息。”
- 为什么重要： 因为它是冻结的，它保留了构建它的那个拥有 300 亿参数的大型模型的原始智能和知识。它永远不会遗忘或感到困惑。
可训练扩散去噪塔（编辑）：
- 角色： 这是一个“快速工作者”。它是一个可以训练的 AI，允许进行改变和学习。
- 职责： 它不再是一个词一个词地写，而是抓取一个包含 16 个单词（或 Token）的块，这些单词目前是空白的（被遮盖的）。它观察“图书管理员”的笔记，并尝试同时填补这 16 个空格。
- 过程： 它不会立即做到完美。它先做一个猜测，然后观察自己的猜测和图书管理员的笔记，并进行精炼。它反复进行这个过程（迭代），直到这一块单词变得清晰且符合逻辑。
- 神奇之处： 因为它可以同时观察这 16 个单词并一起修正它们，所以它比一个一个写要快得多。

它们如何协同工作（流水线）

想象一个工厂的流水线：

图书管理员阅读提示词和已经写好的单词，建立起故事的完美心理地图。
编辑抓取一块空白的纸张（长度为 16 个单词）。
编辑询问图书管理员：“接下来是什么？”图书管理员提供上下文。
编辑填补 16 个空格，检查自己的工作并修正错误。
一旦编辑确信这 16 个单词没问题，它们就会被“提交”（定稿）。
图书管理员将这新增的 16 个单词纳入其心理地图中。
编辑抓取下一个 16 个单词的空白块，循环往复。

结果：速度 vs. 质量

论文声称该系统实现了“甜点位”（最佳平衡点）：

速度： 它比传统的逐词生成法快了 2.42 倍。这就像是从蜗牛变成了短跑运动员。
质量： 它保留了原模型 98.7% 的质量。它并没有为了获得速度而损失太多智能。
效率： 他们用大约 2.1 万亿个单词训练了这个新的“编辑”塔，这只是原始“图书管理员”所用 25 万亿个单词的一小部分。这意味着他们不需要从头开始重新教整个 AI；他们只是教会了“编辑”如何使用“图书管理员”的笔记。

为什么这与以往的尝试不同

以前的尝试试图使用单个大脑来同时完成两项工作：既要记住故事，又要修正单词。论文指出，这就像要求一个人同时成为完美的记忆守护者和快速的编辑——这工作量太大了，结果导致两项工作都做不好。

通过将角色拆分为两个塔（一个冻结，一个可训练），他们让每个部分都去做自己最擅长的事。图书管理员保持完美，而编辑变得快速。

“置信度”技巧

该系统还有一个聪明的特性——“置信度解掩码”（confidence unmasking）。

如果编辑对某个单词非常有信心，它会立即锁定该词。
如果它不确定，它会留下该单词的空白（遮盖），并在下一轮编辑中尝试修复它。
这意味着它不会在已经确定正确的单词上浪费时间重复检查，而是会将额外的时间花在打磨那些棘手的环节上。

总结

Nemotron-TwoTower 是一种让 AI 文本生成速度更快且不降低智能的方法。它通过将“记住故事”（保持缓慢且完美）的任务与“编写下一块内容”（变得快速且具有迭代性）的任务分离来实现。其结果是，该系统编写故事的能力几乎与旧的缓慢方法持平，但速度快了近两倍。

技术摘要：Nemotron-TwoTower

问题陈述

自回归（AR）语言模型在文本生成领域占据主导地位，但由于其逐个 token 的顺序解码过程，面临着吞吐量瓶颈。离散扩散语言模型提供了一种极具前景的替代方案，通过并行生成 token 并进行迭代细化，有望实现更高的吞吐量和更精细的控制力。然而，现有的扩散方法通常依赖单个网络同时承担两个不同的角色：代表纯净上下文 token 和迭代去噪损坏的 token。这种角色纠缠迫使同一组权重去服务于相互冲突的目标，从而限制了模型在任一角色上达到卓越水平的能力。此外，目前尚不清楚在结合了 Mamba、注意力机制和混合专家（MoE）层的现代混合架构上，完全解耦这些角色是否在大型规模下仍然可行。

方法论

作者提出了 TwoTower，这是一种块状（block-wise）自回归扩散模型，它将上下文表示和去噪角色解耦为两个独立的塔，并基于 Nemotron-3-Nano-30B-A3B 混合架构（该架构交替使用 Mamba-2、自注意力层和 MoE 层）构建。

架构

该模型利用了两个具有互补角色的预训练网络副本：

冻结的 AR 上下文塔（Frozen AR Context Tower）： 该塔保持冻结状态，作为提示词（prompt）和已提交 token 的因果自回归模型。它处理纯净的 token 以生成层对齐的键值（KV）缓存和 Mamba 状态。它保留了骨干网络原始的自回归能力，并维持了跨块的可重用状态。
可训练的扩散去噪器塔（Trainable Diffusion Denoiser Tower）： 该塔经过训练以细化噪声块。它在当前的噪声块内采用双向块注意力（允许 token 相互关注），同时与冻结的上下文塔的相应层保持因果交叉注意力。这种层对齐的交叉注意力为去噪器提供了对骨干网络表示的多尺度访问，这与以往仅广播最终隐藏状态的方法形成了对比。

生成过程

生成过程按块进行（算法 1）：

上下文塔因果地处理提示词和已提交的 token。
一个新的块由 $S$ 个掩码（masked）token 初始化。
扩散去噪器塔在 $T$ 个步骤内迭代细化该块。在每一层，去噪器都会关注上下文塔中过去块的 KV 缓存以及当前块自身的双向 KV 缓存。
一旦一个块被充分细化（基于置信度阈值），它就会被提交。随后，上下文塔使用新的纯净 token 更新其状态，并重复此过程以处理下一个块。

训练与条件化

掩码扩散（Masked Diffusion）： 去噪器使用掩码扩散目标进行训练。根据线性调度，通过将 token 替换为 [MASK] 来损坏纯净块。模型最小化掩码位置处原始 token 的负对数似然。
时间条件化（Time Conditioning）： 去噪器通过自适应层归一化（adaLN）受扩散时间步 $t$ 的约束，这是图像扩散中常用的技术，但在掩码语言模型中较少使用。这增加了极少的参数（约 150 万个），但提升了采样质量。
课程学习（Curriculum）： 训练遵循两阶段课程学习。第一阶段使用多样化的数据混合以实现广泛覆盖；第二阶段则增加高质量和 STEM 领域来源的权重。训练块大小逐渐从 $S=32$ 减少到 $S=16$ ，以平衡并行性与条件化质量。
采样： 作者采用了置信度去掩码（confidence-unmasking）策略（算法 2）。在每一步中，去噪器并行预测所有掩码 token；置信度高于阈值 $\gamma$ 的 token 会被提交，而不确定的 token 则保持掩码状态以进行进一步细化。这使得提交的 token 数量是自适应的。

关键结果

在 Nemotron-3-Nano-30B-A3B 骨干网络（30B 参数，混合 MoE）上进行的评估显示，Nemotron-TwoTower 表现如下：

质量与吞吐量的权衡： 该模型保留了自回归基准模型 98.7% 的综合基准测试质量，同时实现了 2.42 倍 的实际运行（wall-clock）生成吞吐量提升。
基准测试性能：
- 通用知识： 性能与 AR 基准模型保持在约 1 个百分点以内。
- 代码与数学： 与基准模型相比表现出轻微下降。
- 常识与多语言： 性能得到了恢复或略有提升。
设计消融实验：
- 双向 Mamba： 测试了一个双向 Mamba 变体，结果显示质量提升微乎其微，但使 SSM 计算量翻倍；最终设计保留了因果 Mamba。
- 时间条件化： 显式的时间条件化显著提升了生成、代码和数学得分。
- 塔解耦： 保持上下文塔冻结并训练独立的去噪器得到了最佳结果。在两个塔之间共享权重或继续对上下文塔进行 AR 训练会导致显著的准确率下降。
块大小敏感性： 较小的训练块大小（ $S=16$ ）比较大的块大小（ $S=128$ ）提升了质量，但降低了吞吐量。作者选择 $S=16$ 作为默认值，以平衡质量保留与速度。使用大于训练块大小的块进行采样会导致严重的性能退化。

意义与主张

论文声称，Nemotron-TwoTower 成功证明了预训练的自回归表示可以支持可训练的扩散去噪器，而无需重新预训练整个模型。通过将上下文和去噪角色解耦，该方法在保留原始 30B 混合 MoE 模型广泛能力的同时，释放了扩散模型的并行生成优势。

作者强调，这种质量恢复是通过在仅为原始预训练数据一小部分（约 2.1T token 对比 25T token）的数据上进行适配实现的。这项工作建立了一个新的帕累托前沿（Pareto frontier），为实现更高吞效生成的语言模型提供了一条可行路径，且无需牺牲大规模自回归模型的底层质量。代码和模型权重已发布，以促进该方向的进一步研究。