Reflected Diffusion Models

Aaron Lou, Stefano Ermon

Abstract: Score-based diffusion models learn to reverse a stochastic differential equation that maps data to noise. However, they can generate values that violate natural boundary constraints, such as images that lie outside of pixel space or categorical probabilities with negative values. Since naive recovery techniques produce unnatural samples, previous methods resort to projecting to the feasible set after each diffusion step, leading to a mismatch between the training and generation processes. To address the boundary problem in a principled manner, we present Reflected Diffusion Models, which learn to reverse a reflected stochastic differential equation that perturbs data on a prescribed domain. Our method learns the score functions of the perturbed densities with a new score matching loss for bounded domains and parameterizes a reverse reflected stochastic differential equation with these scores. Additionally, to train with and evaluate likelihoods, we extend Girsanov's theorem to derive an equivalence between the new score matching loss and the ELBO of the diffusion model. We also bridge the theoretical gap with the aforementioned "projection after each step" trick; the method is precisely a sampling method for reflected stochastic differential equations and can be improved with our proposed score matching objective. Finally, on standard image generation benchmarks our method remains competitive with or surpasses the state of the art without any architectural changes.