0%||0/0
MolSSAI x Bio Seminar

Strong Stochastic Flow Maps

Sam McCallum1* · Zander W. Blasingame2* · Timothy Herschell1
Niklas Rindtorff2 · Alexander Tong2† · James Foster1†
*Equal authorship †Equal supervision
1University of Bath 2AITHYRA
Aithyra Aithyra
We want to learn the Itô integral.
02 · motivation

Solving SDEs is very slow

$ density — 1D reaction coordinate
$ chignolin — 3D conformation
$\mathrm{d}\boldsymbol X_t=\boldsymbol f(t,\boldsymbol X_t)\,\mathrm{d}t+\boldsymbol g(t,\boldsymbol X_t)\,\mathrm{d}\boldsymbol W_t\qquad t:0\to 1$ $\boldsymbol\Psi_{s,t}(\boldsymbol x)=\boldsymbol x+\displaystyle\int_s^{\,t}\boldsymbol f(\tau,\boldsymbol X_\tau)\,\mathrm{d}\tau+\displaystyle\int_s^{\,t}\boldsymbol g(\tau,\boldsymbol X_\tau)\,\mathrm{d}\boldsymbol W_\tau$
solving SDE0000/2000steps

Integrating the SDE from a unit Gaussian ($t=0$) to a conformation of chignolin ($t=1$).

Strong Stochastic Flow Maps

learn the Itô integral!

03 · Flow maps

Solving ODEs lightning fast

$\boldsymbol X_s$
$\boldsymbol X_t$
$\boldsymbol f(s,\boldsymbol X_s)$
$\boldsymbol X_u$
 t -->
$s$
$t$
$u$
The flow map teleports our particle in both space and time!
$\dfrac{\mathrm{d}\boldsymbol X_t}{\mathrm{d}t}=\boldsymbol f(t,\boldsymbol X_t)$
vector field
$\boldsymbol X_u=\boldsymbol X_s+\displaystyle\int_s^{\,u} \boldsymbol f(\tau,\boldsymbol X_\tau)\,\mathrm{d}\tau$
integrate the field
$\boldsymbol\Phi_{s,t}(\boldsymbol x)=\boldsymbol x+\displaystyle\int_s^{\,t} \boldsymbol f(\tau,\boldsymbol X_\tau)\,\mathrm{d}\tau$
the flow map
04 · Stochastic Differential Equations

From ODEs to SDEs

$t$
$X$
$\boldsymbol X_{k+1}=\boldsymbol X_k+\class{sde-t-drift}{\boldsymbol f(t,\boldsymbol X_k)\,\Delta t}+\class{sde-t-kick}{\boldsymbol g(t)\,\xi_k}$
$\mathrm{d}\boldsymbol X_t=\class{sde-t-drift}{\boldsymbol f(t,\boldsymbol X_t)\,\mathrm{d}t}+\class{sde-t-kick}{\boldsymbol g(t,\boldsymbol X_t)\,\mathrm{d}\boldsymbol W_t}$
Shrink the step: $\Delta t\to\mathrm{d}t$ — the random walk becomes a stochastic differential equation.
 Δt large dt (Δt→0)
05 · Strong vs weak convergence

Two notions of "learning" an SDE

Strong
Learn the Itô integral→ paths converge

Drive the model and SDE with the same noise. Their individual trajectories should agree.

Weak
Learn the distribution, not the path

Individual trajectories differ. Across many samples, the distribution of outcomes should agree.

strong ⊋ weak

06 · What we want to learn

The strong stochastic flow map (Itô map)

Recall the deterministic flow map: integrate the vector field from $s$ to $t$ and jump there in one shot

$\boldsymbol\Phi_{s,t}(\boldsymbol x_s)=\boldsymbol x_s+\displaystyle\int_s^{\,t}\boldsymbol f(\tau,\boldsymbol x_\tau)\,\mathrm{d}\tau$.

The strong stochastic flow map (also known as an Itô map) integrates against the driving Brownian motion $\{\boldsymbol W_t\}$:

$\class{fm-t-param}{\boldsymbol\Psi_{s,t}}(\boldsymbol X_s)=\boldsymbol X_s+\class{fm-t-drift}{\displaystyle\int_s^{\,t}\boldsymbol f(\tau,\boldsymbol X_\tau)\,\mathrm{d}\tau}+\class{fm-t-diff}{\displaystyle\int_s^{\,t}\boldsymbol g(\tau)\,\mathrm{d}\boldsymbol W_\tau}.$

We parameterize the model as the mean drift and diffusion, i.e.,

$\class{fm-t-param}{\boldsymbol\Psi^\theta_{s,t}}(\boldsymbol X_s,\boldsymbol W_{[s,t]})=\boldsymbol X_s+\class{fm-t-drift}{(t-s)\boldsymbol f_{s,t}^\theta(\boldsymbol X_s, \boldsymbol W_{[s,t]})}+\class{fm-t-diff}{(\boldsymbol W_t - \boldsymbol W_s)\boldsymbol g_{s,t}^\theta(\boldsymbol W_{[s,t]})}.$
Learn a neural Itô map $\boldsymbol\Psi^{\theta}_{s,t}$ that reproduces the true flow strongly (pathwise), so we can leap $s\to t$ in a single step.
07 · two obstacles to learning the stochastic map

Learning SDEs is much harder than ODEs

one Brownian sample pathtime
01
Brownian path input
$\boldsymbol W_{[s,t]}\in \mathcal C([s,t];\mathbb R^d)$
A function-valued input, infinite-dimensional!
02
Nowhere differentiable
the limit does not exist
No derivatives defined anyhwere!
ODE-based flow maps use time derivative based flow-map losses $\boldsymbol{\mathcal L}_{\mathrm{LSD}}(\theta)\propto\mathbb E\!\left[\left\|\partial_t\boldsymbol\Psi^\theta_{s,t}(\boldsymbol X_s)-\boldsymbol v_\theta\!\left(t,\boldsymbol\Psi^\theta_{s,t}(\boldsymbol X_s)\right)\right\|_2^2\right]$ $\boldsymbol{\mathcal L}_{\mathrm{ESD}}(\theta)\propto\mathbb E\!\left[\left\|\partial_s\boldsymbol\Psi^\theta_{s,t}(\boldsymbol X_s)+\nabla_{\boldsymbol X_s}\boldsymbol\Psi^\theta_{s,t}(\boldsymbol X_s)\,\boldsymbol v_\theta(s,\boldsymbol X_s)\right\|_2^2\right]$
The representation and the training objective must respect the path.
08 · how do you feed a random path to a network?

The increment alone isn't sufficient information

$\class{brown-t-a}{\boldsymbol W_1-\boldsymbol W_0}$
$\displaystyle\int_0^1 \class{brown-t-a}{(\boldsymbol W_r-\boldsymbol W_0)\otimes\mathrm{d}\boldsymbol W_r}=\tfrac12\!\left[\Delta\boldsymbol W_{0,1}^{\otimes2}-\mathbf I_d\right]+\class{brown-t-area}{\boldsymbol A_{0,1}}$
$W^1$
$W^2$
$\boldsymbol W_t(\omega_1)$
$\boldsymbol W_t(\omega_2)$
$\boldsymbol W_0(\omega_i)$
$\boldsymbol W_1(\omega_i)$
09 · a tractable representation of the path

The Brownian polynomial

L = 16 L
$x$
$y$
N = 4
On one segment $[s,t]$ of the $L$ coefficients
$$\boldsymbol I_{\class{poly-t-seg}{s,t}}^{(\class{poly-t-deg}{n})}=\int_{\class{poly-t-seg}{s}}^{\class{poly-t-seg}{t}}\widetilde P_{\class{poly-t-deg}{n}}\!\left(\tfrac{u-s}{t-s}\right)\mathrm{d}\boldsymbol W_u$$
Degree-N reconstruction on that segment
$$\boldsymbol W_{u,v}^{(\class{poly-t-deg}{N})}=\sum_{\class{poly-t-deg}{n}=0}^{\class{poly-t-deg}{N}-1}\tfrac{2n+1}{t-s}\,\boldsymbol I_{\class{poly-t-seg}{s,t}}^{(\class{poly-t-deg}{n})}\int_u^v\widetilde P_{\class{poly-t-deg}{n}}\!\left(\tfrac{r-s}{t-s}\right)\mathrm{d}r,\quad[u,v]\subseteq[\class{poly-t-seg}{s,t}]$$
Tractable coefficients
$$\boldsymbol I_{\class{poly-t-seg}{s,t}}^{(\class{poly-t-deg}{n})}\sim\mathcal N\!\left(\boldsymbol 0,\tfrac{t-s}{2\class{poly-t-deg}{n}+1}\boldsymbol I\right)$$
Chen relations · merge $[s,u]$, $[u,t]$ into $[s,t]$
$$\boldsymbol I_{\class{poly-t-seg}{s,t}}^{(\class{poly-t-deg}{n})}=\sum_{\class{poly-t-deg}{m}=0}^{\class{poly-t-deg}{n}}c_{n,m}\Big(\boldsymbol I_{\class{poly-t-seg}{s,u}}^{(\class{poly-t-deg}{m})}+(-1)^{\class{poly-t-deg}{n}+\class{poly-t-deg}{m}}\,\boldsymbol I_{\class{poly-t-seg}{u,t}}^{(\class{poly-t-deg}{m})}\Big)$$
Pathwise convergence
$$\big\|\boldsymbol W^{(\class{poly-t-deg}{N})}-\boldsymbol W\big\|_{\alpha\mathrm{-Höl}}\xrightarrow[\;N\to\infty\;]{}0,\quad\alpha\in[0,\tfrac12)$$
10 · how do we train the map?

Two losses: match locally, compose globally

$\class{loss-t-match}{\mathcal L_{\boldsymbol f,\boldsymbol g}}=\mathbb E\big[\lVert\boldsymbol f(t,\boldsymbol X_t)-\hat{\boldsymbol f}_{t,t}(\boldsymbol X_t)\rVert^2+\lVert\boldsymbol g(t)-\hat{\boldsymbol g}_{t,t}\rVert^2\big]$
$\class{loss-t-match}{\mathcal L_{\boldsymbol f,\boldsymbol g}}=(t-s)^{-1}\big\lVert\underbrace{\boldsymbol X_s+\boldsymbol f(s,\boldsymbol X_s)(t-s)+\boldsymbol g(s)\class{loss-t-sg}{(\boldsymbol W_t-\boldsymbol W_s)}}_{\text{small EM step}}-\boldsymbol\Psi^\theta_{s,t}(\boldsymbol X_s,\boldsymbol I^{(N)}_{s,t})\big\rVert^2$
$\class{loss-t-sg}{\mathcal L_{\mathrm D}}=\mathbb E\big[\lVert\hat{\boldsymbol\Psi}_{s,t}(\boldsymbol X_s,\boldsymbol W_{[s,t]})-\mathrm{sg}\big[\hat{\boldsymbol\Psi}_{u,t}(\hat{\boldsymbol\Psi}_{s,u}(\boldsymbol X_s,\boldsymbol W_{[s,u]}),\boldsymbol W_{[u,t]})\big]\rVert^2\big]$
$\mathcal L_{\mathrm{SD}}=\class{loss-t-match}{\mathcal L_{\boldsymbol f,\boldsymbol g}}+\class{loss-t-sg}{\mathcal L_{\mathrm D}}\;\Rightarrow\;\hat{\boldsymbol\Psi}=\text{Itô map (unique global min)}$
time
$X$
11 · interlude

We can learn the Itô map — now what?

✓ representation — the Brownian polynomial $\boldsymbol I^{(N)}_{s,t}$ ✓ objective — EM matching + semigroup distillation
Both obstacles cleared — the network can carry a strong solution map $\boldsymbol\Psi^\theta_{s,t}$.
So… what do we do with it?
12 · experiments · image generation

Strong maps: same path, same horse — at any step count

n = 2n = 4n = 8n = 16n = 32 same $\boldsymbol W$ CIFAR-10 sample, 2 steps, shared Brownian path 4 steps, shared path 8 steps, shared path 16 steps, shared path 32 steps, shared path different $\boldsymbol W$ 2 steps, independent Brownian path 4 steps, independent path 8 steps, independent path 16 steps, independent path 32 steps, independent path
Fixed $\boldsymbol X_0$. Top: one Brownian path ⇒ the same image at every step count (strong convergence). Bottom: fresh $\boldsymbol W$ each time ⇒ stochasticity retained.
13 · experiments · image generation

Competitive visual fidelity in 2–16 steps

SSFMGLASSDiamond MapPSD-MPSD-U
14 · experiments · molecular systems

Alanine dipeptide: 10× speedup

ground truth
100 NFE10 NFE2 NFE diffusion SSFM
PMF error ↓2201001000
Diffusion16.7283.4000.0840.066
SSFM0.2350.0890.0670.062
JS × 10−2 ↓2201001000
Diffusion48.73516.2930.7870.618
SSFM1.9900.8600.6400.590

SSFM at 100 NFE is competitive with diffusion at 1000 NFE — a 10× reduction in network evaluations.

15 · experiments · molecular systems

Chignolin: all-atom generation in under 8 steps

Chignolin · 10 residues · all atoms
$\mathbb T$-W$_2$ ↓1248163264100
Diffusion4.7375.8825.5872.7161.7931.7141.6991.735
SSFM2.4392.0601.8241.7641.7331.7071.6941.670
tICA-W$_2$ ↓1248163264100
Diffusion1.97619.1473.2360.7120.3230.2620.2660.330
SSFM0.7580.5300.4140.3680.3350.3160.2900.304

SSFM captures the backbone torus and slow tICA modes where the diffusion baseline fails.

16 · experiments · molecular systems

Extended Chignolin results

pooled backbone Ramachandran density
truth1248163264100 NFE diffusion SSFM
slow conformational modes · tICA projection
truth1248163264100 NFE diffusion SSFM
We learned the Itô integral.
✓Brownian polynomial encoding
✓Matching & semi-group loss
Questions?

Strong Stochastic Flow Maps

Paperarxiv.org/abs/2606.01086Codegithub.com/sammccallum/ssfm