Functional Adjoint Sampler: Scalable Sampling on Infinite Dimensional Spaces

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.

• The paper introduces FAS as a scalable sampler that uses stochastic optimal control and adjoint matching to efficiently sample from infinite-dimensional Gibbs-type distributions.
• It employs spectral expansions, neural operator parameterizations, and local matching objectives to overcome computational challenges in high-dimensional path sampling.
• Empirical results on synthetic potentials and molecular systems show FAS achieves high transition-hit probabilities, lower energy transition points, and zero-shot discretization independence.

Functional Adjoint Sampler: Scalable Sampling on Infinite Dimensional Spaces

Introduction and Problem Formulation

The paper introduces the Functional Adjoint Sampler (FAS), a scalable sampler for Gibbs-type distributions on infinite-dimensional Hilbert spaces. The primary challenge addressed is efficient and reliable sampling from distributions over trajectory/path spaces, particularly transition path sampling (TPS) for systems with rare events or rigid boundary constraints. Prior work has demonstrated notable progress for finite-dimensional spaces, but there has been a lack of practical and generalizable approaches for infinite-dimensional functional spaces, especially regarding direct control-based methods for sampling in such settings.

FAS formulates path sampling as stochastic optimal control (SOC) over functions, leveraging the infinite-dimensional generalization of the adjoint matching principle. The approach rests on the stochastic maximum principle (SMP), enabling scalable local matching objectives suitable for functional representations. This is in contrast to global path-wise classical control objectives, which typically suffer from prohibitive computational cost and scalability issues.

Theoretical Foundations

Hilbert Space and Infinite-Dimensional SDEs

FAS works in a separable Hilbert space $H$ with norm and inner product $\langle \cdot, \cdot \rangle_H$. Sampling objectives utilize Gaussian reference measures with trace-class covariance operators to ensure well-posedness, and stochastic processes are defined via $Q$-Wiener processes yielding appropriately smooth, square-integrable noise.

An infinite-dimensional SDE of the form:

$dX_t = A X_t dt + \sigma_t dW_t^Q$

with $A$ generating a $C_0$-semigroup and $Q$ self-adjoint and trace-class, leads to a well-defined Ornstein–Uhlenbeck semigroup with explicit mean/covariance structure for the process.

Non-Equilibrium Sampling via Change of Measure

FAS implements non-equilibrium sampling (NES) by importance weighting with explicit Radon–Nikodým derivatives (RND) between Gaussian measures, using Ornstein–Uhlenbeck reference dynamics. The existence and closed-form of RNDs ensures proper density transformation in infinite-dimensional space—a critical theoretical advance, as prior literature primarily addresses compatibility and singularity of measures in such contexts.

Stochastic Optimal Control and Adjoint Matching

Sampling from the target measure is cast as minimizing the Kullback–Leibler divergence between path measures, which is equivalent to the SOC problem:

$$\min_\alpha \mathbb{E}_{P^\alpha}\left[\int_0^T \frac{1}{2} \|Q^{1/2} \alpha_t\|^2 dt + g(X_T^\alpha)\right]$$

where $g(x) = U(x) + \log q_T(x_0, x)$ encodes the unnormalized energy and RND corrections. The SMP yields a dual local optimality condition, transforming global minimization to pointwise convex matching, dramatically improving scalability and permitting efficient parallelized training.

Control policies $\alpha^\theta$ are optimized by minimizing the local matching objective:

$$\int_0^T \mathbb{E}_{P^{\bar{\alpha}}} \left[ \frac{1}{2}\|\alpha^{\theta}(X_t, t) + \sigma_t Q^{1/2} Y_t\|^2 \right] dt$$

with adjoint process $Y_t$ solving backward SDEs or approximated by conditional expectation for tractable computation.

Numerical Implementation

Operator Parameterizations

FAS utilizes spectral expansions of the Laplacian (with Dirichlet boundary conditions) for both the reference dynamics and the covariance operator. The eigenbasis (sine functions for $[0, L]$) and trace-class basis truncation enable efficient Galerkin approximations for simulation and learning. Boundary conditions are enforced by function decomposition: a fixed reference path with zero boundary residual, ensuring sampled paths respect prescribed endpoints without explicit penalties or projection steps.

Neural Operator Parameterization

The control $\alpha^{\theta}$ is parameterized as a neural operator, tailored by functional architectures (modified U-shaped Neural Operator), embedding both time and spatial coordinates. The operator incorporates spectral convolutions modulated by sinusoidal time embedding to achieve grid and resolution invariance, permitting zero-shot transfer across discretization levels.

Training and Sampling Algorithms

Replay buffers are employed to amortize sampling over recent trajectories. Training alternates between control updates and path sampling, using Monte Carlo samples of the adjoint process to compute unbiased gradient estimates. In practice, post-processing (gradient flow correction, IDPP refinement) is applied, especially in molecular cases, to improve path physical plausibility and suppress artifacts.

Empirical Results

Synthetic Potential

On the Müller–Brown potential benchmark, FAS attains 100% transition-hit probability (THP) and lowest energy of transition states (ETS), outperforming prior DL-based and MCMC methods. Log-likelihood (LLK) scores for the sampled transition path distribution also surpass prior approaches, indicating high probabilistic fidelity and reliability.

Figure 1: TPS for synthetic potential—sampled reactive path traversing energy landscape valleys, demonstrating low ETS and robust THP.

Molecular Systems

Alanine Dipeptide

FAS achieves perfect THP and zero endpoint RMSD; Dirichlet enforcement of boundaries guarantees endpoint correctness. The sampled paths reliably cross near the free-energy saddle, producing physically plausible transitions. FAS yields smoother transitions and lower ETS compared to state-of-the-art TPS methods, PIPS, TPS-DPS, TR-LV, and steered/MD baselines.

Figure 2: TPS on alanine dipeptide—projected paths cross near the conformational energy saddle between C5 and C7ax states, visualized on $[\phi, \psi]$ CVs.

Chignolin Protein Folding

For chignolin, a problem with substantially higher dimension and a rugged landscape, FAS demonstrates lowest ETS and perfect THP across sampled trajectories, whereas gradient-based or steered baselines exhibit reduced transition rates and larger RMSD. FAS’s enforcement of boundary conditions ensures all paths are reactive, removing endpoint drift and enhancing transition smoothness.

Figure 3: TPS on Chignolin—energy plot illustrates smoother folding transitions under FAS, with corresponding conformational trajectory visualizations.

Discretization Invariance

FAS’s function-space parameterization enables generation of paths at arbitrarily fine temporal resolutions without retraining or additional energy/force computations. ETS and LLK scores remain stable across orders-of-magnitude grid refinement ([x1], [x10], [x100]), confirming zero-shot discretization independence—critical for scientific applications demanding high-resolution inference.

Figure 4: Sampled transition path on alanine dipeptide for various discretization steps, underscoring resolution-agnostic behavior.

Ablation Studies

Investigations into diffusion magnitude, control network scaling, and buffer size confirm that larger models and higher early exploration improve final ETS. All training objectives are highly parallelizable due to the pointwise local structure, further supporting scalability for large and high-dimensional domains.

Figure 5: Ablation studies on alanine dipeptide show decreasing ETS with increased $\beta_{max}$ and improved statistics with larger model scaling.

Implementation Guidelines and Limitations

• Computational Complexity: Simulation cost is dictated by the number of basis modes, path resolution, and neural operator size; spectral truncation balances fidelity and efficiency.
• Boundary Handling: Dirichlet eigen-decompositions must be tailored to task/domain geometry (e.g., graph Laplacians for molecular graphs, DCT basis for images).
• Adjoint Computation: In complex domains, approximating conditional expectations for the adjoint process may require additional sampling or post-processing.
• Path Initializations: For highly stiff potentials, initialization via IDPP or gradient flow refinement avoids non-physical bond stretching or clashing, essential in molecular contexts.
• Noise Scheduling: Geometric schedules increase exploration; hyperparameters should be adjusted proportionally for resolution changes.

Implications and Future Directions

FAS provides a general framework for direct sampling of functionals over infinite-dimensional spaces, extending the reach of SOC-based diffusion samplers beyond finite dimensions. It enables exact enforcement of boundary conditions, discretization independence, and high computational scalability—characteristics desired in computational physics, chemistry, and scientific machine learning.

The practical impact spans molecular dynamics (protein folding, reaction rates), data assimilation in scientific computing, inverse problems for PDEs, and generative modeling over domains where sample-to-energy settings motivate functional approaches.

Areas for future development include multi-dimensional and irregular domains, adaptive basis selection for complex geometries, domain-specific neural operator architectures, and integration with adjoint-based variational inference for uncertainty quantification.

Conclusion

The Functional Adjoint Sampler represents a principled advance in scalable non-equilibrium sampling over infinite-dimensional spaces. By generalizing adjoint matching via stochastic optimal control and implementing spectral-enforced boundary constraints, FAS achieves high-efficiency, physically reliable sampling for complex path distributions. Its discretization independence and local matching structure increase its applicability to broad scientific domains, with theoretical guarantees supported by strong empirical performance. Continued extension to multidimensional and heterogeneous spaces will further enhance its role in AI-driven scientific discovery.