Improved Training of Strong Stochastic Flow Maps
Abstract
Stochastic flow maps promise few-step sampling of stochastic differential equations (SDEs) by learning to directly jump between times along the solution path, yet training these models to learn this solution map remains difficult. The primary obstacle is the Brownian path itself: it is an infinite-dimensional object which needs to be passed to the neural network and is nowhere differentiable, which makes the implementation of derivative-based losses that power the state-of-the-art deterministic flow maps just out of reach. We identify two core design principles which address these issues: a) represent the Brownian path with a basis which carries higher-order local information, and b) train with an objective that forces the network to use the full representation of the path, not just the first order. To address this, we leverage the shifted Legendre polynomial expansion with a Wong-Zakai styled objective, unlocking derivative-based losses. Finally, we propose a recipe for fine-tuning large pre-trained flow models into strong stochastic flow maps. We ablate through this design space on several toy problems and use our findings to transform a large flow model into a strong stochastic flow map on ImageNet, achieving roughly a three-fold improvement in FD-DINOv2 over SSFM at four NFE.
Related
- Model reduction in time: Learning flow maps for reducing stochastic dynamical systems in time
- Numerical approximation of probabilistically strong solutions to the stochastic harmonic map heat flow
- Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows
- Strong Positive Recurrence for Strongly Entropy Expanding Maps
- Improved weighted bounds for the strong maximal function
BibTeX
@misc{blasingame2026improved,
title = {{Improved Training of Strong Stochastic Flow Maps}},
author = {Zander W. Blasingame and Antonio Franca and Raul Minan and Christian Koke and Niklas Rindtorff and Alexander Tong},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.29605},
eprint = {2610.29605},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.29605},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Improved Training of Strong Stochastic Flow Maps
AU - Zander W. Blasingame
AU - Antonio Franca
AU - Raul Minan
AU - Christian Koke
AU - Niklas Rindtorff
AU - Alexander Tong
PY - 2026
DA - 2026/10/21
PB - arXiv
DO - 10.48550/arxiv.2610.29605
UR - https://arxiv.org/abs/2610.29605
AB - Stochastic flow maps promise few-step sampling of stochastic differential equations (SDEs) by learning to directly jump between times along the solution path, yet training these models to learn this solution map remains difficult. The primary obstacle is the Brownian path itself: it is an infinite-dimensional object which needs to be passed to the neural network and is nowhere differentiable, which makes the implementation of derivative-based losses that power the state-of-the-art deterministic flow maps just out of reach. We identify two core design principles which address these issues: a) represent the Brownian path with a basis which carries higher-order local information, and b) train with an objective that forces the network to use the full representation of the path, not just the first order. To address this, we leverage the shifted Legendre polynomial expansion with a Wong-Zakai styled objective, unlocking derivative-based losses. Finally, we propose a recipe for fine-tuning large pre-trained flow models into strong stochastic flow maps. We ablate through this design space on several toy problems and use our findings to transform a large flow model into a strong stochastic flow map on ImageNet, achieving roughly a three-fold improvement in FD-DINOv2 over SSFM at four NFE.
ER -