Nemotron-TwoTower：基于预训练自回归上下文的扩散语言建模

# Nemotron-TwoTower：利用预训练自回归上下文的扩散语言建模
来源：https://arxiv.org/html/2606.26493
John Kamalu∗Roger WaleffeMostofa PatwaryMohammad ShoeybiBryan Catanzaro

###### 摘要

扩散语言模型因其并行和迭代生成的潜力，成为自回归模型的一种有前景的替代方案。然而，现有方法使用单个网络同时负责上下文表示和迭代去噪，迫使一个模型承担双重角色，限制了其在任一角色上的能力。我们提出TwoTower，一种分块自回归扩散模型，将这两个角色解耦为两个塔：一个**冻结的**AR上下文塔，以因果方式处理干净词元；以及一个可训练的扩散去噪塔，具有双向分块注意力，通过对上下文的交叉注意力来精炼噪声块。该模型基于Nemotron-3-Nano-30B-A3B（一个开源的30B混合Mamba-Transformer MoE模型），并在约2.1T词元上训练，Nemotron-TwoTower保留了自回归基线98.7%的质量，同时提供了2.42倍的墙钟生成吞吐量。我们在https://huggingface.co/collections/nvidia/nemotron-twotower上发布代码和模型权重。

**核心贡献者。††项目负责人。联系方式：[email protected]

## 1 引言

自回归语言模型是文本生成的主导范式(radford2019language; grattafiori2024llama; liu2024deepseek)。然而，逐词元解码过程造成了吞吐量瓶颈。离散扩散语言模型采取了不同的方法：它们并行生成词元并迭代精炼，提供更高的吞吐量和更细粒度的可控性(austin2021structured; sahoo2024simple; nie2025large; arriola2025block)。在这种范式中，现有模型通常在每一步去噪中使用单个解码器扮演两个不同角色：表示干净词元和对损坏词元去噪。这种纠缠使得同一组权重被拉向不同方向，限制了它们在任一角色上表现出色的能力。

arriola2025encoder观察到这两个角色可以由独立的模块处理，提出了一个编码器-解码器架构，其中编码器表示干净词元，轻量级解码器迭代地对每个块进行去噪。他们最大的实验在1.7B规模上，使用共享权重训练两个模块。完全解耦这两个角色是否能在更大规模上成立，特别是在结合了Mamba、注意力和混合专家的现代混合架构上，仍有待观察。

我们提出TwoTower，构建于Nemotron-3-Nano-30B-A3B(blakeman2025nvidia)之上，这是一个开源的混合模型，交错着Mamba-2、注意力和混合专家层。我们取预训练网络的两个副本，并分配互补的角色：AR上下文塔保持**冻结**，作为干净词元上的因果模型，保留了主干的自回归能力；扩散去噪塔通过掩码扩散目标进行训练，在每个噪声块内具有双向自注意力，并对过去的干净上下文具有因果交叉注意力。

参照图注

(a)双塔架构：冻结的AR上下文塔在每个噪声块上调节扩散去噪塔。
参照图注准确率 vs 吞吐量（相对于基线AR）吞吐量（×基线）准确率（基线百分比）

(b)跨置信度阈值和去噪预算的质量-吞吐量曲线，相对于AR基线归一化。

图1：Nemotron-TwoTower概览。(a) 冻结的AR上下文塔在提示和已提交词元上因果运行，保留预训练主干并维护可重用的KV缓存和Mamba状态。可训练的扩散去噪塔使用层对齐的上下文注意力和上下文种子的Mamba状态，迭代地解析每个噪声块。(b) 相对于逐词元AR基线的质量-吞吐量权衡。两个塔逐层连接。每个去噪层交叉注意力到上下文塔的对应层，为去噪器提供对主干表示的多尺度访问。这与先前仅广播最后隐藏状态的方法形成对比。去噪器还通过自适应层归一化(peebles2023scalable)受扩散时间步调制，这在图像扩散Transformer中很常见，但在掩码扩散语言模型中较少见(ou2024your)。

Nemotron-TwoTower在约2.1T词元上训练——这是用于预训练主干的25T总词元的一小部分——在30B混合MoE架构上，保留了自回归基线的大部分质量，同时提供了显著更高的吞吐量。在默认操作点，Nemotron-TwoTower在解码早期每步精炼中提交多个词元，这有助于解释观察到的相对于逐词元AR解码的墙钟加速。我们在huggingface.co/collections/nvidia/nemotron-twotower (https://arxiv.org/html/2606.26493v1/huggingface.co/collections/nvidia/nemotron-twotower)上发布训练方法、代码和模型权重。

## 2 方法

在本节中，我们描述双塔架构（第2.1节）、掩码扩散公式（第2.2节）和训练方法（第2.3节）。

### 2.1 双塔架构

TwoTower是一种通用方法，可以应用于任何预训练的自回归语言模型。在本工作中，我们在Nemotron-3-Nano-30B-A3B(blakeman2025nvidia)上实例化它，该模型包含52层：23个Mamba-2层、6个自注意力层和23个混合专家层。

我们创建该网络的两个副本并分配互补的角色（见图1(a)）。给定一个提示，**AR上下文塔**作为提示和先前已提交词元上的因果自回归模型运行，产生逐层KV对和最终Mamba状态。然后逐块生成：每个新块初始化为[MASK]词元，并通过**扩散去噪塔**在T步去噪步骤中迭代精炼。在去噪器的每一层，注意力层交叉注意力到对应的上下文塔KV缓存，而Mamba-2层从对应的上下文塔Mamba状态中种子化其初始状态。一旦一个块中的所有词元都是干净的，该块就被提交；上下文塔因果处理已提交的词元以更新其缓存，生成继续下一个块。算法1概述了这个分块自回归生成过程。

算法1 TwoTower分块生成。0: 提示 x_prompt，块大小 S，去噪步数 T
1: (KV, States) ← ContextTower(x_prompt) ▷ 构建AR上下文
2: 对于每个块 b = 1,2,... 执行
3:     z_b ← ([MASK], ..., [MASK])
4:     对于 t = 1,...,T 执行
5:         z_b ← SampleStep(z_b; DenoiserTower(z_b, t, KV, States))
6:     结束
7:     x_b ← z_b ▷ 提交块
8:     (KV, States) ← ContextTower(x_b, KV, States) ▷ 更新AR上下文
9: 结束

在推理时，Nemotron-TwoTower为上下文塔权重引入了固定的内存占用，同时保留单个前缀缓存。只有上下文塔跨块维护KV和Mamba状态，随着块被提交而更新它们，因此序列长度相关的缓存内存规模与AR基线类似。

我们保持上下文塔主体不变。其最终的词汇投影/LM头是可选的：可以在默认的扩散生成路径中省略，其中上下文塔只需要产生状态；当上下文塔用于投机解码、验证、似然评估或AR评分时，则保留它。相比之下，我们对去噪器塔进行了一些架构修改以适应扩散训练和推理。我们在下面描述每个修改，并在第3.2节中提供消融实验。

**双向注意力。**在正在精炼的块内，去噪器放宽了因果掩码：噪声词元双向注意力到其他噪声词元，同时相对于过去的干净块保持因果。这没有添加参数，并且对于标准密集注意力核，不会改变每层FLOPs。在去噪器层i，来自噪声块的查询注意力到上下文塔的层i KV（用于过去的块0，...，b-1）和去噪器自身的层i KV（用于块b）的连接键/值序列：
Attn(Q_b^(i), [K_<b^(ctx,i); K_b^(den,i)], [V_<b^(ctx,i); V_b^(den,i)])。 (1)
交叉注意力是**层对齐的**：去噪器层i注意力到上下文塔层i。由于两个塔都从相同的预训练检查点初始化，相同索引的层在可比的表示级别上操作，使层对齐交叉注意力成为一种自然的配对。

**双向Mamba。**我们还测试了一个无参数的双向Mamba-2变体，通过从零状态从左到右和从右到左运行预训练的Mamba-2权重，然后平均输出。这大致使每层状态空间模型FLOPs翻倍；消融实验显示只有边际质量增益，因此最终设计保持Mamba因果。

**时间条件。**我们使用adaLN-single(peebles2023scalable; chen2023pixart)将去噪器以时间步t为条件。一个全局MLP将t映射到共享的尺度、移位和门控参数，并具有每层学习的嵌入用于层特定调制。在30B主干上，这仅增加了1.5M参数；我们在每个张量并行秩上复制这些小型模块而不是分片它们，避免了额外的通信成本。时间条件提高了采样质量（第3.2节）。

**专家路由。**我们使用主干现有的MoE路由机制，包括序列级负载均衡。词元通过adaLN调制感知噪声。我们不做出明确的路由选择，并允许专家从去噪目标中学习专门化。

### 2.2 块扩散语言建模

我们在掩码扩散框架(sahoo2024simple; shi2024simplified; ou2024your)下训练去噪器，应用于每个自回归块。令x_b = (x_b^1, ..., x_b^S)表示第b个包含S个词元的块，令c_<b = (KV_<b, States_{b-1})表示上下文塔在前一个干净块上的注意力KV缓存和块b-1后的逐层Mamba边界状态。我们将序列分布建模为块自回归：
log p_θ(x) = Σ_{b=1}^B log p_θ(x_b | x_<b)， (2)
其中每个条件块分布由以c_<b为条件的掩码扩散过程参数化。

对于扩散时间t ∈ (0,1]，前向过程通过独立地将每个词元替换为[MASK]来将干净块x_b损坏为噪声块z_t^b。调度α_t是干净词元在噪声水平t下保持未掩码的概率，因此1-α_t是掩码概率；在我们的实验中，我们使用线性调度α_t = 1 - t。使用δ_y表示词元y上的点质量，m表示[MASK]点质量，每个词元的边缘分布为
q(z_t^{b,ℓ} | x_b^ℓ) = Cat(z_t^{b,ℓ}; α_t δ_{x_b^ℓ} + (1-α_t) m)， (3)
其中α_t随着t增加从接近1减小到0。小的t对应于轻度损坏的块，而大的t接近完全掩码的块。

去噪器预测掩码位置的干净词元，以噪声块、扩散时间和上下文缓存为条件。掩码扩散ELBO激发了掩码词元上的时间加权负对数似然(sahoo2024simple; shi2024simplified)；对于线性调度α_t = 1 - t，理论权重简化为1/t。在训练中，为了稳定性我们省略了这个重要性权重，并优化掩码位置上的平均负对数似然：
L_MD = E_{t, z_t} [1/|M_t| Σ_{(b,ℓ)∈M_t} -log p_θ(x_b^ℓ | z_t^b, t, c_<b)]。 (4)
这里M_t = {(b,ℓ): z_t^{b,ℓ} = [MASK]}是掩码词元位置的集合。

**采样。**算法1描述了外层的块自回归循环，而算法2总结了我们主要结果中使用的单块采样器。采样器从上下文塔接收前缀缓存(KV_<b, States_{b-1})，并在精炼块b时保持它们固定；缓存只有在块提交后才被推进。我们还考虑了块扩散模型(arriola2025block; arriola2025encoder)中使用的标准预测-和-噪声采样器，它预测一个干净块，然后根据噪声调度重新掩码低置信度位置。我们的主要结果使用置信度-去掩码变体，带有阈值γ：在每一步，去噪器并行预测所有当前掩码的词元，高于γ的预测被提交，剩余的不确定位置保持掩码以供后续精炼。这使得提交的词元数量适应模型的置信度，同时确保块在T步内完成。

算法2 单块置信度去掩码。0: 前缀缓存 (KV_<b, States_{b-1})，块大小 S，步数 T，置信度阈值 γ
1: z ← ([MASK], ..., [MASK])
2: 对于 r = 1,...,T 执行
3:     τ ← |{ℓ: z^ℓ = [MASK]}| / S
4:     p_θ ← DenoiserTower(z, τ, KV_<b, States_{b-1})
5:     z ← ConfidenceSample(z, p_θ, γ, T - r + 1)
6: 结束
7: 返回 z

### 2.3 训练

我们