🤖 machine learning

First-Order Trajectory Matching: Fast Ensemble Predictions of Chaotic, Turbulent, Stochastic Systems

Original authors: Shreya Jha, Timo Schorlepp, Nicholas Geissler, Jules Berman, Benjamin Peherstorfer

Published 2026-06-10

📖 5 min read🧠 Deep dive

Original authors: Shreya Jha, Timo Schorlepp, Nicholas Geissler, Jules Berman, 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 predict the weather. You have a supercomputer that can simulate the atmosphere perfectly, but running it takes days. You need a prediction now for a flight plan. So, you build a "surrogate" model—a fast, simplified version of the weather system.

The problem is that weather is chaotic. If you just ask a standard AI, "What will the temperature be in an hour?" it will give you the average answer. It might say "70 degrees." But in reality, the temperature could be 60 or 80. If you run this "average" prediction 100 times, you get the same boring 70 degrees every time. You lose the variety of possible outcomes.

This paper introduces a new method called First-Order Trajectory Matching (FTM) to solve this. Here is how it works, using simple analogies.

The Problem: The "Average" Trap

Most current AI models for chaotic systems (like turbulence or financial markets) act like a lazy tour guide.

The Old Way: If you ask the guide, "Where will the crowd be in 10 minutes?" they point to the center of the crowd. They ignore the fact that some people are running left and others right. They just give you the "average" location.
The Result: If you try to simulate a whole crowd moving, everyone ends up in the exact same spot. The simulation collapses into a single point, losing all the excitement and movement of the real world.

Other methods try to fix this by matching the "shape" of the crowd at the end of the day (the marginals).

The "Shape" Fix: This is like a traffic controller who ensures the right number of cars are in the right neighborhoods at 5:00 PM. But they don't care how the cars got there. The cars might be driving in circles, or they might be stuck in a traffic jam that never moves. The controller only sees the final map, not the movement.

The Solution: FTM (The "Flow" Detective)

FTM is different. Instead of just asking "Where will the crowd be?" or "What is the average path?", FTM asks: "How is the crowd actually moving right now?"

Think of a river.

Old Models might tell you where the water is at 2:00 PM and 2:01 PM, but they might miss the fact that the water is swirling in a whirlpool.
FTM looks at the water from two angles:
1. Forward: "If I drop a leaf here, where does it go next?"
2. Backward: "If I see a leaf here, where did it come from?"

By combining these two views, FTM calculates the symmetric motion. It learns the "current velocity"—the actual speed and direction the water is flowing, including the swirls and eddies.

How It Works (The Magic Trick)

Usually, to learn how a chaotic system moves, you have to do a lot of heavy math to figure out the "drift" (the push) and the "diffusion" (the random jitter). This is slow and hard.

FTM uses a clever shortcut (a mathematical identity called the Stratonovich identity).

The Analogy: Imagine you want to know how fast a car is driving. Instead of measuring the engine (drift) and the bumpy road (diffusion) separately, FTM just looks at the car's position a split second ago and a split second ahead. It calculates the speed directly from the movement itself.
The Benefit: It doesn't need to guess the hidden rules of the system. It just learns the "flow" directly from the history of the trajectories.

Why It's Fast and Accurate

The paper claims FTM is a "sweet spot" between speed and accuracy:

It's Fast: It acts like a deterministic model. Once trained, you can run it 1,000 times in a second. You don't need to run complex simulations for every single step. It takes one step to predict the next moment.
It's Accurate: Unlike the "lazy tour guide" (which collapses to the average), FTM keeps the swirls and currents alive. It predicts that the crowd is moving in a circle, not just sitting in a pile.
It's Stable: The authors proved mathematically that if you have enough data points and the time steps aren't too tiny, the method stays stable. It doesn't blow up or get confused.

What It Can and Cannot Do

The paper is very specific about what this tool is good for:

It IS good for: Predicting the "flow" of things. If you want to know how much "stuff" is crossing a bridge, or how energy circulates in a storm, FTM is excellent. It captures the currents and circulations.
It is NOT good for: Predicting the exact path of a single, random particle if that path depends on pure luck (like a "hitting time" or when a specific coin flip happens). It predicts the ensemble (the group behavior) perfectly, but it simplifies the random jitter of a single individual.

Summary

FTM is like a smart traffic camera system.

Old AI: "Everyone is at the intersection." (Boring, wrong).
Old "Shape" AI: "There are 100 cars at the intersection, but we don't know if they are moving or parked." (Incomplete).
FTM: "There are 100 cars, and they are all swirling in a roundabout at 30 mph." (Accurate, fast, and captures the movement).

It allows scientists to run fast simulations of chaotic systems (like weather or turbulence) that don't collapse into a boring average, but still move fast enough to be useful for real-time decisions.

Technical Summary: First-Order Trajectory Matching (FTM)

1. Problem Statement

Many scientific dynamical systems are effectively stochastic due to chaos, coarse-graining, partial observation, or unresolved degrees of freedom. While high-fidelity numerical simulators exist, repeated traditional simulation is computationally prohibitive for outer-loop tasks such as design, control, and uncertainty quantification. This necessitates fast surrogate models trained on trajectory data.

Existing approaches face specific limitations in this context:

Operator Learning: Typically trained with mean-squared error, these models predict the conditional mean. This collapses the distribution of possible future states into a single point, failing to capture the transport of probability mass required for ensemble statistics.
Population Dynamics (Marginal Matching): Methods that match time marginals ensure the correct probability mass exists in the correct regions at each time step. However, they do not specify how mass flows over time. Consequently, they miss critical trajectory-dependent phenomena such as circulations, fluxes, and barrier-crossing currents in statistically stationary nonequilibrium systems.
Conditional Generative Models (Diffusion/Flow): These models capture full conditional transition laws and stochastic variability. However, they typically require multiple neural network evaluations or solver steps per physical time step to generate the next state, leading to high inference costs during autoregressive rollouts.

The core challenge is to develop a surrogate model that rapidly propagates ensembles, preserves time marginals, captures first-order trajectory transport (currents), and maintains low deterministic rollout costs (ideally one neural network evaluation per time step).

2. Methodology: First-Order Trajectory Matching (FTM)

FTM addresses these challenges by learning the probability current velocity directly from stochastic trajectories, rather than learning drift, diffusion, or score functions separately.

2.1 Theoretical Foundation

FTM seeks a deterministic velocity field $u(t, x)$ such that the induced Ordinary Differential Equation (ODE) $\dot{\hat{x}}(t) = u(t, \hat{x}(t))$ preserves the time marginals of the underlying Stochastic Differential Equation (SDE) and matches the local transport of probability mass.

The method relies on the probability current velocity $v(t, x)$ , defined as:
$v(t, x) = b(t, x) - \frac{1}{2}\Sigma(t)\nabla \log \rho(t, x)$
where $b$ is the drift, $\Sigma$ is the diffusion matrix, and $\rho$ is the probability density. This velocity satisfies two key properties:

Global Consistency: Its flow preserves the time marginals $\rho(t)$ (solving the Fokker-Planck equation).
Local Consistency: It represents the symmetric first-order motion of trajectories. By combining forward increments (where mass goes) and backward increments (where mass came from), FTM captures the net flux of mass through a space-time point.

2.2 Scalable Loss Function

Directly regressing $v$ is impossible because the target velocity is unknown in the data. FTM derives a simulation-free, pathwise loss function using the Stratonovich identity:
$E\left[\int_0^T \langle v_\theta(t, X(t)), v(t, X(t)) \rangle dt\right] = E\left[\int_0^T v_\theta(t, X(t)) \circ dX(t)\right]$
This identity allows the linear term in the regression objective to be estimated using only observed trajectory increments and the candidate velocity field $v_\theta$ , without needing to know the true current velocity or estimate the score function.

2.3 One-Step Loss and Stability Analysis

To make the loss scalable for long trajectories, the authors introduce a chunked loss approach. They analyze the trade-off between the time-step size $h$ (discretization error) and the chunk length $\tau$ (variance reduction).

Stability Bound: Theoretical analysis (Proposition 1) shows that the variance of the loss is governed by the product of the sample size $N$ and the time-step $h$ .
One-Step FTM: When $N$ and $h$ are balanced, the authors propose a one-step loss where $\tau \approx h$ . This loss depends only on adjacent states ( $X(t-h), X(t), X(t+h)$ ) and requires a single neural network evaluation per data point.
Result: This approach avoids the computational cost of full-trajectory averaging while maintaining a stable training signal, provided the time resolution and sample size are sufficiently balanced.

3. Key Contributions

Formulation: The paper formulates ensemble prediction as learning local probability mass transport (the probability current velocity) rather than conditional means or marginal-matching flows. This identifies the current velocity as the deterministic surrogate that preserves both marginals and first-order trajectory transport.
Algorithm: It derives a scalable, one-step, simulation-free loss function that learns the current velocity directly from trajectories. This bypasses the difficult and often unreliable separate estimation of drift, diffusion, and score functions.
Theoretical Analysis: The authors provide a stability analysis separating discretization error from sampling variance, justifying the use of the simple one-step loss when time resolution and sample sizes are balanced.
Empirical Performance: Across stochastic dynamical systems and Partial Differential Equation (PDE) examples, FTM demonstrates:
- Accurate matching of ensemble statistics (time marginals).
- Superior capture of trajectory-based quantities of interest (QoIs) such as fluxes, circulations, and barrier-crossing currents compared to marginal-only methods.
- Low Inference Cost: It achieves these results with only one neural network evaluation (NFE) per physical time step, avoiding the collapse to the mean seen in operator learning and the high inference costs of autoregressive generative models.

4. Experimental Results

The authors evaluated FTM on four examples:

Duffing Oscillator: A double-well potential system with Brownian forcing.
Chaotic Rayleigh–Bénard Convection: A stochastically forced nine-mode model.
Stochastically Forced Burgers: A simplified turbulence model.
Stochastically Forced Turbulence: 2D incompressible Navier–Stokes equations.

Key Findings:

Distribution Accuracy: FTM matches the evolving ensemble distribution (measured by Wasserstein distance) competitively with marginal-only flows and explicit marginal-matching methods (e.g., DICE).
Trajectory QoIs: FTM significantly outperforms marginal-only methods in capturing trajectory-based QoIs. For the Duffing oscillator, it minimized barrier-crossing current errors; for Rayleigh–Bénard, it minimized rotational current errors.
Efficiency vs. Accuracy: In turbulence examples, FTM matched energy and enstrophy statistics with one NFE per time step. In contrast, autoregressive diffusion models (ARDM) required up to 100 NFEs per time step to achieve similar accuracy. Distilled conditional flow matching models (AR-CFM) offered lower costs but resulted in errors an order of magnitude higher than FTM.
Robustness: The one-step loss was shown to be stable across 100 random neural network initializations.

5. Significance and Limitations

Significance:
The paper claims that FTM offers a practical route to fast ensemble predictions for stochastic systems. By learning the probability current velocity, it bridges the gap between the low cost of deterministic operator learning and the high fidelity of generative models. It allows for "trajectory-aware" predictions (capturing fluxes and circulations) at the computational cost of a deterministic rollout, avoiding the "collapse to the mean" inherent in standard operator learning.

Limitations:

Scope of Statistics: FTM is designed to preserve time marginals and first-order trajectory transport. It is not intended to reproduce martingale-dominated path statistics, such as hitting-time distributions or temporal autocorrelations.
Theoretical Scope: The theoretical guarantees are established for finite-dimensional additive-noise SDEs. While the method is empirically applied to high-dimensional PDEs (Burgers, Navier-Stokes), the theoretical analysis does not fully extend to these infinite-dimensional settings.
Trade-off: The method trades off the ability to model full stochastic path laws for the benefit of deterministic rollout efficiency. It is most effective when the quantities of interest are ensemble statistics or current-like transport observables.