Schrödinger bridges on Lie groups cut protein-torsion NLL from TDM's 0.69 to 0.59

Schrödinger Bridges on Lie Group Manifolds for Probabilistic Intrinsic Generation

Shizhe Zhang, Mingyang Zhao, Lei Ma

stat.ML, cs.AI, cs.LG

2026-09-02

WKBC and RCCBM run Schrödinger bridges on kinetic Lie groups. Protein-torsion NLL drops from TDM's 0.69 to 0.59; paths are 95% more direct than MD.

What problem this solves

Molecular torsion angles are periodic. Rigid orientations live on rotation groups. Flattening those quantities into Euclidean space and projecting back creates coordinate inconsistency and geometric error. Diffusion and flow matching usually run from a simple prior to data. In scientific problems both endpoint populations are observed, and the random path between them is itself the object of interest, for example how a protein walks from one conformation to another.

A Schrödinger bridge is entropy-regularized stochastic transport: stay close to a reference process in relative entropy while matching two prescribed endpoint laws. This paper puts that construction on Lie groups and lets endpoints constrain only what is measured. Configuration lives on the group, velocity lives in the Lie algebra, and unobserved velocities are filled in by the entropy projection.

Method

The state is Xt=(gt, ξt). gt moves on a Lie group G; velocity ξt lives in the fixed Lie algebra g. The reference process is kinetic Ornstein-Uhlenbeck: the group coordinate is integrated by left translation of the velocity, while the velocity is damped and noised. Stochastic forcing and learned controls act in that fixed vector space, so the configuration never leaves the manifold.

Endpoints constrain only observed coordinates. The entropy projection yields a unique bridge and the conditional law of latent endpoint velocities. Forward generation needs a calibrated initial law and one Doob controller, a drift correction that steers the process toward the terminal observation.

The same bridge has two solvers. On compact Abelian groups (tori, protein and RNA torsions) the periodized kinetic kernel is closed-form, so Wrapped-Kernel Bridge Calibration (WKBC) computes the propagated factors and the velocity-space log gradient directly. Compact non-Abelian groups (SO(3), U(n)) lack that kernel. Reciprocal Conditional-Control Bridge Matching (RCCBM) calibrates a two-sided reciprocal endpoint law, then regresses one Markov forward control from mollified finite-energy conditional teachers. Rigid frames are first stripped of global rotation and learned on a compact reduced group. Errors from endpoints, control regression, initialization, and discretization are separated in bounded-Lipschitz path distance.

Results

On the TDM-aligned prior-to-data likelihood protocol, held-out negative log-likelihood (lower is better):

DataRDMRFMTDMWKBC
General 2D1.041.010.690.59
Glycine1.971.901.040.84
Proline0.120.15-0.60-1.02
Pre-Pro1.241.180.520.17
RNA 7D-3.70-5.20-6.86-8.86

On the General torsion task, ValidRate rises from TDM's 0.595 to 0.873 and Ramachandran JSD falls from 0.277 to 0.125. RNA geometric RMSE is one of the few cells WKBC loses: 0.854 versus TDM's 0.843.

On SO(3) log-likelihood (higher is better) RCCBM leads all four rows: GMM32 0.328 (TDM 0.292), GMM64 0.378 (0.174), GMM128 0.200 (0.112), RingBand 0.834 (0.428). The GMM32 baselines are cited from the TDM paper; the other three are retrained under the same protocol. Against a high-accuracy WKBC reference on T2, RCCBM's source/target Sinkhorn values are 0.158/0.159 versus 0.104/0.111. Same direction, visible gap.

Protein conformational transitions use 15 held-out pairs from mdCATH domain 1jvmB00. Relative to molecular dynamics, RCCBM paths are 95.06% more direct and 99.21% less rough. That is regularized transport. It is not a claim of more physical thermal motion.

Why it matters

For manifold generation, molecular conformation, and rotation distributions, this is a computable way to run a two-observed-endpoint Schrödinger bridge in kinetic Lie-group coordinates. Use the kernel when the torus gives you one; learn a controller when it does not. Errors are split by module. Code is public.

It is not a drop-in replacement for generic manifold diffusion. Unconditional image generation has no paired endpoints. The ground that currently holds is SO(3) and low-dimensional torsions. Read the 95% directness number as "a less wander-prone random bridge," not as better molecular dynamics.

Limitations

The authors say the experiments stay on moderate-dimensional compact groups and selected molecular tasks. Higher dimension, larger systems, noncompact Lie groups, and homogeneous spaces are out of scope. Cartesian reconstruction of rigid frames is postprocessing; the stochastic guarantee lives on the reduced group.

GMM32 baselines are cited rather than rerun, so that row is not fully matched to the other three. RNA geometric RMSE is slightly worse, so a likelihood win is not a sweep of every geometric metric. The pathway comparison is against noisy MD trajectories; almost any smooth interpolant would look straighter. RCCBM has many moving parts (endpoint calibration, mollified teachers, importance correction, terminal refinement). Ablations show that dropping endpoint calibration or terminal replay hurts relative MSE and MMD. The engineering cost is real.

Terms

Source

What people are saying

Related papers

All paper explainers