Parameterization of Lorenz-96 and Testing Online Stability Against Training Data Volume
01 Mathematical Formulation
Atmospheric circulation models parameterize subgrid-scale physical processes due to computational limits. The two-tier Lorenz-96 system provides a standard benchmark for studying turbulence closures:
where $k = 1, \dots, K$ represents resolved slow variables with cyclic boundary conditions ($X_{k+K} = X_k$), $F$ is constant external forcing, and subgrid tendency $U_k$ is governed by coupling to fast variables $Y_{j,k}$:
When reducing the system to resolved variables only, we train a neural parameterization $\hat{U}_k(X) \approx U_k$. While offline metrics ($R^2$, MSE) assess one-step error, they do not guarantee stability once $\hat{U}_k$ is coupled back into the differential equations.
02 Project Motivation
Core Question: How much training data does a small neural network need to parameterize subgrid coupling $U_k$ without drifting or becoming unstable when integrated online in a differential equation?
03 Numerical Setup & Experiments
I set up a numerical experiment using the two-tier Lorenz-96 equations. The baseline model, offline training, and online coupling pipeline are reproduced from the M²LInES notebooks, Learning Machine Learning with Lorenz-96.
- Sample sizes: Generated training subsets across 10 dataset sizes, logarithmically spaced across simulation time.
- Architectures: Tested two activation functions to compare spectral properties and long-term stability.
- Ensembles: Trained 20 neural networks across the data sizes to compare offline fit against online behavior.
- Online testing: Integrated each trained model forward for $t > 10,000$ steps with RK4, tracking kinetic energy, autocorrelation, and trajectory boundedness.
04 Key Observations
Integrating the models online showed clear differences between offline accuracy and operational stability:
- Offline metrics are not enough: Models trained on small datasets achieved offline $R^2 > 0.94$, but quickly exploded or experienced unphysical energy drift when integrated closed-loop.
- Data volume threshold: A distinct threshold appeared where sufficient trajectory coverage allowed the neural parameterization to maintain energy balance and stay on the attractor.
- Activation function impact: Smooth activations demonstrated better stability over long rollouts compared to piece-wise linear ones, reducing high-frequency accumulation in the resolved variables.
05 Code & Implementation
Integration notebooks and model training scripts are available on GitHub: