NVIDIA 发布了 Nemotron-Labs-TwoTower，这是一个基于预训练自回归骨干网络构建的扩散语言模型。该模型以开放权重形式发布，采用 NVIDIA Nemotron 开放模型许可证。此次发布旨在解决文本生成中的吞吐量瓶颈问题。
自回归（AR）模型每次解码一个 token。这种串行处理方式限制了生成吞吐量。离散扩散语言模型则另辟蹊径。它们并行生成 token，并通过迭代方式逐步优化。
大多数扩散语言模型使用单一网络承担两项任务。它在每一步既要表示干净 token，又要对受损 token 进行去噪。TwoTower 将这两项任务分离到两个独立模块中。它保留了自回归基线模型 98.7% 的综合基准质量，同时实现了 2.42 倍的挂钟生成吞吐量提升。
摘要
- TwoTower 将扩散过程拆分为一个冻结的自回归上下文模块和一个经过训练的去噪器模块。
- 它以 2.42 倍的吞吐量（γ=0.8, S=16, 2×H100）保留了自回归模型 98.7% 的质量。
- 去噪器在约 2.1T token 上训练；骨干网络使用了 25T token。
- 单个检查点支持扩散、模拟自回归和自回归解码三种模式。
Nemotron-Labs-TwoTower
TwoTower 是一种逐块自回归扩散模型。它基于 Nemotron-3-Nano-30B-A3B 构建，这是一个开放权重的混合骨干网络。该骨干网络交错使用了 Mamba-2、自注意力机制和混合专家（MoE）层。
每个模块有 52 层：23 层 Mamba-2、6 层自注意力机制和 23 层 MoE。发布的检查点包含两个模块，总参数量约为 600 亿。每个模块每个 token 的活跃参数量约为 30 亿。MoE 使用 128 个可路由专家，其中 6 个被激活，外加 2 个共享专家。
两个模块初始时均为同一骨干网络检查点的副本。只有去噪器模块经过训练。自回归上下文模块保持冻结状态。去噪器在约 2.1T token 上训练，仅占骨干网络 25T token 预训练数据的一小部分。
双模块工作原理
自回归上下文模块对提示词和已提交 token 进行因果处理。它生成每层的 KV 缓存和最终的 Mamba-2 状态。该模块保留了骨干网络的自回归能力。
扩散去噪器模块对噪声块进行优化。在块内部，它使用块内双向注意力机制。对于过去的干净块，它保持因果性。
这些塔层逐层连接。去噪器第 i 层交叉关注上下文塔第 i 层。这种按层对齐的交叉注意力机制使得去噪器能够多尺度地访问骨干网络的表征。此前的方法仅广播最后一个隐藏状态。
另外两项去噪器改进也很重要。Mamba-2 层从其初始状态由上下文塔的 Mamba 状态播种。扩散时间步通过 adaLN-single 时间条件调节每一层。该 adaLN 模块仅增加约 150 万参数。
生成过程逐块进行。每个块以 S 个 [MASK] token 开始。去噪器经过 T 步对其进行精炼，然后将其提交。随后，上下文塔处理已提交的 token 以更新其缓存。
这解释了为什么多次去噪步骤仍能胜过单 token 解码。自回归解码每步仅提交一个 token。而 TwoTower 在精炼早期每步提交多个 token。
基准测试
评估在 2×H100 GPU 上使用 BF16 精度进行。默认运行点为置信度去掩码，阈值 γ=0.8，块大小 S=16。该表格比较了自回归基线方法与 TwoTower 扩散解码。
| 任务 | Nemotron-3-Nano-30B-A3B（自回归） | Nemotron-Labs-TwoTower（扩散） |
|---|---|---|
| MMLU（5-shot，准确率） | 78.56 | 78.24 |
| MMLU-Pro（5-shot，CoT 精确匹配） | 62.59 | 60.93 |
| ARC-Challenge（25-shot，归一化准确率） | 91.72 | 92.66 |
| WinoGrande（5-shot，准确率） | 76.09 | 76.09 |
| RACE（0-shot，准确率） | 88.90 | 88.90 |
| HumanEval（0-shot） | 79.27 | 75.58 |
| MBPP-Sanitized（3-shot） | 74.71 | 74.28 |
| GSM8K（8-shot，准确率） | 92.49 | 90.14 |
| MATH-500（4-shot） | 84.40 | 80.60 |
| MMLU Global Lite（5-shot） | 73.97 | 73.94 |
| MGSM（8-shot，平均准确率） | 80.80 | 80.40 |
| 质量保持 | 100% | 98.7% |
| 生成吞吐量（× 自回归） | 1.0× | 2.42× |
通用知识得分与自回归基线相差约一个百分点以内。代码和数学任务表现出适度下降。常识和多语言得分已恢复或略有提升。降低 γ 值会使每步提交更多 token 并提高吞吐量，但会降低质量。
运行方式：三种生成模式
该检查点提供了三种推理路径。完整的双塔扩散模式使用 2 块 GPU，每块 GPU 在 BF16 精度下约占用 59GB。仅自回归模式可在单块 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 以换取速度。
第三个用例是即插即用的适配。上下文塔保留其语言模型头，用于推测解码、验证或自回归评分。团队可以从同一个检查点运行自回归和扩散模型。
优势与不足
优势：
- 采用 NVIDIA Nemotron 开放模型许可证开放权重；可用于商业用途
- 在默认运行点，以 2.42 倍吞吐量保留 98.7% 的自回归质量
- 单个检查点支持扩散、模拟自回归和自回归解码
- 去噪器在约 2.1T token 上训练，并非完整重新预训练
- 序列长度缓存内存规模与自回归基线相当
不足：
- 完整的双塔扩散需要 2 块 GPU，每块 GPU 约 59GB（BF16 精度）
- 代码和数学能力下降幅度大于通用知识（HumanEval 从 79.27 降至 75.58）
- 同时驻留两个塔会增加固定模型权重的内存占用
- 发布的检查点是基础模型，尚未经过指令微调或对齐
- 吞吐量超过 3 倍时，质量损失更大