Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
作者: Heli Ben-Hamu, Itai Gat, Daniel Severo, Niklas Nolte, Brian Karrer
分类: cs.LG
发布日期: 2025-05-30
💡 一句话要点
提出基于熵界非掩蔽的EB-Sampler，加速掩码扩散模型采样过程。
🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)
关键词: 掩码扩散模型 采样算法 熵界非掩蔽 语言建模 加速采样
📋 核心要点
- 现有掩码扩散模型采样效率低，单个预测包含的额外信息未被充分利用。
- 提出EB-Sampler，通过熵界非掩蔽过程，一次性动态非掩蔽多个token。
- 实验表明，EB-Sampler在编码、数学推理、迷宫导航和数独等任务上加速采样，且性能无损。
📝 摘要（中文）
最近的掩码扩散模型(MDMs)在语言建模方面表现出与自回归模型(ARMs)相媲美的性能。虽然大多数文献都集中在提高性能的采样过程上，但对MDMs的有效采样却鲜有探索。我们观察到，通常一个给定的部分掩码token序列能够确定多个未知token的值，这意味着掩码模型的单个预测包含标准采样过程未使用的额外信息。基于此，我们引入了EB-Sampler，这是一个简单的即插即用替代现有采样器的方法，它利用熵界非掩蔽过程，在预定义的近似误差容限下，通过一次函数评估动态地非掩蔽多个token。我们将EB-Sampler构建为广泛的自适应采样器系列的一部分，并为其提供误差分析，从而证明我们算法选择的合理性。EB-Sampler在标准编码和数学推理基准测试中，将当前最先进的MDMs的采样速度提高了约2-3倍，且性能没有损失。我们还验证了相同的程序在较小的推理任务（包括迷宫导航和数独）中也能很好地工作，而这些任务通常是ARMs难以处理的。
🔬 方法详解
问题定义：论文旨在解决掩码扩散模型（MDMs）采样效率低下的问题。现有的采样方法通常一次只预测一个被掩码的token，忽略了模型单次预测中可能包含的关于多个token的信息。这种低效的采样方式限制了MDMs在实际应用中的潜力，尤其是在需要快速生成序列的任务中。
核心思路：论文的核心思路是利用模型单次预测中包含的关于多个token的信息，通过一次性非掩蔽多个token来加速采样过程。具体而言，论文提出了一种基于熵界的非掩蔽（Entropy Bounded Unmasking）策略，该策略根据模型预测的不确定性（熵）来动态地决定每次非掩蔽的token数量，从而在保证采样质量的同时提高采样效率。
技术框架：EB-Sampler可以作为现有采样器的即插即用替代品。其主要流程如下： 1. 输入：部分掩码的序列。 2. 预测：使用掩码扩散模型预测被掩码的token的概率分布。 3. 熵计算：计算每个被掩码token的预测概率分布的熵。 4. 非掩蔽：根据熵界，选择一批token进行非掩蔽。熵越低的token，表示模型对其预测越确定，越有可能被非掩蔽。 5. 重复步骤1-4，直到所有token都被非掩蔽。
关键创新：论文的关键创新在于提出了熵界非掩蔽策略。该策略能够动态地调整每次非掩蔽的token数量，从而在采样效率和采样质量之间取得平衡。与传统的每次只非掩蔽一个token的方法相比，EB-Sampler能够更有效地利用模型预测中的信息，从而加速采样过程。此外，论文还提供了对EB-Sampler的误差分析，为算法的选择提供了理论依据。
关键设计：EB-Sampler的关键设计包括： 1. 熵的计算方式：论文使用预测概率分布的熵来衡量模型预测的不确定性。 2. 熵界的设定：熵界决定了每次非掩蔽的token数量。熵界的设定需要根据具体的任务和模型进行调整，以在采样效率和采样质量之间取得平衡。 3. 误差容限：EB-Sampler允许一定的近似误差，以换取更高的采样效率。误差容限的设定也需要根据具体的任务和模型进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，EB-Sampler在标准编码和数学推理基准测试中，将当前最先进的MDMs的采样速度提高了约2-3倍，且性能没有损失。此外，该方法在迷宫导航和数独等推理任务上也表现出良好的性能，验证了其在不同任务上的泛化能力。
🎯 应用场景
该研究成果可应用于各种需要快速序列生成的领域，如代码生成、数学推理、自然语言生成等。通过加速掩码扩散模型的采样过程，可以提高这些应用的效率和响应速度，使其更具实用性。此外，该方法在迷宫导航和数独等推理任务上的成功应用，也表明其具有解决复杂问题的潜力。
📄 摘要（原文）
Recent masked diffusion models (MDMs) have shown competitive performance compared to autoregressive models (ARMs) for language modeling. While most literature has focused on performance enhancing sampling procedures, efficient sampling from MDMs has been scarcely explored. We make the observation that often a given sequence of partially masked tokens determines the values of multiple unknown tokens deterministically, meaning that a single prediction of a masked model holds additional information unused by standard sampling procedures. Based on this observation, we introduce EB-Sampler, a simple drop-in replacement for existing samplers, utilizing an Entropy Bounded unmasking procedure that dynamically unmasks multiple tokens in one function evaluation with predefined approximate error tolerance. We formulate the EB-Sampler as part of a broad family of adaptive samplers for which we provide an error analysis that motivates our algorithmic choices. EB-Sampler accelerates sampling from current state of the art MDMs by roughly 2-3x on standard coding and math reasoning benchmarks without loss in performance. We also validate the same procedure works well on smaller reasoning tasks including maze navigation and Sudoku, tasks ARMs often struggle with.