Functional Adjoint Sampler: Scalable Sampling on Infinite Dimensional Spaces

TL;DR

The paper introduces the Functional Adjoint Sampler, a scalable method for sampling from Gibbs distributions in infinite-dimensional spaces, enabling efficient simulation of diffusion process trajectories with boundary constraints.

Contribution

It extends adjoint sampling to infinite-dimensional Hilbert spaces using stochastic optimal control, providing a scalable and effective sampling algorithm.

Findings

Achieves superior transition path sampling in synthetic and molecular systems.

Generalizes adjoint sampling to infinite-dimensional function spaces.

Demonstrates effectiveness on Alanine Dipeptide and Chignolin.

Abstract

Learning-based methods for sampling from the Gibbs distribution in finite-dimensional spaces have progressed quickly, yet theory and algorithmic design for infinite-dimensional function spaces remain limited. This gap persists despite their strong potential for sampling the paths of conditional diffusion processes, enabling efficient simulation of trajectories of diffusion processes that respect rare events or boundary constraints. In this work, we present the adjoint sampler for infinite-dimensional function spaces, a stochastic optimal control-based diffusion sampler that operates in function space and targets Gibbs-type distributions on infinite-dimensional Hilbert spaces. Our Functional Adjoint Sampler (FAS) generalizes Adjoint Sampling (Havens et al., 2025) to Hilbert spaces based on a SOC theory called stochastic maximum principle, yielding a simple and scalable matching-type objective for a functional representation. We show that FAS achieves superior transition path sampling performance across synthetic potential and real molecular systems, including Alanine Dipeptide and Chignolin.