Efficient Diffusion LLMs via Temporal-Spatial Parallel Decoding and Confidence Extrapolation

📄 arXiv: 2605.30753v1 📥 PDF

作者: Zekai Li, Ji Liu, Yiqing Huang, Ziqiong Liu, Dong Li, Emad Barsoum

分类: cs.CL

发布日期: 2026-05-29


💡 一句话要点

提出时空并行解码与置信度外推方法,加速扩散语言模型的推理。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 扩散模型 语言模型 并行解码 置信度外推 动态控制 文本生成 推理加速 轨迹分析

📋 核心要点

  1. 扩散语言模型推理速度慢,主要原因是冗余的去噪步骤和重复的token重掩码操作。
  2. 提出Trace-aware解码框架,包含时空并行解码(TSPD)和置信度外推(CE)两个模块,加速推理过程。
  3. TSPD和CE能够减少不必要的去噪迭代,同时保持输出质量,并且可以与KV缓存等系统优化结合。

📝 摘要(中文)

基于扩散的大型语言模型(dLLMs)通过迭代去噪支持并行文本生成,但由于许多步骤花费在冗余的细化以及对最终值已经确定的token的重复重掩码上,因此推理仍然存在高延迟。先前的加速方法主要依赖于步骤局部置信度启发式方法或固定调度,这些方法对prompt和任务变化敏感,并且忽略了序列中强大的位置效应。我们将扩散解码视为一个动态控制问题,并表明token级别的去噪轨迹为可靠控制提供了关键信号。我们提出了一个具有两个组件的trace-aware解码框架。首先,时空并行解码(TSPD)使用一个轻量级的时空控制器,该控制器消耗每个token的轨迹特征,包括置信度、熵和动量,以及token位置,以决定何时token已经收敛并且可以安全地固定。其次,我们引入置信度外推(CE),这是一个无需训练的状态空间模块,它预测未来logit趋势及其不确定性,以支持主动决策,包括安全的前瞻和当轨迹振荡或置信度不足时的有针对性的稳定。TSPD和CE共同减少了不必要的去噪迭代,同时保持了输出质量,并且它们可以与诸如KV缓存之类的系统优化干净地结合。

🔬 方法详解

问题定义:扩散语言模型(dLLMs)在文本生成任务中表现出色,但其推理过程依赖于迭代去噪,计算成本高昂。现有加速方法,如基于置信度的启发式方法或固定调度,对不同的prompt和任务的适应性较差,并且忽略了序列中token的位置信息,导致效率提升有限。因此,如何减少不必要的去噪迭代,同时保持生成质量,是本文要解决的关键问题。

核心思路:本文将扩散解码过程视为一个动态控制问题,通过分析token级别的去噪轨迹,提取置信度、熵和动量等特征,并结合token的位置信息,来判断token是否已经收敛,从而避免不必要的去噪迭代。此外,通过置信度外推,预测未来logit趋势,实现主动决策,进一步提升效率。

技术框架:该方法提出了一个Trace-aware解码框架,包含两个主要模块:时空并行解码(TSPD)和置信度外推(CE)。TSPD模块利用轻量级的时空控制器,根据token的轨迹特征和位置信息,动态决定何时停止对该token的去噪。CE模块则通过状态空间模型预测未来logit趋势,支持安全的前瞻和有针对性的稳定。整体流程是,在每个去噪步骤中,TSPD和CE模块协同工作,判断哪些token可以停止去噪,哪些token需要进行稳定处理,从而减少总的去噪迭代次数。

关键创新:该方法的核心创新在于将扩散解码过程视为一个动态控制问题,并利用token级别的去噪轨迹信息进行决策。与以往依赖于步骤局部信息或固定调度的方法不同,该方法能够根据token的动态变化情况,自适应地调整去噪策略,从而更有效地减少冗余计算。此外,置信度外推模块通过预测未来logit趋势,实现了主动决策,进一步提升了效率。

关键设计:TSPD模块中的时空控制器是一个轻量级的神经网络,输入包括token的置信度、熵、动量以及位置信息,输出是该token是否需要停止去噪的概率。CE模块采用状态空间模型,对logit趋势进行建模和预测,并输出预测的不确定性。损失函数的设计目标是最大化生成质量,同时最小化去噪迭代次数。具体的网络结构和参数设置在论文中有详细描述,但此处未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文提出的TSPD和CE方法能够显著减少扩散语言模型的推理时间,同时保持甚至提高生成质量。具体的性能提升数据未知,但摘要表明该方法能够减少不必要的去噪迭代,并且可以与KV缓存等系统优化方法结合,进一步提升效率。与现有方法相比,该方法在效率和质量之间取得了更好的平衡。

🎯 应用场景

该研究成果可应用于各种需要快速文本生成的场景,例如机器翻译、文本摘要、对话系统等。通过加速扩散语言模型的推理速度,可以降低计算成本,提高用户体验,并促进扩散模型在实际应用中的普及。此外,该方法也可以推广到其他基于迭代优化的生成模型中。

📄 摘要(原文)

Diffusion-based large language models (dLLMs) support parallel text generation via iterative denoising, yet inference remains latency-heavy because many steps are spent on redundant refinement and repeated remasking of tokens whose final values are already determined. Prior acceleration methods mainly depend on step-local confidence heuristics or fixed schedules, which are sensitive to prompt and task variation and ignore strong positional effects within a sequence. We cast diffusion decoding as a dynamic control problem and show that token-wise denoising trajectories provide the key signal for reliable control. We propose a trace-aware decoding framework with two components. First, Temporal-Spatial Parallel Decoding (TSPD) uses a lightweight temporalspatial controller that consumes per-token trajectory features, including confidence, entropy, and momentum, together with token position, to decide when a token has converged and can be safely fixed. Second, we introduce Confidence Extrapolation (CE), a training-free state-space module that forecasts future logit trends with uncertainty to support proactive decisions, including safe look-ahead and targeted stabilization when trajectories are oscillatory or underconfident. Together, TSPD and CE reduce unnecessary denoising iterations while preserving output quality, and they compose cleanly with system optimizations such as KV caching.