Rough Stochastic Flow Maps

Learning finite-time flow maps can amortize the cost of repeatedly solving differential equations, but extending this idea to pathwise stochastic dynamics is challenging when the diffusion is state-dependent and noncommuting: Brownian increments alone do not determine the finite-time response. We introduce Rough Stochastic Flow Maps (RSFM), a framework for learning pathwise solution flows of Stratonovich SDEs using rough-path structure. RSFM combines a prescribed second-order rough-flow expansion with a neural remainder scaled according to its short-time rough-path order, and enforces consistency across time intervals through a cocycle objective. We show that, under suitable regularity assumptions, superlinear local agreement with the exact flow together with exact cocycle consistency identifies the true pathwise solution flow. For computation, we use a finite-dimensional Legendre representation of Brownian motion that determines the first two rough-path levels required by the analytic backbone, and train by cocycle self-distillation without ground-truth flow-map targets. On a nonlinear multiplicative SDE with noncommuting diffusion, RSFM reduces same-driver flow-map error by factors of $2.5$-$6.7$ relative to the second-order backbone and $7.9$-$21.9$ relative to an adapted state-aware SSFM-style baseline. RSFM also improves end-to-end simulation accuracy over the backbone at all $20$ tested resolution step-size settings, while suppressing the learned remainder by roughly four orders of magnitude when the analytic backbone is already exact.