Intrinsic Schrödinger Bridge Diffusion
Implementation & benchmark of arXiv:2608.25838 — Hard-Constrained Sampling on Embedded Riemannian Manifolds via Adjoint Schrödinger Bridges
Source paper: Hard-Constrained Sampling on Embedded Riemannian Manifolds via Adjoint Schrödinger Bridges arXiv:2608.25838 ⓘ · analyzed Sep 1, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace an unconstrained Euclidean diffusion sampler plus projection with a controlled diffusion whose state always lies on an embedded manifold \(\mathcal M\). The neural controller predicts a tangent vector, while the stochastic forcing is also tangent; this preserves constraints during every intermediate denoising step and avoids the bias caused by repeatedly projecting off-manifold states.
Formulas
Mathematical statement
Let \(\mathcal M\subset\mathbb R^d\) be a smooth embedded manifold, \(\Pi_x=J(x)J(x)^\dagger\) the orthogonal projector onto the tangent space \(T_x\mathcal M\), and \(E:\mathcal M\to\mathbb R\) the target energy. The intrinsic Langevin reference process is \(dX_t=-\operatorname{grad}_{\mathcal M}E(X_t)\,dt+\sqrt{2\beta^{-1}}\Pi_{X_t}\circ dW_t\), where \(\operatorname{grad}_{\mathcal M}E=\Pi_x\nabla \widetilde E\), \(W_t\) is ambient Brownian motion, \(\beta\) is inverse temperature, and \(\circ\) denotes Stratonovich integration. Add a learned tangent control \(u_\theta(x,t)\in T_x\mathcal M\): \(dX_t=[-\operatorname{grad}_{\mathcal M}E(X_t)+u_\theta(X_t,t)]dt+\sqrt{2\beta^{-1}}\Pi_{X_t}\circ dW_t\). Schrödinger-bridge optimality states that, for reference diffusion generator \(\mathcal L\), the minimum-control solution has \(u^*(x,t)=2\beta^{-1}\operatorname{grad}_{\mathcal M}\log\phi(x,t)\), where \(\phi\) solves the backward Kolmogorov equation \(\partial_t\phi+\mathcal L\phi=0\) with endpoint conditions determined by the source and target densities. The control cost is \(\frac{\beta}{4}\mathbb E\int_0^T\|u_\theta(X_t,t)\|^2dt\), and the endpoint mismatch plus this cost estimates the path-space relative entropy.
Implementation notes
1. Exact integration point: use this as the sampler inside a diffusion or energy-based generative model whose outputs must lie on a known manifold. Start with the sphere \(\mathbb S^{d-1}\), where \(\Pi_x=I-xx^\top\), or the Stiefel manifold \(\mathrm{St}(n,k)=\{X:X^\top X=I_k\}\), where \(\Pi_X(V)=V-X\operatorname{sym}(X^\top V)\). Let an MLP or transformer controller output an ambient vector \(v_\theta(x,t,c)\), then set \(u_\theta=\Pi_xv_\theta\). 2. Pseudocode: initialize \(x_0\sim q_0\) on \(\mathcal M\); for \(k=0,\ldots,N-1\), compute \(g=\Pi_x\nabla E(x)\), \(u=\Pi_xv_\theta(x,t_k)\), sample \(z\sim\mathcal N(0,I_d)\), set \(z_T=\Pi_xz\), and update the intrinsic Euler-Maruyama state \(x\leftarrow\operatorname{Retr}_x[h(-g+u)+\sqrt{2h/\beta}\,z_T]\). Use the exact exponential map for the sphere when feasible; for Stiefel use a QR or polar retraction, and measure its numerical orthogonality error. Train with the displayed control-energy plus terminal-energy/endpoint loss, or first use supervised score/control matching from short reference trajectories. 3. Computed quantities: \(\Pi_x\), the manifold gradient, tangent noise, and retraction are geometric operations; \(\phi\) and its gradient are estimated by the controller, not analytically computed. Estimate endpoint KL with held-out energy samples or importance weighting, and log \(\int\|u\|^2dt\). 4. First cheap experiment: sample a von Mises-Fisher target or a multimodal energy on \(\mathbb S^{d-1}\) using a two-layer MLP, comparing intrinsic Schrödinger-bridge diffusion, Euclidean diffusion, and Euclidean diffusion with final projection. Sweep step size \(h\), noise temperature \(\beta^{-1}\), and control capacity. The predicted signature is maximum constraint violation \(|\|x\|_2-1|\) at every intermediate step below numerical tolerance for the intrinsic method, while the unconstrained method has violations scaling approximately as \(O(\sqrt h)\) before projection. As control capacity and horizon increase, endpoint KL should decrease monotonically until discretization error dominates, while the minimum control energy should increase when the source-target discrepancy is increased. A practical acceptance criterion is at least a 10-fold reduction in intermediate constraint violation and endpoint KL no worse than projection-based diffusion at equal network evaluations.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Confirmed ✓
Stage 2 · Practical benchmark result: Significantly worse ✗
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built and numerically verified a sphere-manifold intrinsic diffusion MVP with tangent projection, tangent noise, and normalization retraction. Post-retraction constraint error stayed at 3.33e-16 across all step sizes, unconstrained one-step violation scaled as h^0.570 versus the predicted h^0.5, and the stated control cost scaled quadratically with amplitude with exponent 2.0. Intrinsic and projection-at-every-step updates are equivalent in this toy setting, so this validates constraint preservation but not a generative-quality win over frequent projection.
- Agent confidence
- 8/10
- Baseline
- Ambient Euler diffusion: median one-step norm violations 0.1596, 0.1043, 0.0706, 0.0487 for h=0.04, 0.02, 0.01, 0.005; fitted exponent 0.570. Short endpoint mean final norm error was 0.4220 and alignment was 0.9739.
- Idea
- Intrinsic/retracted diffusion: maximum intermediate norm error was 3.33e-16 for every tested step size; short endpoint mean final norm error was 5.05e-17 and alignment was 0.6095. Control-cost sweep exponent was 2.0, with costs 0.02738, 0.10953, 0.43812, 1.75249, and 7.00997.
Limitations: Only the sphere manifold and a simple linear energy were tested. No learned neural Schrödinger-bridge controller, endpoint KL estimation, multimodal target, Stiefel manifold, likelihood evaluation, or wall-clock/FLOP comparison was implemented. The sphere intrinsic update uses normalization retraction and is equivalent to projection after every Euler step, so the result validates the geometric mechanism rather than superiority over a well-designed frequent-projection baseline.
How to run: python3 experiment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.25838)
Significantly WORSE than the tuned baseline (+183.8%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Manifold_pendulum
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0062
- Idea mean
- 0.0177
- Effect (Δ)
- 0.0115 (+183.8%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±107.2%
- Mechanism
- Confirmed ✓
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
Tested one custom S1-embedded pendulum prediction task, one rnn_small architecture, 18 epochs, and a three-value learning-rate grid. This benchmark did not train a neural Schrödinger-bridge controller, estimate endpoint KL, test Stiefel manifolds, compare sampling speed, or evaluate multimodal generative targets.
How to run: python3 stage2_bench.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 9/10
Built an end-to-end intrinsic S1 retraction intervention with the shared rnn_small architecture, shared learning-rate grid, baseline sweep, eight paired seeds, and permutation testing. The trained-model mechanism signature was confirmed: intrinsic maximum S1 norm violation was 1.19e-7 versus 0.1689 for baseline. However, test MSE was significantly worse for the idea: 0.0177347 versus 0.0062498, paired delta +0.0114849, p=0.0081, with 0/8 paired wins.
- Baseline
- Best lr=0.01; full 8-seed test MSE mean 0.0062498191, std 0.0014967029.
- Idea
- Best lr=0.01; full 8-seed test MSE mean 0.0177347455, std 0.0075728509; paired delta +0.0114849264, p=0.0081, idea wins 0/8. Signature confirmed: baseline max S1 violation 0.1689367294, idea max violation 1.1920928955e-7.
Benchmark result
Significantly WORSE than the tuned baseline (+183.8%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Manifold_pendulum
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0062
- Idea mean
- 0.0177
- Effect (Δ)
- 0.0115 (+183.8%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±107.2%
- Mechanism
- Confirmed ✓
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
Tested one custom S1-embedded pendulum prediction task, one rnn_small architecture, 18 epochs, and a three-value learning-rate grid. This benchmark did not train a neural Schrödinger-bridge controller, estimate endpoint KL, test Stiefel manifolds, compare sampling speed, or evaluate multimodal generative targets.
How to run: python3 stage2_bench.py
Artifacts
- bench_report.json 3.9 KB View Raw
- experiment.py 4.1 KB View Raw
- manifold_dynamics_track.py 1.2 KB View Raw
- report.md 1.8 KB View
- report_bench_2026-09-01T232540.md 4.1 KB View
- results.json 3.1 KB View Raw
- stage2_bench.py 5.8 KB View Raw