英伟达提出“思考用扩散,说话用自回归”:实现语言模型效率与质量的双赢!

近年来,大型语言模型(LLMs)在自然语言处理任务中取得了突破性进展,但其生成效率仍面临严峻挑战。主流的自回归(AR)模型在生成文本时只能逐词输出,导致 GPU 计算资源利用率低,生成速度受内存带宽限制。另一方面,扩散语言模型(dLMs)支持并行生成多个词,理论上能大幅提升吞吐量,但往往以牺牲生成质量为代价。这种"效率与质量不可兼得"的困境,成为制约语言模型实际应用的关键瓶颈。

本文提出的 TiDAR (Think in Diffusion, Talk in Autoregression)架构,首次在单一模型前向传播中实现了扩散模型的并行"思考"(草稿生成)与自回归模型的"对话"(高质量采样)。通过巧妙设计的注意力掩码,TiDAR 在保持自回归模型生成质量的同时,将吞吐量提升了 4.71 倍至 5.91 倍,突破了传统推测解码与扩散模型的局限,为高效语言模型生成提供了全新范式。

背景与相关工作

扩散语言模型(dLMs)的潜力与挑战

扩散模型通过多步去噪过程生成文本,支持并行解码多个词,但其质量往往随并行词数增加而下降。例如,Dream 和 Llada 等开源模型在生成多个词时,数学推理任务(如 GSM8K)准确率会显著降低。根本原因在于并行解码引入了"词间独立性假设",破坏了语言建模中的因果依赖关系。

推测解码(Speculative Decoding)的局限

推测解码使用一个小型模型快速生成候选词序列,再由大型模型验证。这种方法虽能加速,但草稿模型容量有限,导致接受率低,整体加速效果受限。如表 1 所示,传统推测解码方法(如 EAGLE-3)在草稿能力、并行性和验证过程上存在明显短板。

现有混合方法的不足

Block Diffusion 等模型尝试在块内并行生成,但依然无法同时实现高质量与高效率。其训练目标与推理过程不一致,导致生成质量下降。TiDAR 通过统一训练目标与推理机制,解决了这一问题。

TiDAR 方法详解

整体架构设计

TiDAR 的核心思想是将序列分为三部分:前缀词、上一步采样的词、下一步预草稿的词。在单一前向传播中,模型同时执行两项任务:

  1. 自回归采样:对上一步草稿的词进行拒绝采样,确保生成质量。

  2. 扩散草稿:基于所有可能的前缀结果,并行预生成下一步的候选词。

这种设计充分利用了 GPU 的"免费令牌槽",即在内存受限时,增加词数不会显著增加延迟(如图 1 所示)。

训练策略:扩散-自回归双模式训练

TiDAR 使用一种结构化注意力掩码,在训练时同时学习自回归与扩散目标。输入序列被分为两部分:

  • 因果注意力部分:用于自回归损失计算,模拟传统 AR 模型。

  • 块内双向注意力部分:用于扩散损失计算,支持并行生成。

训练目标公式如下:

其中,
是损失平衡因子,

分别是自回归与扩散的交叉熵损失。关键创新在于将扩散部分全部设为掩码词,简化训练策略并增强损失信号。

推理过程:并行草稿与采样

在推理时,TiDAR 通过一步扩散生成草稿词,再通过自回归拒绝采样确保质量。如图 2 所示,模型在每一步同时验证上一步草稿并预生成下一步草稿,形成无缝衔接的生成流程。

关键优化:注意力掩码复用与 KV 缓存

TiDAR 支持精确的 KV 缓存,避免了传统扩散模型因双向注意力导致的缓存效率低下问题。通过预初始化注意力掩码并动态切片(如图 3),TiDAR 在推理时无需重复计算掩码,进一步降低开销。

实验与结果分析

实验设置

TiDAR 在 1.5B 和 8B 参数规模下进行测试,基线包括 AR 模型(Qwen2.5、Qwen3)、扩散模型(Dream、Llada)和推测解码方法(EAGLE-3)。任务涵盖编码(HumanEval、MBPP)、数学(GSM8K、Minerva Math)、知识(MMLU)和常识推理(ARC、Hellaswag)。

生成任务性能

如下表 所示,TiDAR 在编码和数学任务中表现优异。例如,TiDAR 1.5B 在 HumanEval 上达到 43.29%,接近 AR 模型水平,同时每前向传播生成 7.45 个词。TiDAR 8B 在"信任扩散"模式下,质量进一步提升,生成效率达 8.25 词/步。

似然任务评估

TiDAR 支持与 AR 模型相同的似然计算方式,无需蒙特卡洛采样,评估效率更高。如表 3 所示,TiDAR 在 MMLU 和常识推理任务中与 AR 模型性能相当,显著优于传统扩散模型。

效率基准测试

如图显示,TiDAR 1.5B 和 8B 的吞吐量分别达到 AR 模型的 4.71 倍和 5.91 倍,且质量损失极小。相比之下,Block Diffusion 和 EAGLE-3 在效率-质量权衡上均不如 TiDAR。

消融研究
  • 帕累托前沿分析:如图 5 所示,TiDAR 在相同训练配方下优于 Block Diffusion 和 AR 模型,接近微调 AR 的质量,同时生成更多词。

  • 解码策略对比:如表 4 所示,TiDAR 的并行草稿与采样策略在生成质量和效率上均优于置信度解码或左到右解码。

  • 掩码策略有效性:全掩码训练(如表 5)显著提升生成质量与效率,减少训练-测试差异。

  • AR 与扩散输出信任度:如图 6 所示,TiDAR 对 AR 或扩散输出的信任度不敏感,说明模型训练均衡,质量由自回归采样保障。

图5 表4 表5 图6

讨论与局限性

TiDAR 的成功源于三大设计:

  1. 并行草稿与采样:充分利用 GPU 计算密度。

  2. 双模式训练:统一 AR 与扩散目标,增强模型能力。

  3. 注意力掩码优化:支持高效推理与缓存管理。

局限性包括:

  • 批处理大小:当前研究聚焦批大小为 1 的场景,大规模批处理性能需进一步验证。

  • 长上下文扩展:训练时序列长度翻倍(因添加掩码词),可能限制长上下文处理。

  • 系统优化:当前实现基于 PyTorch 与 Flex Attention,定制化内核可进一步提升性能。

结论

TiDAR 是首个在单一前向传播中实现扩散并行草稿与自回归采样的混合架构,成功解决了语言模型生成中的效率-质量权衡问题。实验表明,TiDAR 在保持 AR 级生成质量的同时,将吞吐量提升近 6 倍,超越了现有推测解码与扩散模型。这一成果不仅为低延迟应用提供了实用解决方案,也为未来混合语言模型架构的设计指明了方向。

相关推荐
zl_vslam2 小时前
SLAM中的非线性优-3D图优化之四元数在Opencv-PNP中的应用(五)
人工智能·算法·计算机视觉
EAIReport2 小时前
企业人力资源管理数据分析:离职因素与群体特征研究
人工智能·数据挖掘·数据分析
Paraverse_徐志斌2 小时前
基于 PyTorch + BERT 意图识别与模型微调
人工智能·pytorch·python·bert·transformer
Mintopia2 小时前
🧠 Claude Code 接入 Context7 MCP:一场上下文扩展的浪漫协作
人工智能·claude·全栈
视觉&物联智能2 小时前
【杂谈】-人工智能浪潮中的网络安全守护者:CISO的战略角色与使命
人工智能·web安全·ai·信息安全·aigc·agi·ciso
x***J3482 小时前
Docker数据挖掘开发
docker·容器·数据挖掘
Mintopia2 小时前
🖥️ 老旧设备适配 AIGC:Web 前端兼容性技术解决方案
人工智能·aigc·trae
泰迪智能科技012 小时前
数据挖掘平台建设案例分享——长春大学
人工智能·数据挖掘
aneasystone本尊2 小时前
学习 LiteLLM 的用户管理体系
人工智能