@akshay_pachaar: NVIDIA可能刚刚解决了LLM中最大的权衡问题。每个LLM都让你在速度和品质之间做出选择。自回归…
摘要
NVIDIA推出了TwoTower,这是一种在扩散语言模型中解耦上下文表示和去噪的方法,在30B MoE主干上实现了2.42倍的吞吐量,同时保留了98.7%的自回归质量。
查看缓存全文
缓存时间: 2026/07/11 15:26
NVIDIA 可能刚刚解决了大语言模型中最大的权衡问题。
每个大语言模型都迫使你在速度和质量之间做出选择。像 GPT 这样的自回归模型质量好,但速度慢,因为它一次只生成一个 token。
扩散模型则相反。它们并行生成整个块,因此速度快,但这种速度始终以牺牲质量为代价。
在讨论 NVIDIA 的解决方案之前,值得先了解扩散模型为何会落后。
扩散语言模型从一个被掩码的 token 块开始,通过几次迭代逐步揭开掩码。每次迭代,它都会读取所有已写好的文本,并同时预测该块中被掩码的 token。
因此,同一个网络必须同时完成两项截然不同的任务:既要理解已有的上下文,又要猜测尚未存在的 token。
这相当于一组权重被向两个方向拉扯,导致它无法在任何一个任务上表现优秀。
TwoTower 通过不让一个网络承担双重任务来解决这个问题。它将一个预训练的自回归模型克隆成两个塔。
第一个是冻结的上下文塔,它读取所有已写好的内容,与原始模型完全相同。它从不参与训练,因此模型的智能得以完美保留。
第二个是可训练的降噪塔,它并行填充每个新 token 块。为了保持正确,它不断交叉注意力回上下文塔。
巧妙之处在于:两塔逐层连接,因此降噪器第5层与上下文塔第5层通信,依此类推。
这赋予了降噪器骨干网络完整的多尺度理解能力,而不仅仅是最终摘要。
结果:
↳ 生成吞吐量提升 2.42 倍 ↳ 保留原始模型 98.7% 的质量 ↳ 基于 30B 混合 Mamba-Transformer MoE 骨干网络 ↳ 仅用约 2.1T token 训练,仅为原始模型预训练所用 25T token 的一小部分
最后一个数字才是真正的亮点。这里没有任何内容是从头训练的。
它是一种适配方案,可以附加到已有的自回归模型上,将缓慢的逐个 token 解码转变为快速的并行块生成,且成本低廉。
因此,如果你正在推出任何延迟敏感的产品,这个蓝图值得研究。你再也不必在信任的模型和所需的速度之间做出选择。
权重、代码和配方均已开源。
论文链接:https://arxiv.org/pdf/2606.26493
话说回来,如果你想深入了解,我写了一份详细文章,从基本原理出发,逐步解释大语言模型推理。
文章摘录如下。
Nemotron-Labs-TwoTower: 利用预训练自回归上下文的扩散语言建模
来源:https://arxiv.org/html/2606.26493 John Kamalu∗Roger WaleffeMostofa PatwaryMohammad ShoeybiBryan Catanzaro
摘要
摘要扩散语言模型因其并行和迭代生成的潜力,成为自回归模型的有前景替代方案。然而,现有方法使用单个网络同时进行上下文表示和迭代降噪,迫使一个模型承担两种角色,限制了它在任一角色上的能力。我们提出了 TwoTower,一种逐块自回归扩散模型,将这两种角色解耦为两个塔:一个冻结的 AR 上下文塔,以因果方式处理干净的 token;一个可训练的扩散降噪器塔,具有双向块注意力,通过交叉注意力与上下文相互作用来细化含噪块。基于 Nemotron-3-Nano-30B-A3B(一个开源的 30B 混合 Mamba-Transformer MoE 模型),并在约 2.1T token 上训练,Nemotron-Labs-TwoTower 保留了自回归基线 98.7% 的质量,同时提供了 2.42 倍更高的实际生成吞吐量。我们将代码和模型权重发布在 Nemotron-Labs-TwoTower 集合中 (https://huggingface.co/collections/nvidia/nemotron-labs-twotower) 。
**footnotetext:核心贡献者。\\dagger\\daggerfootnotetext:项目负责人。联系方式:[email protected]## 1 引言
自回归(AR)语言模型是文本生成的主要范式(radford2019language; grattafiori2024llama; liu2024deepseek)。然而,一次一个 token 的解码过程造成了吞吐量瓶颈。离散扩散语言模型采用不同的方式:它并行生成 token 并迭代细化,从而提供更高的吞吐量和更细粒度的可控性(austin2021structured; sahoo2024simple; nie2025large; arriola2025block)。在这种范式中,现有模型通常在每一步降噪中使用单个解码器承担两种截然不同的角色:表示干净 token 和对被破坏的 token 降噪。这种纠缠将同一组权重向不同方向拉扯,限制了它们在任何一项任务中表现出色的能力。
arriola2025encoder 观察到这两种角色可以由不同的模块处理,提出了一种编码器-解码器架构,其中编码器表示干净 token,轻量级解码器迭代地对每个块降噪。他们最大的实验在 1.7B 规模,使用共享权重训练两个模块。在更大规模上,特别是结合了 Mamba、注意力和混合专家层(MoE)的现代混合架构上,这两种角色是否完全解耦仍有待观察。
我们提出 TwoTower,并在 Nemotron-3-Nano-30B-A3B (blakeman2025nvidia) 上实例化为 Nemotron-Labs-TwoTower。Nemotron-3-Nano-30B-A3B 是一个开源的混合模型,交织了 Mamba-2、注意力和混合专家层。我们取预训练网络的两份副本,并分配互补的角色:AR 上下文塔保持冻结,作为干净 token 上的因果模型,保留了骨干网络的自回归能力;扩散降噪器塔通过掩码扩散目标进行训练,在每个含噪块内部使用双向自注意力,并对过去干净上下文使用因果交叉注意力。
参见标题
(a) 双塔架构:一个冻结的 AR 上下文塔在每个含噪块上为扩散降噪器塔提供条件。 参见标题准确度 vs 吞吐量(相对于基线 AR)吞吐量(×基线)准确度(基线百分比)
(b) 质量-吞吐量曲线,涵盖不同置信度阈值和降噪预算,以 AR 基线为标准化。
图 1:Nemotron-Labs-TwoTower 概述。(a) 冻结的 AR 上下文塔因果地运行在提示和已提交 token 上,保留预训练骨干网络的能力,并维护可重用的 KV 缓存和 Mamba 状态。可训练的扩散降噪器塔利用层对齐的上下文注意力和上下文种子的 Mamba 状态,迭代地解析每个含噪块。(b) 相对于逐个 token 的 AR 基线的质量-吞吐量权衡。两塔逐层连接。每个降噪器层交叉注意力到上下文塔的对应层,使降噪器能够多尺度地访问骨干网络的表示。这与之前仅广播最后隐藏状态的方法形成对比。降噪器还通过自适应层归一化 (peebles2023scalable) 受扩散时间步调制,这在图像扩散 Transformer 中常见,但在掩码扩散语言模型中不太常见 (ou2024your)。
在約 2.1T token 上训练(仅为预训练骨干网络所用 25T token 的一小部分),Nemotron-Labs-TwoTower 保留了自回归基线的大部分质量,同时在 30B 混合 MoE 架构上提供了显著更高的吞吐量。在默认工作点,Nemotron-Labs-TwoTower 在解码早期每个细化步骤提交多个 token,这有助于解释观察到的相对于逐个 token AR 解码的实际加速。我们将训练配方、代码和模型权重发布在 Nemotron-Labs-TwoTower 集合中 (https://huggingface.co/collections/nvidia/nemotron-labs-twotower)。
2 方法
在本节中,我们描述双塔架构(第 2.1 节 (https://arxiv.org/html/2606.26493#S2.SS1))、掩码扩散公式(第 2.2 节 (https://arxiv.org/html/2606.26493#S2.SS2))以及训练配方(第 2.3 节 (https://arxiv.org/html/2606.26493#S2.SS3))。
2.1 双塔架构
TwoTower 是一种通用方法,可应用于任何预训练的自回归语言模型。在本工作中,我们在 Nemotron-3-Nano-30B-A3B (blakeman2025nvidia) 上实例化它,该模型包含 52 层:23 个 Mamba-2 层、6 个自注意力层和 23 个混合专家(MoE)层。
我们创建该网络的两份副本,并分配互补的角色(参见图 1(a) (https://arxiv.org/html/2606.26493#S1.F1.sf1))。给定一个提示,AR 上下文塔作为因果自回归模型作用于提示和先前已提交的 token,生成每层的 KV 对和最终的 Mamba 状态。然后,生成逐块进行:每个新块用 [MASK] token 初始化,并由扩散降噪器塔在 T 个降噪步骤中迭代细化。在降噪器的每一层,注意力层交叉注意力到对应的上下文塔 KV 缓存,而 Mamba-2 层则从对应的上下文塔 Mamba 状态种子化其初始状态。一旦块中所有 token 变为干净状态,该块即被提交;上下文塔因果地处理已提交的 token 以更新其缓存,然后继续生成下一个块。算法 1 (https://arxiv.org/html/2606.26493#alg1) 概述了这种逐块自回归生成过程。
算法 1 TwoTower 逐块生成。0: 提示
xprompt\mathbf{x}_{\text{prompt}}, 块大小
SS, 降噪步数
TT 1:
(KV,States) ← ContextTower(xprompt)(\text{KV},\text{States})\leftarrow\text{ContextTower}(\mathbf{x}_{\text{prompt}}) ⊳ 构建 AR 上下文
2: 对于每个块
b=1,2,…b=1,2,\dots 执行
3:
zb ← ([MASK],…,[MASK])\mathbf{z}^b\leftarrow(\texttt{[MASK]},\dots,\texttt{[MASK]}) 4: 对于
t=1,…,Tt=1,\dots,T 执行
5:
zb ← SampleStep(zb; DenoiserTower(zb,t,KV,States))\mathbf{z}^b\leftarrow\textsc{SampleStep}\bigl(\mathbf{z}^b;\,\text{DenoiserTower}(\mathbf{z}^b,t,\text{KV},\text{States})\bigr) 6: 结束对于
7:
xb ← zb\mathbf{x}_b\leftarrow\mathbf{z}^b ⊳ 提交块
8:
(KV,States) ← ContextTower(xb,KV,States)(\text{KV},\text{States})\leftarrow\text{ContextTower}(\mathbf{x}_b,\text{KV},\text{States}) ⊳ 更新 AR 上下文
9: 结束对于
在推理时,TwoTower 为上下文塔权重引入固定的内存占用,同时保持单个前缀缓存。只有上下文塔跨块维护 KV 和 Mamba 状态,并在块提交时更新它们,因此与序列长度相关的缓存内存规模类似于 AR 基线。
我们保持上下文塔主体不变。其最终的词汇投影/LM 头部是可选的:在默认的扩散生成路径中可以省略(上下文塔只需要生成状态),而在上下文塔用于推测性解码、验证、似然评估或 AR 评分时则保留。相反,我们对降噪器塔进行了一些架构修改,以适应扩散训练和推理。我们在下文描述每项修改,并在第 3.2 节 (https://arxiv.org/html/2606.26493#S3.SS2) 提供消融研究。
双向注意力。在被细化的块内,降噪器放宽了因果掩码:含噪 token 双向注意力到其他含噪 token,同时相对于过去的干净块保持因果性。这不会增加参数,对于标准稠密注意力核,也不会改变每层的 FLOPs。在降噪器第 ii 层,来自含噪块的查询与上下文塔第 ii 层的 KV(针对过去的块 0,…,b-1)以及降噪器自身第 ii 层的 KV(针对块 b)的连接序列进行注意力计算:
Attn(Qb(i),[K<bctx,(i);Kbden,(i)],[V<bctx,(i);Vbden,(i)]).\text{Attn}\big(\mathbf{Q}_b^{(i)},\;[\mathbf{K}_{<b}^{\text{ctx},(i)};\,\mathbf{K}_b^{\text{den},(i)}],\;[\mathbf{V}_{<b}^{\text{ctx},(i)};\,\mathbf{V}_b^{\text{den},(i)}]\big). (1) 交叉注意力是层对齐的:降噪器第 ii 层注意力到上下文塔第 ii 层。由于两个塔都从相同的预训练检查点初始化,相同索引的层在可比的表示级别上操作,使得层对齐的交叉注意力成为一种自然的配对。
双向 Mamba。我们还测试了一种无参数的双向 Mamba-2 变体,方法是从零状态从左到右和从右到左运行预训练的 Mamba-2 权重,然后对输出进行平均。这大致使每层状态空间模型(SSM)FLOPs 加倍;消融实验显示质量提升甚微,因此最终设计保持 Mamba 的因果性。
时间条件。我们使用 adaLN-single (peebles2023scalable; chen2023pixart) 使降噪器以时间步 tt 为条件。一个全局 MLP 将 tt 映射到共享的尺度、偏移和门控参数,并带有每层可学习的嵌入以实现层特定调制。在 30B 骨干网络上,这仅增加 1.5M 参数;我们将这些小模块在每个张量并行 rank 上复制而不是分片,避免了额外的通信开销。时间条件改善了采样质量(第 3.2 节 (https://arxiv.org/html/2606.26493#S3.SS2))。
专家路由。我们使用骨干网络现有的 MoE 路由机制,包括序列级负载均衡。Token 通过 adaLN 调制对噪声感知。我们不做出显式路由选择,而是允许专家从降噪目标中学习专长。
2.2 块扩散语言建模
我们在掩码扩散框架 (sahoo2024simple; shi2024simplified; ou2024your) 下训练降噪器,将其应用于每个自回归块。设 xb = (xb1,…,xbS)\mathbf{x}_b = (x_b^1,\dots,x_b^S) 表示第 b 个包含 S 个 token 的块,c<b = (KV<b,Statesb-1)\mathbf{c}_{<b} = (\text{KV}_{<b},\text{States}_{b-1}) 表示上下文塔在之前干净块上的注意力 KV 缓存以及块 b-1 之后每层 Mamba 的边界状态。我们将序列分布建模为块自回归形式:
log pθ(x) = Σb=1^B log pθ(xb | x<b), (2)
其中每个条件块分布由一个以 c<b 为条件的掩码扩散过程参数化。
对于扩散时间 t∈(0,1],前向过程将干净块 xb 破坏为含噪块 zt^b,方法是用 [MASK] 独立替换每个 token。调度 αt 是在噪声水平 t 下干净 token 保持未被掩码的概率,因此 1-αt 是掩码概率;在我们的实验中,我们使用线性调度 αt = 1 - t。使用 δy 表示 token y 的点质量,m 表示 [MASK] 点质量,每个 token 的边缘分布为:
q(zt^{b,ℓ} | xb^ℓ) = Cat(zt^{b,ℓ}; αt δ_{xb^ℓ} + (1-αt) m), (3)
其中 αt 从接近 1 减小到 0,随着 t 增加。小的 t 对应于轻微破坏的块,大的 t 则接近完全掩码的块。
降噪器预测掩码位置上的干净 token,条件为含噪块、扩散时间和上下文缓存。掩码扩散的 ELBO 激发了针对掩码 token 的时间加权负对数似然 (sahoo2024simple; shi2024simplified);在线性调度 αt = 1 - t 下,
相似文章
@LiorOnAI:现在你可以将任何LLM转换成更快的版本,而无需从头重新训练。NVIDIA刚刚在他们30B的模型上实现了这一点。她…
NVIDIA提出了一种方法,将任何LLM转换为更快的版本,方法是将模型拆分为两个副本:一个冻结用于上下文,另一个训练用于并行生成多个token,实现了2.4倍加速,且质量保留约99%,仅使用了8%的训练数据。
@PavloMolchanov: 我们发布了Nemotron-Labs-Diffusion - 首个三模式语言模型系列(3B/8B/14B),可在自回归……之间切换
NVIDIA发布了Nemotron-Labs-Diffusion,这是首个三模式语言模型系列(3B/8B/14B),通过改变注意力模式在自回归、扩散和自推测解码之间切换,实现高达4倍的实际吞吐量。
Nemotron-TwoTower:基于预训练自回归上下文的扩散语言建模
本文提出了Nemotron-TwoTower,一种扩散语言模型,通过冻结的自回归塔和可训练的扩散去噪器解耦上下文表示与去噪过程,以2.42倍吞吐量达到了基线质量98.7%的水平。
@NVIDIAAI:我们将一个300亿参数模型一分为二,并行生成token,而非逐个生成。推出Nemotron-Labs…
NVIDIA Research推出Nemotron-Labs-TwoTower,这是一种扩散语言模型,它将一个300亿参数的模型分成两半,实现并行token生成,生成速度提升2.42倍,同时保持了原始质量的98.7%。
@simplifyinAI: 研究人员刚刚通过零精度损失将 LLM 的速度提升了 8.5 倍。这项技术被称为 DFlash。它取代了缓慢的自回归…
研究人员提出了 DFlash,这是一种用块扩散模型替代自回归草稿模型的方法,在零精度损失的情况下实现了 8.5 倍的 LLM 推理加速。