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

📄 arXiv: 2512.22737v1 📥 PDF

作者: Aiwei Liu, Minghua He, Shaoxun Zeng, Sijun Zhang, Linhao Zhang, Chuhan Wu, Wei Jia, Yuan Liu, Xiao Zhou, Jie Zhou

分类: cs.CL

发布日期: 2025-12-28

备注: 23 pages, 8 figures, project page: https://wedlm.github.io/

💡 一句话要点

WeDLM：通过拓扑重排序和因果注意力，加速扩散语言模型的并行推理。

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

关键词: 扩散语言模型 并行推理 因果注意力 拓扑重排序 流式解码 prefix KV缓存 大语言模型

📋 核心要点

  1. 现有扩散语言模型依赖双向注意力，破坏了prefix KV缓存，导致重复上下文计算，推理效率低。
  2. WeDLM通过拓扑重排序，在保持因果关系的同时，将已观测token置于prefix，实现prefix缓存兼容的并行解码。
  3. 实验表明，WeDLM在保证模型质量的同时，推理速度显著提升，在推理任务上加速3倍，低熵生成任务上加速10倍。

📝 摘要（中文）

自回归(AR)生成是大语言模型(LLM)的标准解码范式，但其逐token的特性限制了推理时的并行性。扩散语言模型(DLLM)通过每步恢复多个被mask的token来提供并行解码；然而，在实践中，它们通常无法将这种并行性转化为超过优化AR引擎(例如，vLLM)的部署速度提升。一个关键原因是许多DLLM依赖于双向注意力，这打破了标准的prefix KV缓存，并强制重复上下文处理，从而损害了效率。我们提出了WeDLM，一个完全建立在标准因果注意力之上的扩散解码框架，使并行生成对prefix缓存友好。核心思想是让每个被mask的位置以所有当前观察到的token为条件，同时保持严格的因果mask，这通过拓扑重排序实现，该方法将观察到的token移动到物理prefix，同时保持它们的逻辑位置。基于此，我们引入了一种流式解码过程，该过程不断地将置信的token提交到不断增长的从左到右的prefix中，并保持固定的并行工作负载，避免了块扩散方法中常见的停止等待行为。实验表明，WeDLM保留了强大AR backbone的质量，同时提供了显著的加速，在具有挑战性的推理基准上接近3倍，在低熵生成机制中高达10倍；重要的是，我们的比较是针对由vLLM提供的AR baseline，在匹配的部署设置下，证明了扩散式解码在实践中可以优于优化的AR引擎。

🔬 方法详解

问题定义：现有的大语言模型通常采用自回归的方式进行生成，即逐个token生成，这限制了推理过程的并行性，导致速度较慢。扩散语言模型虽然可以并行生成多个token，但现有方法通常采用双向注意力机制，这与标准的prefix KV缓存不兼容，需要重复计算上下文信息，反而降低了推理效率。因此，如何在保证模型质量的前提下，充分利用扩散模型的并行性，提高推理速度，是本文要解决的核心问题。

核心思路：WeDLM的核心思路是设计一种与prefix KV缓存兼容的扩散语言模型解码框架。具体来说，通过拓扑重排序，将已观测到的token移动到物理上的prefix位置，同时保持其逻辑位置不变，从而使得每个被mask的token可以依赖于所有已观测到的token，并且仍然满足因果关系。这样，就可以利用标准的因果注意力机制，并充分利用prefix KV缓存，避免重复计算。

技术框架：WeDLM的整体框架包括以下几个主要步骤：1) Masking：随机mask掉一部分token。2) Topological Reordering：将已观测到的token移动到prefix位置，同时保持逻辑位置不变。3) Denoising：利用扩散模型恢复被mask的token。4) Streaming Decoding：不断将置信度高的token提交到prefix中，并重复上述步骤，直到生成完整的序列。

关键创新：WeDLM的关键创新在于拓扑重排序和流式解码。拓扑重排序保证了因果关系和prefix KV缓存的兼容性，流式解码则避免了传统块扩散方法中的停止等待行为，提高了生成效率。

关键设计：WeDLM的关键设计包括：1) 使用标准的因果注意力机制，保证与现有模型的兼容性。2) 设计合适的拓扑重排序算法，保证逻辑位置的正确性。3) 设计合适的流式解码策略，平衡生成质量和速度。具体的参数设置、损失函数和网络结构与backbone模型保持一致，不做特殊修改。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

WeDLM在多个基准测试中取得了显著的加速效果。在具有挑战性的推理基准测试中，WeDLM实现了接近3倍的加速。在低熵生成机制中，WeDLM实现了高达10倍的加速。这些结果表明，WeDLM在保证模型质量的前提下，能够显著提高推理速度，并且优于经过优化的AR引擎（vLLM）。

🎯 应用场景

WeDLM具有广泛的应用前景，可以应用于各种需要快速文本生成的场景，例如机器翻译、文本摘要、对话生成、代码生成等。通过提高推理速度，WeDLM可以降低部署成本，并提升用户体验。此外，WeDLM还可以应用于低延迟要求的场景，例如实时对话系统。

📄 摘要（原文）

Autoregressive (AR) generation is the standard decoding paradigm for Large Language Models (LLMs), but its token-by-token nature limits parallelism at inference time. Diffusion Language Models (DLLMs) offer parallel decoding by recovering multiple masked tokens per step; however, in practice they often fail to translate this parallelism into deployment speed gains over optimized AR engines (e.g., vLLM). A key reason is that many DLLMs rely on bidirectional attention, which breaks standard prefix KV caching and forces repeated contextualization, undermining efficiency. We propose WeDLM, a diffusion decoding framework built entirely on standard causal attention to make parallel generation prefix-cache friendly. The core idea is to let each masked position condition on all currently observed tokens while keeping a strict causal mask, achieved by Topological Reordering that moves observed tokens to the physical prefix while preserving their logical positions. Building on this property, we introduce a streaming decoding procedure that continuously commits confident tokens into a growing left-to-right prefix and maintains a fixed parallel workload, avoiding the stop-and-wait behavior common in block diffusion methods. Experiments show that WeDLM preserves the quality of strong AR backbones while delivering substantial speedups, approaching 3x on challenging reasoning benchmarks and up to 10x in low-entropy generation regimes; critically, our comparisons are against AR baselines served by vLLM under matched deployment settings, demonstrating that diffusion-style decoding can outperform an optimized AR engine in practice.