Exact Conditional Score-Guided Generative Modeling for Amortized Inference in Uncertainty Quantification

TL;DR

This paper introduces a novel two-stage framework combining exact conditional score-guided diffusion models and neural networks to enable fast, accurate, and scalable amortized inference for complex uncertainty quantification problems.

Contribution

It develops an analytical method for constructing exact conditional score functions and trains a neural network for direct posterior sampling, improving efficiency over traditional methods.

Findings

Achieves fast inference without iterative sampling.

Handles high-dimensional, multi-modal posteriors effectively.

Demonstrates superior performance in physical system parameter estimation.

Abstract

We propose an efficient framework for amortized conditional inference by leveraging exact conditional score-guided diffusion models to train a non-reversible neural network as a conditional generative model. Traditional normalizing flow methods require reversible architectures, which can limit their expressiveness and efficiency. Although diffusion models offer greater flexibility, they often suffer from high computational costs during inference. To combine the strengths of both approaches, we introduce a two-stage method. First, we construct a training-free conditional diffusion model by analytically deriving an exact score function under a Gaussian mixture prior formed from samples of the underlying joint distribution. This exact conditional score model allows us to efficiently generate noise-labeled data, consisting of initial diffusion Gaussian noise and posterior samples conditioned on various observation values, by solving a reverse-time ordinary differential equation. Second, we use this noise-labeled data to train a feedforward neural network that maps noise and observations directly to posterior samples, eliminating the need for reversibility or iterative sampling at inference time. The resulting model provides fast, accurate, and scalable conditional sampling for high-dimensional and multi-modal posterior distributions, making it well-suited for uncertainty quantification tasks, e.g., parameter estimation of complex physical systems. We demonstrate the effectiveness of our approach through a series of numerical experiments.