用KL散度加速掩码扩散模型采样，提升速度同时保持生成质量。
KLASS: KL-Guided Fast Inference in Masked Diffusion Models
- 基于令牌级KL散度识别高置信度预测，每轮解码多个标记。
- 推理速度提升最高达2.78倍，且优于贪心解码性能。
- 无需额外训练，适用于文本、图像、分子生成等多种场景。
掩码扩散模型在语言生成等任务中表现优异，但其迭代优化过程导致推理速度慢且固定。为此，我们提出一种名为KL-自适应稳定性采样（KLASS）的快速高效采样方法，利用令牌级别的KL散度识别稳定且高置信度的预测结果。通过在每轮迭代中无需额外训练即可解码多个标记，该方法显著加快生成速度，同时保持样本质量。在推理基准测试中，KLASS实现高达2.78倍的墙时加速，性能优于标准贪心解码，并在扩散模型采样器中达到最先进水平。我们在文本、图像和分子生成等多个领域验证了其有效性，证明其作为通用采样器在不同模型上的广泛适用性。
原文摘要 · Abstract (English)
Masked diffusion models have demonstrated competitive results on various tasks including language generation. However, due to its iterative refinement process, the inference is often bottlenecked by slow and static sampling speed. To overcome this problem, we introduce `KL-Adaptive Stability Sampling' (KLASS), a fast yet effective sampling method that exploits token-level KL divergence to identify stable, high-confidence predictions. By unmasking multiple tokens in each iteration without any additional model training, our approach speeds up generation significantly while maintaining sample quality. On reasoning benchmarks, KLASS achieves up to $2.78\times$ wall-clock speedups while improving performance over standard greedy decoding, attaining state-of-the-art results among diffusion-based samplers. We further validate KLASS across diverse domains, including text, image, and molecular generation, showing its effectiveness as a broadly applicable sampler across different models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。