让离散扩散模型在生成时可反复修正错误,提升文本与分子生成质量。
Remasking Discrete Diffusion Models with Inference-Time Scaling
- 引入重掩码机制,允许生成后重新调整已输出的标记。
- 增加采样步数可逼近自回归模型效果,计算资源少时更稳定。
- 适合需要高精度生成的科学建模与自然语言任务。
扩散模型的成功部分源于其迭代优化能力,即在生成过程中持续修正输出。然而,现代掩码离散扩散模型缺乏这一特性:一旦生成某个标记,便无法再更新,即使产生错误。为此,我们提出重掩码扩散模型(ReMDM)采样器,一种可应用于预训练掩码扩散模型的合理方法,源自具有定制重掩码反向过程的离散扩散模型。最令人关注的是,ReMDM为离散扩散模型带来了推理时计算量扩展的能力。通过增加采样步数,ReMDM生成的自然语言输出质量接近自回归模型;而在计算预算受限时,仍能更好保持性能。ReMDM还提升了离散图像掩码扩散模型的样本质量,并在分子设计等科学领域,增强了扩散引导能力,使可控性优于传统掩码和均匀噪声扩散。代码及项目博客已公开:https://guanghanwang.com/remdm。
原文摘要 · Abstract (English)
Part of the success of diffusion models stems from their ability to perform iterative refinement, i.e., repeatedly correcting outputs during generation. However, modern masked discrete diffusion lacks this capability: when a token is generated, it cannot be updated again, even when it introduces an error. Here, we address this limitation by introducing the remasking diffusion model (ReMDM) sampler, a method that can be applied to pretrained masked diffusion models in a principled way and that is derived from a discrete diffusion model with a custom remasking backward process. Most interestingly, ReMDM endows discrete diffusion with a form of inference-time compute scaling. By increasing the number of sampling steps, ReMDM generates natural language outputs that approach the quality of autoregressive models, whereas when the computation budget is limited, ReMDM better maintains quality. ReMDM also improves sample quality of masked diffusion models for discretized images, and in scientific domains such as molecule design, ReMDM facilitates diffusion guidance and pushes the Pareto frontier of controllability relative to classical masking and uniform noise diffusion. We provide the code along with a blog post on the project page: https://guanghanwang.com/remdm
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。