Improved strong stochastic flow maps: a smooth stand-in for the Brownian path

A deterministic flow map learns to jump an ODE from time $s$ to time $t$ in one network call. The strongest recipes for training one, MeanFlow and its improved version, differentiate the map with respect to its start time. A strong stochastic flow map (SSFM) does the same for an SDE: it jumps from $s$ to $t$ along one specific realisation of the Brownian motion. That is where the derivative breaks. The map depends on the Brownian path between $s$ and $t$, and a Brownian path has no derivative anywhere.

The original Strong Stochastic Flow Maps paper (McCallum, Blasingame, Herschell, Rindtorff, Tong and Foster) sidestepped this by training with a semigroup loss: one big step must equal two half steps. The follow-up, Improved Training of Strong Stochastic Flow Maps by Zander W. Blasingame, Antonio Franca, Raúl Miñán, Christian Koke, Niklas Rindtorff and Alexander Tong (AITHYRA, TU Wien, TU München, MCML), gets the derivative losses back. Blasingame's thread calls the first paper a NeurIPS 2026 Oral; its arXiv listing says only "Preprint", so I have not confirmed the venue. Code is at zblasingame/improved-strong-stochastic-flow-maps (MIT) and the five ImageNet checkpoints at zblasingame/improved-strong-stochastic-flow-maps on Hugging Face.

The same group's Langevin Flow Maps used the SSFM machinery for molecular dynamics. This piece is about the training method underneath. I read the paper and its LaTeX source, read the code, pulled the checkpoint headers, and ran the authors' toy-1D losses on a CPU. Numbers are labelled measured (I computed it from a file or a run), reported (the authors' figure) or reasoned (my arithmetic on the other two).

The short version: a classical theorem, used well. Replace the rough path with a smooth polynomial that converges to it, train against the random ODE that polynomial drives, and the derivative exists. On ImageNet-256 that gives about three times lower FD-DINOv2 than SSFM at 4 and 8 steps, against the authors' own port of SSFM. The "fifteen times fewer steps" is read off an interpolated curve of the weakest available baseline.

Flow maps, deterministic first

Flow matching learns a velocity field $u_t(x)$ whose ODE carries noise ($t = 0$) to data ($t = 1$). The interpolant is $X_t = \alpha_t X_1 + \sigma_t X_0$; on ImageNet the schedule is linear, $\alpha_t = t$, $\sigma_t = 1 - t$. Sampling integrates the ODE, which costs one network call per solver step.

A flow map skips the solver. It learns the two-time map

where $u_{s,t}$ is the mean velocity over $[s, t]$. On the diagonal, $u_{t,t} = u_t$, the ordinary flow-matching field. Three conditions characterise the true map, and each becomes a loss (reported, Appendix A.2). Write $h = t - s$ and $v_s$ for the instantaneous field:

• Eulerian: $\partial_s \Phi_{s,t}(x) + v_s(x)^\top \nabla_x \Phi_{s,t}(x) = 0$. Moving the start point along the flow does not change where you land.
• Lagrangian: $\partial_t \Phi_{s,t}(x) = v_t(\Phi_{s,t}(x))$. The endpoint moves with the flow.
• Semigroup: $\Phi_{s,t} = \Phi_{u,t} \circ \Phi_{s,u}$. One big step equals two smaller ones. This is the shortcut model loss, and it needs no derivative.

Improved MeanFlow (iMF) rewrites the Eulerian residual in terms of $u_{s,t}$. With the material derivative $D_s = \partial_s + v_s^\top \nabla_x$, it becomes

with $\mathrm{sg}$ a stop-gradient: regress onto the instantaneous field, corrected by a forward-mode JVP of the network itself. The paper calls this Eulerian family the most performant at scale.

The stochastic version, and why the derivative was missing

Every flow-matching ODE has an SDE with the same marginals (the paper's Equation 1):

For the linear schedule, $\nu_t^2 = 2(1-t)/t$, which blows up as $t \to 0$ (reported, Appendix B.4.1).

A weak stochastic flow map learns to sample $X_t$ given $X_s$: right in distribution, with fresh noise every call. A strong one learns the pathwise solution: given $X_s$ and the Brownian path $W_{[s,t]}$, return the $X_t$ that this path produces. Step 4 times or 16 times along the same path and you should land on the same sample.

SSFM writes the map as an Euler-Maruyama-shaped step with learned averages:

where $\mathcal I$ is a finite summary of the path on $[s, t]$ and $g$ was a learned diffusion head.

The summary is the shifted Legendre expansion (Foster, Lyons and Oberhauser). Its coefficients are Itô integrals of the Legendre polynomials $\widetilde P_n$ against the path:

$I^{(0)}$ is the increment $W_t - W_s$. The higher ones describe how the path bends away from the straight chord. They are independent, and two intervals' coefficients combine into the parent's by a closed-form Chen relation, which is all the semigroup loss needs (the Langevin Flow Maps piece covers both properties).

Now try the Eulerian loss. It needs $\partial_s \Psi_{s,t}$: slide the start point and see how the output changes. But $\Psi$ takes the path on $[s, t]$ as input, so sliding $s$ changes the path the map sees, and the rate of that change is $dW/ds$, which does not exist. The paper also cites Itô maps for any-step SDEs (Pan et al.), whose Proposition 2 shows that a naively applied Eulerian loss places no pathwise constraint on the map: it can be satisfied while ignoring the path.

The fix: train against a random ODE

Step one: give the path a derivative. The degree-$N$ Legendre polynomial built from the first $N$ coefficients is a smooth curve. For a sub-interval $[u, v]$ of $[s, t]$ it is (the paper's Equation 3)

and its time derivative is an ordinary polynomial, in closed form (Lemma B.1):

Step two: replace the SDE with the ODE this derivative drives. Write $b_r$ for the SDE's drift. The Wong-Zakai random ODE on $[s, t]$ is

Wong and Zakai showed in 1965 that ODEs driven by smooth approximations of Brownian motion converge to the Stratonovich SDE, and for additive noise Itô and Stratonovich coincide. The paper's Proposition C.3 makes it quantitative for this basis. If $b$ is $L$-Lipschitz, then almost surely

and the last factor goes to zero as $N$ grows. The proof is two lines of integration by parts and Grönwall.

The random ODE is a plain ODE for any fixed draw of coefficients, so its flow map obeys the Eulerian condition with an ordinary derivative. The residual (Equation 103) is

Two details make it work. First, $\mathfrak D_s$ is a fixed-path derivative: hold the polynomial curve still and slide $s$ along it. The coefficients of the restricted curve change smoothly, because a degree-$(N-1)$ polynomial restricted to a sub-interval is still exactly $N$ coefficients, so one forward-mode JVP gives the whole thing. Second, the transport field $v^{(N)}$ contains $\nu_s \dot W^{(N)}_s$, which depends on all $N$ coefficients. A network that ignores the higher ones cannot satisfy the loss, which is the paper's second design principle in one term.

The widget runs the theorem on one path. It draws 2,048 Brownian increments on $[0, 1]$, computes the first 32 Legendre coefficients from them, and solves $dX = 4(X - X^3)\,dt + dW$ twice from $X_0 = 0.3$: by Euler-Maruyama on the rough path, and by Heun on the random ODE driven by $\dot W^{(N)}$.

What it shows (measured, in-browser, deterministic seeds). On seed 7, the endpoint error $|X^{(N)}(1) - X(1)|$ falls from 0.5192 at $N = 1$ to 0.0921 at $N = 2$ and 0.0089 at $N = 32$. $N = 1$ is the chord: the ODE sees only the increment, which is all a map conditioned on the increment alone can know. The sup gap between the two Brownian curves falls more slowly, from 1.296 to 0.209, and not monotonically: on seed 11 the endpoint error at $N = 4$ is worse than at $N = 3$. Both curves end at $W(1)$, because every higher Legendre polynomial integrates to zero. The cost is in the last column: the polynomial's slope has RMS 1.98 at $N = 3$ and 6.23 at $N = 32$, while the rough path's finite-difference slope at this grid is 45.8. A higher $N$ is a better approximation and a rougher driver.

expand

Two alternatives lose (reported, Section 3.2). Discretised Eulerian and Lagrangian losses, a small Euler-Maruyama step in place of the derivative, beat the semigroup loss on the toy and slightly lose to it at scale. The Doss-Sussmann transform, $Z_t = X_t - M_t$, makes the SDE an ODE in $Z$, but its derivative never involves the higher coefficients: on the 1D toy at 1 step its strong error sits near 6.5 at every $N$ (reported, Table 4).

Three changes to the map itself

The second half of the paper is a fine-tuning recipe. It starts from a pre-trained Decoupled MeanFlow model, DMF-XL/2+, which is itself distilled from SiT-XL/2+REPA. DMF turns a one-time flow model into a two-time map by feeding $s$ to the first 20 of the 28 transformer blocks and $t$ to the last 8 (reported, Appendix G.1), so at $s = t$ the map is the pre-trained velocity field. iSSFM adds the Brownian coefficients through a zero-initialised residual scaled by $h$, and freezes the backbone under rank-16 LoRA adapters. On the diagonal the coefficients vanish and the network is exactly the teacher at initialisation.

expand

The SDE's drift is $2u_t(x) - (\dot\alpha_t/\alpha_t)x$. The linear term is stiff and diverges as $t \to 0$, and a learned map would have to integrate it on top of the teacher. The three changes:

1. Exponential lift. Multiply through by $\alpha_t$. Itô's product rule gives $d[\alpha_t X_t] = 2\alpha_t u_t(X_t)\,dt + \alpha_t \nu_t\,dW_t$: no linear term. The map becomes an exponentially weighted step,

2. Closed-form martingale. The noise term $M^\alpha_t = \int \alpha_r \nu_r\,dW_r$ is a time-changed Brownian motion (Dambis-Dubins-Schwarz), run on the clock of its quadratic variation. For the linear schedule that clock is $\langle M^\alpha\rangle_t = t^2 - \tfrac23 t^3$ as $t_0 \to 0$, bounded where $\nu_t$ is not (reported, Appendix B.4.1). SSFM learned a diffusion head for this term; iSSFM computes it exactly and builds the Legendre coefficients on the lifted clock, so the map's noise and the coefficients the network sees come from the same path.
3. Drift matching. With the noise exact, the diagonal can regress directly onto the teacher's drift instead of onto a small Euler-Maruyama step.

The paper says the order matters: the closed-form martingale without the lift has infinite variance near $t = 0$, and their run of it "diverged completely into pure noise".

The training loss in the lifted state is iMF: regress $2\alpha_s u_{s,t} - 2\mathcal A_{s,t}\,\mathrm{sg}[D_s u_{s,t}]$ onto $2\alpha_s\,\mathrm{sg}[u^{\text{teacher}}_s]$, with the JVP's tangent along the smoothed lifted field (reported, Appendix G.2). The code matches (measured, dmf_wz_imf.py): the compound is lift_s * u - a_st * material_derivative.detach(), the target is lift_s * v_field, and the Wong-Zakai driver enters only the JVP tangent b_y = lift_s * v_field + driver_s. A comment says why: the driver "cancels in the compound identity, so the regression target stays the deterministic start field". The path information reaches the network through the material derivative, which is exactly where the semigroup loss had nothing.

The recipe (reported, Appendix G): 75% of each batch on the distillation term and 25% on the diagonal, $t \sim \mathcal U[0.06, 1]$, step at least 0.05, AdamW at $10^{-4}$, effective batch 128, 100k steps, EMA 0.9999, $N = 3$ coefficients, one H100 NVL per run. The released checkpoint holds only what was trained (measured, safetensors header of issfm/model.safetensors): 38,571,008 float32 parameters, of which 12,003,584 are LoRA, 26,547,840 sit in a module named fixed_residual.levy_net (the coefficient pathway) and 19,584 in the coefficient embedder. The SSFM checkpoint has 41,245,968, the extra 2.67 million in its learned diffusion layer. Both load on top of the DMF-XL/2+ backbone.

Checking the ImageNet claims

All numbers in this section are FID-50k and FD-DINOv2-50k on class-conditional ImageNet-256, fine-tuned for 100k steps from the same DMF-XL/2+ (reported, Tables 1, 3 and 5). Lower is better.

"+ param." is SSFM with all three parameterisation changes, still trained with the semigroup loss. iSSFM adds the Wong-Zakai iMF loss.

"~3x lower FD-DINOv2 than SSFM at 4 NFE"

Holds at 4 and 8 steps: 475.1 / 143.0 = 3.3 and 229.1 / 75.2 = 3.0 (reasoned). It is 2.1 at 1 step, 1.8 at 2 and 2.5 at 16. The split between the two contributions is uneven. The parameterisation alone gives 1.65 to 1.92 at 4 to 16 steps, more than the "roughly 1.5x" the paper says. The loss alone gives 2.0 at 4 steps, 1.7 at 8 and 1.3 at 16 (reasoned). At 16 steps the improved parameterisation with the old loss has the better FID, 2.03 against 2.07. The new loss matters most where steps are few, which is where it should.

The baseline needs a note. "SSFM" here is the authors' port of SSFM into the fine-tuning setting, "the most faithful replication of SSFM in our setting". The original SSFM paper trained from scratch and reported CIFAR-10 and CelebA-64 only: FID 3.49 at 4 steps on CIFAR-10, against deterministic flow maps at 3.34 (reported, SSFM Table 2). There is no head-to-head on a shared benchmark between the published SSFM and iSSFM.

"4-step iSSFM matches 16-step SSFM"

In FD-DINOv2, roughly: 143.0 against 135.6, 5% worse (reasoned). In FID, no: 5.96 against 4.82.

"~15x fewer NFE than the SiT SDE sampler at matched FD"

This one is read off a curve. The SiT-XL/2+REPA baseline is the SDE integrated by Euler-Maruyama, measured at 1, 3, 8, 17, 85, 170 and 426 steps, with FD-DINOv2 of 1288.1 at 17 steps and 88.5 at 85 (reported, Table 5). The dashed "15x" line in Figure 1 runs from iSSFM's 2-step point, 583.5, to the SiT curve between those two samples. Interpolating log-linearly between them puts SiT at 583.5 near 27 steps, about 14x (reasoned). That is a matched FD of 583.5, where FID is 53.11: an unusable sample quality. At usable quality the table gives a direct comparison: iSSFM at 8 steps (FD 75.2, FID 2.37) beats SiT at 85 (FD 88.5, FID 3.65), about 10.6x fewer steps, and SiT at 170 steps is better than any iSSFM row in FID (1.77 against 2.07).

Euler-Maruyama is also the weakest SDE sampler one could pick, and iSSFM starts from DMF-XL/2+, a deterministic few-step model whose own scores the paper does not tabulate. The limitations section says it plainly: iSSFMs "still struggle to match the performance of deterministic few-step models". What iSSFM buys over a deterministic map is not quality per step. It is a sampler whose noise you control pathwise.

expand

Is it still strong?

The figure above is the qualitative test. The quantitative one is Figure 5b: latent MSE of a $k$-step rollout against the same model's 32-step rollout on the same path, over 1,000 images. Here the source and the picture disagree. The text says the semigroup-trained map "exhibits stronger pathwise consistency, at the cost of worse metrics". The plotted legend says the opposite: the lower, more consistent curve is labelled iSSFM. The LaTeX resolves it (measured, tikz/strong-consistency.tikz and its .dat): the olive curve is drawn from the column psd (semigroup) and the navy from imf, but the legend lists iSSFM first. The legend is swapped. The data say: the iMF-trained iSSFM goes from 0.956 at 1 step to 0.432 at 4 and 0.188 at 16; the semigroup map from 0.230 to 0.158 to 0.091; the weak baseline, which redraws the path every rollout, stays near 1.0 to 1.3. So iSSFM is strong, its error falls with steps where a weak map's does not, but it is about twice as far from its own 32-step answer as the semigroup map at 4 and 16 steps. The derivative loss bought sample quality with some pathwise consistency, as the text says.

The toy losses, rerun on a CPU

The paper's loss comparison is a 1D non-linear SDE, $dX = (X - X^3)\,dt + \sqrt{\beta_t}\,dW$ with $\beta$ linear from 20 to 0.1, and strong error against a 512-step Euler-Maruyama reference on 1,024 shared paths (reported, Table 4). At one step, WZ-iMF falls from 0.65 at $N = 1$ to 0.0763 at $N = 4$, an 8.5x drop as coefficients are added. The semigroup loss goes the other way, 0.513 to 2.44, and at $N = 4$ blows up at 4 and 8 steps (289 and 67,000). The discrete Eulerian loss is competitive at most step counts and best at 8. Doss-Sussmann iMF is the worst at 1 step and the best at 32 (0.00224). The win is a few-step win.

expand

I ran two cells of that grid myself (measured, shared 4-core cloud VM, CPU only, PyTorch 2.14.1, 2 threads). The release needs Lightning and Hydra to train; I wrote a short loop that imports the repo's own Toy1DSSFMLoss, SSFM_MLP and strong-error helpers at commit cb7b92d and mirrors configs/experiment/toy1d.yaml: a 256-wide, 4-deep MLP, AdamW at $10^{-3}$, EMA 0.999, $N = 3$. The difference is the budget. Six authors shared the machine, so each run got 1,200 steps at batch 1,024, about 6% of the paper's 5,000 steps at 4,096, and no learning-rate scheduler. Evaluation copies the repo's callback: 1,024 paths, seed 12345, 512-step reference.

At this budget I do not reproduce the ranking. WZ-iMF is slightly better at one step and worse from two steps on, and both are an order of magnitude or more from the paper's numbers. The 32-step column says why: it measures the diagonal, and the iMF run's diagonal loss ended at 0.32 against the semigroup run's 0.14. My reading (reasoned): the iMF compound puts a JVP-corrected target on a quarter of each batch, and in the first 1,200 steps that slowed the drift fit the whole map rests on. It is a statement about early training, not about the converged method. One full-budget cell is about 17 times the samples of mine, and the paper's grid is 36 cells; neither fit.

What holds

• The mechanism. Training the flow map of the Wong-Zakai random ODE is a sound way to get an Eulerian loss for a pathwise SDE map. The bound is standard and correctly proved, the derivative is closed-form, and the transport term makes the network read every coefficient.
• The ImageNet gain over SSFM. About 3x in FD-DINOv2 at 4 and 8 steps, roughly half from the parameterisation and half from the loss, against the authors' port of SSFM. The loss's share falls to 1.3x at 16 steps.
• The strong property. Yes, at a cost visible in the authors' own data: the iMF map is less self-consistent than the semigroup map, and the published legend says the reverse.
• The step savings against SiT. About 10x at matched quality by the table, against an Euler-Maruyama SDE sampler. The 15x is an interpolation at a quality nobody samples at.
• The toy ranking. Not reproduced in my 6%-budget runs; not refuted either.

What I find most useful is the framing: a strong stochastic flow map now trains like a deterministic one, on an ODE, with the noise turned into an input the derivative can see. The Feynman-Kac result in Section 4.3, where the map carries a particle's log-weight along its own path so resampling works at 4 steps, is something only a strong map can do, and the better reason to want one. The data-space iteration piece covers the deterministic few-step side of this family; the diffusion transformer page covers the backbone under DMF and SiT.