Covariance-Aware Sampling in Diffusion Models

Updated 17 May 2026

The paper introduces covariance-aware sampling techniques that leverage posterior covariance to mitigate mode collapse and enhance sample quality in diffusion models.
Methods include analytic, neural, and low-rank approximations of covariance, enabling unbiased expectation estimation and improved calibration for inverse problems.
Experiments demonstrate superior FID scores and effective sample sizes in low-step regimes, backed by rigorous theoretical guarantees on convergence and variance control.

Covariance-aware sampling for diffusion models refers to a family of techniques that explicitly estimate and utilize the conditional covariance structure of the reverse diffusion process for improved sample quality, unbiased expectation estimation, effective control, or enhanced calibration in generative and inverse problem settings. Unlike classical samplers relying solely on the predicted mean of the denoising distribution, covariance-aware methods leverage analytic, parametric, or data-driven estimates of the posterior covariance—full, diagonal, or structured—to address issues such as mode collapse, poor uncertainty quantification, high variance in importance weights, or insufficient statistical control, especially under constrained computational budgets or few-step sampling regimes. This article traces the mathematical foundations, algorithmic methodologies, and key variants of covariance-aware sampling, summarizing rigorous developments and empirical evaluations across the literature.

1. Mathematical Foundations of Covariance Estimation

The theoretical basis for covariance-aware sampling lies in the explicit characterization of the posterior $p(x_{t-1} \mid x_t)$ or $p(x_0 \mid x_t)$ in the reverse Markov chain or SDE, which, under Gaussian assumptions, is fully determined by its mean and covariance. The mean is obtainable via Tweedie’s formula: for the forward marginal $p(x_t)$ with score $s_t = \nabla_{x_t} \log p_t(x_t)$ , the posterior mean and covariance are: $\mu_t(x_t) = x_t + \sigma_t^2 s_t(x_t), \quad \Sigma_t(x_t) = \sigma_t^4 \nabla_{x_t}^2 \log p_t(x_t) + \sigma_t^2 I,$ as established in multiple works (Hamidi et al., 2024, Rissanen et al., 2024, Ou et al., 2024, Schioppa et al., 13 May 2026). These depend on the score network and the Hessian (second-order score), which is generally intractable for high-dimensional $x_t$ . To circumvent this, recent approaches employ either analytic diagonalization (Ou et al., 2024), finite-difference approximations (Hamidi et al., 2024), low-rank or Fourier-space decomposition (Rissanen et al., 2024, Schioppa et al., 13 May 2026), or amortized neural approximators (Ou et al., 2024).

2. Core Covariance-Aware Sampling Algorithms

Several algorithm families implement covariance awareness in diffusion-based generative modeling:

Variance-Tuned Diffusion Importance Sampling (VT-DIS): VT-DIS learns per-step reverse process covariances $\Sigma_\phi(n)$ (isotropic, diagonal, or low-rank variants) for a frozen pretrained score network, optimizing a trajectory-level $\alpha$ -divergence objective ( $\alpha=2$ ). This yields unbiased expectations via a trajectory-wise importance weight and achieves high effective sample size (ESS) by minimizing sample weight variance (Zhang et al., 27 May 2025).
Likelihood Matching (LM): LM explicitly parameterizes both the score and Hessian, maximizing a quasi-likelihood along the reverse path that matches both conditional mean and covariance. The resulting sampler stochastically draws from a Gaussian proposal at each reverse step, leveraging both learned score and curvature, with theoretical guarantees on TV distance and consistency (Qian et al., 5 Aug 2025).
Fourier-Structured and Low-Rank Covariance Estimation: Practical samplers for image generation employ structured covariance approximations to achieve computational tractability. For instance, (Schioppa et al., 13 May 2026) leverages a block-diagonal model in the DCT or ConvDCT basis, estimating per-block variances via Hutchinson’s estimator and integrating this structured noise within DDIM-based sampling steps via one extra Jacobian-vector product per step.
Optimal Covariance Matching (OCM): OCM amortizes the diagonal of the analytic Tweedie covariance via a dedicated neural regressor, trained with an unbiased loss involving Hessian-vector products. At test time, only a single forward pass of the covariance prediction network is required, enabling efficient, diagonal-covariance-informed sampling (Ou et al., 2024).
Covariance-Aware Posterior Sampling for Inverse Problems: In posterior inference for linear inverse problems, closed-form formulas for the covariance allow improved marginal likelihood estimation, superior posterior calibration, and enhanced reconstruction without tuning new hyperparameters. The covariance is approximated via finite differences or training-free updates, and is directly incorporated in the generation of conditional or guided samples (Hamidi et al., 2024, Rissanen et al., 2024, Zhang et al., 7 Oct 2025).

3. Structural and Computational Aspects

Several architectural choices are motivated by computational considerations:

Covariance parameterizations: Isotropic (scalar), diagonal, and block-diagonal (frequency-structured) forms balance expressivity and scalability. Full-rank parametrizations entail $O(N d^2)$ parameters and can become infeasible for high $p(x_0 \mid x_t)$ 0 (Zhang et al., 27 May 2025). Diagonal and block-structured alternatives reduce this to $p(x_0 \mid x_t)$ 1 or $p(x_0 \mid x_t)$ 2.
Estimation cost:
- Hessian-vector products (HVPs) or Jacobian-vector products (JVPs) are used for diagonal or trace estimation with minimal overhead.
- Structured estimators utilize the fast transforms (DCT, ConvDCT) for efficient projection and averaging in the frequency domain (Schioppa et al., 13 May 2026).
- Time and space updates in low-rank settings leverage online BFGS-style rank-one updates to combine data-driven and locally propagated covariance information (Rissanen et al., 2024).
Integration with existing samplers: Covariance-aware correction can be post-hoc (as in test-time trajectory adjustment (Zhang et al., 27 May 2025)), amortized in training (as in OCM (Ou et al., 2024) or LM (Qian et al., 5 Aug 2025)), or realized via lightweight controllers adapting initial noise covariances to achieve output constraints (Song et al., 7 Feb 2025).

4. Empirical Performance and Ablations

Covariance-aware samplers demonstrate notable improvements in domains with limited sampling steps or for high-stakes inverse problems:

Sampling quality at low NFE: In few-step regimes, covariance-aware extensions to DDIM surpass the best second-order samplers (Heun, DPM-Solver++) on FID for both ImageNet-128 and ImageNet-512, especially as NFE drops below 32–64 (Schioppa et al., 13 May 2026).
Posterior calibration in inverse problems: Covariance-aware diffusion posterior sampling (CA-DPS) improves FID, LPIPS, and SSIM metrics across various restoration tasks on FFHQ and ImageNet at $p(x_0 \mid x_t)$ 3 (Hamidi et al., 2024). Inpainting, deblurring, and super-resolution tasks all benefit with statistically tighter posteriors.
Unbiased expectation estimation and ESS: VT-DIS achieves effective sample sizes (ESS) up to 80% for DW-4 and 35–40% for molecular clusters (LJ-13), a significant improvement over vanilla diffusion + IS (Zhang et al., 27 May 2025).
Robustness in high dimensions: Structured and low-rank covariance approximations stably outperform purely diagonal or fixed-covariance baselines in settings with strong pixel correlations or where high-rank structure is essential (Rissanen et al., 2024).

5. Theoretical Guarantees and Limitations

The mathematical analysis of covariance-aware sampling yields rigorous guarantees and identifies challenges:

Convergence and estimation error: LM (Qian et al., 5 Aug 2025) proves non-asymptotic TV bounds on sampling error as a function of step size, score, and Hessian estimation quality, while establishing consistency of quasi-maximum-likelihood estimation.
Bias-variance tradeoffs: KL-based objectives or variational lower bound (VLB) training can lead to degenerate ESS or mode collapse in high-dimensions; alternative objectives ( $p(x_0 \mid x_t)$ 4 in VT-DIS, OCM loss) provide more robust variance control (Zhang et al., 27 May 2025, Ou et al., 2024).
Computational constraints: Full-rank adaptive covariance learning is memory- and computation-limited for high $p(x_0 \mid x_t)$ 5; low-rank, block-diagonal, or diagonal methods are preferred in practical settings (Schioppa et al., 13 May 2026, Ou et al., 2024).
Coverage limitations: Support mismatch between forward and reverse processes can persist, particularly in high dimensions or for misspecified covariance schedules, limiting forward ESS (Zhang et al., 27 May 2025). Covariance tuning must often be repeated for each sampling schedule.

6. Applications and Extensions

Covariance-aware sampling is relevant to a range of applications and can be combined or extended as follows:

Unbiased sampling for physical and molecular systems: Efficient expectation estimation in Boltzmann statistics, with applications to molecular dynamics and cluster formation simulations (Zhang et al., 27 May 2025).
Inverse problems and controlled generation: Posterior-inference reconstructions in image inpainting, deblurring, and super-resolution, with covariance-aware guidance achieving state-of-the-art PSNR and error metrics (Hamidi et al., 2024, Zhang et al., 7 Oct 2025, Rissanen et al., 2024).
Statistical control and constrained sampling: The CCS paradigm shows that carefully shaped initial noise, via linearization through the reverse ODE, allows precise control of sample output covariance, enabling targeted generation with minimal sample quality loss (Song et al., 7 Feb 2025).
Classifier(-free) guidance calibration: Locally-adapted covariances reduce over-sharpening/over-smoothing artifacts in classifier-guided or classifier-free setups, by balancing likelihood terms according to posterior uncertainty (Rissanen et al., 2024).

7. Open Directions and Practical Considerations

Research on covariance-aware sampling continues to address the following:

Adaptive and data-driven structures: Potential improvements include sample-dependent low-rank covariances, Kronecker or E(3)-equivariant structures, and neural predictors for optimal grouping or local variance scales (Schioppa et al., 13 May 2026, Zhang et al., 27 May 2025).
Scalability: Joint schedule and covariance tuning techniques, amortized across time and datasets, are a focus for future extension (Zhang et al., 27 May 2025).
Extension to latent diffusion models: Current approaches for pixel-space models do not always transfer directly; noise addition in latent space may degrade output quality unless the latent encoder/decoder is co-trained for such covariances (Schioppa et al., 13 May 2026).
Limitations: Diagonal-only models may fail under strong intrinsic correlations; high accuracy requires high-fidelity score and Hessian estimation, and quality depends critically on the baseline model (Ou et al., 2024, Schioppa et al., 13 May 2026).

By leveraging tractable analytic identities, fast structured estimators, and stable learning objectives, covariance-aware sampling establishes a unifying framework that advances both theoretical understanding and empirical performance in diffusion model generation, statistical inference, and scientific posterior estimation.