Itô maps for any-step SDEs

The field of deep generative modeling has experienced rapid growth, driven by models such as Variational Autoencoders, Generative Adversarial Networks, and more recently, diffusion and score-based models. Diffusion models, exemplified by DDPM (Denoising Diffusion Probabilistic Models) and score-based models like Song et al.'s work, have achieved state-of-the-art results in image synthesis, leveraging iterative denoising processes that approximate complex data distributions. These models typically rely on discretized SDEs or ODEs to define the evolution from noise to data, with the sampling process involving multiple sequential steps. While effective, this approach is computationally intensive, limiting real-time applications.

Recent advances have focused on distilling multi-step samplers into fewer steps or even single-step mappings, such as flow-matching and flow-map distillation techniques. These methods aim to compress the continuous-time dynamics into direct maps, often deterministic, which can generate samples rapidly but at the expense of losing stochasticity and the ability to model residual uncertainty. Score-based models inherently involve stochasticity through their SDE formulation, but most existing one-step models approximate the stochastic transition as a deterministic map, thus failing to capture the full posterior distribution conditioned on intermediate states.

The challenge remains in modeling the residual uncertainty and diversity in the generative process, especially for applications like posterior sampling, stochastic control, and conditional generation. Existing methods struggle to represent the full conditional law p_{1|t}(· | xt), which encodes the distribution of the endpoint conditioned on an intermediate noisy state. This gap motivates the development of a stochastic, pathwise approach that can accurately model the evolution of data along the entire stochastic trajectory, supporting both diversity and control.

The core problem addressed in this work is the inability of existing one-step generative models to accurately represent the full conditional distribution of the data endpoint given an intermediate noisy state. Deterministic flow maps, while fast, do not encode the residual uncertainty, limiting their use in posterior sampling and stochastic control. Conversely, stochastic models based on SDEs retain uncertainty but are computationally expensive and difficult to implement efficiently at arbitrary steps. The key challenge is to develop a method that can perform accurate, efficient, and flexible path prediction conditioned on partial information, while preserving the stochastic nature of the underlying dynamics. This is particularly important for applications requiring diverse samples, precise control, and real-time inference, such as image editing, data augmentation, and autonomous decision-making. The problem is compounded by the high dimensionality of data like images and videos, which makes path modeling computationally challenging. Therefore, a scalable, low-dimensional representation of the stochastic paths, combined with a theoretically sound predictive model, is essential to address these issues.

The paper introduces the Itô map, a novel stochastic flow map that predicts the future state of an SDE conditioned on both the current state and a low-dimensional Brownian path representation. Unlike traditional deterministic flow models, the Itô map explicitly models the residual stochasticity by conditioning on the Brownian path, enabling exact sampling of the conditional law p_{1|t}(· | xt). The key innovations include:

• �� Low-dimensional Brownian path features: Using Karhunen–Loève expansion and Haar wavelets to compress the infinite-dimensional Brownian motion into manageable feature vectors (e.g., 5 modes), facilitating scalable learning.
• �� Pathwise learning via self-distillation: Training the model to enforce path consistency across multiple time points, ensuring the learned stochastic flow respects the SDE dynamics.
• �� Control and inference: Deriving Bayesian estimators (BEL series) and gradient-based control estimators (Itô-G, Itô-GF) that leverage the path structure for reward-guided sampling.
• �� Theoretical guarantees: Ensuring the learned maps are pathwise consistent and respect the stochastic calculus foundations, enabling accurate, diverse, and controllable sampling.

This approach fundamentally shifts the paradigm from deterministic transport to stochastic pathwise modeling, supporting both high-quality sample generation and flexible control.

• �� Path feature extraction: Use Karhunen–Loève expansion to represent Brownian paths with a small set of Gaussian coefficients (e.g., 5 modes). Additionally, employ Haar wavelet decomposition to capture multi-scale path fluctuations, forming structured, finite-dimensional features.
• �� Itô map definition: Construct Φs,t as a function taking current state xt and Brownian feature vector W, predicting the future state xu along the same path.
• �� Training objectives: Minimize path consistency loss (LSD) by matching the predicted path derivatives with the true SDE drift and diffusion terms, and enforce pathwise evolution via LPSD, which ensures the learned map respects the semigroup property.
• �� Pathwise training: Sample pairs of times (s, t), generate intermediate states, simulate Brownian paths, and optimize the model to produce path-consistent predictions across multiple time points.
• �� Control estimation: Derive Bayesian estimators (BEL series) based on pathwise Jacobians and Brownian increments, enabling reward-guided sampling without explicit reward gradient dependence (Itô-GF). When gradients are available, use Itô-G for more accurate control.
• �� Implementation details: Use neural networks to parameterize G_{s,t} and G_{t,t}, train with stochastic gradient descent, and incorporate low-dimensional Brownian features for scalability. Regularize to avoid singularities at t=1 by setting σt = √2(1−t).

The authors conduct comprehensive experiments across synthetic and real datasets. In 1D Gaussian mixtures, they compare true trajectories with predictions, achieving errors around 6.92×10^-2. In 2D posterior sampling, the model outperforms baselines like MFM-G and DPS, with SW2 distances of 0.16 and MMD of 0.024, demonstrating high-quality conditional samples. On MNIST, the model accurately reproduces image paths with pixel MSE of 0.05, validating path prediction in high dimensions. For control tasks, the model successfully modulates class proportions, with KL divergence below 0.025, indicating effective reward-guided sampling. The experiments also include ablation studies on the number of Brownian modes, path feature types, and training strategies, confirming the robustness and scalability of the approach.

Quantitative results show that the proposed Itô map achieves a path prediction RMSE of 6.92×10^-2 in synthetic tasks, significantly outperforming deterministic flow models. In posterior sampling, SW2 and MMD metrics demonstrate clear advantages over baseline methods, with SW2 reduced from 1.95 (unsteered) to 0.16 (Itô-G). In image generation, pixel MSE drops to 0.05, confirming accurate path modeling. The control experiments show the model can steer class proportions with KL divergence as low as 0.024, outperforming existing control estimators. These results collectively validate the framework’s effectiveness in diverse tasks, highlighting its potential for scalable, flexible, and high-quality stochastic sampling and control.

This framework can be directly applied to high-dimensional image synthesis, posterior inference, and reward-guided generation. Its ability to perform exact conditional sampling from intermediate states makes it suitable for tasks like image editing, data augmentation, and personalized content creation. The low-dimensional Brownian feature extraction allows for scalable training and inference, even in large datasets like ImageNet. In the long term, integrating this approach with reinforcement learning and autonomous decision systems could enable real-time path planning and control in robotics, autonomous vehicles, and complex simulation environments. Its flexibility also supports multi-modal data modeling, such as video and 3D point clouds, broadening its impact across AI domains.

Despite promising results, the current model’s scalability to extremely high-resolution data remains computationally demanding, especially in extracting and optimizing low-dimensional path features. The fixed number of Brownian modes (e.g., 5) may limit expressiveness for highly complex or non-Gaussian paths, requiring adaptive feature selection. Additionally, robustness under severe noise or non-Gaussian dynamics has not been fully explored, and performance may degrade in such scenarios. Further research is needed to improve the efficiency of training in large-scale settings, extend the theoretical guarantees to more general stochastic processes, and develop adaptive path feature mechanisms to handle diverse data distributions.