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 without architectural modifications and, for classifier-free guidance, our approach enables fast exact sampling with ODEs and produces more faithful samples under high guidance weight.
翻译:基于得分的扩散模型学习逆向映射数据到噪声的随机微分方程。然而,对于复杂任务,数值误差会累积并导致高度不自然的样本。以往工作通过阈值化缓解这种漂移——在每次扩散步骤后将结果投影到自然数据域(如图像像素空间),但这会导致训练过程与生成过程之间的不匹配。为以原理性方式纳入数据约束,我们提出反射扩散模型,该模型逆向演化定义在数据支撑集上的反射随机微分方程。我们的方法通过广义得分匹配损失学习扰动得分函数,并扩展了标准扩散模型的关键组件,包括扩散引导、基于似然的训练及常微分方程采样。我们还弥合了与阈值化方法之间的理论鸿沟:此类方案实质是反射随机微分方程的离散化形式。在标准图像基准上,本方法无需架构修改即可与当前最优方法竞争甚至超越,且对于无分类器引导,我们的方法能通过常微分方程实现快速精确采样,并在高引导权重下生成更保真的样本。