Likelihood Matching for Diffusion Models
Abstract
We propose a Likelihood Matching approach for training diffusion models by first establishing an equivalence between the likelihood of the target data distribution and a likelihood along the sample path of the reverse diffusion. To efficiently compute the reverse sample likelihood, the equivalence, a quasi-likelihood is considered to approximate each reverse transition density by a Gaussian distribution with matched conditional mean and covariance, respectively. The score and Hessian functions for the diffusion generation are estimated by maximizing the quasi-likelihood, ensuring a consistent matching of both the first two transition moments between every two time points. A stochastic sampler is introduced to facilitate the computation that leverages both the estimated score and Hessian information. We establish consistency of the quasi-maximum likelihood estimation, and provide non-asymptotic convergence guarantees for the proposed sampler, quantifying the rates of the approximation errors due to score and Hessian estimation, dimensionality, and the number of diffusion steps. Empirical and simulation evaluations demonstrate the effectiveness of the proposed Likelihood Matching and validate the theoretical results.
Lay Summary
Many modern AI systems create images and other data by starting with random noise and gradually turning it into a realistic example. A central challenge is teaching this noise-removal process accurately. Most existing methods learn the direction in which the data should move at each step, but they do not directly train the model to make the full sequence of steps as likely as possible under the data. This can leave useful information unused, especially information about how uncertainty changes during generation. We propose Likelihood Matching, a new way to train these models by looking at the whole path from noisy data back to clean data. Instead of only learning where to move next, the method also learns how uncertain each move should be, which helps the model take better reverse steps. We provide mathematical evidence that the method becomes accurate under suitable conditions, and we design an efficient version that can be used with image models. Experiments on synthetic and image datasets show that Likelihood Matching can produce better samples than a standard training approach, especially when fewer generation steps are used.