Technical Summary: Controlling Transient Amplification Improves Long-horizon Rollouts

1. Problem Statement

Autoregressive neural simulators have achieved state-of-the-art performance in short-horizon predictions for physical systems, matching or surpassing classical numerical solvers. However, a critical limitation persists: when these models are rolled out iteratively over long horizons (tens to thousands of steps), their accuracy degrades rapidly, often diverging from ground truth despite high single-step accuracy.

The prevailing explanation for this phenomenon is distributional shift: during training, models receive clean ground-truth states, but during inference, they consume their own imperfect predictions, causing errors to accumulate. While recent methods attempt to mitigate this via multi-step training or inference-time corrections (e.g., iterative denoising), these approaches treat the symptom rather than the underlying structural cause. This paper posits that the root cause is transient amplification of perturbations driven by the structural properties of the model's latent Jacobians, specifically their non-normality and non-commutativity.

2. Theoretical Analysis

The authors analyze error propagation by linearizing the autoregressive operator $F_\theta$ around a predicted trajectory. Let $\varepsilon_t$ be the error at step $t$ and $J_t$ be the Jacobian of the operator at that step. The error evolves as $\varepsilon_{t+1} \approx J_t \varepsilon_t$, leading to a total error after $T$ steps of $\varepsilon_T \approx \Phi_T \varepsilon_0$, where $\Phi_T = J_{T-1} \cdots J_0$.

Standard stability analysis focuses on the spectral radius $\rho(J_t) < 1$, ensuring asymptotic decay. However, the paper highlights that for non-normal matrices (where $J^\top J \neq J J^\top$), the spectral norm $\|J\|_2$ can be significantly larger than the spectral radius $\rho(J)$. This allows for transient growth, where errors amplify significantly before eventually decaying.

Furthermore, in a time-varying system (LTV), even if individual steps are contractive, the product of non-commuting matrices can exhibit massive transient amplification. The paper establishes that:

1. Non-normality of individual Jacobians allows errors to grow transiently within a single step.
2. Non-commutativity between Jacobians across steps ($[J_s, J_t] \neq 0$) causes error components to rotate and reinforce one another, compounding transient growth over the rollout.

The authors derive a propagator bound (Theorem 2.1) showing that if Jacobians are approximately normal and approximately commuting, the error growth is bounded by the spectral radius $\rho^T$ rather than the spectral norm $\alpha^T$. This represents a strict improvement for non-normal systems, as $\rho < \alpha$.

3. Methodology: Commutativity Regularization

To directly control transient amplification, the authors propose commutativity regularization, a training-time objective that penalizes structural defects in the latent space Jacobians without incurring inference-time costs.

Regularization Terms

The method introduces two auxiliary losses estimated via Jacobian-vector products (JVPs) and vector-Jacobian products (VJPs), avoiding the storage of explicit $n \times n$ Jacobian matrices:

1. Normality Penalty ($L_{norm}$): Penalizes the deviation of individual Jacobians from normality.
   $$L_{norm} = \mathbb{E}_v \left[ \| J_t^\top J_t v - J_t J_t^\top v \|^2 \right]$$
   This discourages the generation of large within-step transients.

2. Commutator Penalty ($L_{comm}$): Penalizes the non-commutativity of consecutive Jacobians.
   $$L_{comm} = \mathbb{E}_v \left[ \| J_{t+1} J_t v - J_t J_{t+1} v \|^2 \right]$$
   This discourages the cross-step compounding of transients.

Implementation

• Latent Space: The penalties are applied to the Jacobians of the "latent advance map" $G_\theta$ (the bottleneck-to-bottleneck dynamics), rather than the full input-output map. For architectures like UNets, this is the bottleneck feature; for FNOs, it is the middle layer features.
• Efficiency: The losses are computed using random probe vectors (Hutchinson estimator) via torch.func.jvp and torch.func.vjp.
• Training Objective: The total loss is $L = L_{pred} + \lambda_c L_{comm} + \lambda_n L_{norm}$, where $L_{pred}$ is the standard single-step MSE.
• Inference: The rollout procedure remains unchanged; the regularization only affects the learned weights.

4. Experimental Results

The method was evaluated across four diverse settings: 1D integrable systems, 2D chaotic PDEs, real-world sea surface temperature, and large-scale weather forecasting.

4.1 Korteweg–de Vries (KdV) Equation

• Setup: Trained on single-soliton dynamics (200 steps) and evaluated on horizons up to 5,000 steps, including out-of-distribution (OOD) multi-soliton interactions.
• Results:
  • In-Distribution: The regularized UNet (UNet+CR) showed a 5x improvement over the baseline at 1,000 steps and maintained coherence where the baseline diverged.
  • OOD: While baselines failed to maintain coherent multi-soliton trajectories, the regularized model remained stable.
  • Spectral Backbones (FNO/U-FNO): Unregularized FNOs collapsed around step 100 despite excellent single-step accuracy. Commutativity regularization converted these divergent rollouts into bounded, accurate trajectories.
  • Comparison: UNet+CR matched the performance of PDE-Refiner (an inference-time correction method) at intermediate horizons but outperformed it at long horizons without PDE-Refiner's 5x inference-time cost.

4.2 Barotropic Vorticity Equation (BVE)

• Setup: A chaotic 2D fluid system.
• Results: The unregularized baseline destabilized within the training window (10s), with RMSE saturating at 6.0. The regularized model remained bounded with an RMSE of 0.66 at the same horizon, representing a 9x improvement in RMSE and an 82x improvement in normalized MSE. The regularizer preserved coherent vortex structures while the baseline amplified vorticity to unphysical scales.

4.3 FourCastNet on ERA5

• Setup: Fine-tuning a pre-trained 74M parameter global weather model (FourCastNet) on a small slice of its own training data (2013–2015) to test structural improvements without new data.
• Results: Plain fine-tuning degraded long-horizon forecasts for 2m temperature ($t2m$), causing a 3x degradation at day 10. Adding commutativity regularization reversed this, improving $t2m$ forecasts by 41% at day 10 compared to the frozen pre-trained model. The method also improved upper-air variables where the baseline drifted.

4.4 Sea Surface Temperature (SST)

• Setup: Training a UNet on weekly NOAA OISST data for a 7.2-year autoregressive rollout.
• Results: The baseline drifted in phase relative to the annual cycle and amplified anomalies. The regularized model tracked the seasonal cycle accurately for the entire 376-week horizon, achieving a 3.0x improvement in average accuracy.

5. Key Contributions

1. Identification of Mechanism: The paper identifies non-normal, non-commuting latent Jacobians as the structural driver of transient amplification in autoregressive neural operators, distinct from simple distributional shift.
2. Commutativity Regularization: Proposes a novel training objective combining commutator and normality penalties, estimated efficiently via JVPs with zero inference-time overhead.
3. Theoretical Bound: Provides a propagator bound showing that under approximate commutativity and normality, error growth is governed by the spectral radius rather than the spectral norm, offering a strict improvement for non-normal systems.
4. Empirical Validation: Demonstrates consistent long-horizon accuracy gains across integrable, chaotic, and real-world geophysical systems, including OOD generalization and large-scale weather models.

6. Significance and Limitations

Significance:
The paper argues that bridging the gap between single-step accuracy and long-horizon reliability requires addressing the structural properties of the learned operator, not just the data distribution. By enforcing normality and commutativity, the method allows models to maintain stability over thousands of steps without expensive inference-time corrections. The results on FourCastNet are particularly notable, showing that structural regularization can rescue fine-tuning efforts that would otherwise degrade performance.

Limitations:

• Linearization Assumption: The theoretical guarantees rely on linearization around the trajectory. Far from the training distribution, where linearization may be invalid, the guarantees become suggestive rather than rigorous.
• Stochastic Estimation: The penalties are estimated using single-probe JVPs, serving as a tractable surrogate for the ideal all-pairs joint diagonalizability condition.
• Architecture Specificity: The method is most natural for architectures with explicit latent bottlenecks. Applying it to purely spectral or attention-based models without such bottlenecks requires further adaptation.
• Complementarity: The authors note that this method targets the structural component of error growth and is complementary to, rather than a replacement for, methods addressing distributional shift (e.g., diffusion-based refinement). They expect the strongest results from combining both approaches.