Learning to Emulate Chaos: Adversarial Optimal Transport Regularization

1. Problem Statement

The paper addresses the fundamental challenge of training data-driven emulators (surrogate models) for chaotic dynamical systems.

• The Core Difficulty: Chaotic systems exhibit extreme sensitivity to initial conditions (the "butterfly effect"). Consequently, exact long-term trajectory prediction is theoretically infeasible. Standard training objectives, such as Mean Squared Error (MSE), fail in this regime because they penalize trajectory divergence. Even with a perfect model, small measurement noise or initial condition errors cause the MSE to grow exponentially over time, leading to catastrophic failure in long-term forecasting.
• The Goal: Instead of matching specific trajectories, the objective is to learn an emulator that preserves the invariant physical measure (the statistical properties) of the chaotic attractor. The emulator should reproduce the correct long-term statistical distribution, energy spectra, and dynamical invariants, even if individual trajectories diverge from the ground truth.
• Limitations of Existing Methods: Previous approaches rely on hand-crafted summary statistics (e.g., specific moments or Lyapunov exponents) or require diverse datasets from multiple dynamical regimes. These methods often lack flexibility, require expert domain knowledge to select features, or fail to generalize from a single noisy trajectory.

2. Methodology

The authors propose a novel framework using Adversarial Optimal Transport (OT) Regularization. The core idea is to jointly learn a high-quality emulator and an adaptive set of summary statistics that best distinguish between the true system's statistics and the emulator's generated statistics.

A. The Framework

The method formulates a min-max optimization problem involving three components:

1. The Emulator ($g$): A neural network (e.g., Fourier Neural Operator) that predicts the next state $u_{t+1} = g(u_t)$.
2. The Summary Map ($f$): A learnable neural network that maps states to a summary space $S$, extracting informative statistics $s = f(u)$.
3. The Critic/Dual Function ($\phi$): (In the WGAN variant) A function that distinguishes between the distribution of summaries from the real data and the emulator.

B. The Objective Function

The training objective combines a standard prediction loss with an adversarial OT term:
$$\min_{g} \max_{f} \mathcal{L}(g, f) = \mathcal{L}_{MSE} + \lambda \cdot \text{OT}_{cost}(f_\#\mu, f_\#\hat{\mu})$$
Where:

• $\mathcal{L}_{MSE}$ ensures short-term one-step fidelity.
• $f_\#\mu$ and $f_\#\hat{\mu}$ are the pushforward measures (distributions) of the summary statistics for the true data and the emulator, respectively.
• The OT cost measures the discrepancy between these distributions.
• The Adversarial Nature: The summary map $f$ is trained to maximize the discrepancy (finding the most discriminative statistics), while the emulator $g$ is trained to minimize it (forcing the emulator to match the true statistics in the most critical dimensions).

C. Two Instantiations

The paper implements two specific formulations:

1. WGAN-style (1-Wasserstein): Uses the Kantorovich-Rubinstein duality. It employs a critic network to approximate the 1-Wasserstein distance. This requires enforcing Lipschitz constraints (e.g., via weight clipping or gradient penalty).
2. Sinkhorn Divergence (2-Wasserstein): Uses entropic regularization to make the OT cost differentiable and computationally efficient via the Sinkhorn algorithm. This avoids explicit Lipschitz constraints on a critic and interpolates between OT and Maximum Mean Discrepancy (MMD).

3. Key Contributions

1. Adversarial Learning of Summary Statistics: Unlike prior work that relies on fixed, hand-crafted features, this method learns the summary statistics jointly with the emulator. This ensures the statistics are optimally discriminative for the specific chaotic system and dataset, removing the need for domain expertise.
2. Theoretical Analysis of Robustness:
   • The authors prove that while MSE loss explodes exponentially with rollout time due to noise (Corollary 5.9), the Wasserstein distance remains bounded and stable (Theorem 5.14).
   • They demonstrate that the OT regularizer is robust to initial condition noise because chaotic systems "forget" initial conditions in a distributional sense (exponential mixing), whereas trajectory-level metrics do not.
3. Single-Trajectory Learning: The method is capable of training high-fidelity emulators using only a single noisy trajectory, a significant improvement over methods requiring diverse multi-environment datasets.
4. Comprehensive Empirical Validation: The approach is tested on three high-dimensional chaotic systems:
   • Lorenz-96: A 60-variable atmospheric model.
   • Kuramoto-Sivashinsky (KS): A 1D PDE modeling reaction-diffusion instabilities.
   • Kolmogorov Flow: A 2D incompressible Navier-Stokes system.

4. Experimental Results

The experiments compare the proposed Learnable OT methods (Sinkhorn and WGAN) against:

• No OT: Standard MSE training.
• Fixed OT: Using hand-crafted summary statistics (Jiang et al., 2023).

Key Findings:

• Statistical Fidelity: The adversarial OT methods significantly outperform the MSE baseline and often outperform the Fixed OT baseline in long-term statistical metrics (L1 histogram error and spectral distance).
• Noise Robustness: In noisy settings ($\sigma=0.3$), the MSE baseline fails to capture the attractor geometry, often collapsing into limit cycles or underestimating the attractor's extent. The adversarial OT methods maintain the correct "butterfly" shape and spatial extent even at high noise levels.
• Dynamical Properties: The learned emulators reproduce the Leading Lyapunov Exponent (LLE) much more accurately. For Lorenz-96, the WGAN method achieved an LLE of 2.336 (vs. ground truth 2.334), whereas the MSE baseline was 1.615.
• Generalization: The method works effectively for both multi-trajectory and single-trajectory training regimes.

5. Significance

This work represents a paradigm shift in learning chaotic dynamics:

• From Trajectory to Distribution: It formally moves the training objective from matching specific time-series points (which is impossible for chaos) to matching the invariant probability measure of the system.
• Automation: By adversarially learning the summary statistics, it removes the bottleneck of manual feature engineering, making the approach applicable to complex systems where physical invariants are unknown or difficult to derive.
• Practical Impact: The ability to learn from a single noisy trajectory makes this method highly relevant for real-world applications like weather prediction, climate modeling, and plasma physics, where data is often scarce and noisy.

In conclusion, the paper demonstrates that Adversarial Optimal Transport provides a theoretically sound and empirically superior regularization technique for training neural emulators of chaotic systems, ensuring long-term statistical consistency where traditional loss functions fail.