One-step generative models produce stochastic trajectories efficiently

One-Step Generative Surrogate Models via Block-Triangular Joint Drifting

Machine Learning

Summary

Modeling random changes from one state to the next usually needs many examples of possible next steps from the same starting point, but typical data only shows one outcome. The authors propose a new method that transforms the entire joint data of consecutive states to generate realistic next steps in just one step. This approach keeps the original state's distribution intact while allowing direct sampling of possible future states. Their experiments show this method can accurately capture the behavior of stochastic systems with less computational cost than other popular techniques.

What this means in practice

  • For simulation engineers: Generate realistic trajectories of stochastic systems in simulations with efficient one-step sampling based on typical trajectory data.
  • For machine learning engineers: Build surrogate models that produce stochastic transitions faster and more accurately without multiple sampling steps.

Authors

Nicholas Geissler, Shreya Jha, Ricardo Baptista, Benjamin Peherstorfer

Abstract

Drifting provides a direct route to one-step generative models, but applying it directly to stochastic transition modeling requires multiple samples of the next state conditioned on the same current state. Standard trajectory data, however, typically provide only one realized next state for each observed current state and therefore do not provide an empirical approximation of the corresponding conditional distribution over possible next states. We introduce block-triangular joint drifting, which instead applies a projected drift field to the empirically accessible joint distribution of consecutive states. Importantly, the block-triangular architecture preserves the current-state marginal while making its second component a direct sampler of the conditional distribution of possible next states. The resulting surrogate generates stochastic trajectories with one model evaluation per time step, without auxiliary generative steps between time steps. Numerical experiments demonstrate accurate marginal and trajectory-dependent statistics and favorable accuracy-cost tradeoffs compared with deterministic, diffusion-, flow-, and distillation-based generative surrogate models.