Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Variance-Reduced Diffusion Sampling via Conditional Score Expectation Identity

Alois Duston, Tan Bui-Thanh

TL;DR

This work introduces the Conditional Score Expectation (CSE) identity for affine diffusion processes, establishing an exact link between time-evolved scores via forward dynamics. It pairs a nonparametric CSE estimator with the classical Tweedie denoising score and proves their Monte Carlo errors are negatively correlated, enabling a variance-minimizing convex blend of the two estimators. The blended score estimator, adaptable to Bayesian inverse problems through likelihood tilting of SNIS weights, improves sampling fidelity and posterior reconstruction quality without changing the underlying sampler or model. Empirical results across a controlled 6D helix manifold and challenging inverse problems (Navier–Stokes and MNIST deblurring) demonstrate reduced estimator variance, better transport properties, and closer alignment to reference posteriors compared with Tweedie or CSE alone. The framework also includes learned local score proxies and a Critic–Gate distillation path to scale the method to high-dimensional settings, with a roadmap toward neural distillation and curvature-aware embeddings for broader applicability.

Abstract

We introduce and prove a \textbf{Conditional Score Expectation (CSE)} identity: an exact relation for the marginal score of affine diffusion processes that links scores across time via a conditional expectation under the forward dynamics. Motivated by this identity, we propose a CSE-based statistical estimator for the score using a Self-Normalized Importance Sampling (SNIS) procedure with prior samples and forward noise. We analyze its relationship to the standard Tweedie estimator, proving anti-correlation for Gaussian targets and establishing the same behavior for general targets in the small time-step regime. Exploiting this structure, we derive a variance-minimizing blended score estimator given by a state--time dependent convex combination of the CSE and Tweedie estimators. Numerical experiments show that this optimal-blending estimator reduces variance and improves sample quality for a fixed computational budget compared to either baseline. We further extend the framework to Bayesian inverse problems via likelihood-informed SNIS weights, and demonstrate improved reconstruction quality and sample diversity on high-dimensional image reconstruction tasks and PDE-governed inverse problems.

Variance-Reduced Diffusion Sampling via Conditional Score Expectation Identity

TL;DR

This work introduces the Conditional Score Expectation (CSE) identity for affine diffusion processes, establishing an exact link between time-evolved scores via forward dynamics. It pairs a nonparametric CSE estimator with the classical Tweedie denoising score and proves their Monte Carlo errors are negatively correlated, enabling a variance-minimizing convex blend of the two estimators. The blended score estimator, adaptable to Bayesian inverse problems through likelihood tilting of SNIS weights, improves sampling fidelity and posterior reconstruction quality without changing the underlying sampler or model. Empirical results across a controlled 6D helix manifold and challenging inverse problems (Navier–Stokes and MNIST deblurring) demonstrate reduced estimator variance, better transport properties, and closer alignment to reference posteriors compared with Tweedie or CSE alone. The framework also includes learned local score proxies and a Critic–Gate distillation path to scale the method to high-dimensional settings, with a roadmap toward neural distillation and curvature-aware embeddings for broader applicability.

Abstract

We introduce and prove a \textbf{Conditional Score Expectation (CSE)} identity: an exact relation for the marginal score of affine diffusion processes that links scores across time via a conditional expectation under the forward dynamics. Motivated by this identity, we propose a CSE-based statistical estimator for the score using a Self-Normalized Importance Sampling (SNIS) procedure with prior samples and forward noise. We analyze its relationship to the standard Tweedie estimator, proving anti-correlation for Gaussian targets and establishing the same behavior for general targets in the small time-step regime. Exploiting this structure, we derive a variance-minimizing blended score estimator given by a state--time dependent convex combination of the CSE and Tweedie estimators. Numerical experiments show that this optimal-blending estimator reduces variance and improves sample quality for a fixed computational budget compared to either baseline. We further extend the framework to Bayesian inverse problems via likelihood-informed SNIS weights, and demonstrate improved reconstruction quality and sample diversity on high-dimensional image reconstruction tasks and PDE-governed inverse problems.
Paper Structure (48 sections, 8 theorems, 146 equations, 13 figures, 4 tables, 3 algorithms)

This paper contains 48 sections, 8 theorems, 146 equations, 13 figures, 4 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Relation to Prior Work
  3. Theory: From Exact Identities to Optimal Estimators
  4. Score-Based Sampling with the Ornstein-Uhlenbeck Process
  5. The Tweedie Identity and Denoising Score Matching
  6. The CSE Identity
  7. Optimal Blending of Complementary Estimators
  8. Opposite growth and decay of the two estimators
  9. Negative correlation of the two estimators
  10. Optimal blending as variance minimization
  11. A learned proxy for the initial score
  12. Application to Bayesian Inverse Problems
  13. Results
  14. Moderate-dimensional manifold: 6D helix GMM
  15. Characterizing the Posterior Sampling Advantage Regime
  16. ...and 33 more sections

Key Result

Lemma 3.1

Let $p_0$ be a distribution on $\mathbb{R}^d$, let $p_{t|0}(\cdot\mid y)$ denote the OU posterior defined in eq:OU_posterior, and let $s(\cdot,t)$ denote the time-$t$ score function defined in eq:OU_score_def. Then, for any $t>0$,

Figures (13)

  • Figure 1: MMD and KSD vs. number of references (lower is better) on the 6D helix GMM. Left: MMD with an RBF kernel, reflecting global mass placement and coverage. Right: KSD with an inverse multiquadric kernel, a score-based discrepancy sensitive to local geometry through the target score. Blend (proxy, dashed blue, using the diagonal learned score proxy from \ref{['sec:learned_proxy']}) is comparable with Tweedie in terms of global mass placement (left), while achieving lower KSD than Tweedie (right). The oracle Blend (solid blue, using exact $s_0$) further approaches the ground-truth floor. While CSE (green) shows high variance, the blended score estimators $\hat{s}_{\textsc{BLEND}}$\ref{['eq:snis_blended_score']} stabilize it.
  • Figure 1: Critic–and–Gate distillation on 48-D GMM (10 steps). Qualitative density projections: left column (truth), middle (Critic–Gate), right (DSM). The distilled critic, trained by \ref{['eq:critic-pop-loss']}, recovers thin filamentary sets that DSM blurs.
  • Figure 1: Anti-correlation between CSE estimator \ref{['eq:nonParametricCSE']} and Tweedie estimator \ref{['eq:nonParametricTweedie']} on the 6D helix GMM. The grey region highlights the regime where anti-correlation is strongest (near $t\approx 10^{-3}$). The diagonal-proxy curve (dashed; Diag from \ref{['sec:learned_proxy']}) preserves the negative-correlation effect that underlies variance reduction in the blended estimator \ref{['eq:nonParametricBlend']}.
  • Figure 2: Relative variance and bias (due to SNIS) of the Tweedie and CSE non-parametric score estimators as a function of time $t$. The former has low variance/bias at large $t$ but diverges at $t=0$, while the latter has low variance/bias at small $t$ but grows exponentially. For both bias and variances, the crossover occurs at the same point as for variance: $t^* = \ln(2)/2 \approx 0.347$.
  • Figure 3: Qualitative comparison on the 6D helix GMM (N = 750). Each panel displays a 2D histogram of samples projected onto the principal directions $(d_1,d_2)$ (top row) and $(d_3,d_4)$ (bottom row), with PCA fitted to the target distribution and held fixed across methods. Columns show the True distribution, Blend score (True), Blend score (Proxy) using the LR+D proxy from \ref{['sec:learned_proxy']}, Tweedie score, and CSE score.
  • ...and 8 more figures

Theorems & Definitions (19)

  • Lemma 3.1: Conditional Score Expectation (CSE)
  • Proof 1: Proof (OU case)
  • Proposition 3.2: Gaussian case: exact anti-correlation
  • Proof 2
  • Remark 3.3
  • Theorem 3.4: Negative correlation for small time $t$ and large $N_{\text{ref}}$
  • Proof 3
  • Proposition 3.5: Variance-optimal blending weight
  • Remark 3.6: On variance vs. MSE
  • Theorem A.1: CSE for Linear/Affine SDEs
  • ...and 9 more