Functional Adjoint Sampler: Scalable Sampling on Infinite Dimensional Spaces

TL;DR

This work addresses the challenge of sampling Gibbs-type distributions on infinite-dimensional spaces, particularly for transition-path sampling in molecular systems. It introduces the Functional Adjoint Sampler (FAS), a diffusion-based method that operates in Hilbert spaces by formulating non-equilibrium sampling (NES) as an infinite-dimensional stochastic optimal control problem and applying the stochastic maximum principle for adjoint matching. A practical training objective, ${\mathcal L}_{\texttt{FAS}}$, is derived to sidestep pathwise backpropagation by using conditional adjoints and a boundary-enforced function-space representation, enabling scalable learning and discretization-invariant path generation. Empirical results on synthetic Müller–Brown potential and real molecular systems (Alanine Dipeptide, Chignolin) show FAS achieves higher transition-hit rates and lower transition-state energies than baselines, with stable performance across discretization refinements. The method offers a principled, scalable way to sample complex trajectory distributions directly in function spaces, with potential applications to diffusion-based simulations, PDE-constrained generation, and other TPS-like tasks.

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.

Figures (6)

• Figure 1: Comparison on synthetic potential. We report the metrics averaged over $64$ paths. Best results are highlighted.
• Figure 2: TPS on alanine dipeptide. (a) Sampled transition path projected on conformational landscape as a function of CVs. The saddle point that separates the two-meta stable states is depicted as $\bigstar$ and the highest energy along paths sampled from FAS depicted as $\bullet$. (b) Potential energy plot. Grey box highlights the transition-state region where paths reach the ETS. (c) Visualization of conformation for given meta-stable states C5 (far left) and C7ax (far right).
• Figure 3: TPS on Chignolin. (a) Potential energy plot. (b) Visualization of protein folding trajectory from unfolded (far left) to folded (far right). Grey box highlights the transition-state region where trajectories reach the ETS.
• Figure 4: Sampled transition path on alanine dipeptide over various discretization steps.
• Figure 5: Sampled transition path on synthetic potential over various discretization steps.

Theorems & Definitions (24)

• Lemma 0: Finite Normalization Constant
• Theorem 1: Explicit RND
• Lemma 1: Stochastic Maximum Principle
• Proposition 1: Adjoint Matching in $\calH$
• Proposition 1: Unbiased Estimator
• Proposition 1: Laplacian Operator
• Lemma 2.4: Stochastic integration by parts in ${\mathcal{H}}$
• proof
• Lemma 2.5: Verification Theorem
• proof