Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models¶

会议: NeurIPS2026
arXiv: 2606.10829
代码: https://github.com/yusufsahin99/ADAS
领域: LLM 效率 / 文本生成
关键词: 掩码扩散语言模型、并行解码、注意力折扣、依赖感知采样、不确定性加权

一句话总结¶

ADAS 在不训练模型、不改变基础采样器停止规则的前提下，用候选位置对本轮已选位置的注意力及后者的不确定性动态折扣置信度，在低去噪调用次数条件下改善并行解码质量；两种模型的任务与方法平均收益分别为 9.11 和 10.46 个百分点。

研究背景与动机¶

掩码扩散语言模型不是每次只写下一个词，而是在部分遮蔽的序列上预测所有待填位置，再一次揭示其中若干位置。减少去噪器调用次数（number of function evaluations，NFE）的直接办法是每轮多填一些词，但同一轮的预测都来自尚未包含这些新词的旧上下文。若两个位置相互依赖，它们各自看起来很可信，并不意味着同时固定两者就能得到一致的文本。数学步骤和代码结构尤其容易受到这种同步承诺的影响。

Top-k、Fast-dLLM 和 EB-Sampler 分别用固定数量、置信度条件和熵预算控制一轮填多少词，通常仍按单位置分数排队。减少一轮的数量能够缓解错误，却不能回答“相同数量下，哪些词适合一起填”。理论上，单位置预测误差与把联合分布当作独立边缘分布乘积产生的误差并不是一回事；挑出最高置信度的几个位置，只直接照顾了前一类问题。另一条路线是训练规划器、搜索未来解码轨迹，或用注意力构造硬冲突图，但这些方法会引入额外训练、计算或新的批次约束。

本文选择只替换候选排序这一小步：同一次前向传播已经产生了注意力，可以用它估计某个候选是否依赖本轮刚选中的其他位置，而不必为每个候选额外运行模型。注意力也不是“必须禁止同时更新”的证据，因此作者使用连续惩罚而非二元冲突。核心 idea：保留基础采样器对“选够了没有”的判断，只让“下一项选谁”取决于它对本轮已选、不够确定位置的有向注意力。

方法详解¶

整体框架¶

ADAS 是推理时的贪心候选选择规则，不是新的去噪网络。输入是当前部分遮蔽序列的一次去噪结果：各位置的词分布、最高概率对应的置信度，以及最后一层自注意力在各头之间的平均值。输出是本轮同时揭示的位置集合；词值仍来自基础模型已有的预测，不由折扣分数重新生成。

一轮开始时，已选集合为空，所有待填位置先使用原始置信度。选择当前得分最高的候选后，先用基础采样器的可接纳条件检查它能否加入；若可以，加入集合并降低其余候选的得分，再重复。若不可以，立即停止构造，将已经选中的位置一起揭示。下一次模型前向传播读取这些新词后，重新计算分布、注意力和本轮集合。

这里的“已选”专指当前轮内选中但尚未一起写入上下文的位置，不是前几轮已经揭示的文本。组内排序过程只更新缓存中的得分，不在每次选中一个位置后重新运行去噪器。正是这种区分，使 ADAS 能估计联合承诺风险，又不把并行生成退化成逐词前向计算。

关键设计¶

1. 置信度起点：先保留基础模型的单位置价值判断

对待填位置，置信度是模型给最可能词的概率，而不是熵、采样温度或折扣之后的分数。设待填位置集合为 \(\mathcal{M}\)，本轮已选集合为 \(S\)，原始置信度为 \(c_i\)。它保留了“模型越确定，越值得先固定”的常规直觉，也让 ADAS 可以直接接入原本使用置信度的采样器。

当 \(S\) 为空时，还没有联合承诺带来的惩罚，因此 ADAS 的第一个选择与置信度排序相同。差异从第二个选择开始：传统排序继续拿固定的原始置信度比较，ADAS 则根据第一个选择改变剩余位置的相对价值。这不是把所有注意力强的位置永久延后，而是随当前选择历史调整下一次决定。

原文中的模型误差与联合依赖误差分解提供了动机，但真实条件联合分布在推理时不可见。ADAS 不计算真实多信息，也不声称折扣就是联合误差的精确数值；置信度与注意力分别只是个体可靠性和条件影响的可计算代理。

2. 有向不确定性折扣：惩罚依赖本轮已选且不够可靠的位置

设 \(A_{is}\) 表示候选位置 \(i\) 对已选位置 \(s\) 的注意力，即查询来自候选、键来自已选位置。ADAS 的核心边际分数是：

\[ u(i\mid S,A)=c_i-\alpha\sum_{s\in S}A_{is}(1-c_s). \]

\(\alpha\) 控制折扣强度，主实验统一使用 \(\alpha=40\)。关键不是只减掉注意力，而是用已选词的 \(1-c_s\) 加权：若候选依赖一个几乎确定的词，与依赖一个尚不可靠的词，不应承受相同惩罚。这里不确定性属于已选词，不能误写成候选自己的 \(1-c_i\)。

方向也不能交换。需要估计的是“已选词被固定之后，候选的分布可能改变多少”，因此读取候选对已选词的注意力，而不是已选词对候选的注意力。算法不将两者对称化，也不把注意力超过阈值的边判成禁止共同更新。一个强依赖候选仍可能因自身置信度足够高而成为下一项。

随着 \(S\) 扩大，惩罚累加；每接受位置 \(s\)，其余候选分数只需再减去 \(\alpha A_{is}(1-c_s)\)。无需重复计算整个求和，最终选出 \(|S|\) 个位置的选择成本为 \(O(|S|\,|\mathcal{M}|)\)。这仍有组内串行选择，但通常远小于大模型前向传播的成本。

这种贪心过程依赖顺序。原文用“向有序选择序列追加元素”的效用解释它，不能改称某个无序集合目标的精确最大化：无序的双向两两惩罚会加入 ADAS 并未使用的反方向项。作者明确没有全局最优或近似比保证；局部边际分数有解释，不代表最终集合经过了全局搜索。

3. 原规则接纳：排序变化与停止条件严格分离

折扣分数只决定提议谁，接纳与停止仍使用基础规则。设提议加入后的集合为 \(T=S\cup\{i^\star\}\)，每个位置的原始预测熵为 \(H_s\)，三种条件分别为：

\[ \begin{aligned} \text{Top-}k &: |T|\leq k,\\ \text{Fast-dLLM} &: |T|\left(1-\min_{s\in T}c_s\right)<f,\\ \text{EB-Sampler} &: \sum_{s\in T}H_s-\max_{s\in T}H_s<\gamma. \end{aligned} \]

Top-k 只是数量上限；Fast-dLLM 检查集合大小乘以最不确定成员的错误概率代理；EB-Sampler 检查总预测熵扣除最大单项熵是否在预算内。置信度和熵仍来自本轮原始分布，不能把折扣分数代入这些公式，也不能说 ADAS 发明了另一种停止阈值。

算法每次找当前折扣分数最高的候选，然后检查上述条件。若它不满足条件，原文算法是直接停止，而不是跳过它、遍历所有更低分候选直至找到可行者。排序改变因此可能间接改变自适应规则最终选出的数量，即便规则和阈值本身完全相同。这也是实验不能简单按相同 \(f\) 或 \(\gamma\) 当作相同计算量比较的原因。

固定 Top-k 能更干净地隔离组内选取的作用：数量固定后，若 NFE 相同而质量变化，变化来自选了哪些位置，而不是靠多运行模型换取质量。反过来，当 \(k=1\) 时，集合为空时的第一个选择完全相同，ADAS 不应产生排序收益；附录两模型的对应结果也相同。

4. 注意力代理检验：验证条件影响，而不是把注意力等同于正确性

作者从普通 Top-8 轨迹中抽取状态，先选最可信的位置，再在不改变原轨迹的副本里只揭示该位置，观察同轮另外七个候选的预测分布变化。每个任务和模型有 100 个样本、每个样本五个状态，共 3,500 对位置；这个额外前向传播是离线诊断，不属于 ADAS 日常解码步骤。

候选对所揭示位置的注意力与干预前后 KL 变化呈正相关。跨所有位置对时，未加权注意力的 Spearman 相关为 0.284–0.545，乘以已选词不确定性后为 0.782–0.841。不过同一状态内已选词及其不确定性固定，这个乘法不会改变注意力的状态内排序；它改善的是跨状态校准，不能拿它证明每个状态内的排序都变得更优。

更具象的证据来自两种排序第一次分歧的 1,098 个状态。置信度排序原本要选、却被 ADAS 延后的词，收到第一个新词后平均 KL 变化为 0.793，替代词为 0.230；平均置信度增量分别为 0.092 和 0.024。它支持“先等上下文再固定”的解释，但自由形式答案没有唯一逐词标签，所以诊断测的是稳定性，真正的正确性由终端任务分数检验。

附录 C 进一步给出局部一阶动机：已选词表示的不确定性可由其预测错误概率代理，注意力作为候选对该表示扰动的相对敏感度代理，二者相乘形成折扣项。这依赖局部线性化与敏感度近似，并非对任意 Transformer 的严格全局风险界，也没有处理所有高阶联合关系。

一个完整示例¶

下面是为解释公式自设的数值例子，不是论文观测样本。假设三个待填位置的置信度依次为 0.90、0.88、0.85，使用 Top-2 和 \(\alpha=40\)。集合为空时，第一个位置以 0.90 被选中，与普通排序一致。

第二个候选对第一个位置的注意力为 0.02，第三个为 0.001；第一个已选词的不确定性是 0.10。两者更新后的分数分别是 \(0.88-40\times0.02\times0.10=0.80\) 与 \(0.85-40\times0.001\times0.10=0.846\)，因此第三个位置成为第二项，尽管它原始置信度更低。

最终本轮一起固定第一和第三个位置，第二个继续遮蔽。下一轮它能够在新上下文中重新预测；这不是保证它原来的词错了，也不是保证第三个一定正确，而是优先推迟更可能受到刚选位置影响的预测。若第一项几乎完全确定，则两者的注意力惩罚都会很小。

损失函数 / 训练策略¶

没有新增损失函数、微调、学习到的规划器或额外验证模型；ADAS 使用现有模型权重和预测。实现读取最后一层各头平均注意力，提取权重时需要仅在最后一层禁用 FlashAttention。完整数据集的计时估计平均每次模型前向传播增加 3.1% 开销，这不等于所有硬件或全部解码设置的端到端延迟都只增加 3.1%。

默认 \(\alpha=40\) 不是零调参结果。作者在 LLaDA-8B 的 HumanEval 上，针对两种 Top-k、两种 EB 和两种 Fast-dLLM 设置比较六个非零折扣强度，再以六设置平均 Pass@1 选定一次，之后不按任务、模型或停止规则重新调参。它在该表的六设置平均值为 22.15%，而 \(\alpha=20\) 为 21.95%；相较逐行选最佳强度的事后最优选择，默认值平均损失 0.72 个百分点、最差损失 2.44 个百分点。证据支持一个默认值的迁移性，尚不能证明对其他模型和任务普遍最优。

实验关键数据¶

主实验¶

使用 LLaDA-8B-Base 与 Dream-7B-Base，数学任务评准确率，代码任务评 Pass@1。GSM8K/MATH500/MBPP/HumanEval 分别为 8/4/3/0-shot；温度为 0，few-shot 种子为 1234，最大生成长度分别为 256/512/512/512。

下表摘自原文表 2，单元格是 LLaDA / Dream 相对对应基础采样器的绝对收益，单位为百分点，不是相对百分比。

数据集	EB + AD	Fast-dLLM + AD	Top-k + AD	三方法平均
GSM8K	+9.22 / +14.31	+6.92 / +6.36	+11.90 / +5.88	+9.35 / +8.85
MATH500	+6.49 / +4.84	+5.70 / +2.15	+10.73 / +6.93	+7.64 / +4.64
HumanEval	+10.45 / +15.63	+6.90 / +12.03	+9.55 / +10.77	+8.97 / +12.81
MBPP	+10.37 / +15.01	+7.39 / +12.04	+13.67 / +19.60	+10.48 / +15.55
数据集平均	+9.13 / +12.45	+6.73 / +8.14	+11.46 / +10.80	+9.11 / +10.46

这些收益以基础采样器平均每步至少揭示四个词的运行点为起点，在对应 ADAS 曲线上按相同 NFE 线性插值；超出插值范围的点排除。自适应规则不按相同阈值直接比较，Fast-dLLM+AD 还额外测了 \(f=20\) 来覆盖更高并行区域。因此表中结果可能是插值估计，并非每一项都有两个 NFE 完全相同的原始运行。

下面摘取附录表 18、19 中固定 Top-8 的直接测量，补充无需插值的证据。

模型	数据集 / 指标	NFE（两者相同）	Top-8（%）	Top-8 + AD（%）
LLaDA-8B-Base	GSM8K / 准确率	32.00	22.67	42.46
LLaDA-8B-Base	HumanEval / Pass@1	64.00	5.49	20.12
Dream-7B-Base	GSM8K / 准确率	32.00	13.80	28.81
Dream-7B-Base	HumanEval / Pass@1	64.00	12.20	21.95

消融实验¶

下表来自附录表 9、10，均为 LLaDA-8B-Base、HumanEval、Top-8，指标为 Pass@1。平均不确定性变体替换的是已选词特定的不确定性，而不是完全删除注意力惩罚。

配置	Pass@1（%）	说明
仅置信度，无 ADAS	5.49	原始 Top-8
首层注意力	9.76	对各注意力头取平均
中间层注意力	15.24	对各注意力头取平均
最后一层注意力 + 已选词不确定性	20.12	默认 ADAS
最后一层注意力 + 待填位置平均不确定性	17.07	用全局平均值替代特定已选词权重

关键发现¶

最后一层信号在该消融设置中优于首层和中间层；已选词特定的不确定性比全局平均权重高 3.05 个百分点。这是单模型、单任务、单并行度证据，不是层选择的普遍定律。
不存在逐点必胜：90 个 matched-NFE 运行点有 80 个提升、10 个回退，点级平均收益为 9.27，95% bootstrap 区间为 [7.74, 10.84]。该平均的权重不同于表 2 的任务/方法分组平均，不能混为同一统计量。
直接测量也有反例：GSM8K 的 Top-4 在 LLaDA 上从 53.90% 降为 50.11%，在 Dream 上从 48.45% 降为 46.47%，各自 NFE 均为 64。方法的主要价值是缓解高并行时的崩溃，而非恢复完整逐词生成质量。

亮点与洞察¶

将“选谁”与“选多少”分开，使改进能复用已有预算控制。自适应采样器仍可能因为选到不同成员而改变实际 NFE，因此模块独立不等于计算量自动相等。
不确定性来自被依赖的已选词，符合扰动来源的方向。只看候选自身置信度，或只看无权重注意力，都会遗漏这一风险来源。
连续折扣允许强依赖但高价值的位置仍被接受，比硬冲突图更柔性。代价是它只是一种可检验的代理启发式，没有独立性证明或全局优化保证。

局限与展望¶

注意力可反映语法、位置或格式，不等同于有害依赖；低注意力也不能保证联合兼容。更细的头/层选择及高阶交互值得验证。
只测了两个基础模型和数学/代码任务，开放式生成、更大模型及不同置信度校准尚未覆盖；全局折扣强度在 HumanEval 上选定，也限定了“无需重新调参”的外推范围。
当前承诺不可逆。重掩码或修订可能纠正早期错误，但会改变解码动力学与 NFE，必须重新进行等计算量比较。
温度为 0 的确定性运行没有重复种子的方差估计；运行点 bootstrap 仅描述跨配置稳定性，不能当作固定配置的逐样本显著性检验。不同 few-shot 排序的敏感性也未测试。

评分¶

新颖性: 4/5 — 将有向、已选词不确定性加权的软折扣嵌入既有停止规则，贡献边界清晰。
实验充分度: 4/5 — 双模型、四任务、三规则及干预诊断较完整，外部任务与提示方差仍缺。
写作质量: 4/5 — 算法与无全局保证说明清楚；代码状态表述仍有摘要与附录的历史差异。
价值: 4/5 — 为高并行掩码语言模型提供轻量推理改进，收益须按 NFE 与适用区域解读。