An actively trained neural operator for near-explicit atmospheric chemistry

Zhi-Song Liu; Wenqing Peng; Ammar Kheder; Risto Makkonen; Michael Boy

Figure 1: One network both emulates near-explicit chemistry and chooses its own training simulations. a , The FNO-MoE emulator maps the full retained chemical state

x_{0}

(4,691 species, shown for a real ARCA state, grouped by chemical family and coloured by mixing ratio from

10^{-14}

(blue) to

10^{2}

ppb (red)) and four physical controls

m

(temperature, relative humidity, short-wave radiation and condensation sink) to all 30 minute-resolved states of the next 30 minutes in one forward pass (stacked sheets: the emulator’s prediction of the change from

x_{0}

for the same case; blue, decrease; red, increase). Inside, the state is lifted and broadcast over the horizon with a time coordinate; four independent Fourier neural operator experts retain 4, 8, 12 and 16 Fourier modes, each applying four spectral blocks (Fourier transform

F

, weights

R_{k}

on the lowest

k

modes, inverse transform

F^{-1}

, a

1\times 1

bypass

W

, FiLM(

m

) and GELU); a static species-wise gate

\pi=\mathrm{softmax}(g)

combines the experts’ increments as

\sum_{e}\pi_{s,e}\Delta^{(e)}

, which is added to

x_{0}

. The last predicted state seeds the next call. b , Single-round active sampling: after a warm-up of 2,500 ARCA simulations, one near each of 2,500 anchor states drawn from seven atmospheric regimes (shown: temperature versus short-wave radiation), 16 Sobol candidates are drawn within

r=0.10

of each anchor (40,000 in total); the frozen experts score each candidate by the spread

\sigma

of their predictions for its first 30 minutes, without running the box model; the three highest-scoring candidates per anchor (7,500) are simulated by ARCA, and the emulator is retrained once on all 10,000 simulations. Blue arrows and glyphs, emulator only; orange, ARCA box-model runs. c , A real six-hour rollout of the deployed emulator drawn as its twelve chained 30-min calls (rows, species; colour, change from the starting state on the same blue–red scale as in a), against the 90 calls of a step-wise emulator with a 4-min chemistry step.

Sampling strategy	All species		Key species (204)
Sampling strategy	Initial $\downarrow$	Post-transient $\downarrow$	Initial $\downarrow$	Post-transient $\downarrow$
Whole-space random	0.2670	0.0108	0.1315	0.0063
Static local	0.1469	0.0085	0.0716	0.0063
Active local (this work)	0.1289	0.0075	0.0641	0.0051

Table 1: Active local sampling gives the lowest one-step error in every metric. Normalized-log RMSE on the 2,500 regime anchors for three strategies with identical architecture and a matched budget of 10,000 one-hour simulations. Initial windows start in the first five minutes, during the stiff initialization transient; post-transient windows start at or after minute five. All species: each dataset’s retained species (4,691 for active, 4,713 for static and 4,725 for random sampling); Key: 204 key species (Supplementary Note Supplementary Note 4: key species ). Lower is better; best values in bold.

Figure 2: Emulator-guided sampling lowers one-step error on most anchors at a fixed budget of 10,000 simulations. Grey, whole-space random; amber, static local; blue, active local sampling (this work). a , Median per-anchor all-species RMSE of the three strategies for all 2,500 evaluation anchors and for each atmospheric regime; right, the median paired change of active relative to static sampling and the share of anchors on which active sampling is more accurate. All regime-level changes have 95% paired-bootstrap intervals that exclude zero (all anchors:

-10.5\%

[

-11.4

,

-10.0

]; improved on 76% [74, 78] of anchors). b , Distribution of the paired per-anchor change in RMSE of active relative to static sampling (filled; blue, improved; grey, worse) and relative to random sampling (outline); dashed lines, medians (

-11\%

and

-50\%

). c , RMSE of the 30-min block versus its starting offset within the reference simulation; shading marks the floor-seeded transient (first 5 min), where errors are 17–25 times larger than later. d , Pre-gate expert disagreement versus prediction error for 500 evaluation states (Pearson

r=0.31

, 95% bootstrap interval

[0.20,0.41]

); circles and bars, median and interquartile range of the error in eight equal-count bins of disagreement. Disagreement ranks candidates usefully but is not a calibrated error estimate.

Model	6-h FT	Transient	In envelope (6 cases)			Envelope edge (8 cases)
Model	6-h FT	Transient	RMSE $\downarrow$	Drift $\downarrow$	Div. $\downarrow$	RMSE $\downarrow$	Drift $\downarrow$	Div. $\downarrow$
1h-Stiff	✗	✓	0.322	0.628	6/6	0.354	0.483	7/8
1h-NoStiff	✗	✗	0.234	0.270	6/6	0.257	0.231	8/8
6h-Stiff	✓	✓	0.0371	0.0057	0/6	0.0269	0.0121	0/8
6h-NoStiff (deployed)	✓	✗	0.0274	0.0051	0/6	0.0161	0.0052	0/8

Table 2: Six-hour fine-tuning removes divergence and lowers final error by 88–94%. Six-hour autoregressive rollouts from post-transient states (minute five). RMSE (normalized log) and elemental drift at the last minute; Div., number of diverged cases (any normalized-log value

|x|>3

or non-finite). 6-h FT, fine-tuned on six-hour trajectories; Transient, stiff initialization windows included in training. Envelope-edge RMSE and divergence are computed over the eight distinct cases; drift is averaged over the 29 runs, in which seven edge cases appear four times (Methods). Lower is better; best values in bold; the deployed model is shaded.

Figure 3: Six-hour fine-tuning keeps chained 30-min predictions stable. a , b , Normalized-log RMSE versus rollout time for the four variants in the six in-envelope (a) and eight envelope-edge (b) cases; lines are means over runs (b: 29 runs, eight distinct cases; Methods), shading the 10th–90th percentiles for 6h-NoStiff; coloured counts give diverged distinct cases. c , Hydroperoxyl radical (HO 2 ) over a six-hour rollout from a post-transient reference state: ARCA reference, 1h-Stiff and 6h-NoStiff (log axis); 1h-Stiff ends about 30 times too high. d , Scorecard of the RMSE of 6h-NoStiff relative to 1h-NoStiff from the same starting states, by chemical family (species counts in brackets) and rollout hour, pooled over the six in-envelope cases; blue, six-hour model more accurate; colours saturate at

\pm 100\%

. The variants in c differ in both training horizon and transient inclusion; Table 2 separates the two effects.

Figure 4: The deployed emulator reproduces the six-hour evolution of all 4,691 retained species. a , Two-dimensional density of emulator (6h-NoStiff) versus reference concentration for all retained species at every minute of the six in-envelope rollouts (

10^{7}

points; log colour scale); dashed, 1:1 line. b , Per-species

R^{2}

(normalized-log space) of all 4,691 retained species, ranked from best to worst; the vertical axis is scaled as

-\log_{10}(1-R^{2})

so that every species is shown. Dashed guides mark the share of species with

R^{2}\geq 0.99

(90.9%) and

R^{2}\geq 0.9

(99.5%); open dot, median (0.998); red dot, the only species with negative

R^{2}

(C 2 H 5 CHO,

-0.97

). c , Per-species relative concentration error by chemical family (median 1.9%, dashed line). Boxes, interquartile range; line, median; whiskers, 5th–95th percentiles; dots, species beyond them. d , Reference change from the starting concentration,

\log_{10}[c(t)/c_{0}]

, for every species (rows, grouped by family) in a typical rollout (the case with the median share within 0.5 dex). e , Error of 6h-NoStiff,

\log_{10}(\hat{c}/c)

, for the same case. f , As e for the one-hour model 1h-NoStiff from the same starting state. In d, blue marks a decrease and red an increase; in e and f, blue marks under- and red over-prediction. Boxed numbers give the share of species–minutes within 0.5 dex in this case; over all six rollouts the shares are 99.6% (6h-NoStiff) and 84.6% (1h-NoStiff). Colours saturate beyond the colour-bar limits.

Figure 5: Skill transfers to novel compositions but not to the most extreme forcing. a , Share of species whose six-hour per-species

R^{2}

exceeds a threshold, for the in-envelope anchors and 25 compositions far from every anchor (physical controls inside the training envelope); the curves separate only in the high-skill tail. b , Median per-species

R^{2}

by chemical family for both sets. c , Direction diagram for the ten extreme scenarios, each a realistic base state with one driver pushed to an extreme: six-hour O 3 change of 6h-NoStiff against that of the ARCA reference (symmetric-log axes, linear within

\pm 1

ppb). Points in the shaded same-sign quadrants capture the direction of the reference response (filled, 5 of 10); open circles mark missed directions; dashed, 1:1 line. Time series are shown in Supplementary Fig. 14 .

Figure 6: The emulator does not yet track time-varying forcing, but one call per 30 minutes makes it fast. a , Short-wave radiation over the six-hour morning ramp (temperature rises in parallel from 288 to 299 K). b – g , Key species for one representative in-envelope composition: reference (solid), emulator supplied with the time-varying forcing for each block (dashed) and a constant-forcing control (dotted). h , Median per-species

R^{2}

by chemical family under constant and time-varying forcing (all-species medians 0.998 and 0.870). i , Cost of a six-hour simulation as time per call against calls per simulation; grey diagonals are constant total time per six-hour case (labelled at 0.1 s and at the 150-s ARCA integration). FNO-MoE needs 12 calls: 0.20 s on one single-threaded CPU core (750 times faster than the 150-s single-core ARCA integration), 0.099 s on one GH200 GPU (1,500 times) and 0.35 ms per case amortized over a batch of 1,024 (

4.3\times 10^{5}

times). A step-wise emulator with a 4-min step needs 90 calls and would match FNO-MoE on the CPU only if each call took less than 2.2 ms. Data transfer and coupling overhead are excluded.

Regime	Sampling weight	Evaluation anchors
Clean biogenic summer day	0.20	550
Urban polluted	0.20	550
Cold low-biogenic winter night	0.15	400
Mixed background	0.15	400
Marine remote background	0.15	300
Biomass-burning plume	0.10	200

Supplementary Table 1: Seven hand-constructed atmospheric regimes define the anchor design. The weight sets each regime’s share of anchors. The anchor counts are fixed quotas from two generation batches, 500 anchors from an earlier four-regime design and 2,000 from the seven-regime design allocated as weight

\times

2,000, and therefore differ from weight

\times

2,500. The same anchors are used for acquisition and evaluation; regime-specific ranges are listed in Supplementary Table 2 .

Physical controls
Parameter	Regime
Parameter	Clean	Urban	Winter	Mixed	Marine	Biomass	Polar
Design weight	0.20	0.20	0.15	0.15	0.15	0.10	0.05
$T$ (K)	290–310	–	260–280	–	–	285–305	260–270
SW (W m -2 )	500–900	–	0–100	–	–	300–900	0–5
RH	0.2–0.55	–	0.6–0.99	–	0.7–0.99	0.2–0.5	0.6–0.99

Supplementary Table 2: Regime-specific parameter ranges for the anchor design. Sampling range within each regime for the four physical controls and nine signature precursors; a dash means the regime inherits the global range. Ranges are hand-constructed from atmospheric-chemistry relationships, not fitted to measurements, and are clamped inside the global bounds. Controls: temperature

T

, short-wave radiation SW, relative humidity RH and condensation sink CS (ARCA inputs TEMPK, SW_RADIATION, REL_HUMIDITY and CONDENS_SINK). Global ranges, used where a regime does not override and for whole-space random sampling:

T

260–310 K; SW 0–900 W m -2 ; RH 0.2–0.99; CS

10^{-4}

–

10^{-2}

s -1 ; and, in ppb, NO 0.01–10, NO 2 0.05–10, O 3 10–80, CO 50–500, SO 2 0.02–5,

\alpha

-pinene 0.01–2, isoprene 0.01–5, toluene 0.01–5 and DMS 0.0022–0.22. The full 146-species precursor vector carries further regime overrides, given in the released configuration.

Dataset	Simulations	Duration	Role
Training and fine-tuning
Whole-space random	10,000	1 h	Random-sampler training
Static local sampling	10,000	1 h	Static-sampler training
Active local sampling	10,000	1 h	1h-Stiff and 1h-NoStiff training
Long-trajectory anchors	2,500	6 h	6h-NoStiff and 6h-Stiff fine-tuning
Evaluation

Supplementary Table 3: Reference-trajectory datasets and their roles.

Control	Training envelope	Direction	Simulated value (all $d$ )
Temperature	260–310 K	low, high	260 / 310 K
Relative humidity	0.2–0.99	low	0.2
Short-wave radiation	0–900 W m -2	high	900 W m -2
Condensation sink	$10^{-4}$ – $10^{-2}$ s -1	low, high	$10^{-4}$ / $10^{-2}$ s -1
Corner (all four)	—	outward	310 K, 0.2, 900 W m -2 , $10^{-2}$ s -1

Supplementary Table 4: Construction of the envelope-edge rollout cases. The generator was designed to push the physical controls beyond the training envelope by a distance

d

(half-ranges beyond the edge for the linear controls, decades for the condensation sink), one control at a time plus an all-control corner, with the 146 initial concentrations at the median-anchor base. Every control was, however, clipped to the sampling bounds, so all four distances of a direction produced the same simulation at the edge. The 29 runs therefore contain eight distinct cases: the centred control point (285 K, RH 0.595, 450 W m -2 ,

5.05\times 10^{-3}

s -1 ) and the seven edge settings listed below.

Acquisition strategy	Overall RMSE $\downarrow$
Whole-space random	0.0528
Single-chain Markov-chain Monte Carlo	0.0335
Static local sampling	0.0308
Active local sampling (best-of- $N$ disagreement)	0.0271

Supplementary Table 5: Coverage-guaranteed active sampling outperforms a single-chain Metropolis strategy. Overall normalized-log RMSE (all species, full trajectory) at a matched 10,000-simulation budget; this metric is shown as vertical lines in Supplementary Fig. 2 . Only the scalar RMSE was retained for the single-chain strategy, which visited about 30 of the 2,500 anchors per round. Lower is better; best value in bold.

	SIMADy Yang et al. (2025)	This work
Reference	PartMC–MOSAIC with MCM	ARCA-box with MCM
Representation	Six latent species	4,691 retained species as direct outputs
Dynamics	Sparse mass-action system, integrated in time	Fourier mixture of experts, 30 states per call
Training data	Latin-hypercube scenarios	Emulator-selected simulations
Stability	Bounded by construction	Empirical six-hour rollout tests
Evaluation	Time-varying multi-day forcing	Constant physical controls; one time-varying test

Supplementary Table 6: Direct prediction and SIMADy make complementary design choices. Methodological distinctions, not a matched performance ranking.

Supplementary Fig. 1: Regime anchors combine physical drivers and precursor mixtures that co-vary in the atmosphere. Six reproducible draws for each of the seven hand-constructed regimes, showing the joint ranges of physical drivers (a,b) and chemically informative initial precursors (c,d). Variables without a regime-specific range keep the global bounds. The panels document the sampling design, not a measurement-derived distribution.

Supplementary Fig. 2: Active sampling shifts the per-species error distribution towards lower error. Histograms of per-species normalized-log RMSE for the three 10,000-simulation strategies; vertical lines mark each strategy’s overall RMSE (Supplementary Table 5 ). Paired species-level differences would be needed to identify individual species that regress.

Supplementary Fig. 3: The emulator reproduces the principal precursor decay at evaluation anchors. Minute-resolved ARCA solutions (solid black) and active-sampling emulator predictions (dashed blue) from the same initial retained state. Columns: clean-biogenic, urban-polluted and biomass-burning anchors; rows: O 3 , isoprene,

\alpha

-pinene and SO 2 (ppb). Deviations include the curvature of O 3 . Illustrative traces, not an additional accuracy metric.

Supplementary Fig. 4: Six-hour fine-tuning prevents divergence in the in-envelope rollouts. Final normalized-log RMSE (a), elemental drift (b) and number of diverging cases out of six (c) for the four variants. Panels a and b use logarithmic axes.

Supplementary Fig. 5: The stability gain persists at the edges of the training envelope. As Supplementary Fig. 4 , for the 29 envelope-edge runs, which comprise eight distinct cases because the intended extrapolation distances were clipped to the sampling bounds (Supplementary Table 4 ). RMSE (a) and divergence counts (c) are over the eight distinct cases (7/8, 8/8, 0/8 and 0/8); the elemental drift (b) is averaged over the 29 runs because per-case drift was not stored.

Supplementary Fig. 6: Rollout error grows gradually over six hours. Per-case normalized-log RMSE of 6h-NoStiff versus rollout time for six in-envelope cases (a) and the 29 envelope-edge runs (b; eight distinct cases). Solid line, median; shading, minimum–maximum range and empirical 10th–90th percentiles. These ranges describe the tested cases and are not confidence intervals; the lower aggregate error at the envelope edges should not be read as evidence that extrapolation is easier.

Supplementary Fig. 7: Six-hour rollout trajectories reproduce chemistry across families. Predicted (6h-NoStiff, dashed) and reference (ARCA, solid) concentrations over a representative in-envelope six-hour simulation for twelve species spanning inorganic, radical, biogenic, aromatic, oxygenated and sulfur families. Axes are logarithmic where the dynamic range exceeds a factor of 40.

Supplementary Fig. 8: Skill is high in every chemical family. a , Per-species

R^{2}

(normalized-log space, six-hour rollouts of 6h-NoStiff) by chemical family (species count in brackets); the axis is scaled as

-\log_{10}(1-R^{2})

as in Fig. 4 b, and values below 0.9 are drawn at the left edge. b , Per-species relative concentration error (% of reference ppb) by family, logarithmic axis. Boxes, interquartile range; line, median; whiskers, 5th–95th percentiles; dotted lines, all-species medians (0.998 and 1.9%). c , Percentage of all 4,691 retained species whose

R^{2}

exceeds a threshold

\tau

(99.5% at 0.9, 90.9% at 0.99). Relative errors are largest for low-abundance families.

Supplementary Fig. 9: Batching amortizes GPU overhead. Emulator speedup relative to the single-core ARCA reference (orange line, 150 s per six-hour case) versus batch size on one GH200 GPU, using the amortized time per case; the dashed line is the emulator on one single-threaded CPU core. Coupled-model acceleration also depends on data movement, other processes and scheduling.

Supplementary Fig. 10: The configured mechanism is a dense network organized around a few inorganic and radical hubs. Species–species graph in which a directed link joins a reactant to a product of the same reaction (6,068 species, 18,490 reactions, 31,156 reactant–product links). (a) Number of links from each reactant family (rows) to each product family (columns), on a logarithmic colour scale; counts of at least 100 are printed. (b) Distribution of the number of links per species; most species have fewer than 20 links, whereas OH (3,686), HO 2 (3,424), NO 2 (2,626), NO (1,623), NO 3 (1,606) and CO (950) connect to a large part of the mechanism. The emulator predicts the filtered retained state.

Supplementary Fig. 11: A low-dimensional linear code reconstructs the retained state with about 8% error at every tested latent dimension. Reconstruction error of the SIMADy-inspired linear autoencoder (% of RMS concentration) versus the number of latent variables, for all species, key species and ozone. Dotted lines mark the 30-min prediction errors of 6h-NoStiff in the same units (0.36% for all species, 0.95% for ozone); the all-species and key-species curves coincide.

Supplementary Fig. 12: The expert gate is used jointly and does not collapse. Experts are coloured light to dark by retained Fourier modes (4, 8, 12, 16), as in Fig. 1 . a , Mean softmax gate weight per expert; dotted line, uniform value 1/4. b , Distribution of per-species gate entropy (logarithmic count axis), normalized so that 0 is single-expert routing and 1 a uniform mix (mean 0.99). c , For each chemical family (species count in brackets), the share of species whose highest-weighted expert is each of the four.

Supplementary Fig. 13: Finite-horizon sensitivities of the emulator. Mean magnitude of the derivative of predicted species

i

with respect to initial species

j

,

|\partial\widehat{X}_{i}/\partial x_{0,j}|

, averaged over 200 states for 22 inorganic, radical and biogenic species; outlined cells mark pairs directly linked by an MCM reaction, and the self-sensitivity diagonal (off scale) is hatched grey; colours use a logarithmic scale clipped at the 5th percentile of the off-diagonal values, so that structure outside the strongly sensitive

\beta

-caryophyllene row remains visible. Linked pairs are only 1.1 times as sensitive as unlinked pairs on average: sensitivities respond broadly to inorganic and radical inputs rather than to individual reaction links (Supplementary Note Supplementary Note 2: learned sensitivities do not map onto individual reaction links ).

Supplementary Fig. 14: Ozone trajectories in the ten extreme scenarios. Six-hour ozone from the reference (ARCA, solid) and the deployed 6h-NoStiff model (dashed) for each scenario, a realistic base state with one driver pushed to an extreme (Methods). Symbols mark whether the emulator captures the direction of the reference ozone change (5 of 10); panels marked “off scale” are clipped to the reference range. Vertical scales differ between panels. The direction diagram in Fig. 5 c summarizes these trajectories.

An actively trained neural operator for near-explicit atmospheric chemistry

Abstract

Figures & tables

Neural Surrogate Modeling

Toward Mechanistic Interpretability of an AI Foundation Model Fine-Tuned for Atmospheric Chemistry

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

AeroChem: Closed-loop Physics-Informed State Space Modeling for Long-term Chemically-Reactive Air Quality Forecasting--Extended Version

Abstract

Neural Surrogate Modeling

Explore similar work

Toward Mechanistic Interpretability of an AI Foundation Model Fine-Tuned for Atmospheric Chemistry

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

AeroChem: Closed-loop Physics-Informed State Space Modeling for Long-term Chemically-Reactive Air Quality Forecasting--Extended Version