WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference

自回归（AR）生成是大语言模型（LLM）的标准解码范式，但其逐个生成token的特性限制了推理阶段的并行性。扩散语言模型（DLLM）通过每步恢复多个被掩码的token实现并行解码；然而在实践中，它们往往难以将这种并行性转化为相比优化后的自回归引擎（例如vLLM）的实际部署速度优势。一个关键原因在于，许多DLLM依赖双向注意力机制，这破坏了标准的前缀KV缓存机制，导致必须反复进行上下文重构，从而损害了效率。为此，我们提出了WeDLM——一种完全基于标准因果注意力的扩散解码框架，旨在使并行生成过程对前缀缓存友好。其核心思想是：在保持严格因果掩码的前提下，让每个被掩码的位置都能基于所有当前已观测到的token进行条件建模，这一目标通过“拓扑重排序”（Topological Reordering）实现，即将已观测到的token移至物理上的前缀位置，同时保留其逻辑上的原始位置。基于这一特性，我们进一步提出了一种流式解码流程，持续将高置信度的token加入不断增长的从左到右的前缀序列中，并维持固定的并行计算负载，从而避免了块状扩散方法中常见的“停等”行为。实验表明，WeDLM在保持强自回归主干模型生成质量的同时，实现了显著的速度提升：在具有挑战性的推理任务上接近3倍加速，在低熵生成场景下甚至可达10倍；尤为关键的是，我们的对比是在与vLLM驱动的自回归基线相同部署条件下进行的，证明了扩散式解码在实际应用中确实可以超越优化后的自回归引擎。

论文试图解决当前大语言模型（LLM）推理过程中自回归（AR）解码因逐token生成而导致的低效问题。尽管扩散语言模型（DLMs）理论上支持并行解码，但由于广泛采用双向注意力机制，破坏了前缀KV缓存，导致实际部署中无法超越高度优化的AR引擎（如vLLM）的速度。这并非一个全新问题，但如何在保持生成质量的同时实现真正高效的并行解码，仍是未被充分解决的挑战。

提出WeDLM，一种基于纯因果注意力的扩散解码框架，通过拓扑重排序（Topological Reordering）将已观测token移至物理前缀位置，同时保持其逻辑位置不变，从而在严格因果掩码下实现每个掩码位置对上下文的条件依赖。这一设计使WeDLM兼容标准前缀KV缓存，避免重复计算。进一步引入流式解码机制，持续提交高置信度token并维持固定并行负载，消除传统块状扩散方法中的停等瓶颈。相比现有DLMs，该方法首次实现了与高效AR系统公平比较下的真实速度优势。

实验表明WeDLM在保持强AR骨干模型生成质量的同时，在复杂推理任务上接近3倍加速，在低熵生成场景下最高达10倍加速，且对比基线为使用vLLM服务的AR模型，具备实际部署意义。流式解码策略提升了硬件利用率。论文强调了工程部署效率与算法设计协同的重要性。目前未提及代码是否开源，值得后续关注。未来工作可探索更优的置信度估计策略、适配更多模型架构及扩展到多模态生成任务。

1. FlowSeq: Generative PreTraining from Autoregressive to Bidirectional Generation 2. Mask-Predict: Parallel Decoding of Conditional Masked Language Models 3. Diffusion-LM Improves Controllable Text Generation 4. Parallel Text Generation with Deterministic Non-Autoregressive Decoding 5. CXLMT: A Continuous Linear Model for Sequence Generation 6. Simple and Effective Parallel Decoding in Diffusion Language Models