强随机流映射的训练改进：用移位 Legendre 多项式与 Wong-Zakai 目标实现基于导数的损失

View PDF

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.

Comments:	preprint
Subjects:	Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Dynamical Systems (math.DS)
Cite as:	arXiv:2610.29605 [cs.LG]
	(or arXiv:2610.29605v1 [cs.LG] for this version)
	https://doi.org/10.48550/arXiv.2610.29605 arXiv-issued DOI via DataCite (pending registration)

Submission history

From: Zander Blasingame [view email]
[v1] Wed, 21 Oct 2026 11:43:19 UTC (5,648 KB)