提出基于注意力的采样器,让扩散语言模型生成更准且更快。
Attention-Based Sampler for Diffusion Language Models

- 用注意力矩阵列和排序决定采样顺序,理论推导更优。
- 在多个数据集上生成质量超越现有方法,支持并行采样。
- 无需训练,适合追求高效高质量生成的研究者使用。
自回归模型(ARMs)在语言建模中占据主导地位,但其严格序列采样限制了推理效率与建模灵活性。为解决这一问题,基于扩散的大型语言模型(dLLMs)被提出,具备并行采样和灵活建模潜力。然而,当前dLLMs主要依赖词元级信息进行采样,忽视全局序列结构,常导致次优结果。本文从对数似然最大化角度研究采样顺序选择问题,证明该问题为NP难,并提出基于采样秩的近似方法,使目标可计算。进一步证明,该可计算目标可通过按注意力矩阵列和降序采样实现最优。此发现为注意力引导采样提供了理论依据,提供了一种替代贪婪搜索的严谨方案。我们据此设计无需训练的新采样算法Attn-Sampler,还引入动态注意力阈值以加速实际应用。多基准实验验证了方法有效性,表明其在保持更高生成质量的同时显著提升采样并行性。
原文摘要 · Abstract (English)
Auto-regressive models (ARMs) have established a dominant paradigm in language modeling. However, their strictly sequential sampling paradigm imposes fundamental constraints on both inference efficiency and modeling flexibility. To address these limitations, diffusion-based large language models (dLLMs) have been proposed, offering the potential for parallel sampling and flexible language modeling. Despite these advantages, current dLLMs sampling strategies rely primarily on token level information, which fails to account for global sequence structure and often yields suboptimal results. In this paper, we study the sampling order selection problem from the perspective of log-likelihood maximization. We show that this problem is NP-hard and propose an optimal sampling-rank-based approximation that makes the objective computationally tractable. We further prove that the tractable objective is optimized by sampling tokens in descending order of their attention-matrix column sums. This finding provides a principled justification for attention-guided sampling and offers a theoretically grounded alternative to greedy search. We instantiate this theoretical insight in a new training-free sampling algorithm, termed Attn-Sampler, and further propose dynamic attention thresholding for practical acceleration. Extensive experiments across multiple benchmarks validate the effectiveness of our proposed method, demonstrating that it achieves superior generation quality while enhancing the sampling parallelism.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。