PaperlayerView exact revision on arXiv

Source-linked AI summary

Hard-Constrained Sampling on Embedded Riemannian Manifolds via Adjoint Schrödinger Bridges

Mattia Mosso, Jaemoo Choi, Heng Yang

arXiv:2608.25838v1eess.SY

TL;DR

The paper addresses sampling from unnormalized Boltzmann distributions supported on manifolds, where Euclidean methods may not enforce exact feasibility. It formulates an intrinsic stochastic-control Schrödinger bridge sampler, develops practical geometric approximations, and reports validation across several physics and engineering applications. The authors conclude that the method is effective while identifying convergence certification as future work.

  • Problem

    Existing adjoint-sampling formulations primarily target Euclidean spaces and do not directly enforce exact support constraints for manifold-valued variables.

  • Method

    The paper extends adjoint Schrödinger bridge sampling to smooth, compact, path-connected embedded Riemannian manifolds using an intrinsic controlled diffusion and geometry-aware approximations.

  • Results

    The methods are empirically validated on spherical and Stiefel-manifold sampling, closed-loop inverse kinematics with obstacle avoidance, and robust Wahba optimization on SO(3).

  • Takeaways & Limitations

    The proposed approach provides a theoretically supported route to sampling unnormalized distributions while keeping the controlled process feasible on the manifold.

  • Takeaways & Limitations

    The practical geometric substitutions have short-time theoretical support but do not provide an exact convergence bound, and expanding convergence certificates remains future work.

Abstract

from arXiv · show

A variety of tasks require sampling from unnormalized Boltzmann distributions supported on manifolds. Building upon the foundations of adjoint matching and adjoint Schrödinger bridge sampling, this paper provides a theoretically justified method, through the lens of stochastic optimal control, to address this problem on smooth, compact, path-connected embedded Riemannian manifolds. As an element of novelty compared to existing literature, feasibility is enforced at the level of the state space, meaning the controlled diffusion is defined intrinsically on the curved space. Empirical validations are provided for several physics applications.

1 INTRODUCTION

The paper extends adjoint Schrödinger bridge sampling to unnormalized Boltzmann distributions on smooth embedded Riemannian manifolds, enforcing feasibility intrinsically. It develops exact manifold formulations, practical geometric approximations, and empirical validations across several applications.

  • Sampling unnormalized Gibbs–Boltzmann distributions matters across computational physics, molecular simulation, Bayesian inference, and statistical mechanics.
  • Euclidean adjoint-sampling formulations do not directly enforce exact support constraints for variables on lower-dimensional geometric spaces.Examples include rotations, orthogonal matrices, conservation-law level sets, and closed-loop robot configurations.
  • The proposed controlled diffusion evolves intrinsically on a smooth, compact, path-connected manifold, with drift and diffusion acting through its tangent bundle.This enforces feasibility at the state-space level rather than through penalties or ambient projections.
  • R–ASBS derives exact denoising- and adjoint-matching identities, then approximates heat-kernel, bridge, and transport terms using local geometric constructions.The method uses intrinsic transition kernels, manifold bridges, Log maps, and parallel transport.
  • Extended R–ASBS supports implicit manifolds by replacing unavailable geometric primitives with nearest-point retraction, projection-as-transport, and a projected-chord corrector.
  • The methods are evaluated on multimodal spherical distributions, Stiefel-manifold Gibbs sampling, closed-loop inverse kinematics with obstacle avoidance, and robust Wahba optimization on SO(3).

2 MATHEMATICAL PRELIMINARIES

The preliminaries formulate diffusion sampling and adjoint Schrödinger bridge methods through controlled path measures, stochastic optimal control, and matching objectives. Alternating control and corrector updates connect these objectives to Schrödinger bridge solutions.

  • A diffusion sampler transports an easy-to-sample source distribution to a target distribution using a base drift, noise schedule, and learned control field.
  • The controlled path measure is the trajectory law induced by the controlled SDE, while the base path measure corresponds to setting the control to zero.
  • Euclidean ASBS represents source-to-target transport as a Schrödinger bridge with an equivalent terminal-cost stochastic optimal control problem.
  • The optimal control satisfies an adjoint-matching identity involving the Schrödinger bridge potential and the noise schedule.
  • Adjoint matching uses endpoint samples from the current controlled process and intermediate points sampled from the base bridge conditioned on the endpoints.
  • Corrector matching estimates the endpoint corrector from the current controlled process without requiring samples from the target distribution.
  • The alternating optimization relaxes the interdependence between control and corrector objectives and is interpreted as iterative proportional fitting toward the Schrödinger bridge.

3 EXTENSION TO EMBEDDED RIEMANNIAN CONSTRAINT MANIFOLDS

The paper extends Schrödinger bridge and stochastic optimal control sampling to smooth embedded Riemannian manifolds, enforcing feasibility intrinsically. It derives exact manifold identities and introduces practical geometric approximations for R–ASBS.

  • 3 Extension to Embedded Riemannian Constraint Manifolds: The controlled diffusion evolves directly on M through tangent drift and diffusion, so feasibility is preserved for all times.This intrinsic construction avoids approximate constraint satisfaction from ambient penalty or projection methods.
  • 3.1 Reference and Controlled Diffusions on M: The framework assumes a smooth, boundaryless, compact, path-connected embedded manifold with dimension n = d − m and tangent spaces TxM.The orthogonal projection Px maps ambient vectors onto TxM, while Riemannian differential operators are defined intrinsically.
  • 3.2 Manifold SOC–SB Equivalence: Theorem 1 establishes equivalence between the manifold Schrödinger Bridge problem and a terminal-cost stochastic optimal control formulation under smoothness, positivity, and well-posedness assumptions.The optimal path measure is induced by the controlled intrinsic diffusion, with admissible tangent controls satisfying the stated martingale and integrability conditions.
  • 3.3 Adjoint Matching on Manifolds: Adjoint matching on curved manifolds replaces the Euclidean derivative-semigroup identity with a formula involving damped stochastic parallel transport and its metric adjoint.The adjoint pulls terminal covectors back along the reference Brownian path through covariant differentiation and Ricci-based damping.
  • 3.4 Practical Geometric Approximations and R–ASBS: Algorithm 1 makes the exact identities implementable by approximating heat-kernel scores, reference bridges, and damped transport with local geometric constructions.It uses short-time heat-kernel or Varadhan approximations, noisy geodesic bridges, and Levi–Civita parallel transport; Extended R–ASBS additionally uses projection-as-transport and projected chords.

4 NUMERICAL EXPERIMENTS

Experiments evaluate R–ASBS on spherical, Stiefel, closed-loop kinematic, and robust Wahba sampling problems, demonstrating applications across constrained geometric spaces.

  • Spherical and semi-spherical distributions: 86% of geometric Langevin MCMC particles became trapped in one energy basin, whereas R–ASBS placed 43.8% in the northern hemisphere versus the theoretical ratio 0.5.The target assigns equal mass to both hemispheres, but MCMC showed slow mixing and initialization bias.
  • Spherical and semi-spherical distributions: R–ASBS was validated on global earthquake distributions and cosmic ray arrivals on S2.The empirical Boltzmann target was defined using a training-data energy landscape.
  • Stiefel manifold and matrix orthogonality condition: The Stiefel experiment studies trace energy E(X) = tr(X⊤HX), whose low-temperature limit concentrates around the eigenspace of H’s p lowest eigenvalues.At high temperature, the target converges to the uniform invariant measure on St(n,p).
  • Stiefel manifold and matrix orthogonality condition: 5000 orthogonal matrices were used to approximate the Stiefel-manifold energy expectation, with results confirming the established properties.The experiment uses Algorithm 2 and observes d⟨E⟩β/dβ = −Varβ(E) ≤ 0.
  • High-dimensional redundant inverse kinematics: Extended R–ASBS generates diverse valid closed-loop robot configurations while avoiding static obstacles in a 10-joint planar manipulator.The configuration space is defined by nonlinear loop-closure constraints, with obstacle repulsion and a weak resting-posture prior.
  • The Wahba problem with outliers: The Wahba experiment reports TLS objectives, rotation errors, recovered true inliers, and clipped true outliers across outlier ratios using 10,000 samples.Results are means with standard deviations over 5 independent repeats.

5 CONCLUSION AND FUTURE WORK

The paper extends ASBS to intrinsically defined unnormalized distributions on non-Euclidean spaces and supports the method theoretically and empirically. Convergence certificates for the practical algorithm remain future work.

  • Conclusion: The paper enlarges ASBS to non-Euclidean spaces through a manifold SOC–SB equivalence and a geometry-aware approximate R–ASBS algorithm.The controlled process is defined intrinsically on the manifold.
  • Conclusion: Extensive numerical results demonstrate the efficacy of the method across the studied applications.
  • Future work: Expanding the theoretical convergence certificates for the practical algorithm is left as future work.

A RELATED WORK

The related work contrasts manifold-aware R–ASBS with sequential constrained samplers, Euclidean adjoint methods, and alternative diffusion approaches. It motivates intrinsic bridge and control formulations while identifying computational and structural limitations of prior methods.

  • Adjoint sampling: Adjoint sampling imposes memorylessness through restricted source distributions or strong noising, which can destroy useful initial–terminal dependence.
  • Schrödinger bridges and SOC: The paper recasts the computationally hard Schrödinger Bridge equations as a stochastic optimal control problem with a terminal cost.The kinetic-optimal drift solves the associated SOC problem.
  • Adjoint matching: Adjoint matching relies on the additive structure of the reference transition kernel, unlike denoising matching.

E.1 PROOF OF THEOREM 1 (MANIFOLD SOC–SB EQUIVALENCE)

The proof establishes feasibility, regularity, Markov structure, entropy identities, and equivalence between the manifold Schrödinger Bridge and terminal-cost stochastic optimal control solutions.

  • Feasibility: Under positivity, smoothness, and compactness assumptions, the static Schrödinger problem is feasible with finite entropy and has a unique dynamic minimizer.The independent endpoint coupling provides a finite-entropy feasible coupling, not the optimizer itself.
  • Regularity and admissibility: Smooth positive Schrödinger potentials yield a continuous bounded tangent optimal control, making the control admissible on the compact space-time domain.
  • Doob transform: The Doob h-transform produces a valid Markov transition kernel, and uniqueness in law identifies it with the diffusion controlled by the optimal tangent field.
  • Entropy identity: The proof derives the entropy identity using the Brownian anti-development and an associated orthonormal stochastic frame.
  • SOC identification: The manifold SOC and Schrödinger Bridge solutions coincide as complete path measures.The optimal control belongs to the admissible class and induces the Schrödinger Bridge minimizer.

E.2 RECIPROCAL PROPERTY

The reciprocal property shows that conditioning the optimal controlled bridge on endpoints yields the same path law as the uncontrolled reference bridge. This supports exact reference-bridge sampling when available and motivates noisy-geodesic approximations otherwise.

  • The optimal controlled bridge conditioned on its endpoints is identical to the uncontrolled reference bridge.
  • Endpoint conditioning cancels the Radon–Nikodym endpoint tilt, so the complete conditional path laws coincide.
  • The reciprocal property follows from endpoint factorization and does not depend on Euclidean geometry.
  • When the exact reference bridge is unavailable, the property motivates the noisy-geodesic approximation.
  • The exact manifold denoising and adjoint-matching lemmas characterize the corrector and controller through transported terminal information.

F ADDITIONAL THEORETICAL ANALYSIS

The additional analysis establishes that normal components of ambient controls do not affect the intrinsic diffusion, while tangent projection makes standard Euclidean-output networks admissible. It also records the extended algorithmic training procedure and its path-measure equivalence.

  • Algorithm 1 trains controller and corrector networks using manifold-valued rollouts, bridge interpolation, terminal adjoints, and denoising-corrector targets.
  • The extended algorithm re-samples trajectories after updating the controller and uses transported adjoints for matching.
  • Tangent projection removes the normal component of an ambient control before it enters the intrinsic dynamics.
  • Ambient controls with the same tangent projection induce identical intrinsic controlled diffusions and the same path measure.
  • Standard networks with outputs in Rd are admissible when their outputs are projected and the control cost uses the projected field.

F.3 LOCALITY OF THE TILTED HEAT-KERNEL APPROXIMATION

The locality analysis shows that bounded terminal tilting preserves short-time heat-kernel behavior on compact manifolds. This supports the local geometric substitutions, but not convergence of the full algorithm.

  • A bounded terminal reward tilt leaves the tilted endpoint at geodesic distance O(√r) from its start with Gaussian tails.
  • For small effective diffusion time, terminal tilting does not destroy the local nature of heat-kernel transitions.
  • The locality result supports the substitutions used in the corrector and controller approximations.
  • Damped transport differs from ordinary parallel transport through a bounded curvature-driven finite-variation term along the same short path.
  • The analysis does not prove convergence of Algorithm 1 because the stochastic reference bridge is replaced by a noisy-geodesic construction.

F.4 COMPUTATIONAL COMPLEXITY

R–ASBS has per-epoch cost O(BN) with fixed network size and closed-form geometric primitives, while extended R–ASBS has larger constants and stores full trajectories. After training, sampling is amortized and linear in generated samples and SDE steps.

  • R–ASBS has per-epoch complexity O(BN) with fixed network size and closed-form geometric primitives.
  • Extended R–ASBS also scales as O(BN), but repeated nearest-point projections and BN-sample controller updates increase its constant.
  • R–ASBS memory is O(Bd + Nu + Nh), whereas extended R–ASBS memory is O(BNd + Nu + Nh).
  • After training, both methods sample by simulating only the learned controlled diffusion; the corrector is unused.
  • Post-training generation costs O(N(Fu(S) + 2PS + ΠS)) and is amortized and linear in generated samples and SDE steps.

G ADDITIONAL NUMERICAL RESULTS

Additional results detail manifold-specific geometry, computational compromises, and parameter sensitivity in the R–ASBS experiments. The results indicate that sample budget strongly affects performance, while β is comparatively robust across the tested range.

  • The Stiefel manifold is defined by XᵀX = I_p, with tangent vectors satisfying VᵀX + XᵀV = 0.
  • Projection onto the Stiefel tangent space separates an ambient matrix into tangent and normal components using sym(XᵀZ).
  • The Riemannian gradient is obtained by orthogonally projecting the ambient Euclidean gradient onto the tangent space.
  • QR retraction replaces the expensive exact Stiefel exponential map, while unavailable analytical transport motivates Extended R–ASBS.The exact map involves 2p × 2p matrix exponentials or ODE solves; QR retraction uses the orthogonal Q factor of X + V.
  • Sample budget is the dominant performance factor, whereas R–ASBS remains comparatively robust to β over the tested range.Training also uses a smooth approximation to min(s, c̄²) to avoid nondifferentiability.
  • Figure 5 depicts Extended R–ASBS samples concentrated near the optimal quaternion on a geodesic slice of S3.Its directional-correspondence panel includes noisy measured endpoints, rotation arcs, and the true rotation axis on S2.
Loading 2608.25838v1…