Assessing Emulator Design and Training for Modal Aerosol Microphysics Parameterizations in E3SMv2

This paper systematically evaluates the design and training of scientific machine learning emulators for the MAM4 aerosol microphysics module in E3SMv2, demonstrating that effective scaling, convergence monitoring, and moderate network complexity are critical for accurately reproducing aerosol concentration changes under cloud-free conditions.

The Gist

Imagine the Earth's atmosphere as a giant, bustling kitchen. In this kitchen, tiny particles called aerosols (like dust, sea salt, and pollution) are constantly being cooked, mixed, and transformed by invisible chefs. These particles are crucial because they affect how clouds form, how much sunlight reaches the ground, and ultimately, our climate.

The scientists at Pacific Northwest National Laboratory are trying to build a super-fast recipe book (a computer program) to predict what happens to these particles.

Here is the simple breakdown of what this paper is about:

1. The Problem: The Kitchen is Too Slow

The current computer model used to simulate the Earth's climate (called E3SM) is like a very detailed, slow-motion movie of this kitchen. It calculates every single chemical reaction and movement of every dust particle using complex math equations. While accurate, it's incredibly slow. If you want to run a climate simulation for 100 years, it might take months of supercomputer time just to do the math on these tiny particles.

The scientists wanted to replace this slow, heavy math with a smart shortcut. They wanted to train an AI (a "neural network") to look at the ingredients and instantly guess the outcome, just like a master chef who can taste a sauce and know exactly what's in it without measuring everything.

2. The Experiment: Training the AI Chef

The team tried to teach a simple AI to mimic the "kitchen logic" of the aerosol particles, but only when the sky is clear (no clouds). They fed the AI millions of examples from the slow computer model:

• Input: "Here are the ingredients right now (temperature, humidity, current dust levels)."
• Output: "Here is what the ingredients will look like 30 minutes later."

They didn't just throw random data at the AI. They had to be very careful about how they taught it.

3. The Big Challenges (and How They Solved Them)

The paper details three main hurdles they had to clear to make the AI work well:

• The "Size" Problem (Scaling):
  Imagine trying to teach a child to count. If you ask them to count "1" and then "1,000,000" in the same breath, they get confused. Similarly, the AI got confused because some particle numbers were tiny (like a single grain of sand) and others were huge (like a mountain of sand).

  • The Fix: The scientists used a special mathematical "translator" (called a power transformation) to shrink the big numbers and stretch the small numbers so they all fit on the same scale. It's like converting all currencies to dollars before comparing them.

• The "Architecture" Problem (How complex should the brain be?):
  They asked: "Does the AI need a giant, 10-story brain, or is a small, 2-story brain enough?"

  • The Discovery: They found that a moderately sized brain (3 layers deep with 256 neurons each) was the "Goldilocks" zone.
    • Too small? It couldn't understand the complex chemistry.
    • Too big? It started memorizing the answers instead of learning the rules, and it was a waste of computer power.
    • Just right? It learned the patterns perfectly.

• The "Patience" Problem (Training Convergence):
  Teaching an AI is like training a dog. If you stop the training too early, the dog hasn't learned the trick. The scientists found that they had to let the AI train for a long time (5,000 rounds) to ensure it truly understood the physics and didn't just guess.

4. The Results: A Fast and Accurate Shortcut

When they tested their final AI model:

• It was extremely accurate, matching the slow, detailed computer model with about 99% accuracy.
• It could handle the tricky "weird" particles (like sea salt and organic matter from the ocean) just as well as the common ones.
• Most importantly, it proved that you don't need a super-complex, mysterious AI to do this job. A simple, well-designed "feedforward" neural network works great if you treat the data correctly.

5. Why This Matters

Think of this paper as the instruction manual for building a better engine.

• Before this, people were trying to build AI climate models and getting mixed results (some worked, some didn't), but no one knew exactly why.
• This paper says: "Here is the recipe. If you normalize your data, pick a medium-sized network, and train it long enough, you can speed up climate simulations without losing accuracy."

In a nutshell: The scientists built a smart, fast AI assistant that can predict how dust and pollution move in the air. They figured out that the secret to making it work wasn't a fancy, complicated AI, but rather cleaning up the data and training it patiently. This paves the way for faster, more accurate climate predictions in the future.

Technical Summary

1. Problem Statement

The paper addresses the computational challenges associated with simulating aerosol microphysics in global Earth System Models (ESMs), specifically the Energy Exascale Earth System Model version 2 (E3SMv2).

• The Challenge: Aerosol microphysics involves complex, non-linear processes (nucleation, condensation, coagulation, aging) that occur at scales much smaller than the model's grid spacing. These are typically solved using computationally expensive systems of coupled Ordinary Differential Equations (ODEs).
• The Gap: While Machine Learning (ML) emulators have shown promise in accelerating these processes, previous studies (e.g., Harder et al., Bai and Rouson) reported highly variable performance (coefficient of determination $R^2$ ranging from 0.48 to 0.99). There is a lack of systematic understanding regarding why this variability occurs, specifically how architectural choices, training strategies, and data scaling affect emulator fidelity.
• Objective: To establish a robust baseline for ML emulation of aerosol microphysics by systematically analyzing design choices (architecture, normalization, convergence) using a standard Feedforward Neural Network (FNN) on the MAM4 (Modal Aerosol Module, 4-mode) parameterization suite under cloud-free conditions.

2. Methodology

2.1 Data Generation and Scope

• Source: Data was generated using an E3SMv2 simulation (F2010 component set) at a horizontal resolution of ~165 km (ne30pg2 grid) with 72 vertical layers.
• Scope: The study focuses on cloud-free interstitial aerosols in the troposphere (excluding the top 26 stratospheric layers and cloudy grid cells).
• Target Variables: The emulator predicts the 30-minute evolution of 20 mixing ratios (18 aerosol species across 4 modes + 2 gas precursors: $H_2SO_4$ and SOAG).
• Inputs: 39 input features, including the initial 20 mixing ratios, atmospheric state variables (temperature, pressure, humidity), aerosol properties (particle diameters, density), and chemical production rates.
• Dataset: Over $6.6 \times 10^6$ samples were randomly split into training (50%), validation (25%), and testing (25%) sets.

2.2 Preprocessing and Variable Transformation

A critical component of the methodology is handling the multi-scale nature of aerosol data (spanning multiple orders of magnitude) and skewness.

• Power Transformation: Instead of standard log-transformation (which fails on negative values and compresses large extremes), the authors applied a power transformation: $T(r) = r^{1/n}$ (where $n$ is an odd integer, specifically $n=3$). This handles both positive and negative changes while compressing the dynamic range.
• Z-Score Normalization: Following the power transformation, variables were standardized to have a mean of 0 and standard deviation of 1.
• Output Definition: The network predicts the change ($\Delta$) in mixing ratios over one timestep, rather than the absolute final value, to leverage the small magnitude of changes relative to the state.

2.3 Neural Network Architecture

• Base Model: A standard Feedforward Neural Network (FNN) with Rectified Linear Unit (ReLU) activation.
• Residual Connections: The architecture incorporates residual connections (identity mappings) between layers. This improves optimization conditioning, prevents vanishing/exploding gradients, and allows for stable training across varying depths and widths.
• Configuration: Input layer size = 39; Output layer size = 20.
• Training:
  • Loss Function: Mean Squared Error (MSE) on transformed variables.
  • Optimizer: Adaptive Nesterov momentum (Adan) with a learning rate of $3 \times 10^{-4}$.
  • Batch Size: 4096.
  • Duration: 5000 epochs (no early stopping) to ensure full convergence.

3. Key Contributions

1. Systematic Design Space Exploration: The paper provides the first controlled, systematic analysis of how network depth, width, and training hyperparameters impact aerosol microphysics emulation, moving beyond "black box" trial-and-error.
2. Validation of Simple Architectures: It demonstrates that a relatively simple FNN (without complex attention mechanisms or operator networks) can achieve high fidelity ($R^2 \approx 0.99$) if training convergence and data scaling are handled correctly.
3. Importance of Variable Transformation: The study highlights that standard linear normalization is insufficient for aerosol data. The specific combination of power transformation + z-score normalization is identified as critical for capturing both small and large magnitude changes accurately.
4. Convergence Analysis: The authors emphasize that many previous failures in ML emulation may stem from premature stopping or suboptimal optimizer settings, rather than inherent model limitations.

4. Results

4.1 Training Convergence

• The selected configuration (Adan optimizer, specific learning rate, 5000 epochs) resulted in monotonic convergence of the loss function.
• The "suboptimality gap" (difference between current loss and minimum loss) decreased by over three orders of magnitude, confirming that the optimizer reached a stable minimum.
• Comparisons showed that configurations from previous studies (e.g., [29]) applied to this data required significantly more epochs to converge, suggesting that previous lower performance scores might have been due to under-training.

4.2 Architecture Complexity (Depth vs. Width)

• Shallow/Narrow Networks: Networks with only 1 hidden layer or very narrow widths (32 neurons) failed to capture the complexity of the physics.
• Diminishing Returns on Depth: Increasing depth beyond 2–3 layers yielded marginal improvements compared to increasing width.
• Optimal Configuration: A 3-layer network with 256 neurons per layer ($3 \times 256$) was identified as the optimal compromise between accuracy, parsimony, and computational cost.

4.3 Performance Metrics

• Overall Accuracy: The $3 \times 256$ architecture achieved an average $R^2$ of ~0.99 across all 20 target variables on the testing dataset.
• Variable-Specific Performance:
  • Most variables were emulated with high fidelity.
  • Challenging Variables: Marine Organic Matter (mom), Sea Salt (ncl), Secondary Organic Aerosols (soa, SOAG), and Aitken mode particle number (num_a2) showed slightly lower $R^2$ scores. This suggests these specific processes may require specialized architectures or further investigation.
• Error Distribution: 2D histograms of predicted vs. true values showed strong agreement along the 1:1 line. The power transformation ensured that the model performed well across the entire dynamic range, avoiding the bias toward large values often seen with linear transformations.

5. Significance and Future Work

• Practical Impact: This work provides a "blueprint" for developing ML emulators for atmospheric physics. It proves that replacing expensive ODE solvers with ML is feasible without sacrificing accuracy, provided the training pipeline is rigorous.
• Generalizability: The insights regarding variable transformation, residual connections, and convergence monitoring are applicable to other multi-scale, non-linear parameterizations in Earth System Models (e.g., cloud microphysics, turbulence).
• Next Steps: The authors identify two immediate future directions:
  1. Seasonal Generalization: Extending the training data to cover full seasonal cycles and varying aerosol source scenarios.
  2. Online Integration: Embedding the emulator directly into the E3SMv2 host model to test numerical stability and long-term climate impacts (spin-up behavior) in a coupled environment.

In conclusion, the paper successfully establishes that scientific machine learning can accurately emulate complex aerosol microphysics, but success is heavily dependent on data preprocessing (power transformation) and rigorous training protocols rather than just architectural complexity.