An actively trained neural operator for near-explicit atmospheric chemistry

What it claims

FNO-MoE emulates 4,691 species of the Master Chemical Mechanism (MCM) and predicts the next 30 minute-resolved states in one forward pass, so a six-hour run takes twelve calls. The model's own experts choose which box-model simulations to run. At a budget of 10,000 simulations, this active sampling gives the lowest one-step error on every metric reported. After fine-tuning on six-hour trajectories, none of 70 six-hour rollouts diverges. The median per-species R² is 0.998, and the emulator runs about 750 times faster than the solver on one CPU core.

How it works

The reference is the ARCA box model with an MCM configuration of 6,068 species and 18,490 reactions. A low-concentration filter reduces this to 4,691 retained species. Each simulation starts from a 150-dimensional vector: four controls (temperature, relative humidity, short-wave radiation, condensation sink) plus 146 precursor concentrations. It runs one or six hours under constant controls and saves every minute.

Training pairs are cut from anywhere along a trajectory at a one-minute stride. The input is the full state plus the four controls, never the seed vector. The state is lifted by an MLP and broadcast over the 30-minute horizon with a time encoding. It then passes to four independent FNO experts that keep 4, 8, 12 and 16 Fourier modes. Each expert is conditioned on the controls through FiLM and predicts an increment. A static species-wise gate sums the increments, and the result is added to the input state. Longer runs feed the last predicted state back in as the next input.

Active sampling is a single round:

flowchart TD
  A[7 regimes: 2,500 anchors] --> W[Warm-up: 2,500 ARCA sims]
  W --> F[Train, then freeze emulator]
  A --> C[16 Sobol candidates per anchor]
  F --> S[Score by expert spread]
  C --> S
  S -->|top 3 per anchor| R[Run ARCA: 7,500 sims]
  S -->|other 13| X[Discarded, never simulated]
  W --> T[Retrain once on 10,000 sims]
  R --> T

Candidates are drawn within radius 0.10 in normalized units. In total, 40,000 candidates are scored and 7,500 are simulated. Scoring one 30-minute block takes about 16 ms on one CPU core. Keeping a fixed number per anchor means every regime gets new simulations.

Evidence

Error is the RMSE of normalized log concentration. A value of 0.01 corresponds to roughly a 10% concentration error. All three strategies use the same architecture, 10,000 one-hour simulations and the same 2,500 evaluation anchors.

• All species, post-transient: active 0.0075, static local 0.0085, whole-space random 0.0108.
• 204 key species, post-transient: active 0.0051; static and random both 0.0063.
• Initial-transient windows, all species: 0.1289, 0.1469 and 0.2670.

The advantage holds in each of the seven regimes. A single-Markov-chain variant performed worse than static sampling. Expert disagreement correlates only moderately with error: Pearson r = 0.31 (95% CI 0.20 to 0.42) across 500 states and r = 0.46 across species.

Error concentrates in the first five minutes, where windows are 17 to 25 times worse. A variant trained without these windows (1h-NoStiff) reaches 0.0083 post-transient against 0.0111 for 1h-Stiff. Its full-trajectory error, however, is 0.2553 against 0.0271.

For rollouts, the six-hour fine-tuned models diverged in 0 of 70 six-hour runs. The two one-hour models diverged in 66 of 70. On compositions far from the training design, the median R² falls to 0.990. Batched GPU inference reaches 4.3×10⁵ times the solver's throughput. The authors cite prior work in which a step-wise emulator 80 to 130 times faster offline sped up a transport model's chemistry by only 2 to 4.6 times once embedded.

Limits

All training uses constant meteorology. A first test with time-varying forcing drops the median R² to 0.870, and the authors name this as the main step before coupled deployment. The model misses the most extreme states under strong forcing. The authors treat expert disagreement as a ranking heuristic, not a calibrated uncertainty. Because the evaluation anchors also guide acquisition, the sampling comparison tests prediction near the design set. The fetched text is cut off mid-Results and mentions no code, weights or licence.