Learning to Emulate Chaos: Adversarial Optimal Transport Regularization

In plain words

A learned stand-in for a chaotic system, such as weather or a turbulent flow, drifts away from the real system. Small errors grow until the predicted path bears no resemblance to the true one. The long-run behaviour of such a system still has a stable shape: over a long time it visits some states often and others rarely. A good stand-in should reproduce that pattern even when it cannot follow the exact path.

The paper trains the stand-in with two goals at once. The first is the usual one: predict the next state correctly. The second is a penalty that compares how the stand-in's states are spread out with how the real states are spread out, measured through a few learned features of each state. A second network learns those features and is trained to make the comparison as sensitive as possible. Because the comparison works on the overall pattern, not on single paths, one noisy recording of the real system can be enough.

The method is tested on simulated systems, including ones with high-dimensional chaos. The gains the authors report are in the long-run pattern, not in short-term accuracy.

TL;DR

Learned emulators of chaotic systems drift from the true trajectory, so a per-step error objective is a poor guide for long rollouts §1. The paper adds an optimal-transport term between the distributions of learned summary features of true and emulated states, and learns those features adversarially §4. Author-reported: long-run histogram and spectral metrics improve on Lorenz-96, while the no-OT baseline has the lowest or tied-lowest one-step RMSE in every block of the results table tab. 2. The MIT code on the default branch implements only the Lorenz-96 path, and several settings differ from the paper (see Paper vs code).

Problem

Training a one-step emulator on noisy chaotic trajectories with a squared-error loss is said to fail over long rollouts, because trajectory errors grow exponentially with the horizon §1 §5.2.1. Approaches that match long-run statistics need either handcrafted summary statistics that an expert must pick, or a diverse multi-environment dataset §3.

Setting §2.1: a compact state space with a true one-step map Φ that has an invariant, ergodic measure μ. The observed data is one stationary trajectory from μ. The emulator g induces the measure ĝ = g#μ.

Assumptions the analysis needs §5.1 §5.2.2: the summary maps are Lipschitz, and the emulator is exponentially mixing in W_p for the noise-forgetting result. The card found no test of the mixing assumption on trained emulators.

Method

Objective §4. The emulator g and the summary map f: U → S are trained together. The emulator minimises a one-step error plus λ times a transport cost between the distributions of summaries of the true and emulated states. The summary map maximises the same cost:

\min_g \max_f \; L_{\mathrm{MSE}}(u_{t+1}, g(u_t)) + \lambda \, W_{d_S^p}(f_\#\mu, f_\#\hat\mu) eq. 4

The cost uses the ground metric d_S^p for some p ≥ 1, and λ > 0 §4. The authors call the one-step term MSE. The released code does not use a squared error for it (see Paper vs code).

Transport term §2.2 eq. 2 eq. 3. The Monge–Kantorovich problem gives the cost. Its entropic version is solved with Sinkhorn iterations eq. 2. The Sinkhorn divergence removes the entropic bias: it subtracts half of each self-transport term from the cross term eq. 3.

WGAN-style variant (p = 1) §4.1 eq. 5 eq. 6. The 1-Wasserstein term is replaced by its Kantorovich–Rubinstein dual, with a 1-Lipschitz critic φ. The expectations are taken under the invariant measure μ, not over i.i.d. samples §4.1. The critic is kept Lipschitz by weight clipping, gradient penalty or spectral normalisation §4.1; the released code uses weight clipping app. D.1.

Sinkhorn variant §4.2 eq. 7 eq. 8. A learned ground cost c_f(u, û) = d_S^p(f(u), f(û)) eq. 7 feeds a Sinkhorn divergence between the two pushforwards. The emulator minimises the one-step error plus λ times this divergence, and f maximises it eq. 8. No critic is needed; the paper enforces Lipschitz behaviour through weight clipping of the summary network app. D.1.

Summary map §4 app. D.1. For Lorenz-96 the map is applied at every time step and every spatial location. It returns d numbers per location, so a trajectory becomes a point cloud of size (time × space) in R^d. The default is a three-layer MLP with two hidden layers of width 128 and ReLU. The handcrafted analogue with d = 3 is in Table 1 tab. 1.

Training schedule app. D.1. Warm start: the emulator is pretrained on the one-step error alone, then the transport term is switched on. Teacher forcing uses windows of length 100 with stride 2. Even time steps come from the data and odd time steps are predicted, one step each. The paper does not state the number of warm-start epochs or the optimiser settings.

Data §6 app. D.1. Lorenz-96 with N = 60 variables. Observation noise is Gaussian, with standard deviation 0.3 times the standard deviation of the clean trajectory. The paper also describes a single-trajectory regime at fixed forcing F = 10 and a multi-trajectory regime with forcing drawn from Unif[10, 18].

Pseudocode, WGAN variant. The paper describes the game in eq. 6. The normalisation step is not in the paper and is marked below.

# Learnable WGAN regulariser  [§4.1] [eq. 4] [eq. 6]
# Inputs: true windows (B, T, N); emulator g; summary f; critic phi
# Runs after the MSE warm start  [app. D.1]

def training_step(true_windows):
    # 1. Emulator rollout with anchors every 2 steps  [app. D.1]
    pred_windows = rollout_with_anchors(g, true_windows)  # (B, T, N)

    # 2. Summary and critic: maximise the dual estimate  [§4.1] [eq. 6]
    for _ in range(summary_steps):  # count not stated in paper
        true_feats = f(true_windows)              # (B, P, d)
        pred_feats = f(pred_windows.detach())     # (B, P, d)
        gap = mean(phi(true_feats)) - mean(phi(pred_feats))
        # gradient ascent on gap for f and phi
        # clip the weights of phi to [-c, c]  [§4.1] [app. D.1]

    # 3. Emulator: minimise one-step error plus lambda * gap  [eq. 4]
    true_feats = f(true_windows)
    pred_feats = f(pred_windows)
    gap = mean(phi(true_feats)) - mean(phi(pred_feats))
    loss = one_step_error(pred_windows, true_windows)
    loss = loss + lam * gap                       # lam = 3  [app. D.1]
    # update only the emulator on this loss
    return loss

Pseudocode, Sinkhorn variant §4.2 eq. 8:

# Learnable Sinkhorn regulariser  [§4.2] [eq. 8]
# The paper stops updating f early  [app. D.1]

def training_step(true_windows):
    pred_windows = rollout_with_anchors(g, true_windows)
    # 1. Summary: maximise the Sinkhorn divergence  [§4.2]
    for _ in range(summary_steps):  # count not stated in paper
        div = sinkhorn_div(f(true_windows), f(pred_windows.detach()))
        # gradient ascent on div for f only
    # 2. Emulator: minimise one-step error plus lam * divergence [eq. 8]
    div = sinkhorn_div(f(true_windows), f(pred_windows))
    loss = one_step_error(pred_windows, true_windows) + lam * div
    return loss

Figure 1 shows the framework: a one-step loss regularised by transport between real and generated trajectory distributions through a learnable summary map fig. 1.

Hyperparameters

Values are from the paper unless the column says code. "Not stated" means the paper does not give the value.

Results

Author-reported. Setup: Lorenz-96 with 60 variables, KS with 256 grid points, 2D Kolmogorov flow with Re = 1e4. Baselines: no OT term ("No OT") and the fixed-statistics OT of the cited prior work ("Fixed OT") §6 §7. Values are medians. Each system has a noisy block (σ = 0.3) and a clean block (σ = 0). The main text does not give the number of runs; the appendix gives percentile ranges for the Lorenz-96 multi-trajectory block only app. E.

Table 2, reproduced from the paper tab. 2:

Leading Lyapunov exponent on single-trajectory Lorenz-96 tab. 3: ground truth 2.334; No OT 1.615; Fixed OT 1.890; Sinkhorn 2.030; WGAN 2.336. Single values, no spread reported. The paper's argument that linear Koopman models cannot reach a positive exponent is stated in §7.

Headline claims, as the paper states them §7: all OT variants beat No OT on the statistical metrics, and the learnable variants beat Fixed OT on spectral distance on all systems. Table 2 has exceptions to both. WGAN with 1D convolution on KS is slightly worse than No OT on clean L1 (0.245 against 0.241). WGAN on single-trajectory Lorenz-96 is worse than Fixed OT on clean spectral distance (0.175 against 0.157). The Lorenz-63 results are qualitative only fig. 3 app. F.

Paper vs code

Code: https://github.com/gabrielmelo00/LearningToEmulateChaos_ICML26 at 7fb61d8, MIT, copyright Gabriel Melo (an author). The paper's footnote links it. The repository's README points to branches l63, ks and kolmogorov-2d for the other systems. The compile step cloned only the default branch, so the KS, Kolmogorov and Lorenz-63 code was not read. The single-trajectory Lorenz-96 regime was not found on the default branch either.

Where the paper and the default-branch code differ:

• Reconstruction loss. The paper writes a squared error eq. 4. The training loop uses the relative L2 norm of each trajectory, averaged over the batch, which is not squared main.py train_utils.py. The balance against λ changes.
• Summary normalisation. The paper does not mention it. The code standardises the predicted and true summaries with the mean and standard deviation of the true batch, adding 1e-6 to the standard deviation main.py. The fixed-OT and learnable-Sinkhorn paths call OT_measure.loss, which floors the standard deviation at 0.01 instead OT_utils.py.
• Lipschitz control of the summary map. The paper says the summary network's weights are clipped app. D.1. The code clips only the critic main.py. The summary-clip argument is parsed but never used configuration.py. The Lipschitz hinge regulariser of app. C has no counterpart in the code: a search for hinge, L_max and Jacobian-penalty terms found nothing.
• Critic clip value. The paper gives [-0.01, 0.01] app. D.1. The WGAN launch script passes 0.1 wgan.sh. The default is 0.01 configuration.py.
• Sinkhorn stopping rule. The paper says early stopping at 350 of 400 out of 500 training steps app. D.1. The code stops summary updates at epoch index 200, hard-coded main.py. The launch scripts do not pass an epoch count, so the default of 701 applies configuration.py.
• Warm start. The paper does not give the length. The code's default is 50 epochs configuration.py.
• Forcing range. The paper uses Unif[10, 18] app. D.1. The data generator's default lower bound is 12 generate_data.py.
• Observation noise. The paper does not say whether noise is resampled. The code draws it once when writing the data files and reuses it every epoch dataloader_l96.py.
• Entropic scale. The paper gives ε = 0.02 app. D.1. The code passes blur = 0.02 to GeomLoss's Sinkhorn loss with p = 2 sinkhorn.sh OT_utils.py. Whether that equals ε = 0.02 depends on GeomLoss's blur convention, which was not checked here.
• Summary network. The paper's description matches: three linear layers, hidden width 128, ReLU, one output per location summary.py. The output layer is sized to the spatial grid, so the weights are not shared across locations.
• Critic. Two linear layers, hidden width 64, ReLU summary.py, averaged over the points summary.py. The paper does not say how points are pooled.
• Emulator. The paper says the emulator is the same FNO as the cited prior work. The code's FNO has six spectral layers, while its docstring says four fno_1d_new.py fno_1d_new.py. The forcing enters as an extra input channel fno_1d_new.py.
• Matching settings. λ = 3 wgan.sh; teacher-forcing window 100 and stride 2 configuration.py train_utils.py; rollout length 1500 main.py; Lyapunov settings eval_LE.py LE.py.
• Runnability on the default branch. scripts/main.py imports models.mlp_l63, which is not in the repository main.py. It also reads args.lambda_contra main.py main.py, which configuration.py never defines. As released, the script fails at import. The card did not run it: no GPU or environment was set up. The README's experiments/CL_l96/eval.sh is absent from the repository [README.md].

Pitfalls

• Unstated training settings. Batch size, learning rates, epochs, critic and summary step counts, and FNO width and modes are missing from the paper. The code fills them in with values that differ from the paper's (see Paper vs code). An implementation that follows only the paper will not match the reported numbers.
• The transport term is applied after per-batch standardisation. Because the true summaries set the mean and standard deviation, the term is insensitive to shifts and scalings of the summaries. A summary map can therefore collapse on most of the attractor. The paper's Jacobian table shows this without early stopping for the Sinkhorn variant app. C.
• Histograms use the handcrafted statistics for every method. The histogram error is computed on the same fixed summary S, regardless of which method trained the emulator §6 tab. 1. Fixed OT optimises those very statistics, so the histogram column favours it by construction. The paper points to spectral distance, which Fixed OT does not optimise, as the clearer test §7.
• Claims about noisy histograms overstate the table. The text says the learnable variants match or beat Fixed OT on noisy histograms §7. Table 2 shows Fixed OT lower on KS noisy L1 (0.290 against 0.339 for WGAN and 0.435 for Sinkhorn), on single-trajectory Lorenz-96 noisy L1 (0.178 against 0.222 for Sinkhorn), and on Lorenz-96 multi-trajectory clean L1 (0.055 against 0.083 and 0.109) tab. 2.
• RMSE trades against the statistics. The no-OT baseline has the lowest or tied-lowest one-step RMSE in every block of Table 2 tab. 2. The paper's remark is that RMSE degrades with horizon regardless §6; the table does not show a benefit for the OT variants on this metric.
• Thin evidence for the differences. The main text reports medians without the number of runs, and single-trajectory rows have no spread at all tab. 2. Several RMSE gaps are a few hundredths. Table 3's Lyapunov values are single numbers tab. 3.
• The theory has assumptions the code does not enforce. The bounds assume L-Lipschitz summary maps §5.1. The released code clips only the critic, so the Lipschitz premise is observed (Jacobian norms in app. C) rather than enforced. The noise-forgetting result assumes exponential mixing of the emulator in W_p §5.2.2, which the paper does not test.
• The MSE blow-up corollary is a sketch. Its proof uses an asymptotic growth argument for the Jacobian of the k-step map, not a bound app. A. Treat it as motivation.
• Cost. The Lipschitz regulariser costs about four times as much training per the paper app. C. The released WGAN script runs five critic updates per emulator update wgan.sh, which the paper does not state.
• Typos. A table in app. E gives "(0.72, 0.111)" where the value is most likely 0.072. The Kolmogorov figure caption in app. G writes α = 0.1 where the text writes χ = 0.1 app. D.3.

Lineage

• Builds on: the fixed-statistics OT regulariser of Jiang et al. (2023, NeurIPS, "Training neural operators to preserve invariant measures of chaotic attractors"), which this paper uses as its Fixed OT baseline §3 §6. The Sinkhorn divergence with a learnable cost follows Genevay et al. (2018) §4.2. The WGAN dual follows Arjovsky et al. (2017) and Gulrajani et al. (2017) §4.1. None of these has a card here, so they are named but not linked, and nothing from them is used beyond what this paper states.
• Compared against: MSE-only training and Fixed OT §6. Related work also names Sobolev-norm losses, DySLIM and Poincaré flow networks §3; the card does not compare results against them.
• Superseded by: not known from the content of this wiki.