Itô Maps for Any-Step SDEs: Single-Pass Stochastic Flow for Posterior Sampling

Paper: 2606.11156 Authors: Zhengkai Pan, Peter Potaptchik, Wenxi Yao, Michael S. Albergo, Jakiw Pidstrigach Categories: stat.ML, cs.LG

The Gap

Recent one-step generative models (Rectified Flow, Flow Matching, Consistency Models) learn deterministic flow maps from ODEs. They accelerate sampling but ignore the stochasticity inherent in score-based diffusion or Langevin dynamics. Prior work cannot exactly distill a multi-step SDE into a single deterministic or stochastic map — they either approximate or keep multiple steps. This paper bridges that gap by defining the Itô map: a stochastic flow that ingests a Brownian path segment and outputs the SDE endpoint in one pass.

+---

![Concept animation](/arxiv-visuals/it-maps-for-any-step-sdes/ConceptScene.gif)

-------------------------------+      +-------------------------------------+
| Problem: Existing one-step       |      | Assumption: A single stochastic     |
| flow maps are deterministic      | ---> | map can exactly integrate           |
| (ODE only)                       |      | any-step SDE dynamics               |
+----------------------------------+      +-------------------------------------+
                                                          |
                                                          v
+----------------------------------+      +-------------------------------------+
| Method: Itô map f(x_t, B) -> x_T |      | Evidence: Conditional validity      |
| trained via conditional flow     | <--- | (marginal/conditional tests)        |
| matching                          |      | + steering experiments on images    |
+----------------------------------+      +-------------------------------------+
                                                          |
                                                          v
+-----------------------------------------------------------------+
| Conclusion: Any-step SDE integration is a useful primitive      |
| for posterior sampling and stochastic control                   |
+-----------------------------------------------------------------+

The Increment

One sentence: Before this paper, SDEs required iterative simulation or approximation for generative sampling; after this paper, any-step SDE integration is possible in a single forward pass via the Itô map.

Core Mechanism

The Itô map is a function $f: (x_t, B_{[t,T]}) \mapsto x_T$, where $x_t$ is an intermediate state and $B_{[t,T]}$ is the full Brownian path from time $t$ to $T$. The key insight: by conditioning on the entire Brownian path, the SDE trajectory becomes deterministic, so learning the endpoint mapping reduces to a conditional regression problem.

Training proceeds via conditional flow matching: sample a random intermediate time $t$, an initial state $x_0$, and a Brownian path $B$. Generate the SDE solution $x_T$ by numerically simulating the SDE with that specific path (or use a known closed form). Then train a neural network $f_\theta(x_t, B_{[t,T]})$ to predict $x_T$, using a simple MSE loss. At inference, given any $x_t$ and a newly sampled Brownian path, $f_\theta$ outputs the corresponding endpoint in one shot.

The neural network architecture typically consists of two encoders: one for the state $x_t$ (e.g., a U-Net backbone) and one for the Brownian path (e.g., a transformer that processes discretized path values). Their outputs are fused, then passed to a prediction head that outputs $x_T$.

Input: (x_t, B sampled path)
        |                |
        v                v
   [State Encoder]  [Path Encoder]
        |                |
        +-----[Fusion]----+
                |
          [Prediction Head]
                |
                v
         Output: x_T (predicted endpoint)

Structural metaphor: Think of the Itô map as a recipe transformer. You have a partially cooked dish (state $x_t$) and a detailed step-by-step recipe (Brownian path $B$). The original recipe requires you to follow each step sequentially (iterative SDE solve). The Itô map is a master chef who, given the current state and the full recipe, can directly plate the final dish in one go. The recipe encodes all randomness (e.g., “add a pinch of salt at step 3” is fixed), so the chef only needs to learn the mapping from (current dish, recipe) → final dish. Different chefs (network weights) can be trained on many (dish, recipe, final) triplets.

Key Concepts

• Stochastic Flow Map: A function that maps an initial or intermediate state (and possibly external randomness) to a final state, preserving the distribution of the underlying SDE. Unlike deterministic flow maps (which collapse noise), stochastic flow maps retain the randomness from the Brownian path, allowing them to generate diverse but conditionally valid samples.

• Path Conditioning: The method conditions on the *entire Brownian path, not just the initial noise. This transforms the SDE into a deterministic mapping from path to endpoint. In practice, paths are discretized into a sequence of Brownian increments; a transformer or RNN encodes this sequence into a fixed-size embedding that is injected into the network.

• Conditional Validity: For a given intermediate state $x_t$, the Itô map’s outputs, when marginalized over the Brownian path, should match the true conditional distribution $p(x_T | x_t)$. The paper verifies this by comparing histograms or using statistical tests (e.g., energy distance) between generated and ground-truth conditional samples.

Framework Shift

The core shift is from iterative local simulation to global single-pass prediction conditioned on path randomness.

Before (mainstream approach):         After (this paper):
  x_0                                    x_0 + sampled path B
   |                                        |
   v                                        v
  [SDE Step 1]                          [Itô map network]
   |                                        |
   v                                        v
  x_1                                      x_T (predicted endpoint)
   |
   v
  [SDE Step 2]
   |
   v
  ... (many steps)
   |
   v
  x_T (sampled after N steps)

One sentence: From multi-step iterative SDE simulation to single-pass conditional prediction, the core shift is replacing temporal iteration with function approximation of the entire flow.

Expert Assessment

Problem choice: Real gap. ODE-based one-step models are popular but ignore stochasticity; this paper directly addresses the missing piece for SDEs. However, the question “do we need exact SDE integration?” is open — many applications use deterministic flows and get away with it.

Method maturity: A clever combination of conditional flow matching and path conditioning, but the idea of conditioning on the Brownian path is not entirely new (used in diffusion bridge literature). The main novelty is framing it as a general-purpose “any-step” primitive. The training requires generating full Brownian paths, which could be costly in high dimensions.

Experimental integrity: The paper provides conditional validity tests (histograms, energy distance) and steering results on synthetic data and CIFAR-10. Baselines are reasonable (multi-step SDE, one-step ODE). Minor red flag: the steering experiments only show qualitative outputs; quantitative comparisons to other stochastic control methods are missing.

Writing quality: The abstract and introduction are sharp. However, section 3 (training details) feels rushed — the network architecture for path encoding is described only vaguely. A more detailed ablation (e.g., transformer vs RNN for path embedding) would strengthen the paper.

Verdict: weak accept — The idea of an any-step SDE map is useful and the method is sound, but the paper lacks thorough comparison to simpler alternatives (e.g., just adding noise to ODE flows) and needs clearer architectural guidance.

Takeaways

• Conditional flow matching on noise paths: This technique can be pulled out and applied to any stochastic dynamical system where you want to jump from an intermediate state to the endpoint. For example, in reinforcement learning with stochastic environments, you could learn an Itô-like map to predict terminal rewards given partial trajectories.

• Path encoding as a first-class input: Treating the noise trajectory as a learnable embedding (instead of using a fixed noise seed) unlocks the ability to do posterior inference and control. Practitioners in diffusion models can borrow this idea to better handle stochasticity in their pipelines.

• Single-pass posterior sampling: If you need to sample from $p(x_T | x_t)$ for many different $t$ or many different observed states, the Itô map amortizes the computation. This is immediately useful for inverse problems (e.g., image restoration) where you condition on partial observations.