Stochastic Optimal Control for Diffusion Bridges in Function Spaces

2024-05-31

Neural Information Processing Systems

AI Summary

This work addresses the lack of interpretability and resolution-invariant modeling capability of diffusion models in infinite-dimensional function spaces—such as images, time series, and probability density functions (PDFs). To this end, we extend stochastic optimal control theory to infinite dimensions. Our method establishes, for the first time, a rigorous equivalence between infinite-dimensional Doob h-transforms and stochastic optimal control; overcomes the fundamental challenge of undefined densities in infinite dimensions by constructing diffusion bridges without explicit density assumptions; and jointly leverages variational inference and functional optimization to directly compute optimal transport paths in function space. Experiments demonstrate that the proposed framework achieves resolution invariance in both distributional bridge learning and sampling tasks. It significantly improves fidelity and interpretability across diverse applications—including image generation, time-series interpolation, and PDF modeling—without dependence on spatial or temporal discretization.

Technology Category

• Reasoning under Uncertainty: Stochastic Optimization
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

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.

Problem

• Extend diffusion-based algorithms to infinite-dimensional function spaces.
• Derive Doob's h-transform from stochastic optimal control perspective.
• Learn diffusion bridges and generative models in infinite dimensions.

Innovation

• Extends stochastic optimal control to infinite-dimensional spaces
• Uses Doob's h-transform for diffusion bridge construction
• Applies theory to generative models and distribution learning