Technical Summary: Metropolis-Adjusted Diffusion Models (MADM)

Problem Statement

Sampling from score-based diffusion models inherently incurs bias arising from two sources: the approximation error of the learned score function and the time-discretisation error introduced by numerical integrators. In the standard Predictor-Corrector framework, a predictor step (often a Probability Flow ODE) moves a sample between noise levels, followed by a corrector step (typically the Unadjusted Langevin Algorithm, ULA) to refine the sample at the current noise level.

While ULA utilizes the score function to propose moves, it is itself a biased sampler because it discretises a continuous Langevin diffusion process. Consequently, the "corrector" in the Predictor-Corrector scheme introduces its own discretisation bias, failing to leave the target marginal distribution invariant. Standard Metropolis-Hastings (MH) or Barker accept-reject mechanisms, which correct for such bias in Markov Chain Monte Carlo (MCMC), cannot be directly applied to diffusion models. This is because MH requires evaluating the ratio of target densities, $p_t(\tilde{x})/p_t(x)$, which is intractable in diffusion models where only the score function $s_\phi(x, t) \approx \nabla_x \log p_t(x)$ is available.

Methodology

The authors propose Metropolis-Adjusted Diffusion Models (MADM), a framework that replaces biased ULA correctors with accept-reject corrected Langevin steps. The core innovation is a method to compute the necessary acceptance probability using only the score function, bypassing the need for the intractable density ratio.

1. Density Ratio Identity

The authors establish that the logarithm of the target density ratio can be expressed as a line integral of the score function along the path between the current state $x$ and the proposal $\tilde{x}$:
$$\log \left( \frac{p_t(\tilde{x})}{p_t(x)} \right) = \int_0^1 \langle s(x + u(\tilde{x} - x), t), \tilde{x} - x \rangle \, du$$
This identity allows the estimation of the density ratio using only score evaluations.

2. Exact Correction via Bernoulli Factory

To achieve an exact correction without evaluating the density, the authors propose an algorithm based on a two-coin Bernoulli factory (Gonçalves et al., 2017).

• Assumption: A bound $C(x, \tilde{x})$ exists such that the absolute value of the integrand in the line integral is bounded by $C$.
• Mechanism: The algorithm constructs an unbiased estimator $W$ of the exponential of the integral (scaled by the bound) using a Poisson-truncated power series.
• Procedure: The algorithm uses a "two-coin" procedure to simulate a Barker accept-reject decision. It first rejects with a fixed probability $\alpha'$, and if not rejected, it flips a coin with probability $W$. If the coin succeeds, the proposal is accepted; otherwise, the round restarts.
• Result: This procedure yields an exact Barker-adjusted Langevin corrector that leaves the target distribution $p_t$ invariant, provided the score is exact and the bound assumption holds.

3. Efficient Approximation via Simpson's Rule

Recognizing that the exact Bernoulli factory method requires a random and potentially large number of score evaluations, the authors propose a scalable approximation.

• Method: The line integral is approximated deterministically using Simpson's 1/3 rule (a Newton-Cotes quadrature rule). This requires evaluating the score at the current state, the proposal, and the midpoint.
• Acceptance: The resulting estimate is plugged into the standard Metropolis-Hastings acceptance probability.
• Properties: While this introduces a small bias due to quadrature truncation, the approximation achieves an accuracy of order $O(h^{5/2})$ in the step size $h$. It is computationally cheap, requiring only one additional midpoint score evaluation per proposal.

Key Contributions

1. MADM Framework: Introduction of a Metropolis-adjusted corrector framework for diffusion models that systematically replaces biased ULA steps.
2. Exact Method: The first exact method for adjusting Langevin corrections in diffusion models using a two-coin Bernoulli factory, requiring only the score function and a bound on the line integral.
3. Scalable Approximation: A practical Simpson's rule approximation that achieves high-order accuracy ($O(h^{5/2})$) with near-zero marginal cost (one extra score evaluation).
4. Empirical Validation: Demonstration of improved sample quality on both synthetic and image datasets.

Results

The authors evaluate MADM on synthetic 2D datasets and standard image benchmarks (CIFAR-10, FFHQ, AFHQv2, ImageNet-64) using pre-trained models from Karras et al. (2022).

• Synthetic Datasets: On 2D datasets (spiral, funnel, Sierpinski triangle, pinwheel), the standard ancestral sampler (equivalent to a predictor with a ULA corrector) produced samples with significant outliers outside the support of the true data distribution. MADM corrector steps successfully moved these samples back to the data support, eliminating outliers.
• Image Datasets:
  • For unconditional models (CIFAR-10, FFHQ, AFHQv2), applying one MADM corrector step per predictor step yielded consistent, albeit modest, improvements in Fréchet Inception Distance (FID) compared to deterministic baselines (Heun and Euler ODE solvers).
  • For the class-conditional ImageNet-64 model, the gains were more pronounced. Using MADM with an Euler predictor yielded larger FID improvements than using a higher-order Heun predictor without correction. This suggests that MH adjustments are more effective at mitigating discretisation errors than simply increasing the order of the ODE solver when the number of steps is large.
• Computational Cost: The exact method (Algorithm 2) was found to be computationally expensive due to the random number of score evaluations required. Consequently, the experiments primarily utilized the Simpson's rule approximation (Algorithm 3), which provided consistent gains with manageable overhead.

Significance and Claims

The paper claims to provide the first exact method for adjusting Langevin corrections in diffusion models, addressing the inherent bias of ULA correctors. By leveraging the score function to estimate density ratios via line integrals, MADM enables the use of rigorous MCMC acceptance criteria in generative modeling without requiring access to the intractable data density.

The authors position MADM as a natural extension of the Predictor-Corrector framework that improves sample quality. They note that while the exact method is theoretically significant, the Simpson's rule approximation offers a practical, scalable solution that consistently improves FID on image benchmarks. The work assumes perfect score estimation, acknowledging this as an inherent limitation of diffusion models rather than a specific flaw of their method. Future work is suggested to focus on developing efficient variants of the exact method or its subroutines.