Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Supervised Guidance Training for Infinite-Dimensional Diffusion Models

Elizabeth L. Baker, Alexander Denker, Jes Frellsen

TL;DR

This work develops a principled approach for conditioning score-based diffusion models in infinite-dimensional function spaces to sample from Bayesian posteriors. It builds an infinite-dimensional Doob’s $h$-transform to obtain the posterior sampler and proves a score decomposition $s^y(t,x)=s(t,x)+C\nabla\log h^y(t,x)$, enabling a practical learning objective. The authors introduce Supervised Guidance Training (SGT), a simulation-free method that learns the intractable guidance term, and connect the framework to stochastic control and Tweedie-type approximations, validating the method on PDE-based inverse problems. The results show that explicit learning of the guidance term yields superior posterior sampling accuracy relative to heuristic baselines, offering a significant advancement for function-space Bayesian inference with diffusion models.

Abstract

Score-based diffusion models have recently been extended to infinite-dimensional function spaces, with uses such as inverse problems arising from partial differential equations. In the Bayesian formulation of inverse problems, the aim is to sample from a posterior distribution over functions obtained by conditioning a prior on noisy observations. While diffusion models provide expressive priors in function space, the theory of conditioning them to sample from the posterior remains open. We address this, assuming that either the prior lies in the Cameron-Martin space, or is absolutely continuous with respect to a Gaussian measure. We prove that the models can be conditioned using an infinite-dimensional extension of Doob's $h$-transform, and that the conditional score decomposes into an unconditional score and a guidance term. As the guidance term is intractable, we propose a simulation-free score matching objective (called Supervised Guidance Training) enabling efficient and stable posterior sampling. We illustrate the theory with numerical examples on Bayesian inverse problems in function spaces. In summary, our work offers the first function-space method for fine-tuning trained diffusion models to accurately sample from a posterior.

Supervised Guidance Training for Infinite-Dimensional Diffusion Models

TL;DR

This work develops a principled approach for conditioning score-based diffusion models in infinite-dimensional function spaces to sample from Bayesian posteriors. It builds an infinite-dimensional Doob’s -transform to obtain the posterior sampler and proves a score decomposition , enabling a practical learning objective. The authors introduce Supervised Guidance Training (SGT), a simulation-free method that learns the intractable guidance term, and connect the framework to stochastic control and Tweedie-type approximations, validating the method on PDE-based inverse problems. The results show that explicit learning of the guidance term yields superior posterior sampling accuracy relative to heuristic baselines, offering a significant advancement for function-space Bayesian inference with diffusion models.

Abstract

Score-based diffusion models have recently been extended to infinite-dimensional function spaces, with uses such as inverse problems arising from partial differential equations. In the Bayesian formulation of inverse problems, the aim is to sample from a posterior distribution over functions obtained by conditioning a prior on noisy observations. While diffusion models provide expressive priors in function space, the theory of conditioning them to sample from the posterior remains open. We address this, assuming that either the prior lies in the Cameron-Martin space, or is absolutely continuous with respect to a Gaussian measure. We prove that the models can be conditioned using an infinite-dimensional extension of Doob's -transform, and that the conditional score decomposes into an unconditional score and a guidance term. As the guidance term is intractable, we propose a simulation-free score matching objective (called Supervised Guidance Training) enabling efficient and stable posterior sampling. We illustrate the theory with numerical examples on Bayesian inverse problems in function spaces. In summary, our work offers the first function-space method for fine-tuning trained diffusion models to accurately sample from a posterior.
Paper Structure (37 sections, 12 theorems, 99 equations, 6 figures, 3 tables, 1 algorithm)

This paper contains 37 sections, 12 theorems, 99 equations, 6 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions
  3. Background and setup
  4. Bayesian Inverse Problems
  5. Conditioning in Finite Dimensions
  6. Infinite-dimensional Score-Based Diffusion Models
  7. Conditioning via Doob's $h$-Transform
  8. Approximating the Guidance Term
  9. Stochastic Optimal Control
  10. Infinite-Dimensional Tweedie Approximation
  11. Learning the Guidance Term
  12. Parametrisation of the Guidance Term
  13. Experiments
  14. Implementation Differences in Function Space
  15. Sparse Observations
  16. ...and 22 more sections

Key Result

Theorem 3.1

Let $\mathbb{P}$ be the path measure of the unconditional time-reversal $Z_t$ in eq:backwardSDE, with $h^y$ as in eq:h-transform. We assume either setting:cameron-martin or setting:gaussian. Then the regular conditional probability $\mathbb{P}(\cdot \mid Y=y)$, denoted as $\mathbb{P}^y$, satisfies Moreover, the conditional process of $Z_t$ given $Y=y$, which we denote by $Z_t^y$, satisfies where

Figures (6)

  • Figure 1: Posterior sampling for the point evaluation. We show the ground truth signal in black, the mean sample in red and individual samples in light blue. Note, that the uncertainty is small on the observed points (green). SGT is evaluated with $k=0$ for the preconditioner.
  • Figure 2: Posterior sampling for the heat equation. We show the ground truth signal in black, the mean sample in red and individual samples in light blue. Note, that due to the boundary condition $f(0)=f(1)=0$ the uncertainty at the boundary of the domain is small.
  • Figure 3: Conditional sampling results on shape data using the first six EFD modes. We show the reconstruction from the first six EFD modes, the ground-truth shape, and conditional samples generated by Conditional Diffusion, FunDPS, and our SGT. Mean predictions are shown in darker colours, with individual samples overlaid in lighter shades.
  • Figure 4: Posterior sampling for the heat equation. We show the ground truth signal in black, the mean sample in red and individual samples in light blue. Note, that due to the boundary condition $f(0)=f(1)=0$ the uncertainty at the boundary of the domain is small.
  • Figure 5: The first $8$ examples of the MNIST test dataset, parametrised using $64$ landmarks.
  • ...and 1 more figures

Theorems & Definitions (23)

  • Remark 2.1
  • Theorem 3.1
  • proof
  • Theorem 3.2
  • Proposition 4.1: Infinite-dimensional Tweedie Estimate
  • Proposition 4.2
  • Remark 4.3
  • Theorem 1.1
  • Proposition 1.2
  • proof
  • ...and 13 more