arXiv:2605.22967cs.LG2026-05

让扩散模型在生成时提前传递信息,提升效率与效果。

Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

  • 引入可学习的信道,跨步骤传递隐状态信息
  • 在编码任务上超越监督微调,推理延迟降低32%
  • 适合追求高效高质文本生成的研究者

当掩码扩散模型通过迭代精炼生成序列时,掩码位置的丰富内部计算被丢弃,导致每一步去噪都需重新计算已有信息。为避免去噪轮次间的硬重置,我们提出学习型中继表示(Relay),使模型在去噪过程中具备前瞻性,显式学习如何将潜在信息传递至后续去噪步骤。Relay引入可微的逐标记通道,在前向传播间传递信息,并通过截断时间反向传播(BPTT)训练。该框架可扩展至当前最先进扩散语言模型(DLM),且与块扩散、键值缓存等技术无缝兼容。我们在一个挑战性的数独规划任务上验证了设计合理性,随后将其应用于Fast-dLLM v2,在编码任务上表现优于标准监督微调,同时推理延迟最多降低32%。实证结果表明,最先进的DLM可通过显式训练实现跨解码步骤的隐状态传递,显著提升性能-延迟权衡。所有实验代码均已公开。

原文摘要 · Abstract (English)

When Masked Diffusion Models (MDMs) generate sequences through iterative refinement, the rich internal computation over masked positions is discarded, forcing every subsequent refinement step to recompute the valuable internal information stored as model representations. To avoid a hard reset between denoising rounds, we propose Learned Relay Representations (Relay), a method that allows MDMs to be forward-thinking when denoising by explicitly learning how to propagate latent information for the benefit of future denoising steps. Relay introduces a differentiable per-token channel that passes information between forward passes and is trained via truncated backpropagation through time (BPTT). We show that this framework can be scaled to state-of-the-art Diffusion Language Models (DLMs), and is seamlessly compatible with techniques like block diffusion and KV caching. We first provide a thorough justification of the design choices in Relay on a challenging Sudoku-based planning task. We then scale Relay to Fast-dLLM v2, a state-of-the-art DLM, outperforming standard supervised finetuning on coding tasks while reducing inference latency by up to 32%. Our empirical results demonstrate that state-of-the-art DLMs can be explicitly trained to relay latent information forward across decoding steps, advancing the performance-latency Pareto frontier. We provide code for all our experiments.

扩散模型文本生成推理加速序列建模

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。