NVIDIA发布Nemotron-Labs-TwoTower：基于冻结自回归骨干网络的开放权重扩散语言模型

NVIDIA正式发布了Nemotron-Labs-TwoTower，这是一种基于预训练自回归骨干网络的扩散语言模型。该模型以开放权重形式提供，采用NVIDIA Nemotron开放模型许可证。此次发布旨在解决文本生成中的吞吐量瓶颈问题。

传统的自回归（AR）模型逐token解码，这种串行过程限制了生成吞吐量。而离散扩散语言模型则采用不同路径：它们并行生成token并迭代优化。然而，大多数扩散语言模型使用单一网络同时处理两项任务：在每个步骤中表示干净token并去噪被污染的token。TwoTower将这两项任务分离到两个独立模块中，分别称为“上下文塔”和“去噪塔”。

TwoTower基于Nemotron-3-Nano-30B-A3B骨干网络构建，该骨干网络是开放权重的混合架构，融合了Mamba-2、自注意力和混合专家（MoE）层。每个塔包含52层：23个Mamba-2层、6个自注意力层和23个MoE层。发布的检查点包含两个塔，总计约600亿参数，每个塔每token有效参数约30亿。MoE使用128个可路由专家，其中6个激活，外加2个共享专家。

两个塔初始化为同一骨干检查点的副本，但仅去噪塔进行训练，AR上下文塔保持冻结。去噪器的训练数据量约为2.1万亿token，仅是骨干网络预训练数据（25万亿token）的一小部分。

在推理过程中，AR上下文塔对提示和已提交token进行因果处理，生成逐层的KV缓存和最终的Mamba-2状态，保留了骨干网络的自回归能力。扩散去噪塔则对噪声块进行迭代优化：在块内使用双向注意力，同时对过去的干净块保持因果性。两个塔逐层连接，去噪塔的每一层都交叉关注上下文塔的对应层，这种逐层对齐的交叉注意力提供了对骨干网络表示的多尺度访问。此外，Mamba-2层从上下文塔的Mamba状态初始化其初始状态，扩散时间步通过adaLN-single时间调节调制每一层，该模块仅增加约150万参数。

生成过程按块进行：每个块初始化为S个[MASK]token，去噪器经过T步迭代后提交该块，然后上下文塔处理已提交的token以更新其缓存。这解释了为什么多次去噪步骤仍能快于逐token解码：自回归解码每步只提交一个token，而TwoTower在早期迭代中每步提交多个token。

在2×H100 GPU上的BF16评估显示，默认操作点（置信度解掩码，阈值γ=0.8，块大小S=16）下，TwoTower在MMLU、ARC-Challenge、WinoGrande等任务上性能接近AR基线，聚合质量保留98.7%，生成吞吐量提升2.42倍。通用知识任务与基线差距在1个百分点以内，代码和数学任务下降稍多，常识和多语言任务则略有恢复或提升。降低γ可提交更多token以提升吞吐量，但会带来质量损失。

该检查点支持三种推理模式：完整双塔扩散（需2块GPU，每块约59GB BF16内存）、模拟自回归和纯自回归（后者可在单块80GB GPU上运行）。最直接的应用场景是加速批量生成，例如合成文本生成任务中，可以以少量质量下降换取大幅吞吐量提升。用户还可调节置信度阈值来权衡质量与速度。此外，由于上下文塔保留了LM head，可用于推测解码、验证或自回归评分，团队可从同一检查点运行多种模式。

TwoTower的主要优势包括：开放权重，可用于商业用途；在默认操作点下以2.42倍吞吐量保留98.7%质量；单个检查点支持多种解码模式；去噪器训练成本较低。其局限包括：完整双塔推理需要2块GPU；代码和数学任务性能下降较多；固定模型权重内存占用较高；发布版本为基础模型，未经过指令微调或对齐。