Inflationary Flows: Calibrated Bayesian Inference with Diffusion-Based Models

TL;DR

This paper introduces inflationary flows, a diffusion-based framework that enables calibrated Bayesian inference on a low-dimensional latent space by mapping data to an asymptotically Gaussian distribution via a deterministic pfODE. It develops two flow classes—dimension-preserving and dimension-reducing—through carefully designed noise schedules and coordinate rescalings, yielding unique, invertible transformations with controllable error. Score functions are learned with diffusion models and used to drive the pfODE, enabling sampling-based posterior estimation in reduced dimensions and thus tractable uncertainty quantification. Empirical results on CIFAR-10 and AFHQv2 demonstrate meaningful generation quality under compression and improved posterior calibration relative to injective flow baselines, highlighting the approach’s potential for principled uncertainty quantification in scientific settings.

Abstract

Beyond estimating parameters of interest from data, one of the key goals of statistical inference is to properly quantify uncertainty in these estimates. In Bayesian inference, this uncertainty is provided by the posterior distribution, the computation of which typically involves an intractable high-dimensional integral. Among available approximation methods, sampling-based approaches come with strong theoretical guarantees but scale poorly to large problems, while variational approaches scale well but offer few theoretical guarantees. In particular, variational methods are known to produce overconfident estimates of posterior uncertainty and are typically non-identifiable, with many latent variable configurations generating equivalent predictions. Here, we address these challenges by showing how diffusion-based models (DBMs), which have recently produced state-of-the-art performance in generative modeling tasks, can be repurposed for performing calibrated, identifiable Bayesian inference. By exploiting a previously established connection between the stochastic and probability flow ordinary differential equations (pfODEs) underlying DBMs, we derive a class of models, inflationary flows, that uniquely and deterministically map high-dimensional data to a lower-dimensional Gaussian distribution via ODE integration. This map is both invertible and neighborhood-preserving, with controllable numerical error, with the result that uncertainties in the data are correctly propagated to the latent space. We demonstrate how such maps can be learned via standard DBM training using a novel noise schedule and are effective at both preserving and reducing intrinsic data dimensionality. The result is a class of highly expressive generative models, uniquely defined on a low-dimensional latent space, that afford principled Bayesian inference.

Figures (13)

• Figure 1: SDE-ODE Duality of diffusion-based models. The forward (noising) SDE defining the DBM (left) gives rise to a sequence of marginal probability densities whose temporal evolution is described by a Fokker-Planck equation (FPE, middle). But this correspondence is not unique: the probability flow ODE (pfODE, right) gives rise to the same FPE. That is, while both the SDE and the pfODE possess the same marginals, the former is noisy and mixing while the latter is deterministic and neighborhood-preserving. Both models require knowledge of the score function $\nabla_\mathbf{x} \log p_t(\mathbf{x})$, which can learned by training either model.
• Figure 2: Dimension-preserving flows for toy datasets. Numerical simulations of dimension-preserving flows for five sample toy datasets. Left sequences of sub-panels show results for integrating the pfODE forward in time (inflation); right sub-panels show results of integrating the same system backwards in time (generation) (Appendix \ref{['ap:deets_pfODE_integration']}). Simulations were conducted with score approximations obtained from neural networks trained on each respective toy dataset (Appendix \ref{['toy_training_deets']}).
• Figure 3: Dimension-reducing flows for toy datasets. Numerical simulations of dimension-reducing flows for the same five datasets as in Figure \ref{['Figure_2']}. For 2D datasets, we showcase reduction from two to one dimension, while 3D datasets are reduced to two dimensions. Colors and layouts are the same as in Figure \ref{['Figure_2']}, with scores again estimated using neural networks trained on each example. Additional results showcasing (1) similar flows further compressing two-dimensional manifolds embedded in $D=3$ space, and (2) effects of adopting different scaling schemes for target data are given in Appendices \ref{['app:embedded_toys']} and \ref{['app:different_scaling']}, respectively.
• Figure 4: Calibration experiments. To assess error in our posterior model estimates, we used Hamiltonian Monte Carlo (HMC) to perform inference in one of our toy datasets (2D circles). Drawing samples from a 3-component Gaussian Mixture Model (GMM) prior, we integrated the generative process backward in time to obtain corresponding data space samples (A, components shown in orange, blue, and purple). We then used HMC to obtain posterior samples from the posterior distribution over the weights of the GMM components. (B, C) Kernel density estimates from the joint posterior samples over the mixture distribution weights in the dimension-preserving and dimension-reducing cases. Dashed vertical and horizontal lines indicate posterior means for each component. Reference ground-truth weights were $\mathbf{w} = [0.5, 0.25, 0.25]$.
• Figure 5: Generation and round-trip experiments for AFHQv2 at IG=1.02 and varying number of preserved dimensions. Top row: Generated samples for select flow schedules (PR-Preserving (PRP), PR-Reducing to 2D ($\approx 0.07\%$), 30D($\approx 1\%$), and 307D($\approx 10\%$), at 1.02 IG. Bottom row: Results for round-trip experiments under same schedules. Leftmost columns are original samples, middle columns are samples mapped to Gaussian latent spaces, and rightmost columns are recovered samples.