通过熵约束解码，让掩码扩散模型采样快2-3倍
Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
- 动态解码多个掩码位置，一次预测获取更多信息
- 在代码与数学推理任务上提速2-3倍，性能不变
- 适合需要高效生成的场景，尤其适配小规模推理任务
近期掩码扩散模型（MDMs）在语言建模方面表现已接近自回归模型（ARMs）。尽管多数研究聚焦于提升性能的采样方法，但对MDMs的高效采样仍鲜有探索。我们观察到：给定部分掩码序列常可确定多个未知词元的值，这意味着单次预测蕴含了标准采样未利用的额外信息。基于此，我们提出EB-Sampler，一种无需修改架构的即插即用采样器，采用熵约束解码机制，在一次函数评估中动态解码多个位置，误差控制在预设范围内。我们将EB-Sampler纳入一类自适应采样器框架，并提供误差分析以支撑算法设计。在主流编码与数学推理基准上，该方法使当前最优MDMs采样速度提升约2-3倍，且性能无损。此外，该方法在迷宫导航、数独等小型推理任务上也表现良好，而这些任务常令ARMs难以应对。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。