Technical Summary: Supervised Guidance Training for Infinite-Dimensional Diffusion Models

Problem Statement

Inverse problems, particularly those arising from partial differential equations (PDEs), often require inferring unknown parameters that exist in infinite-dimensional function spaces. In the Bayesian framework, the goal is to sample from a posterior distribution over these functions, conditioned on noisy observations. While score-based diffusion models (SDMs) have recently been extended to infinite-dimensional function spaces to serve as expressive priors, the theory for conditioning these models to sample from the posterior remains an open problem.

Existing finite-dimensional conditioning techniques, such as Doob's $h$-transform, rely on density decompositions and transition operators that do not directly translate to infinite dimensions. Challenges include the non-existence of Lebesgue densities, the potential singularity of measures (Feldman-Hájek theorem) when shifting Gaussian measures outside their Cameron-Martin space, and the difficulty of ensuring differentiability of transition operators in Hilbert spaces. Furthermore, discretizing data immediately to apply finite-dimensional methods introduces resolution-dependent errors and limits the flexibility of the framework.

Methodology

Theoretical Foundation: Infinite-Dimensional $h$-Transform

The authors propose a principled framework for conditioning infinite-dimensional diffusion models using an extension of Doob's $h$-transform. They establish the existence of a conditional diffusion process under two settings:

1. Setting 1: The prior distribution is supported within a ball of the Cameron-Martin space of the Wiener process used in the diffusion SDE.
2. Setting 2: The prior is absolutely continuous with respect to a Gaussian measure, with specific bounds on the potential function and its gradient.

Under these settings, the authors prove that the time-reversed conditional process satisfies a stochastic differential equation (SDE) where the drift is modified by a guidance term. Specifically, they demonstrate that the conditional score function, $s_y(t, x)$, decomposes into the unconditional score, $s(t, x)$, and a guidance term, $C\nabla \log h_y(t, x)$:
$$s_y(t, x) = s(t, x) + C\nabla \log h_y(t, x)$$
Here, $h_y$ is the infinite-dimensional $h$-transform defined via the conditional expectation of the likelihood potential.

Supervised Guidance Training (SGT)

Since the guidance term $C\nabla \log h_y$ is intractable due to its dependence on the full transition law of the infinite-dimensional process, the authors introduce Supervised Guidance Training (SGT). This is a simulation-free training objective designed to learn the guidance term directly.

• Objective: The method minimizes a denoising score matching loss. Instead of relying on simulation-based stochastic optimal control (which suffers from high variance and requires sampling trajectories), SGT leverages the known unconditional score $s(t, x)$ and the joint distribution of data and observations $(X_0, Y)$.
• Loss Function: The loss is formulated to minimize the difference between the learned guidance term $u(t, x, y)$ and the true guidance term. The objective is:
  $$\mathbb{E} \left[ \left\| \frac{X_t - e^{-t/2}X_0}{1 - e^{-t}} + s(t, X_t) + u(t, X_t, Y) \right\|^2_K \right]$$
  where $K$ is a chosen Hilbert space (e.g., the Cameron-Martin space).
• Parametrization: The guidance term is parametrized using a neural network. The authors explore preconditioning the network output with the covariance operator $C$ to different degrees ($k \in \{0, 1/2, 1\}$) to enforce regularity and align with the noise structure of the SDE.

Approximations and Comparisons

The paper also analyzes existing heuristic approaches:

• Stochastic Optimal Control (SOC): While theoretically sound, SOC requires simulation-based training which is computationally expensive and unstable in early training stages.
• Tweedie Approximation (FunDPS): The authors extend the FunDPS approximation (which uses Tweedie's formula to estimate the clean data) to Setting 2 (absolute continuity), relaxing the restrictive assumptions of previous work. However, they note this approach still incurs computational overhead due to the need for backpropagation through the score model at every discretization step.

Key Contributions

1. Theoretical Extension: The paper provides the first rigorous proof that infinite-dimensional diffusion models can be conditioned via Doob's $h$-transform under general settings (Cameron-Martin support or absolute continuity), establishing the decomposition of the conditional score into an unconditional score and a guidance term.
2. Supervised Guidance Training (SGT): The authors propose a novel, simulation-free training algorithm to learn the intractable guidance term. This approach avoids the high variance of simulation-based methods and the computational cost of differentiating through the score model at every step (as required by Tweedie-based approximations).
3. Resolution Invariance: By treating the problem in infinite dimensions and only discretizing at the final step, the framework offers a resolution-invariant approach to conditioning, avoiding the performance degradation often seen when applying finite-dimensional architectures to higher resolutions.
4. Generalization of FunDPS: The work proves that Tweedie-based approximations (FunDPS) hold under the less restrictive Setting 2 (absolute continuity), broadening their applicability.

Experimental Results

The authors evaluate SGT on three Bayesian inverse problems in function spaces:

1. Sparse Observations: Recovering a 1D function from noisy point observations.
2. Heat Equation: Recovering the initial temperature distribution of a heat equation from noisy observations at a later time.
3. Shape Inpainting: Reconstructing shapes (digit "3" from MNIST) from partial Fourier coefficients (elliptic Fourier descriptors).

Findings:

• Performance: SGT consistently outperforms the heuristic FunDPS baseline in terms of Root Mean Square Error (RMSE) and Energy Score (ES).
• Comparison to Conditional Diffusion: SGT achieves accuracy comparable to a "conditional diffusion" model trained from scratch (which serves as an upper bound), despite the conditional model being significantly larger and trained for more steps.
• Robustness: In experiments with imperfect (randomly initialized or partially trained) unconditional base models, SGT demonstrates superior robustness, recovering realistic samples where FunDPS fails completely.
• Efficiency: SGT requires fewer parameters and training steps than training a full conditional model from scratch, as it only trains the guidance network while reusing a pre-trained unconditional model.

Significance and Claims

The paper claims to offer the first function-space method for fine-tuning trained diffusion models to accurately sample from a posterior distribution. The significance lies in bridging the gap between the theoretical requirements of infinite-dimensional Bayesian inference and the practical implementation of diffusion models.

The authors emphasize that their framework allows for:

• Resolution Invariance: The ability to generate samples at resolutions different from the training data without performance degradation.
• Stability: A training regime that is stable even when the underlying unconditional model is imperfect.
• Flexibility: Alternative options for discretization, training, and loss functions that are not constrained by finite-dimensional density assumptions.

The work concludes that explicitly learning the infinite-dimensional guidance term yields superior accuracy compared to heuristic approximations, providing a principled path forward for applying diffusion models to complex inverse problems in science and engineering.