NVIDIA 发布 Nemotron-Labs-TwoTower 开放权重扩散语言模型

NVIDIA 发布了 Nemotron-Labs-TwoTower，这是一个基于预训练自回归骨干网络构建的扩散语言模型。它以开放权重形式发布，采用 NVIDIA Nemotron Open Model License 许可。此次发布旨在解决文本生成中的吞吐量瓶颈。

自回归（AR）模型一次解码一个 token。这种串行过程限制了生成吞吐量。离散扩散语言模型则另辟蹊径。它们并行生成 token，并迭代地对其进行精炼。

大多数扩散语言模型使用一个网络承担两项任务。它既表示干净的 token，又在每一步对损坏的 token 进行去噪。TwoTower 将这两项任务分离为两个塔。它保留了 AR 基线综合基准质量的 98.7%。同时报告了 2.42 倍的挂钟生成吞吐量提升。

TL;DR

• TwoTower 将扩散拆分为一个冻结的 AR 上下文塔和一个训练好的去噪器塔。
• 它在 2.42 倍吞吐量下保留了 98.7% 的 AR 质量（γ=0.8，S=16，2×H100）。
• 去噪器在约 2.1T token 上训练；骨干网络使用了 25T。
• 一个检查点即可运行扩散、模拟 AR 和 AR 解码模式。

Nemotron-Labs-TwoTower

TwoTower 是一种分块自回归扩散模型。它基于 Nemotron-3-Nano-30B-A3B 实例化，后者是一个开放权重的混合骨干网络。该骨干网络将 Mamba-2、自注意力和混合专家（MoE）层交错排列。

每个塔有 52 层：23 层 Mamba-2、6 层自注意力和 23 层 MoE。发布的检查点包含两个塔，总参数量约为 60B。每个 token 的激活参数量约为每个塔 3B。MoE 使用 128 个可路由专家，其中 6 个激活，另有 2 个共享专家。

两个塔最初都是同一骨干网络检查点的副本。只有去噪器塔被训练。AR 上下文塔保持冻结。去噪器在约 2.1T token 上训练，仅为骨干网络 25T token 预训练量的一小部分。

双塔如何工作

AR 上下文塔以因果方式在提示词和已提交 token 上运行。它生成逐层 KV cache 和最终的 Mamba-2 状态。它保留了骨干网络的自回归能力。

扩散去噪器塔对噪声块进行精炼。在一个块内，它使用双向块内注意力。相对于过去的干净块，它保持因果性。

两个塔逐层连接。去噪器层 i 交叉注意力到上下文塔层 i。这种层对齐的交叉注意力提供了对骨干网络表示的多尺度访问。此前的方法仅广播最后一个隐藏状态。

还有两处去噪器修改值得关注。Mamba-2 层从上下文塔的 Mamba 状态中获取其初始状态。扩散时间步通过 adaLN-single 时间条件对每一层进行调制。该 adaLN 模块仅增加约 1.5M 参数。

生成逐块进行。每个块以 S 个 [MASK] token 开始。去噪器在 T 步内对其进行精炼，然后将其提交。随后上下文塔处理已提交的 token 以更新其缓存。

这解释了为什么多步去噪仍然可以胜过单 token 解码。自回归解码每步恰好提交一个 token。TwoTower 在精炼早期每步提交多个 token。

基准测试

评估在 2×H100 GPU 上使用 BF16 进行。默认工作点为置信度解掩码，阈值 γ=0.8，块大小 S=16。该表将 AR 基线 与 TwoTower 扩散解码进行对比。

通用知识保持在 AR 基线约一个百分点以内。代码和数学出现小幅退化。常识和多语言得分得到恢复或略有提升。降低 γ 会使每步提交更多 token 并提高吞吐量，但质量有所下降。

运行方式：三种生成模式

该检查点暴露了三条推理路径。完整双塔扩散使用 2 块 GPU，BF16 下每块 GPU 约 59GB。仅 AR 模式可在单块 80GB GPU 上运行。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "nvidia/Nemotron-Labs-TwoTower-30B-A3B-Base-BF16"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.bfloat16, trust_remote_code=True,
)
# context tower -> GPU 0, denoiser tower -> GPU 1
model.place_towers_on_devices("cuda:0", "cuda:1")
model.eval()

prompt = "France is a country "
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")

outputs = model.generate_mask_diffusion(
    inputs["input_ids"], max_new_tokens=128,
    block_size=16, steps_per_block=16, mask_token_id=3,
    temperature=0.1, confidence_threshold=0.8,
    eos_token_id=tokenizer.eos_token_id,
)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

这三种模式分别是 generate_mask_diffusion()、generate_mock_ar() 和 generate_ar()。掩码扩散每步最多提交 block_size 个 token。Mock-AR 和 AR 每步提交一个 token。

适用场景：用例

最直接的用例是更快的批量生成。生成合成文本的数据团队可以用小幅质量下降换取吞吐量。在 γ=0.8 时，这一权衡是 1.3% 的质量换取 2.42× 的速度。

第二个用例是调节质量与吞吐量之间的权衡。根据 NVIDIA 的论文，提高 γ 可以保留更多质量。降低 γ 则每步提交更多 token 以换取速度。

第三个用例是即插即用式适配。上下文塔保留其 LM head，用于投机解码、验证或 AR 评分。团队可以从一个检查点同时运行 AR 和扩散。

优势与劣势

优势：

• 在 NVIDIA Nemotron Open Model License 下开放权重；可直接用于商业用途
• 在默认工作点下，以 2.42× 吞吐量保留了 98.7% 的 AR 质量
• 单个 checkpoint 同时支持扩散、mock-AR 和 AR 解码
• 去噪器在约 2.1T tokens 上训练，而非完整的重新预训练
• 序列长度缓存内存的扩展方式与 AR 基线一致

劣势：

• 完整的双塔扩散需要 2 块 GPU，BF16 下每块 GPU 约需 59GB
• 代码和数学能力下降幅度大于通用知识（HumanEval 79.27 → 75.58）
• 同时常驻两个塔会提高固定的模型权重内存占用
• 发布的 checkpoint 是基础模型，尚未经过指令微调或对齐
• 吞吐量超过 3× 后会伴随更大的质量损失

交互式讲解