Strong Stochastic Flow Maps

Abstract: Flow and diffusion models generate high-quality samples in many modalities; however, many network evaluations are required during inference due to numerical integration of an underlying differential equation. Flow maps alleviate this problem by learning the solution map of the differential equation directly, enabling few-step sampling. Yet, current methods are restricted to approximating the solution map of ODEs. These methods can be used to learn the transition kernel of an SDE, thereby obtaining a solution map that recovers the marginal distributions of the process (weak convergence) rather than the solution path (strong convergence). We propose Strong Stochastic Flow Maps (SSFMs) as a novel framework for learning the strong solution map of additive-noise SDEs, directly generalizing deterministic flow maps to the stochastic setting. Further, a polynomial approximation to Brownian motion is introduced and shown to converge pathwise. These results enable a simulation-free training objective for the solution map of diffusion models. We demonstrate that SSFMs outperform previous stochastic flow map methods on image generation and enable few-step sampling of molecular systems.

• The paper introduces a neural SDE generative modeling framework that achieves strong pathwise convergence and simulation-free training.
• It leverages a polynomial encoding of Brownian motion to enable efficient few-step inference with controlled error via adjustable polynomial degree.
• Empirical results on image and molecular datasets demonstrate superior performance and consistent pathwise sampling compared to weak flow map approaches.

Strong Stochastic Flow Maps: A Framework for Pathwise Neural SDE Generative Modeling

Introduction and Motivation

The paper "Strong Stochastic Flow Maps" (2606.01086) introduces a new framework—Strong Stochastic Flow Maps (SSFMs)—designed to learn pathwise solution maps for additive-noise stochastic differential equations (SDEs) in generative modeling. Unlike weak stochastic flow map approaches, which approximate SDE marginal distributions (i.e., weak solutions), SSFMs provide strong (pathwise) convergence, recovering sample trajectories for fixed realizations of Brownian noise. This aligns with deterministic flow map models but extends their capacity into the stochastic regime, unlocking new capabilities for few-step and path-consistent sampling in both image and molecular data domains.

Theoretical Foundations

The authors formalize the SSFM framework as a neural approximation to the Itô map for additive-noise SDEs. The Itô map, $\bm{\Psi}_{s,t}(\bx_s, \bW_{[s,t]})$, produces the solution $\bx_t$ at time $t$ from initial state $\bx_s$ by leveraging the full path of Brownian motion $\bW_{[s,t]}$. Crucially, this map captures strong convergence—meaning for a fixed initial condition and Brownian path, the mapping is deterministic and consistent across temporal decompositions:

(Figure 1)

Figure 1: Strong versus weak stochastic flow maps—only the strong (top) formulation produces pathwise-consistent solutions under the same Brownian realization, while weak (bottom) approaches sample new stochastic paths independently at each timestep.

The SSFM architecture parameterizes the Itô map using a neural network that receives two inputs: the initial condition and a path encoding of the driving Brownian motion. The authors develop a polynomial (shifted Legendre) basis expansion for Brownian motion, which converges in $\alpha$-Hölder distance to the true path as the polynomial degree increases. The expansion coefficients are independent and normally distributed, enabling efficient sampling and the construction of semigroup (composition) properties via closed-form analogues of the Chen relations. This enables consistent pathwise conditioning over discretized intervals.

Figure 2: Left: Ground-truth Brownian path $\bW_t$ and 4th degree polynomial pathwise approximation; Right: The faithful tracking of SDE solution by the SSFM when driven by the polynomial expansion.

Mathematically, the neural SSFM implements a generalized Euler-Maruyama update:

$\bm{\Psi}_{s,t}(\bx_s, \bW_{[s,t]}) = \bx_s + \bm{f}_{s,t}(\bx_s, \bm{I}_{s,t}^{(N)})(t-s) + \bm{g}_{s,t}(\bm{I}_{s,t}^{(N)})(\bW_t - \bW_s),$

where $\bm{I}_{s,t}^{(N)}$ are the polynomial Brownian coefficients, and $\bm{f}_{s,t}$, $\bx_t$0 are neural network-parameterized drift and diffusion mappings.

Training Objective and Algorithm

To train SSFM models, the authors introduce a simulation-free, self-distillation-based loss. The loss enforces:

1. Local Consistency: For infinitesimal intervals ($\bx_t$1), recovered drift and diffusion must match the true SDE coefficients.
2. Semigroup Property: For any partition $\bx_t$2, the flow map satisfies

$\bx_t$3

which ensures compositional consistency across the time axis.

The training is performed using simulated Brownian coefficients, with polynomial degree $\bx_t$4 controlling the fidelity-computation tradeoff. As $\bx_t$5 increases, the SSFM converges strongly to the Itô map, with diminishing truncation error per coefficient added.

Figure 3: True Brownian path $\bx_t$6 and successively refined polynomial path approximations $\bx_t$7 over increasing degrees $\bx_t$8.

Empirical Results

Ablation on Nonlinear SDEs

When applied to nonlinear additive-noise SDEs, increasing the polynomial degree $\bx_t$9 systematically decreases strong convergence error, particularly in the large timestep regime—a key requirement for few-step generative inference.

Image Generation

On CIFAR-10 and CelebA-64, SSFM models outperform all previous weak stochastic flow map approaches, as measured by FID at all tested step counts. Notably, with as few as 2–4 network function evaluations (NFEs), SSFM approaches match or surpass deterministic few-step flow map models.

Molecular Conformation Sampling

In equilibrium molecular generation benchmarks (Alanine-Dipeptide and all-atom Chignolin), SSFM yields orders-of-magnitude smaller pathwise errors and JS divergence for very low NFE regimes compared to diffusion model baselines, while converging to equivalent accuracy as NFE increases.

Figure 4: Ramachandran plots for Alanine-Dipeptide: ground truth (left), diffusion mixture baseline (center), and SSFM predictions (right) across NFE step counts highlight the superior convergence of SSFM.

Figure 5: SSFM predictions closely match the ground truth SDE dynamics across all temporal resolutions.

Figure 6: Pooled backbone Ramachandran plots for Chignolin—ground truth (first column), diffusion baseline (upper rows), and SSFM (lower rows) across decreasing step counts indicate SSFM’s superior high-fidelity sampling at low NFE.

Figure 7: Chignolin tICA projections—SSFMs remain closer to ground truth across all step counts compared to the diffusion baseline.

Implications and Future Directions

The introduction of SSFMs establishes a rigorous, tractable, and efficient pathwise learning paradigm for neural SDE generative models. The pathwise consistency under fixed Brownian realizations enables estimation of pathwise observables, critical for downstream tasks such as reward alignment, trajectory optimization, and structured conditional generation. The polynomial path encoding circumvents the intractability of direct path input, allowing scaling to complex domains.

Strong claims are established: for additive-noise SDEs, the SSFM architecture achieves strong solution map learning, enabling simulation-free, few-step, pathwise-correct generative modeling. Existing weak-flow generative models are fundamentally limited by inability to produce pathwise-consistent solutions, a deficiency addressed by the SSFM approach.

Extensions to state-dependent SDEs appear theoretically tractable under the presented framework but remain empirically untested. Increasing the polynomial degree and scaling to higher-dimensional, longer time horizon domains will be essential for applications in high-resolution video, protein folding, and scientific simulation acceleration.

Conclusion

The SSFM framework generalizes the flow map paradigm to the stochastic domain, providing strong pathwise convergence with practical simulation-free training. Empirical evidence demonstrates state-of-the-art efficiency and accuracy for both image and molecular sample generation in the few-step regime. The theoretical contributions and practical training strategy pave the way for future advances in pathwise generative modeling, with broad implications for sample-efficient, dynamics-consistent real-world AI systems.