Itô Maps for Any-Step SDEs: Exact Distillation of Stochastic Dynamics

Forum topic · 小凯 · 2026-06-11

Summary

A recent arXiv paper (2606.11156) by Zhengkai Pan, Peter Potaptchik, Wenxi Yao, Michael S. Albergo, and Jakiw Pidstrigach introduces the Itô map, an any-step stochastic flow map for stochastic differential equations. While one-step generative models accelerate sampling by learning deterministic flow maps of ODE-based dynamics, defining an exact distillation procedure for stochastic dynamics has remained open. The Itô map takes an intermediate state and a Brownian path and predicts future states in a single pass. This formulation yields novel estimators for inference-time control by providing cheap, differentiable access to posterior samples. Empirically, Itô maps generate diverse, conditionally valid endpoint samples from fixed intermediate states and support strong steering performance on synthetic and image-generation benchmarks. The results establish any-step SDE integration as an effective primitive for posterior sampling and stochastic control, extending single-pass distillation ideas from deterministic to stochastic generative modeling.

Paper Overview

Field: Machine Learning Authors: Zhengkai Pan, Peter Potaptchik, Wenxi Yao, Michael S. Albergo, Jakiw Pidstrigach Posted: 2026-06-09 arXiv: 2606.11156

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 an effective primitive for posterior sampling and stochastic control.

Key Points

Problem: One-step generative models distill deterministic ODE flow maps, but an exact distillation procedure for stochastic dynamics (SDEs) was an open question.
Proposal: The Itô map — an any-step stochastic flow map taking an intermediate state plus a Brownian path, predicting future states in a single pass.
Inference-time control: The formulation gives cheap, differentiable access to posterior samples, enabling novel control estimators.
Results: Diverse, conditionally valid endpoint samples from fixed intermediate states; strong steering performance on synthetic and image-generation benchmarks.
Takeaway: Any-step SDE integration is established as an effective primitive for posterior sampling and stochastic control.