Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stochastic Optimal Control for Diffusion Bridges in Function Spaces

Byoungwoo Park, Jungwon Choi, Sungbin Lim, Juho Lee

TL;DR

This work extends stochastic optimal control to infinite-dimensional Hilbert spaces to enable diffusion-bridge methods in function spaces. By deriving an explicit infinite-dimensional Doob's $h$-transform via a Radon-Nikodym framework on Gaussian references, it creates diffusion-bridge-based sampling and Bayesian learning in function spaces. The paper introduces two learning paradigms: Bridge Matching for transporting between infinite-dimensional distributions and Bayesian Learning for functional posterior sampling, with training losses based on cross-entropy and relative-entropy in the Hilbert-space setting. It demonstrates resolution-free functionality on 1D/2D data and functional Bayesian tasks, highlighting potential for efficient, resolution-agnostic generative modeling in applications like image translation and stochastic-process inference. Limitations include time-invariant coefficient constraints and computational demands, pointing to avenues for scalable algorithms and broader domain extensions in future work.

Abstract

Recent advancements in diffusion models and diffusion bridges primarily focus on finite-dimensional spaces, yet many real-world problems necessitate operations in infinite-dimensional function spaces for more natural and interpretable formulations. In this paper, we present a theory of stochastic optimal control (SOC) tailored to infinite-dimensional spaces, aiming to extend diffusion-based algorithms to function spaces. Specifically, we demonstrate how Doob's $h$-transform, the fundamental tool for constructing diffusion bridges, can be derived from the SOC perspective and expanded to infinite dimensions. This expansion presents a challenge, as infinite-dimensional spaces typically lack closed-form densities. Leveraging our theory, we establish that solving the optimal control problem with a specific objective function choice is equivalent to learning diffusion-based generative models. We propose two applications: (1) learning bridges between two infinite-dimensional distributions and (2) generative models for sampling from an infinite-dimensional distribution. Our approach proves effective for diverse problems involving continuous function space representations, such as resolution-free images, time-series data, and probability density functions.

Stochastic Optimal Control for Diffusion Bridges in Function Spaces

TL;DR

This work extends stochastic optimal control to infinite-dimensional Hilbert spaces to enable diffusion-bridge methods in function spaces. By deriving an explicit infinite-dimensional Doob's -transform via a Radon-Nikodym framework on Gaussian references, it creates diffusion-bridge-based sampling and Bayesian learning in function spaces. The paper introduces two learning paradigms: Bridge Matching for transporting between infinite-dimensional distributions and Bayesian Learning for functional posterior sampling, with training losses based on cross-entropy and relative-entropy in the Hilbert-space setting. It demonstrates resolution-free functionality on 1D/2D data and functional Bayesian tasks, highlighting potential for efficient, resolution-agnostic generative modeling in applications like image translation and stochastic-process inference. Limitations include time-invariant coefficient constraints and computational demands, pointing to avenues for scalable algorithms and broader domain extensions in future work.

Abstract

Recent advancements in diffusion models and diffusion bridges primarily focus on finite-dimensional spaces, yet many real-world problems necessitate operations in infinite-dimensional function spaces for more natural and interpretable formulations. In this paper, we present a theory of stochastic optimal control (SOC) tailored to infinite-dimensional spaces, aiming to extend diffusion-based algorithms to function spaces. Specifically, we demonstrate how Doob's -transform, the fundamental tool for constructing diffusion bridges, can be derived from the SOC perspective and expanded to infinite dimensions. This expansion presents a challenge, as infinite-dimensional spaces typically lack closed-form densities. Leveraging our theory, we establish that solving the optimal control problem with a specific objective function choice is equivalent to learning diffusion-based generative models. We propose two applications: (1) learning bridges between two infinite-dimensional distributions and (2) generative models for sampling from an infinite-dimensional distribution. Our approach proves effective for diverse problems involving continuous function space representations, such as resolution-free images, time-series data, and probability density functions.
Paper Structure (44 sections, 8 theorems, 73 equations, 8 figures, 5 tables, 3 algorithms)

This paper contains 44 sections, 8 theorems, 73 equations, 8 figures, 5 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Notation.
  3. A Foundation on Stochastic Optimal Control for Diffusion Bridges
  4. Preliminaries
  5. Gaussian Measure and Cameron-Martin Space.
  6. Stochastic Differential Equations in Hilbert Spaces.
  7. Stochastic Optimal Control in Hilbert Spaces
  8. Doob's h-transform for Diffusion Bridges in Hilbert Spaces
  9. Approximating path measures
  10. Simulating Diffusion Bridges in Infinite Dimensional Spaces
  11. Infinite Dimensional Bridge Matching
  12. Objective Functional for Bridge Matching.
  13. Bayesian Learning in Function Space
  14. Objective Functional for Bayesian Learning.
  15. Related Work
  16. ...and 29 more sections

Key Result

Lemma 2.1

Let $\mathcal{V}$ be a solution of HJB equation eq:PDE HJB with $R(\alpha) := \frac{1}{2}\left\lVert\alpha\right\rVert^2_{\mathcal{H}}$ satisfying the assumptions in assumptions value function. Then, we have $\mathcal{V}(t, \mathbf{x}) \leq \mathcal{J}(t, x, \alpha)$ for every $\alpha \in \mathcal{U for almost every $s \in [t, T]$ and $\mathbb{P}$-almost surely. Then $(\alpha^{*}, \mathbf{X}^{\al

Figures (8)

  • Figure 1: (Top) Diffusion Bridge $\mathbb{P}^{\star}$ evaluated on $32^2$. (Bottom) Learned process $\mathbb{P}^{\alpha^{\star}}$ evaluated on $256^2$
  • Figure 2: Results on 1D function generation. (Left) Real data and (Right) generated samples from our model.
  • Figure 3: Results on Unpaired image transfer task. (Up) EMNIST $\to$ MNIST (Down) AFHQ-64 Wild $\to$ Cat. (Left) Real data and (Right) generated samples from our model. For generation at unseen resolutions, the images within the red and blue boxed initial conditions were upsampled (using bi-linear transformation) from the observed resolution ($32^2$) for EMNIST and ($64^2$) for AFHQ-64 Wild, respectively.
  • Figure 4: Sampled functions from a learned stochastic process $\mathbf{X}_t^{\alpha}$ evaluated on $[0, I]$ for $t \in [0, \frac{T}{2}, T]$. The grey line represents the mean function $\mathbb{E}[\mathbf{X}^{\alpha}_t]$ and the blue-shaded region represents the confidence interval. (Left) GP with RBF kernel. (Right) Physionet.
  • Figure : Bridge Matching of DBHS
  • ...and 3 more figures

Theorems & Definitions (16)

  • Lemma 2.1: Verification Theorem
  • Theorem 2.2: Hopf-Cole Transform
  • Theorem 2.3: Explicit Representation of $h$
  • Example 2.4: Diffusion Bridge in $\mathcal{H}$
  • Theorem 3.1: Mixture of Bridges in $\mathcal{H}$
  • Theorem 3.2: Exact sampling in $\mathcal{H}$
  • Theorem A.2
  • proof
  • proof
  • Corollary A.3: Markov Control fabbri2017stochastic
  • ...and 6 more