让语言模型更安全地一次猜多个词,提升推理速度。
Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models
- 根据注意力关系动态降低候选词信心分,避免冲突
- 在低步数下显著提升推理准确率,平均提高10个百分点
- 无需训练,适合快速部署到各类扩散模型
掩码扩散语言模型可通过每轮去噪同时揭示多个词来减少推理步数,但这种并行性脆弱:单个位置预测置信度高时,若词间存在耦合,一起确定可能不安全。现有无训练采样器如Top-$k$、Fast-dLLM和EB-Sampler主要控制揭示词数,常基于逐词得分排序,忽略选中集合内的交互。我们提出ADAS,一种无训练重排序规则,保持基采样器停止规则不变,贪婪地根据已选位置的注意力加权其预测不确定性,对每个词的置信度进行折扣。在LLaDA-8B-Base和Dream-7B-Base模型上,于GSM8K、MATH500推理基准和HumanEval、MBPP代码基准测试中,将ADAS接入三种采样器后,在匹配去噪器评估下,平均分别提升9.11%和10.46%的准确率,仅增加3.1%前向运行时间开销。
原文摘要 · Abstract (English)
Masked diffusion language models can reduce inference steps by revealing multiple tokens per denoising iteration, but this parallelism is fragile: positions that are individually confident may be unsafe to commit together when their predictions are coupled. Existing training-free samplers such as Top-$k$, Fast-dLLM, and EB-Sampler mainly control how many tokens to reveal, while often ranking candidates by token-wise scores that ignore interactions within the selected set. We propose ADAS, a training-free reranking rule that leaves the base sampler's stopping rule unchanged and greedily discounts each token-wise confidence score according to its attention to already selected positions, weighted by their prediction uncertainty. Across LLaDA-8B-Base and Dream-7B-Base on the reasoning benchmarks GSM8K and MATH500 and the code benchmarks HumanEval and MBPP, plugging ADAS into all three samplers improves low-NFE performance at matched denoiser evaluations by $9.11$ and $10.46$ percentage points on average, respectively, with $3.1\%$ per-forward runtime overhead.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。