Likelihood Matching for Diffusion Models

Technical Summary: Likelihood Matching for Diffusion Models

Problem Statement
Current mainstream diffusion models, particularly Score-Based Generative Models (SGMs), rely on Score Matching (SM) as their primary training objective. While empirically effective, SM minimizes the discrepancy between a parameterized score function and the true score function at various noise levels. Theoretically, this approach only minimizes an upper bound on the negative log-likelihood of the original data distribution rather than the likelihood itself. This indirect connection can lead to a loss of statistical efficiency compared to Maximum Likelihood Estimation (MLE), which is known for producing consistent and efficient estimators. Furthermore, standard SM focuses exclusively on first-order moment information (the score/gradient), potentially neglecting higher-order structural information contained in the transition densities.

Methodology
The authors propose Likelihood Matching (LM), a framework that directly targets the likelihood of the data distribution by leveraging the equivalence between the likelihood of the target data and the path likelihood of the reverse diffusion process.

1. Theoretical Foundation:
   The method is grounded in the property that the joint likelihood of the reverse diffusion path is equivalent to the likelihood of the original data distribution (up to constants related to the forward process). The authors establish that maximizing the exact path likelihood of the reverse process is equivalent to maximizing the data log-likelihood.

2. Quasi-Maximum Likelihood Estimation (QMLE):
   Since the exact reverse transition densities are intractable, the authors approximate them using a Quasi-Maximum Likelihood approach. They model each reverse transition density as a Gaussian distribution where the conditional mean and covariance are matched to the true reverse process.

   • Key Insight: The conditional mean and covariance of the reverse process depend not only on the score function ($\nabla \log q_t$) but also on the Hessian function ($\nabla^2 \log q_t$).
   • Parameterization: The authors parameterize both the score and the Hessian using neural networks (specifically U-Nets). To manage computational complexity, the Hessian is often approximated using a diagonal-plus-low-rank structure.

3. The Likelihood Matching Objective:
   The proposed objective function, Likelihood Matching (LM), is derived by minimizing the negative quasi-log-likelihood of the reverse transitions.

   • The objective unifies score matching and covariance matching.
   • It incorporates an implicit likelihood weighting mechanism via the matched covariance matrix $(I + \sigma^2 H_t)^{-1/2}$, which automatically adjusts the weight of the score matching term based on local curvature.
   • The training utilizes a stochastic sampler that randomly selects time steps along the diffusion trajectory to compute the loss, ensuring efficient use of evolutionary information from the forward SDE.

4. Stochastic Sampler:
   A new sampling procedure is introduced that utilizes both the learned score and Hessian information. The reverse step is defined as:
   $$\tilde{Y}_{t-1} = \hat{\mu}_{t-1|t}(\tilde{Y}_t) + \hat{\Sigma}_{t-1|t}^{1/2}(\tilde{Y}_t) Z_t$$
   where the mean $\hat{\mu}$ and covariance $\hat{\Sigma}$ are functions of the estimated score and Hessian, respectively.

Key Contributions

• Novel Training Objective: The paper introduces a quasi-likelihood-based training objective (LM) that approximates the reverse path log-likelihood, effectively combining score matching and covariance matching with implicit likelihood weighting.
• Hessian-Enhanced Sampling: A stochastic sampler is derived that leverages both score and Hessian information, offering a more accurate approximation of the reverse transition density than standard DDPM-type samplers.
• Theoretical Guarantees:
  • Convergence: Non-asymptotic convergence guarantees are provided in total variation distance. The error bound scales as $O(d^3 \log^{4.5} T / T)$ for the reverse step error, with additional terms for score and Hessian estimation errors.
  • Consistency: The paper proves the consistency of the proposed quasi-MLE under an oracle parametric setting, showing that the estimator converges to the true data parameters as sample size and diffusion steps increase.
• Efficient Implementation: The use of low-rank Hessian approximations and the Sherman-Morrison-Woodbury formula allows for efficient computation of the likelihood and sampling steps.

Experimental Results
The authors validate the method on both synthetic and real-world datasets:

• Synthetic Data: On Gaussian and t-distribution mixtures, LM consistently outperformed standard Score Matching (SM), particularly in capturing well-separated modes. On a highly coiled Swiss-roll distribution (representing curved manifolds), LM demonstrated superior performance, validating the benefit of Hessian-informed covariance.
• Parameter Estimation: In a parametric Gaussian mixture setting, LM achieved lower Mean Absolute Error (MAE) and standard error in parameter estimation compared to SM, supporting the theoretical consistency claims.
• Image Datasets (MNIST, CIFAR-10, CelebA):
  • LM achieved lower Fréchet Inception Distance (FID) and better Negative Log-Likelihood (NLL) scores compared to SM baselines.
  • Performance improved with the rank of the Hessian approximation ($r$), peaking around $r=20-30$.
  • Speed-Quality Trade-off: LM achieved lower FID for comparable wall-clock time and required fewer reverse sampling steps to reach high visual quality. For instance, on MNIST, LM produced recognizable digits after only 20 reverse iterations, whereas SM samples remained blurred.

Significance and Claims
The paper claims that Likelihood Matching offers a more statistically rigorous alternative to Score Matching by directly targeting the data likelihood through a quasi-likelihood framework. By incorporating second-order information (the Hessian), the method captures richer transition densities and improves sampling efficiency, particularly in low-step regimes and on complex, curved data manifolds. The authors emphasize that while LM incurs higher computational costs per step due to Hessian modeling, the reduction in the number of required sampling steps and the improvement in generation quality (FID and NLL) validate its effectiveness for training high-fidelity generative models. The work bridges the gap between the statistical efficiency of MLE and the practical success of diffusion models.