Quantifying Epistemic Uncertainty in Diffusion Models

AI summary

Overview

Research area: Uncertainty quantification (UQ) for diffusion generative models, specifically separating epistemic (parameter-driven) uncertainty from aleatoric (sampling-noise) uncertainty. The paper sits at the intersection of Bayesian deep learning, randomized numerical linear algebra, and score-based generative modeling. It is posted under stat.ML.

Technical level: Advanced. The paper derives closed-form covariance projections, a multi-step recursion, and a randomized trace-norm error bound, and assumes familiarity with DDPM-style diffusion, Laplace approximations, Gauss–Newton/GGN curvature, and Jacobian-based error propagation.

Scope: The paper proposes a Fisher-information-based method called FLARE that isolates and propagates only the epistemic component of a diffusion model's uncertainty through the reverse denoising trajectory, and validates it on synthetic time-series generation tasks.

What This Paper Is About

Diffusion models generate samples by a stochastic reverse process, so a single generated sample always varies for reasons unrelated to whether the model actually "knows" anything — the sampling noise itself is a source of variability. Existing uncertainty scores for diffusion models, such as BayesDiff, aggregate both that intrinsic randomness and uncertainty about the model's parameters into one number, making it hard to tell whether a sample is unreliable because the model lacks knowledge or merely because sampling is random.

The goal of this work is to isolate the parameter-driven part alone, propagate it through the denoising chain, and turn it into a per-sample score that can be used to filter out low-confidence generations.

Key Contributions

1. Fisher–Laplace projection. The authors derive a closed-form projection of parameter uncertainty into data space using the Jacobian of the denoiser, producing an interpretable epistemic uncertainty map that separates parameter uncertainty from diffusion noise.

2. FLARE (Fisher–Laplace Randomized Estimator). They introduce a scalable randomized approximation that subsamples parameters uniformly across the network instead of restricting to the final layer, and prove it preserves epistemic structure with rapidly decaying relative error.

3. Experiments. They demonstrate improved uncertainty-aware sample filtering on synthetic time-series tasks, reporting up to 100% gap closure, and report consistently outperforming BayesDiff and last-layer Laplace baselines.

4. Evidence against last-layer Laplace. They provide analytic and empirical evidence that last-layer Laplace approximations (LLLA) are insufficient for this task, because restricting uncertainty to the final affine layer discards sensitivity information from earlier representations.

Main Findings

• Epistemic uncertainty is separable from diffusion noise. Using the law of total covariance on the DDPM reverse update, the authors decompose the conditional covariance of a reverse step into an aleatoric term equal to the diffusion noise β̃ₜI plus an epistemic term bₜ²JₜΣ_θJₜᵀ (Proposition 1). The epistemic term is reducible with more data; the aleatoric term is not.

• Epistemic uncertainty accumulates additively along the trajectory. Under a first-order linearization around the MAP estimate and a "local decoupling" condition, the general propagation identity (which contains a cross-covariance term 2aₜbₜCₜ(η)) simplifies to the additive recursion Σᵉᵖ_{t−1}(η) = aₜ²Σᵉᵖ_t(η) + bₜ²JₜΣ_θJₜᵀ, and unrolls into a weighted sum over all reverse steps.

• The cross-covariance term is argued to be negligible. The authors state in Appendix G that the cross term is (i) uniformly bounded and absorbable into the leading variance terms, (ii) asymptotically negligible under posterior concentration and local smoothness of the reverse trajectory, and (iii) numerically insignificant under a full-Hessian Laplace posterior, as verified by Monte Carlo evaluation.

• Randomized subnetwork approximation has a provable rate. Under mild regularity assumptions, the approximation error of the randomized projection in trace norm decays as O(1/√m) in the subnetwork size m, giving a principled cost-accuracy trade-off. The authors state that last-layer Laplace approximations do not admit comparable guarantees under the same assumptions.

• Last-layer Laplace suppresses real epistemic structure. Figure 1's 2D Gaussian-mixture illustration (adapted from Aithal et al. (2024); Jazbec et al. (2025)) has nine Gaussian modes on a square grid. BayesDiff assigns low uncertainty to samples lying between modes, and LLLA suppresses uncertainty further because of its restriction to the final layer, whereas the proposed method assigns high epistemic uncertainty in low-density regions between modes.

• FLARE tracks the full Fisher–Laplace projection. The paper states that the estimator "closely tracks the behavior of the full Fisher–Laplace projection," while scaling compute and memory costs from the number of parameters p to the subnetwork size m.

• Higher uncertainty appears where the model must extrapolate. The authors report that across all tasks, the method recovers structured epistemic uncertainty in regions where models switch modes, traverse sparse areas, or move off distribution.

• Damping constant used in experiments. The regularization parameter in the Laplace approximation is reported as λ = 10⁻⁶ in the authors' experiments.

• Reported quantitative coverage. The abstract and contributions report improved uncertainty estimation on synthetic time-series generation and up to 100% gap closure in uncertainty-aware sample filtering. The truncated text does not report specific sample counts, wall-clock timings, or per-dataset error tables.

Methodology in Plain English

The authors start from a standard discrete-time score-based diffusion model with a fixed noise schedule and a neural denoiser ε_θ(xₜ, t). They treat the trained parameters θ as a random variable and place a Laplace approximation around the MAP estimate θ̂, so the parameter posterior covariance is approximated as (H + λI)⁻¹, where H is the generalized Gauss–Newton (GGN) matrix of the training loss — for the mean-squared-error objective, H ≈ (1/n)ΣᵢJᵢᵀJᵢ over training pairs.

Their key move is to condition on a single realized reverse trajectory (a fixed draw of all non-parametric randomness η, which includes training-side randomness and sampling-side diffusion noise). Once that trajectory is fixed, the reverse state is a deterministic function of θ, so all remaining variability comes from the parameters. They linearize the one-step conditional mean around θ̂ (a delta-method step) and push the parameter covariance through the denoiser's Jacobian at each step, which gives a one-step epistemic covariance bₜ²JₜΣ_θJₜᵀ.

They then accumulate these one-step contributions along the reverse chain using a linear recursion induced by the DDPM update, ending with a final epistemic covariance at x̂₀. Scalar diagnostics — the trace tr(Σ₀ᵉᵖ(η)) and the normalized trace tr(Σ₀ᵉᵖ(η))/d — give a per-sample score; samples with the lowest trace are retained as the most confidently generated. For efficiency, the needed quadratic forms gᵀΣ_θg are computed via conjugate gradient solves of (H + λI)z = g using Hessian-vector products, avoiding explicit inversion.

Because forming and inverting the full p × p posterior covariance is infeasible, they replace it with a posterior over a randomly chosen subset of m ≪ p parameter coordinates: they sample a uniform index set I, form the submatrix H_{I,I}, invert it to get Σ_sub, and keep only the corresponding columns J_{t,I} of the Jacobian in the recursion. This is the FLARE algorithm. The design deliberately avoids the alternative shortcut of keeping only the final affine layer, because the final layer ignores parameter sensitivity that arises in earlier layers and travels through the network.

Why This Matters

Impact on research. The paper draws a sharp conceptual line between posterior predictive uncertainty and epistemic uncertainty in diffusion models, and argues that a widely used reference point, BayesDiff, targets total predictive uncertainty in data space rather than isolating parameter-induced uncertainty. It also provides a concrete theoretical and empirical case against last-layer Laplace approximations in this setting, and contributes a randomized-estimation rate (O(1/√m)) that connects diffusion UQ to randomized numerical linear algebra.

Real-world applications (potential, based on the settings discussed):

• Filtering synthetic time-series generated for forecasting or simulation pipelines, discarding samples that fall in sparse or off-distribution regions.
• Image synthesis quality control, since the introduction names image synthesis as a dominant application of diffusion models.
• Detecting when a generative model is extrapolating beyond its training distribution, which matters for any downstream decision made on generated data.
• Mode-coverage diagnostics for multimodal generative models, illustrated by the nine-mode Gaussian mixture toy example.

Industry relevance. Any deployed diffusion system that needs to flag unreliable outputs — rather than just produce them — needs a per-sample uncertainty signal. FLARE is presented as a simple, sample-level diagnostic with a tunable accuracy-versus-runtime trade-off through the subnetwork size m, which is an attractive property for practitioners who cannot afford ensembles or multiple forward passes.

Future Directions

• Extending beyond synthetic time series. The evaluation is limited in the provided text to synthetic time-series generation tasks probing multimodality, extrapolation, and low-density regions; how the method scales to image-scale diffusion models is not reported.
• Tightening the treatment of the cross-covariance term. The additive recursion relies on a local decoupling approximation that drops the 2aₜbₜCₜ(η) term; further work could characterize regimes where this approximation breaks down.
• Choosing the subnetwork size m. The O(1/√m) bound establishes a trade-off between cost and accuracy, but the paper does not report a procedure for selecting m in practice from the truncated content.
• Comparing against a wider set of UQ baselines. The paper benchmarks against BayesDiff, HyperDM-style parameter perturbation, and a last-layer variant of FLARE; the broader classes of ensembles, dropout, and variational methods discussed in related work are not evaluated here in the truncated text.

Target Audience

Readers who will benefit most are machine learning researchers working on uncertainty quantification, Bayesian deep learning, or diffusion models; statisticians interested in separating epistemic from aleatoric variability in generative processes; and practitioners building generative time-series or image systems who need a per-sample reliability score. The paper assumes comfort with diffusion model notation, Laplace approximations, Jacobians, and randomized linear algebra, so it is best suited to readers with graduate-level ML or statistics background rather than beginners.

Authors’ abstract

To ensure high quality outputs, it is important to quantify the epistemic uncertainty of diffusion models. Existing methods are often unreliable because they mix epistemic and aleatoric uncertainty. We introduce a method based on Fisher information that explicitly isolates epistemic variance, producing more reliable plausibility scores for generated data. To make this approach scalable, we propose FLARE (Fisher-Laplace Randomized Estimator), which approximates the Fisher information using a uniformly random subset of model parameters. Empirically, FLARE improves uncertainty estimation in synthetic time-series generation tasks, achieving more accurate and reliable filtering than other methods. Theoretically, we bound the convergence rate of our randomized approximation and provide analytic and empirical evidence that last-layer Laplace approximations are insufficient for this task.