← 返回信息流
技术博客arXiv cs.CL·3 小时前

Nvidia发布Nemotron-TwoTower扩散语言模型

原标题:Nemotron-TwoTower: Diffusion Language Modeling with Pretrained Autoregressive Context

速览

Nvidia发布了Nemotron-TwoTower,这是一种基于Nemotron-3-Nano-30B-A3B模型的块状自回归扩散语言模型。该模型将上下文表示与迭代去噪角色解耦,采用冻结的自回归上下文塔和可训练的扩散去噪塔。实验显示,该模型在保留98.7%自回归基线质量的同时,将实际生成吞吐量提高了2.42倍。

AI 深度解读

Nemotron-TwoTower:利用预训练自回归上下文实现扩散语言建模

背景

在自然语言处理(NLP)领域,自回归(Autoregressive, AR)模型长期以来占据主导地位。这类模型通过逐个预测下一个 token 来生成文本,虽然质量高,但受限于串行生成的特性,推理速度往往成为瓶颈。相比之下,扩散语言模型(Diffusion Language Models)提供了一种极具潜力的替代方案。扩散模型最初在图像生成领域取得巨大成功,其核心优势在于能够并行生成和迭代优化,理论上可以显著提升生成吞吐量。

然而,现有的扩散语言模型面临着一个关键的架构瓶颈:它们通常使用单一网络同时承担两个截然不同的角色——既负责表示上下文(Context Representation),又负责迭代的去噪过程(Iterative Denoising)。这种“一石二鸟”的设计迫使同一个模型在两种功能间妥协,限制了其在任一角色上的表现能力。上下文建模需要严格的因果顺序以捕捉语言依赖,而去噪过程则需要强大的全局视野和并行处理能力。将两者强行耦合,往往导致模型容量无法被充分利用,进而影响最终的生成质量或效率。

核心内容

为了解决上述架构限制,NVIDIA 的研究团队提出了 Nemotron-TwoTower 架构。这是一种基于块(block-wise)的自回归扩散模型,其核心创新在于将上下文表示和去噪任务解耦为两个独立的“塔”(Towers):

  1. 冻结的 AR 上下文塔(Frozen AR Context Tower): 该塔负责因果地处理干净的 token(clean tokens)。由于它基于预训练的自回归模型构建,因此能够高效且准确地捕捉语言的因果结构和上下文语义。在训练过程中,该塔的参数是冻结的,不参与反向传播更新,从而保留了预训练模型强大的语言理解能力。

  2. 可训练的扩散去噪塔(Trainable Diffusion Denoiser Tower): 该塔专门负责通过迭代过程细化(refine)带有噪声的文本块。它采用双向块注意力机制(bidirectional block attention),允许模型在去噪过程中查看整个块的内容,而不仅仅是过去的 token。更重要的是,该塔通过交叉注意力机制(cross-attention)与上下文塔进行交互,从而利用上下文塔提供的丰富语义信息来指导去噪过程。

模型基础与训练数据: Nemotron-TwoTower 构建在 Nemotron-3-Nano-30B-A3B 模型之上。这是一个开源权重的混合架构模型,结合了 Mamba 和 Transformer 技术,采用混合专家(MoE)结构,总参数量为 30B,激活参数量为 3B。该模型在约 2.1 万亿(2.1T)个 token 的数据集上进行了训练。

性能表现: 实验结果表明,Nemotron-TwoTower 在保持生成质量的同时,显著提升了推理效率:

  • 质量保留:它保留了自回归基线模型 98.7% 的生成质量。这意味着在大多数实际应用场景中,其文本流畅度和逻辑连贯性与传统的自回归模型几乎无异。
  • 吞吐量提升:在墙钟时间(wall-clock time,即实际物理时间)生成吞吐量方面,Nemotron-TwoTower 比自回归基线提高了 2.42 倍。这一提升主要得益于扩散模型并行生成的特性以及去噪塔的高效并行计算能力。

此外,研究团队已公开了该模型的代码和权重,供社区进一步研究和应用。

关键要点

  • 架构解耦:Nemotron-TwoTower 的核心创新是将上下文建模和去噪任务分离为两个独立的网络塔,避免了单一网络在多任务间的容量妥协。
  • 因果与双向结合:利用冻结的自回归塔处理因果上下文,确保语义准确性;利用可训练的双向扩散塔进行并行去噪,确保生成效率。
  • 交叉注意力机制:去噪塔通过交叉注意力机制从上下文塔获取信息,实现了上下文信息与去噪过程的无缝融合。
  • 高效基座模型:基于 Nemotron-3-Nano-30B-A3B(30B 总参数,3B 激活参数)的混合 Mamba-Transformer MoE 架构,兼顾了模型容量和推理效率。
  • 显著的性能增益:在仅损失 1.3% 生成质量的前提下,实现了 2.42 倍的生成吞吐量提升,证明了扩散语言模型在大规模实际应用中的可行性。
  • 开源贡献:代码和模型权重的公开促进了扩散语言建模领域的进一步研究和创新。

意义与影响

Nemotron-TwoTower 的提出标志着扩散语言模型向实用化迈出了重要一步。长期以来,扩散模型在文本生成领域的应用受到生成速度和质量平衡难题的制约。Nemotron-TwoTower 通过创新的 TwoTower 架构,成功地在保持自回归模型高质量的同时,发挥了扩散模型并行生成的速度优势。

这一进展对 AI 基础设施和开发范式具有深远影响:

  1. 降低推理成本:2.42 倍的吞吐量提升意味着在相同硬件资源下,可以处理更多的请求或降低延迟,从而显著降低大规模部署的经济成本。
  2. 推动架构多样化:它证明了混合架构(如 Mamba-Transformer)与扩散建模结合的有效性,为未来语言模型的设计提供了新的思路,不再局限于单一的自回归范式。
  3. 促进开源生态:通过开源 Nemotron-3-Nano 系列模型及 TwoTower 的实现,NVIDIA 降低了研究人员和企业进入扩散语言建模领域的门槛,有助于加速该技术的迭代和应用落地。

总之,Nemotron-TwoTower 不仅是一个技术突破,更是语言模型从“纯自回归”向“多范式融合”演进的重要里程碑,为构建更快、更高效的下一代 AI 系统奠定了坚实基础。

查看原文 →arxiv.org