DAWN: Dependency-Aware Fast Inference for Diffusion LLMs¶
会议: NeurIPS 2026
arXiv: 2602.06953
代码: https://github.com/lizhuo-luo/DAWN
领域: LLM 其他
关键词: 扩散大语言模型, 并行解码, 注意力依赖图, 锚点引导解码, 冲突调度
一句话总结¶
DAWN 提出了一种免训练的依赖感知快速推理方法,通过过滤注意力汇偏置构建稀疏依赖图,结合锚点引导解码与冲突感知贪心调度,在保持生成质量的前提下将扩散大语言模型的推理速度提升 1.80 至 8.06 倍。
研究背景与动机¶
扩散大语言模型(diffusion LLMs, dLLMs)采用全注意力机制对所有位置进行双向建模,通过迭代去噪逐步恢复完整的文本序列。与自回归(AR)模型严格的因果逐 token 生成不同,dLLM 在每个推理步骤都能同时输出所有位置的概率分布,天然具备并行解码的能力,并有望缓解自回归架构中的长程生成延迟与反转诅咒。然而,在现有实际部署中,dLLM 往往面临严重的质量与速度权衡难题。若每一步仅去噪单个置信度最高的位置,推理速度极度迟缓;若粗暴地同时解除多个位置的掩码,生成文本的质量又会剧烈崩溃。
这种质量崩塌的核心矛盾在于:并行解码通常建立在各位置预测条件独立的假设之上,但自然语言 token 之间存在极强的语义耦合。例如在生成“full house”或“straight flush”等复合词组时,两个位置的 token 必须成对呼应;若在缺乏相互约束的情况下依赖边际概率独立解码,极易拼凑出“high house”等语法与语义无效的畸形组合。现有加速方案(如 Fast-dLLM、EB-Sampler、KLASS)多采用严格的置信度或熵阈值进行保守筛选,由于割裂了位置间的内在关联,过度限制了并行的步长与范围,使得 dLLM 固有的并行生成潜力未能得到充分释放。
本文的切入角度是直接利用推理过程中天然伴随的前向注意力图来近似 token 间的依赖关系,以此打破孤立评估各位置的保守限制。作者发现:注意力图中广泛存在与语义无关却异常聚集注意力的“注意力汇(attention sinks)”,会严重扭曲依赖关系的估计;同时,强烈依赖于高置信度已解码 token 的被诱导位置,即便自身瞬时置信度偏低,依然能保持与最终输出的高度一致性。核心 idea:通过过滤注意力汇构建稀疏有向依赖图,一方面将已确定的高置信度 token 作为锚点放宽强相关被诱导位置的解码阈值,另一方面基于依赖冲突图通过贪心独立集调度禁止强耦合的不确定位置在同一步解除掩码,从而安全最大化并行度。
方法详解¶
DAWN 是一种免训练(training-free)的依赖感知解码加速方案,其核心思想是在每个去噪迭代步骤中,通过注意力机制捕捉 token 之间的耦合结构,动态筛选出尽可能大且互不冲突的安全并行解码集合。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入: 提示词 X 与掩码序列 y^(t)"] --> B["前向传播与注意力汇过滤<br/>平均深层多头注意力并剔除汇偏置"]
B --> C["构建稀疏有向依赖图<br/>保留高于阈值 tau_edge 的边"]
C --> D["锚点引导解码<br/>筛选高置信度独立位置与锚点诱导位置 U_anchor"]
C --> E["冲突感知调度<br/>基于依赖图贪心求解最大非冲突独立集 U_conflict"]
D --> F["合并并行解掩码集合<br/>U^(t) = U_anchor ∪ U_conflict"]
E --> F
F --> G["并行更新输出 y^(t+1)"]
整体框架¶
在 dLLM 的离散掩码扩散范式中,模型从全掩码序列 \(y^{(0)} = (\text{[MASK]}, \ldots, \text{[MASK]})\) 出发,经过多轮去噪步骤逐步揭示 token。在去噪步骤 \(t\) 下,DAWN 接收当前序列状态 \(y^{(t)}\) 并执行前向传播。整体流程分为三个协同阶段:首先由依赖图构建模块计算深层注意力图,通过入度异常检测剔除注意力汇并保留显著边,构建出反映 token 依赖的有向图;随后,锚点引导解码(Anchor-Guided Decoding, AGD)识别独立的高置信度位置以及被已确认锚点单向约束的诱导位置,形成第一批安全更新候选 \(\mathcal{U}^{(t)}_{\text{anchor}}\);紧接着,冲突感知调度(Conflict-Based Scheduling, CBS)在依赖图上建立无向冲突边,对剩余满足较低置信度阈值的候选 token 贪心构造最大非冲突独立集 \(\mathcal{U}^{(t)}_{\text{conflict}}\);最终将两部分互不冲突的集合取并集,在当前步完成多位置并行解掩码。
关键设计¶
1. 过滤注意力汇的稀疏依赖图构建:消除注意力异常聚集以精确度量耦合 在利用注意力图表征 token 依赖时,dLLM 普遍存在类似因果模型中注意力汇(attention sinks)的系统性偏差——少量标点符号或特殊 token 会无视上下文语义吸收极高比例的注意力权重,且聚集位置随去噪轮次动态漂移。若直接使用原始注意力得分,这些语义贫瘠的汇节点会被错误判定为掌控全序列语义的核心中枢。为此,DAWN 提取最后 4 层全部注意力头的注意力矩阵并计算均值 \(A^{(t)} \in \mathbb{R}^{L \times L}\),定义位置 \(j\) 的平均输入注意力质量为: $\(\bar{A}^{(t)}_{j} = \frac{1}{L} \sum_{i=1}^{L} A^{(t)}_{i,j}\)$ 当 \(\bar{A}^{(t)}_{j} > \tau_{\text{sink}}\) 时将位置 \(j\) 标记为注意力汇并剔除其入边干扰,同时忽略主对角线上的自注意力。在此基础上,设置边缘过滤阈值 \(\tau_{\text{edge}}\),仅当 query 位置 \(i\) 对 key 位置 \(j\) 的注意力强度满足 \(A^{(t)}_{i,j} \ge \tau_{\text{edge}}\) 时,才在图中添加有向边 \(j \rightarrow i\),代表位置 \(i\) 的生成强烈依赖于位置 \(j\) 提供的上下文。
2. 锚点引导解码:利用可靠上下文约束放宽安全并行门槛 实验表明,一旦序列中某个关键 token 的决策已经以极高置信度确定,与之存在强依赖的关联位置(如介词短语补足、语义连带词)即使在当前步的边际置信度较低,其预测方向也已基本收敛且与最终输出高度一致。DAWN 将置信度达到高保守阈值 \(\tau_{\text{high}}\)(通常设为 0.9)的已解掩码位置以及 Prompt 位置定义为“锚点(Anchors)”。通过在依赖图中追溯从锚点出发的有向边,圈定当前被掩码的受控集合 \(\mathcal{I}^{(t)}\)。对于普通掩码位置仍执行严格的 \(\tau_{\text{high}}\) 筛选以确保独立性,而对于属于 \(\mathcal{I}^{(t)}\) 的被诱导位置,则将其置信度放宽至更激进的阈值 \(\tau_{\text{induced}}\)(如 0.70-0.75),得到第一组解掩码集合: $\(\mathcal{U}_{\text{anchor}}^{(t)} = \left\{ i \in M^{(t)} \mid c_i \ge \tau_{\text{high}} \;\lor\; (i \in \mathcal{I}^{(t)} \land c_i \ge \tau_{\text{induced}}) \right\}\)$ 该设计使模型不必等待边际概率缓慢爬升即可提前解除被锚定位置的掩码,显著扩大了早期去噪的并行规模。
3. 冲突感知贪心调度:阻断耦合位置同步决策以消除边际独立误差 针对未能被锚点覆盖的大量中间置信度候选位置,若简单降低全局阈值会导致成对耦合的 token 同时基于错误的独立假设被硬性解掩码,引发像“high house”一样的逻辑冲突。DAWN 将依赖图视为冲突判定器:只要位置 \(i\) 与位置 \(j\) 之间存在任意方向的有向边(\(i \rightarrow j\) 或 \(j \rightarrow i\)),即表明两者强相关,不能在同一步内同时解掩码。CBS 模块首先将 \(\mathcal{U}_{\text{anchor}}^{(t)}\) 及其在依赖图中的全部冲突邻居统一划入禁区 \(\mathcal{X}\);随后在剩余候选集 \(\mathcal{R} = \{i \in \mathcal{C}^{(t)} \mid c_i \ge \tau_{\text{low}}\} \setminus \mathcal{X}\) 中按置信度由高到低贪心挑选最优位置 \(i^\star\) 纳入 \(\mathcal{U}_{\text{conflict}}^{(t)}\),并在候选集中即时剔除 \(i^\star\) 及其全部冲突邻域 \(\mathcal{N}(i^\star)\)。该策略本质上在强耦合子图上求解带权最大独立集,既充分利用了低至 \(\tau_{\text{low}}\)(如 0.70-0.80)的潜在可用 token,又在结构上物理阻断了联合误差的传播。
一个完整示例¶
假设当前掩码片段包含 4 个位置 \([y_1, y_2, y_3, y_4]\),当前预测置信度分别为 \(c_1=0.92, c_2=0.72, c_3=0.76, c_4=0.78\)。系统设定 \(\tau_{\text{high}}=0.90, \tau_{\text{induced}}=0.70, \tau_{\text{low}}=0.75\)。 1. 依赖图构建:经注意力汇清洗后,识别出已解掩码的锚点 token 对 \(y_2\) 存在强注意力依赖(有向边 \(\text{Anchor} \rightarrow y_2\));同时 \(y_3\) 与 \(y_4\) 之间存在强互注意力边 \(y_3 \leftrightarrow y_4\)。 2. 锚点引导解码:\(y_1\) 满足 \(c_1 \ge 0.90\) 直接入选;\(y_2\) 作为诱导位置满足 \(c_2 \ge 0.70\) 成功破格入选。此时得到 \(\mathcal{U}_{\text{anchor}} = \{y_1, y_2\}\)。 3. 冲突调度:剩余位置为 \(\{y_3, y_4\}\),置信度均高于 \(\tau_{\text{low}}=0.75\)。CBS 先挑选置信度最高的 \(y_4\)(0.78)加入 \(\mathcal{U}_{\text{conflict}}\),紧接着将与其发生冲突的邻居 \(y_3\) 从候选集合中剔除。 4. 最终解掩码:本轮并行解掩码集合为 \(\{y_1, y_2, y_4\}\)。被冲突阻塞的 \(y_3\) 留待下一步在 \(y_4\) 确定形成的全新条件下完成条件解码,完美规避了 \(y_3, y_4\) 同时生成的幻觉风险。
实验关键数据¶
主实验¶
在 NVIDIA H100 80G 环境下,使用标准 lm-eval 评测库对 LLaDA-8B-Instruct、LLaDA-1.5、Dream-v0-Base-7B 及 Dream-v0-Instruct-7B 四款代表性扩散大模型进行端到端评估。固定生成序列长度为 256、块大小为 32。主实验涵盖数学推理(GSM8K、MATH)与代码生成(HumanEval、MBPP)四大主流基准。
| 模型 | 方法 | GSM8K (Acc. / TPS / 加速比) | MATH (Acc. / TPS / 加速比) | HumanEval (Acc. / TPS / 加速比) | MBPP (Acc. / TPS / 加速比) |
|---|---|---|---|---|---|
| LLaDA-8B-Instruct | Original (Top-1) | 77.94 / 10.32 / 1.00× | 33.20 / 14.10 / 1.00× | 40.24 / 26.46 / 1.00× | 29.60 / 9.46 / 1.00× |
| LLaDA-8B-Instruct | Confidence | 78.39 / 33.44 / 3.24× | 32.98 / 36.88 / 2.62× | 40.85 / 86.73 / 3.28× | 30.00 / 34.81 / 3.68× |
| LLaDA-8B-Instruct | KLASS | 76.12 / 23.23 / 2.25× | 32.32 / 25.27 / 1.79× | 36.59 / 56.91 / 2.15× | 27.00 / 22.62 / 2.39× |
| LLaDA-8B-Instruct | LocalLeap | 77.33 / 44.19 / 4.28× | 32.42 / 47.12 / 3.34× | 39.63 / 109.80 / 4.15× | 30.80 / 44.13 / 4.66× |
| LLaDA-8B-Instruct | DAWN (本文) | 77.94 / 44.72 / 4.33× | 32.36 / 48.16 / 3.42× | 40.24 / 108.99 / 4.12× | 30.80 / 45.17 / 4.77× |
| LLaDA-1.5 | Original (Top-1) | 81.12 / 9.65 / 1.00× | 33.36 / 12.75 / 1.00× | 43.90 / 9.26 / 1.00× | 38.80 / 3.45 / 1.00× |
| LLaDA-1.5 | Confidence | 80.74 / 32.49 / 3.37× | 33.38 / 32.87 / 2.58× | 42.68 / 23.39 / 2.53× | 38.80 / 20.40 / 5.91× |
| LLaDA-1.5 | KLASS | 77.63 / 22.63 / 2.35× | 31.80 / 23.29 / 1.83× | 39.63 / 16.49 / 1.78× | 30.00 / 13.25 / 3.84× |
| LLaDA-1.5 | LocalLeap | 80.14 / 42.60 / 4.41× | 32.44 / 42.36 / 3.32× | 42.07 / 30.21 / 3.26× | 39.20 / 27.00 / 7.83× |
| LLaDA-1.5 | DAWN (本文) | 80.82 / 43.24 / 4.48× | 31.80 / 42.84 / 3.36× | 42.07 / 29.53 / 3.19× | 37.60 / 27.80 / 8.06× |
| Dream-v0-Instruct-7B | Original (Top-1) | 76.35 / 7.30 / 1.00× | 37.80 / 17.73 / 1.00× | 53.66 / 22.66 / 1.00× | 54.20 / 8.37 / 1.00× |
| Dream-v0-Instruct-7B | Confidence | 75.20 / 28.81 / 3.95× | 38.24 / 37.92 / 2.14× | 57.32 / 52.11 / 2.30× | 55.60 / 39.09 / 4.67× |
| Dream-v0-Instruct-7B | KLASS | 71.72 / 12.93 / 1.77× | 34.66 / 18.74 / 1.06× | 54.88 / 26.71 / 1.18× | 56.80 / 18.53 / 2.21× |
| Dream-v0-Instruct-7B | LocalLeap | 73.24 / 32.94 / 4.51× | 38.10 / 43.66 / 2.46× | 54.27 / 59.28 / 2.62× | 55.00 / 42.71 / 5.10× |
| Dream-v0-Instruct-7B | DAWN (本文) | 73.16 / 32.99 / 4.52× | 38.22 / 44.18 / 2.49× | 54.88 / 60.23 / 2.66× | 55.80 / 44.03 / 5.26× |
消融实验¶
在 GSM8K 与 HumanEval 任务上测试去除各核心组件对推理效率(TPS)、模型前向评估次数(NFE)以及最终准确率(Acc.)的影响。由于依赖图构建是后续两者的拓扑底座,消融重点针对位置选择策略。
| 模型 | 配置 | GSM8K Acc. | GSM8K TPS | GSM8K NFE | HumanEval Acc. | HumanEval TPS | HumanEval NFE |
|---|---|---|---|---|---|---|---|
| LLaDA-8B-Instruct | DAWN (完整方案) | 77.94 | 44.72 | 55.76 | 40.24 | 109.00 | 61.63 |
| LLaDA-8B-Instruct | w/o AGD (去除锚点引导) | 76.80 | 22.31 | 112.90 | 41.46 | 51.51 | 131.00 |
| LLaDA-8B-Instruct | w/o CBS (去除冲突调度) | 78.47 | 33.83 | 74.69 | 41.46 | 92.52 | 72.93 |
| LLaDA-8B-Instruct | Original (单步Top-1) | 77.94 | 10.32 | 256.00 | 40.24 | 26.46 | 256.00 |
| Dream-v0-Instruct-7B | DAWN (完整方案) | 73.16 | 32.99 | 51.92 | 54.88 | 60.23 | 77.12 |
| Dream-v0-Instruct-7B | w/o AGD (去除锚点引导) | 73.31 | 29.33 | 58.13 | 54.88 | 54.79 | 83.34 |
| Dream-v0-Instruct-7B | w/o CBS (去除冲突调度) | 75.28 | 27.63 | 62.19 | 57.31 | 51.97 | 93.01 |
| Dream-v0-Instruct-7B | Original (单步Top-1) | 76.35 | 7.30 | 256.00 | 53.66 | 22.66 | 256.00 |
关键发现¶
- AGD 是吞吐量跃升的核心引擎:消融数据显示,移除 AGD 模块后,LLaDA-8B-Instruct 在 GSM8K 上的 TPS 从 44.72 骤降至 22.31,NFE 从 55.76 倍增至 112.90。这表明通过锚点上下文合法放宽阈值,大幅增加了每步可解掩码的 token 数量。
- CBS 以极微小的精度代价换取高额并发:去掉 CBS 后模型精度虽略有微升(例如 Dream-v0-Instruct 在 GSM8K 从 73.16% 升至 75.28%),但 TPS 从 32.99 掉至 27.63,NFE 从 51.92 增至 62.19。这证明 CBS 在阻断破坏性冲突的前提下,成功将低置信度候选池转化为安全的加速红利。
- 长序列与块长度扩展性优异:在生成长度 \(L \in \{128, 256, 512, 1024\}\) 的测试中,DAWN 始终保持对原始采样器的速度倍率优势;块长度在 32 附近形成最佳的 Pareto 峰值平衡。
亮点与洞察¶
- 洞察到扩散模型内部的“注意力汇”会伪造虚假的全局依赖:将异常入度均值大于阈值 \(\tau_{\text{sink}}\) 的节点识别并剔除,避免了将无意义标点或特殊符号误当作强因果前置。
- 确立了“锚点诱导一致性”假设并形式化:证实了只要前置依赖项以 \(\ge 0.9\) 置信度固化,后置从属 token 在较低置信度下依然具备极高真值保真度,打破了唯全局置信度论。
- 图论独立集与去噪调度的优雅映射:将并行采样下的复合短语冲突建模为依赖图上的冲突图,以极低开销的贪心算法化解了因边际概率独立性假设引发的文本错乱。
局限与展望¶
- 依赖图估算的近似误差:以深层自注意力权重直接等价于语义因果依赖属于经验性近似,在极端长文本或复杂语篇逻辑中可能遗漏跨段落深层因果。
- 超参数敏感度需针对模型校准:注意力汇阈值 \(\tau_{\text{sink}}\) 和边缘阈值 \(\tau_{\text{edge}}\) 在不同架构(LLaDA 与 Dream)间存在最佳取值漂移,尚未实现自适应动态调节。
- 未来方向:可进一步将依赖感知框架与动态 KV-Cache 剪枝算法或推测解码(Speculative Decoding)结合,实现多级复合加速。
相关工作与启发¶
- vs Fast-dLLM (Confidence-Aware Parallel): Fast-dLLM 仅依据孤立的标量置信度 \(\tau_{\text{high}}\) 筛选更新位置,对未达标位置完全保守闭门;DAWN 引入拓扑图解耦,利用锚点机制成功召回大量被保守规则遗漏的有效 token,实现更高吞吐。
- vs LocalLeap: LocalLeap 依赖局域决定性假说(只在已知 token 的物理空间邻近窗口放宽阈值);DAWN 摆脱了绝对物理距离束缚,通过注意力图捕获全局任意跨距的语义依存,调度更加精确。
- vs KLASS: KLASS 需计算多轮去噪步骤间的分布 KL 散度来判定稳定性,存在额外计算耗时与步序滞后;DAWN 仅依赖单步内的注意力矩阵与贪心图遍历,零额外推理训练开销。
评分¶
- 新颖性: ⭐⭐⭐⭐☆ 敏锐抓住了 dLLM 并行解掩码与自然语言强耦合的本质矛盾,用拓扑依赖图重构了解码调度。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 款不同基座扩散模型、4 个涵盖数学与代码的权威基准,包含深度消融、NFE 与参数敏感性分析。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑链条严密,动机图、形式化算法伪代码与实证分析相互印证,结构清晰明了。
- 价值: ⭐⭐⭐⭐⭐ 免训练、即插即用,显著缩减了扩散语言模型与自回归工业级推理吞吐量之间的代差。