SSFMs: Strong Stochastic Flow Maps

• Strong Stochastic Flow Maps (SSFMs) are a framework that learns pathwise approximations of stochastic differential equations by modeling the entire Brownian path.
• They employ truncated shifted-Legendre polynomial expansions to represent driving noise, ensuring convergence in the Hölder metric and preserving stochastic calculus semantics.
• SSFMs enable simulation-free training with self-distillation losses that balance drift and diffusion components, driving robust performance in high-dimensional generative modeling and molecular simulations.

A Strong Stochastic Flow Map (SSFM) is a mathematical and algorithmic framework that extends the concept of deterministic flow maps to the stochastic setting, with the aim of learning or approximating strong (pathwise) solution maps for stochastic differential equations (SDEs), particularly those with additive noise. These maps enable consistent, pathwise realizations of SDE transitions for generative modeling, sampling, and system simulation, offering both theoretical guarantees and empirical advantages over weakly consistent or purely deterministic approaches. SSFMs are foundational to recent developments in high-dimensional generative modeling, robust molecular dynamics, and unbiased posterior sampling.

1. Mathematical Formulation and Pathwise Consistency

Consider an additive-noise Itô SDE on $\mathbb{R}^d$: $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$ where $f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$ is a drift, $g: \mathbb{R} \to \mathbb{R}^{d \times w}$ is a diffusion coefficient, and $W_t$ is standard $\mathbb{R}^w$ Brownian motion. Under Lipschitz and linear growth conditions, this SDE admits a unique strong solution in the Itô sense.

The "Itô map" $\Psi_{s, t}: (X_s, W_{[s,t]}) \mapsto X_t$ describes the pathwise evolution of the SDE from state $X_s$ using the entire Brownian path over $[s, t]$. A Strong Stochastic Flow Map is a learned or constructed neural or analytical map $\Psi^\theta_{s, t}$ that, given $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$0 and a suitably parameterized representation of $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$1, outputs $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$2 so that the map converges to the true Itô map in supremum norm over paths (i.e., "strong" or pathwise convergence) (McCallum et al., 31 May 2026).

This extends deterministic flow maps, where $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$3 and the map $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$4 defines deterministic ODE evolution. SSFMs generalize to handle the entire admissible randomness in SDEs, not just their marginal transition kernels.

2. Polynomial Approximation and Representation of Driving Noise

To make the driving Brownian path $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$5 amenable to neural network input, SSFMs employ truncated shifted-Legendre polynomial expansions: $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$6 where $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$7 are shifted Legendre polynomials. These coefficients are independent Gaussian vectors: $dX_t = f(t, X_t)\,dt + g(t)\,dW_t$8 and parameterize a degree-$dX_t = f(t, X_t)\,dt + g(t)\,dW_t$9 polynomial path $f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$0. This approximation converges almost surely to the real Brownian path in the $f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$1-Hölder metric for any $f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$2: $f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$3 Chen relations provide consistent concatenation of increments, preserving correct stochastic calculus semantics (McCallum et al., 31 May 2026).

This construction enables neural SSFMs to condition on a finite-dimensional summary of the latent randomness, ensuring pathwise consistency and facilitating strong approximation.

3. Simulation-Free Training Objectives for Solution Maps

SSFMs are trained by minimizing "self-distillation" losses that enforce both local and global SDE consistency: $f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$4 with

$f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$5

enforcing local tangent (SDE coefficient) matching, and

$f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$6

enforcing the semigroup property ($f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$7).

Notably, these losses are simulation-free: for known SDEs (including generative diffusion models), analytic or flow-matching expressions for $f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$8 and $f: \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$9 are used, and training does not require sampling long SDE trajectories (McCallum et al., 31 May 2026).

In small-step limits, these losses balance drift and diffusion components, as the expected Euler–Maruyama loss is

$g: \mathbb{R} \to \mathbb{R}^{d \times w}$0

Polynomial truncation error in representing $g: \mathbb{R} \to \mathbb{R}^{d \times w}$1 decays as $g: \mathbb{R} \to \mathbb{R}^{d \times w}$2, with empirical convergence often achieved with $g: \mathbb{R} \to \mathbb{R}^{d \times w}$3 (McCallum et al., 31 May 2026).

4. Theoretical Guarantees, Contractivity, and Comparison

Under standard regularity (Lipschitz, differentiability), the unique minimizer $g: \mathbb{R} \to \mathbb{R}^{d \times w}$4 of the SSFM loss coincides with the true Itô solution map. As $g: \mathbb{R} \to \mathbb{R}^{d \times w}$5, the model achieves almost sure convergence to the strong solution.

The contractivity of SSFM-based flows has been established in the context of stochastic interpolation frameworks. For instance, when sampling from a Gaussian base $g: \mathbb{R} \to \mathbb{R}^{d \times w}$6 to a strongly log-concave target $g: \mathbb{R} \to \mathbb{R}^{d \times w}$7, the flow map satisfies the sharp Lipschitz bound (Daniels, 14 Apr 2025): $g: \mathbb{R} \to \mathbb{R}^{d \times w}$8 where $g: \mathbb{R} \to \mathbb{R}^{d \times w}$9 is the uniform lower bound on $W_t$0, matching Caffarelli’s constant for optimal transport. The flows admit exponential contractivity for appropriate schedules of interpolation coefficients, and these results extend to non-Gaussian endpoints using log-smoothness constants.

This property ensures both stability for numerical discretization and rapid convergence for high-dimensional sampling, a critical practical advantage for modern generative modeling (Daniels, 14 Apr 2025).

5. Empirical Evaluation Across Image and Molecular Domains

Comprehensive experiments demonstrate the distinct benefits of SSFMs in a range of generative settings (McCallum et al., 31 May 2026):

• Toy SDEs: For nonlinear SDEs such as $W_t$1 with rapidly varying $W_t$2, SSFMs show fast decay of strong pathwise error with increasing polynomial degree $W_t$3, achieving nearly exact solution at $W_t$4 with only 16 integration steps.
• Image Generation: On datasets such as CIFAR-10 and CelebA-64 using EDM2 U-Net architectures, SSFMs (with $W_t$5) achieve FID (Fréchet Inception Distance) competitive with deterministic flow-maps and significantly outperform weak stochastic flow baselines like "Diamond Maps" and "GLASS" in the low-step regime.

| Method | NFE=2 | NFE=4 | NFE=8 | NFE=16 | |--------------------|-------|-------|-------|--------| | SSFM | 4.93 | 3.49 | 3.29 | 3.35 | | Diamond Map (weak) | 5.80 | 5.80 | 6.73 | – | | GLASS (weak) |157.5 | 39.5 |11.6 | – |

• Molecular Sampling: For Alanine-Dipeptide and Chignolin, SSFMs using graph transformers or DiT architectures achieve parity with fine-grained SDE simulation at an order of magnitude fewer steps, with precise match in equilibrium distribution metrics (e.g., PMF squared-error, τ-$W_t$6, tICA-$W_t$7). At even a single step, SSFMs yield plausible molecular conformations, and consistently outperform diffusion baselines at low computational cost.

Empirical analysis establishes that SSFMs confer both high sample fidelity and pathwise consistency for fixed Brownian paths, uniting the advantages of deterministic map-based methods and SDE simulation.

6. Extensions, Generalizations, and Related Approaches

SSFMs provide a unifying language for deterministic flows (including normalizing flows, rectified flows, and probability-flow ODEs), stochastic samplers (reverse-time SDEs), and broader classes of stochastic interpolation flows (Singh et al., 2024, Daniels, 14 Apr 2025).

Key extensions and variations include:

• Adaptive and Optimal Noise Schedules: Control parameters (e.g., $W_t$8, $W_t$9) can be tuned post-training to modulate bias-variance-diversity tradeoffs without retraining the underlying flow; optimal schedules may be learned via FID or likelihood on validation sets (Singh et al., 2024).
• Reward Alignment and Posterior Estimation: Weak stochastic flow maps ("Diamond Maps") and posterior mapping approaches facilitate efficient value function estimation and reward-guided sampling, but only SSFMs guarantee pathwise (strong) consistency (Holderrieth et al., 5 Feb 2026).
• Strong Measurable Modifications: The theoretical foundations for strong modifications of stochastic flows on general metric spaces, including coalescing flows and flows on graphs, have been formalized for pathwise and measurable consistency (Raimond et al., 2023).

SSFMs thus form a flexible and theoretically grounded foundation for generative modeling, high-dimensional sampling, efficient amortized inference, and molecular simulation, enabling robust translation between deterministic and stochastic paradigms.

7. Theoretical and Practical Significance

The advent of SSFMs resolves a central limitation in earlier stochastic flow maps, which were only weakly consistent (i.e., preserved marginal laws but not sample paths). By directly learning solution maps that are strong in the sense of Itô, SSFMs:

• Enable consistent, memory-efficient few-step sampling for complex SDEs without reliance on costly simulation.
• Achieve dimension-free sampling performance and rapid convergence under log-concavity, as substantiated by sharp Lipschitz bounds matching those for optimal transport maps (Daniels, 14 Apr 2025).
• Provide a practical engine for simulation-free training of generative SDEs and diffusion models at scale, as confirmed empirically in imaging and computational chemistry (McCallum et al., 31 May 2026).
• Unify diverse lines of theoretical work in probability, transport, and machine learning, and open avenues for further research in adaptive control, guided sampling, and measure-valued flows.

These advances render SSFMs an essential construct in the modern toolkit of probabilistic modeling, simulation, and generative inference.