🤖 machine learning

Stochastic Lifting for Generating Trajectories of Stochastic Physical Systems

This paper introduces Stochastic Lifting, a method that models smooth stochastic physical systems by augmenting state transitions with independent high-dimensional random labels to learn a deterministic map capable of generating diverse, non-collapsing trajectories through autoregressive inference.

Original authors: Jules Berman, Tobias Blickhan, Benjamin Peherstorfer

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jules Berman, Tobias Blickhan, Benjamin Peherstorfer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to predict the future of a chaotic system, like a swirling storm, a wave crashing on a beach, or a robot arm pushing a box. The problem is that these systems are stochastic, meaning they are full of randomness. If you show the robot the exact same starting point twice, the outcome might be slightly different each time because of tiny, unpredictable factors (like a gust of wind or a microscopic vibration).

The Problem: The Robot's "Average" Trap

Most standard AI models try to learn by looking at past examples and finding the "best guess." If you show a robot a ball rolling down a hill, and sometimes it rolls left and sometimes it rolls right, a standard model will get confused. To minimize its error, it will simply predict the average: a ball rolling straight down the middle.

But in reality, balls don't roll straight down the middle; they roll left or right. By predicting the average, the robot creates a blurry, fake reality that never actually happens. It "collapses" into a boring, mean prediction and loses all the exciting, diverse possibilities.

Other methods (like complex diffusion models) can generate diverse outcomes, but they are slow. To predict just one second of the future, they might need to take 40 tiny, slow steps, like a snail inching forward.

The Solution: Stochastic Lifting

The authors propose a clever trick called Stochastic Lifting. Think of it as giving the robot a "secret cheat sheet" or a random label for every single move it learns.

Here is how it works, using a simple analogy:

The Analogy: The Dance Class
Imagine a dance instructor trying to teach students a routine.

The Problem: If the instructor says, "When the music hits this beat, you jump," but sometimes the students jump left and sometimes right, a strict robot might just tell everyone to "jump straight up" (the average).
The Stochastic Lifting Fix: The instructor gives every student a unique, random colored sticker (a "label") before the music starts.
- Student A (Red Sticker): "When the music hits, jump Left."
- Student B (Blue Sticker): "When the music hits, jump Right."
- Student C (Green Sticker): "When the music hits, jump Forward."

Now, the instructor isn't asking, "What is the average jump?" Instead, they are learning a rule: "If you have a Red Sticker, jump Left. If you have a Blue Sticker, jump Right."

Because every student has a unique sticker, the robot can learn a smooth, clear rule for every specific outcome without getting confused. It doesn't have to guess the average anymore; it just needs to know which "sticker" (random label) goes with which move.

How It Works in Practice

Training: The AI looks at data (like a video of a wave). For every single frame transition (from time $t$ to $t+1$ ), it attaches a random, high-dimensional "label" (like a unique ID code). It then learns a map: "Current State + This Specific Label = Next State."
The Magic of "Lifting": By adding these extra labels, the AI moves the problem into a higher dimension. This makes it mathematically easier to draw a smooth line connecting all the dots without the lines getting messy or crossing over each other.
Inference (The Future): When the AI wants to predict the future, it takes the current state, picks a brand new random label (a new sticker), and runs it through the map.
- Pick a new random label $\rightarrow$ Get a plausible "Left" wave.
- Pick a different random label $\rightarrow$ Get a plausible "Right" wave.

Why It's a Big Deal

Speed: Unlike the "snail" methods that take 40 steps to predict one second, this method does it in one single step. It's like snapping your fingers instead of inching forward.
Diversity: It generates many different, realistic possibilities (trajectories) instead of just one blurry average.
Physics-First: The authors tested this on real physical systems (waves in random media, oil flowing through porous rock, and even robot videos). They found it was more accurate and faster than existing state-of-the-art methods for these specific types of problems.

The Catch (Limitations)

The paper is very clear about where this trick works and where it doesn't:

It needs a "Next Step": This method relies on the fact that the "next state" is closely related to the "current state" (like a wave moving a tiny bit forward). It works great for predicting the next frame of a video or the next second of a storm.
It doesn't work for "Magic from Nothing": You cannot use this to generate an image from pure noise (like creating a picture of a cat out of static). Because there is no "current state" to build upon, the smooth connection the method relies on breaks down.

In short, Stochastic Lifting is a way to teach AI to embrace randomness by giving it a unique "key" for every possible outcome, allowing it to predict diverse futures in a single, lightning-fast step.

Technical Summary: Stochastic Lifting for Generating Trajectories of Stochastic Physical Systems

1. Problem Formulation

The paper addresses the challenge of generating diverse, plausible trajectories from stochastic physical systems. In many computational physics problems, the system state $X_t$ evolves according to a conditional transition law $X_{t+1} \sim \rho(\cdot | X_t = x_t)$ . The goal is not point prediction (estimating the mean next state) but rather sampling from this conditional distribution to generate full trajectory ensembles.

Standard regression approaches fail here because fitting a map $f: x_t \to x_{t+1}$ using mean-squared error (MSE) collapses the solution to the conditional expectation $E[X_{t+1} | X_t = x_t]$ , thereby suppressing the intrinsic variability of the stochastic process. Conversely, general-purpose generative models (e.g., diffusion or flow-based models) can represent distributions but typically require multi-step sampling at inference. Generating a trajectory of length $T$ with an $s$ -step sampler requires $T \times s$ neural network evaluations, which is computationally expensive. While recent "one-step" samplers exist, they often struggle with training stability or are designed for static generation (noise-to-image) rather than the time-dependent regularity found in trajectory data.

2. Methodology: Stochastic Lifting

The authors propose Stochastic Lifting, a method that exploits the smoothness premise of physical systems: for sufficiently small time steps, the transition law varies regularly with the state, and next states can be viewed as small perturbations of the current state.

Core Mechanism

Stochastic Labeling: The training data pairs $(x_t, x_{t+1})$ are augmented with an independently drawn, high-dimensional auxiliary stochastic label $\xi_t \in \mathbb{R}^d$ (typically sampled from a standard normal distribution $\mathcal{N}(0, I_d)$ ). This creates an augmented dataset $D_\xi = \{(x_t, x_{t+1}, \xi_t)\}$ .
Lifting the Map: Instead of learning a map $f: x_t \to x_{t+1}$ $f : x_{t} \to x_{t + 1}$ , the method learns a map $F: (x_t, \xi_t) \to x_{t+1}$ $F : (x_{t}, ξ_{t}) \to x_{t + 1}$ .
- The label $\xi_t$ acts as an auxiliary coordinate that parametrizes the stochasticity.
- By assigning a unique random label to each realized transition, the map $F$ can interpolate multiple possible next states for similar current states without collapsing to the mean.
Training: The map $F_\theta$ (parameterized by a neural network) is trained using a standard empirical regression loss (e.g., $L_2$ or LPIPS) on the labeled pairs:
$\mathcal{L}(\theta) = \frac{1}{MT} \sum_{i,t} \| F_\theta(x_t^i, \xi_t^i) - x_{t+1}^i \|_2^2$
Inference: To generate a trajectory, the model starts with an initial state $x_0$ . At each time step $t$ , a fresh label $\xi_t$ is sampled from the label distribution $\nu$ , and the next state is computed via a single network evaluation: $x_{t+1} = F_\theta(x_t, \xi_t)$ . This process is rolled out autoregressively.

Theoretical Underpinnings

The paper provides finite-sample theoretical analysis based on two key insights:

Smoothness and Lipschitz Regularity: If the learned map $F$ is smooth (has a low Lipschitz constant) and interpolates the training data, the generated samples are close to the true conditional distribution in the Wasserstein-2 metric.
Role of Label Dimension: High-dimensional labels ( $d$ ) increase the separability of training points in the lifted space. Theoretically, increasing $d$ reduces the minimal Lipschitz constant required for any interpolant of the data. This allows the network to fit the data (interpolation) while maintaining the smoothness necessary for generalization to fresh labels.
Finite-Sample Regime: The method is explicitly designed for the finite-sample regime. In the infinite-data limit ( $M \to \infty$ ), the labels become independent of the states, and the minimizer would collapse to the conditional expectation. However, with finite data and high-dimensional labels, the network can achieve zero training error (interpolation) while remaining smooth, thus avoiding the collapse.

3. Key Contributions

Introduction of Stochastic Lifting: A novel framework that assigns stochastic labels to state transitions to learn a one-step generator. This enables trajectory generation with a single network evaluation per time step.
Finite-Sample Theory: The authors derive Wasserstein bounds demonstrating that interpolation combined with Lipschitz regularity controls sampling error. They show that higher-dimensional labels improve the finite-sample geometry, making low-Lipschitz interpolation possible.
Empirical State-of-the-Art: The method is demonstrated to outperform existing approaches on stochastic, time-dependent physical systems and video generation tasks.

4. Experimental Results

The paper evaluates Stochastic Lifting on three domains:

A. Stochastic Physical Systems

Duffing Oscillator: A simple SDE example used to validate the smoothness premise. Shuffling training pairs (breaking the $x_t \to x_{t+1}$ coupling) degraded performance, confirming the method relies on the smooth transition structure. Increasing label dimension reduced training loss and gradient norms, confirming the promotion of smoothness.
Traveling Wave in Random Media: A 2D wave equation with random wave speed. The initial condition is deterministic, but the medium introduces stochasticity.
Two-Phase Flow in Porous Media: An incompressible flow problem with random permeability.
Comparison: In both physics problems, Stochastic Lifting (1 step) matched or exceeded the accuracy of Autoregressive Diffusion Models (ARDM) that required 40 diffusion steps. It accurately captured non-Gaussian distributions (e.g., heavy tails in crossing times) where ARDMs failed even with many steps.

B. Video Generation

BAIR Robot Pushing: A standard benchmark for video prediction. Stochastic Lifting achieved the lowest Fréchet Video Distance (FVD) among all one-step methods and was competitive with multi-step diffusion models that use significantly more compute.
CLEVRER: A synthetic dataset for reasoning about collisions. The model was trained on 16-frame videos but successfully rolled out for 500 frames, generating stable interactions and object accumulations not seen in training.

5. Significance and Limitations

Significance:
The paper argues that Stochastic Lifting offers a computationally efficient alternative to multi-step generative models for trajectory data. By leveraging the structural regularity of time-dependent systems (where $x_{t+1}$ is a small perturbation of $x_t$ ), it achieves high-quality sampling with $O(1)$ network evaluations per time step. The theoretical framework provides actionable guidance: increasing label dimension is a practical lever to improve the smoothness and separability of the learned map.

Limitations (as stated by the authors):

Coupling Requirement: The method critically depends on the strong coupling between current and next states. It fails for unconditional generation (e.g., generating images from pure noise) where the transition structure is absent.
Label Dimension vs. Data: While high dimensions guarantee the existence of a linear interpolant, they may require proportionally more data to train accurately.
Finite-Sample Focus: The theory applies to finite-sample regimes and time marginals. A meaningful population limit for Stochastic Lifting remains an open question.
Manifold Assumption: The quality of generated samples relies on the network learning a meaningful embedding where new labels map to the data manifold; the conditions under which this fails are not fully characterized.