Strong Stochastic Flow Maps

2026-05-31

AI Summary

Existing stochastic flow methods guarantee only weak convergence, which limits their ability to accurately model diffusion trajectories and consequently hinders the quality and applicability of few-step sampling. This work proposes the Strong Stochastic Flow Mappings (SSFMs) framework, which for the first time extends deterministic flow models to the stochastic setting by directly learning the strong solution map of additive-noise stochastic differential equations. By introducing a polynomial approximation of Brownian motion, SSFMs achieve pathwise strong convergence and enable a training objective that does not rely on numerical simulation. Empirical results demonstrate that the proposed method outperforms existing stochastic flow approaches in image generation and achieves high-quality few-step sampling in molecular systems.

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.

Problem

• stochastic differential equations
• strong convergence
• flow maps
• diffusion models
• Brownian motion

Innovation

• Strong Stochastic Flow Maps
• additive-noise SDEs
• pathwise convergence
• simulation-free training
• few-step sampling