Itô maps for any-step SDEs
2026-06-09 • Machine Learning
Machine Learning
AI summaryⓘ
The authors introduce a new method called the Itô map, which helps predict future states in complex systems that involve randomness, like those modeled by stochastic differential equations (SDEs). Unlike previous methods that focused on deterministic processes, their approach works in a single step using random Brownian motion paths. This makes it easier and faster to sample possible outcomes and control these random systems during inference. They tested the method on synthetic data and images, showing it can produce diverse and accurate results while allowing for effective control.
Stochastic differential equations (SDEs)Itô mapBrownian motionDeterministic flow mapsPosterior samplingStochastic controlGenerative modelsInference-time controlStochastic flows
Authors
Zhengkai Pan, Peter Potaptchik, Wenxi Yao, Michael S. Albergo, Jakiw Pidstrigach
Abstract
Recent one-step generative models accelerate sampling by learning deterministic flow maps of the underlying dynamics. These methods rely on learning from ordinary differential equations, leaving open how to define an exact distillation procedure for stochastic dynamics. We introduce the Itô map, an any-step stochastic flow map that takes an intermediate state and Brownian path and predicts future states in a single pass. The Itô map formulation yields novel estimators for inference-time control by providing cheap, differentiable access to posterior samples. Empirically, Itô maps produce diverse, conditionally valid endpoint samples from fixed intermediate states and support strong steering performance on synthetic and image-generation benchmarks. These results establish any-step SDE integration as a useful primitive for posterior sampling and stochastic control.