Score-based diffusion models learn to reverse a stochastic differential equation that maps data to noise. However, for complex tasks, numerical error can compound and result in highly unnatural samples. Previous work mitigates this drift with thresholding, which projects to the natural data domain (such as pixel space for images) after each diffusion step, but this leads to a mismatch between the training and generative processes. To incorporate data constraints in a principled manner, we present Reflected Diffusion Models, which instead reverse a reflected stochastic differential equation evolving on the support of the data. Our approach learns the perturbed score function through a generalized score matching loss and extends key components of standard diffusion models including diffusion guidance, likelihood-based training, and ODE sampling. We also bridge the theoretical gap with thresholding: such schemes are just discretizations of reflected SDEs. On standard image benchmarks, our method is competitive with or surpasses the state of the art and, for classifier-free guidance, our approach enables fast exact sampling with ODEs and produces more faithful samples under high guidance weight.
翻译:基于得分的扩散模型通过学习逆转将数据映射到噪声的随机微分方程来生成样本。然而,对于复杂任务,数值误差会累积并导致高度不自然的样本。先前的工作通过阈值截断(thresholding)缓解这种漂移,即在每个扩散步骤后将结果投影到自然数据域(如图像的像素空间),但这会导致训练与生成过程之间的不匹配。为以原理性方式融入数据约束,我们提出反射扩散模型(Reflected Diffusion Models),该方法逆转的是在数据支撑集上演化的反射随机微分方程。我们的方法通过广义得分匹配损失学习扰动得分函数,并扩展了标准扩散模型的关键组件,包括扩散引导、基于似然的训练和常微分方程采样。同时,我们弥合了与阈值截断的理论差距:此类方案仅是反射SDE的离散化形式。在标准图像基准测试上,我们的方法可与当前最优方法竞争甚至超越,对于无分类器引导,该方法通过ODE实现快速精确采样,并在高引导权重下生成更保真的样本。