Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Control Variate Score Matching for Diffusion Models

Khaled Kahouli, Romuald Elie, Klaus-Robert Müller, Quentin Berthet, Oliver T. Unke, Arnaud Doucet

TL;DR

This paper tackles the persistent variance issues in score-based diffusion methods when sampling from unnormalized densities by unifying two classic score estimators, DSI and TSI, under a control variate framework. It introduces CVSI, an unbiased estimator with an optimal time-dependent coefficient c*(t) that minimizes variance across the entire diffusion time, and shows how CVSI naturally interpolates between DSI and TSI. The authors derive tractable forms for c*(t) using perturbation-kernel scores and demonstrate CVSI as a robust plug-in estimator that improves both data-free training (iDEM) and diffusion sampling without extra computational cost. Empirically, CVSI achieves near-ground-truth performance on high-dimensional, multi-modal targets with far fewer energy evaluations, enhances mode recovery, and scales to challenging energy landscapes like the Double-Well potential.

Abstract

Diffusion models offer a robust framework for sampling from unnormalized probability densities, which requires accurately estimating the score of the noise-perturbed target distribution. While the standard Denoising Score Identity (DSI) relies on data samples, access to the target energy function enables an alternative formulation via the Target Score Identity (TSI). However, these estimators face a fundamental variance trade-off: DSI exhibits high variance in low-noise regimes, whereas TSI suffers from high variance at high noise levels. In this work, we reconcile these approaches by unifying both estimators within the principled framework of control variates. We introduce the Control Variate Score Identity (CVSI), deriving an optimal, time-dependent control coefficient that theoretically guarantees variance minimization across the entire noise spectrum. We demonstrate that CVSI serves as a robust, low-variance plug-in estimator that significantly enhances sample efficiency in both data-free sampler learning and inference-time diffusion sampling.

Control Variate Score Matching for Diffusion Models

TL;DR

This paper tackles the persistent variance issues in score-based diffusion methods when sampling from unnormalized densities by unifying two classic score estimators, DSI and TSI, under a control variate framework. It introduces CVSI, an unbiased estimator with an optimal time-dependent coefficient c*(t) that minimizes variance across the entire diffusion time, and shows how CVSI naturally interpolates between DSI and TSI. The authors derive tractable forms for c*(t) using perturbation-kernel scores and demonstrate CVSI as a robust plug-in estimator that improves both data-free training (iDEM) and diffusion sampling without extra computational cost. Empirically, CVSI achieves near-ground-truth performance on high-dimensional, multi-modal targets with far fewer energy evaluations, enhances mode recovery, and scales to challenging energy landscapes like the Double-Well potential.

Abstract

Diffusion models offer a robust framework for sampling from unnormalized probability densities, which requires accurately estimating the score of the noise-perturbed target distribution. While the standard Denoising Score Identity (DSI) relies on data samples, access to the target energy function enables an alternative formulation via the Target Score Identity (TSI). However, these estimators face a fundamental variance trade-off: DSI exhibits high variance in low-noise regimes, whereas TSI suffers from high variance at high noise levels. In this work, we reconcile these approaches by unifying both estimators within the principled framework of control variates. We introduce the Control Variate Score Identity (CVSI), deriving an optimal, time-dependent control coefficient that theoretically guarantees variance minimization across the entire noise spectrum. We demonstrate that CVSI serves as a robust, low-variance plug-in estimator that significantly enhances sample efficiency in both data-free sampler learning and inference-time diffusion sampling.
Paper Structure (31 sections, 5 theorems, 51 equations, 4 figures, 1 table)

This paper contains 31 sections, 5 theorems, 51 equations, 4 figures, 1 table.

Table of Contents

  1. Introduction & Background
  2. Control Variate Score Matching
  3. Variance Reduction via Control Variates
  4. Optimal Control Coefficient
  5. The Interpolated Estimator
  6. Applications in Data-Free Learning and Diffusion Sampling
  7. Data-Free Training Methods
  8. Iterated Denoising Energy Matching (iDEM)
  9. Variance Reduction for Diffusion Sampling
  10. Experiments
  11. Sampling
  12. Training
  13. Implementation Details & Computational Cost
  14. Gaussian mixture models
  15. Double-Well potential
  16. ...and 16 more sections

Key Result

Proposition 1

For any time-dependent scalar coefficient $c(t) \in \mathbb{R}$, the following estimator is an unbiased estimator of the marginal score $\nabla_{\mathbf{x}(t)} \log q_t(\mathbf{x}(t))$:

Figures (4)

  • Figure 1: Top Left: Variance of the score estimators (log scale) as a function of diffusion time $t$. DSI (blue) and TSI (orange) exhibit variance explosions at $t \to 0$ and $t \to 1$ respectively. Our CVSI (red) consistently achieves the theoretical minimum variance, staying orders of magnitude lower than the baselines in the critical regimes. We use the VP-ISSNR schedule kahouli2025tv_snr with $\eta=1.0$ and $\kappa=0$. Top Right: The derived optimal interpolation weight $\tilde{c}^*(t)$ smoothly transitions from 0 (TSI-dominant) to 1 (DSI-dominant), automatically selecting the stable estimator, dependent on the distribution. Bottom: Kernel Density Estimates of samples generated from a 20-component 2D GMM. TSI fails to resolve modes (high NLL), while CVSI recovers the ground truth structure with higher fidelity than DSI (NLL 2.492 vs. Ground Truth 2.358).
  • Figure 2: We evaluate the Negative Log-Likelihood (NLL) of samples generated from GMMs of varying complexity, using CVSI, TSI, DSI and TSM estimators. Top Row: NLL as a function of the number of components for increasing dimensionality $d \in \{2, \dots, 100\}$ (fixed at 10 MC samples). Bottom Row: NLL versus component count for varying Monte Carlo (MC) samples per step (fixed at $d=15$).
  • Figure 3: Performance comparison of iDEM with TSI (original) and with our CVSI on the 2D GMM task, using two different schedules, the original geometric schedule akhound2024iterated and the well-established KVE schedule by diff_edm. Metrics include the Wasserstein-2 distance ($\mathcal{W}_2$) between generated samples and ground truth samples, and the Negative Log-Likelihood (NLL) as a function of the number of energy function evaluations (NFE) used per training sample. The two rightmost plots show the distributions of samples generated by iDEM with TSI (left) and iDEM with CVSI (right) after training with only 8 NFE per training sample. The ground truth distribution is shown as contour lines for reference, problematic modes (deviating strongly from the ground truth) are highlighted with red circles.
  • Figure 4: Performance comparison of iDEM with TSI (original) and with our CVSI on the Double-Well potential with 4 particles (DW-4), using two different schedules, the original geometric schedule akhound2024iterated and the well-established KVE schedule by diff_edm. The left-most plot shows the Wasserstein-2 distance ($\mathcal{W}_2$) between the distributions, generated and ground truth, of the interatomic distances as a function of the number of energy function evaluations (NFE) used per training sample. The two rightmost plots show histograms representing the distributions of the interatomic distances of the generated and reference test data, using the original iDEM with TSI (left) and iDEM with our CVSI estimator (right) after training with 128 NFE per training sample.

Theorems & Definitions (9)

  • Proposition 1: Unbiased Control Variate Estimator
  • proof
  • Proposition 2: Optimal Control Coefficient
  • proof
  • Proposition 3: Tractable Optimal Coefficient
  • proof
  • Corollary 1: Boltzmann Target Distribution
  • Corollary 2: Optimal Score Interpolation
  • proof