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

AI summary

Overview

Research area: Machine learning for spatio-temporal environmental forecasting — specifically physics-informed deep learning applied to chemically reactive air quality prediction over multi-day horizons.

Technical level: Advanced. The paper assumes familiarity with neural differential equations, state space models (Mamba), and Kalman-style recursive filtering, though the core idea is explained here in plain language.

Scope: The paper proposes AeroChem, a closed-loop framework that repeatedly corrects a physics-based atmospheric simulation with data-driven "virtual measurements" inside the latent state, and introduces a new regional dataset (HanoiAir) for benchmarking long-horizon pollutant forecasting.

What This Paper Is About

Forecasting pollutants such as PM2.5 and ozone days in advance is hard because pollution is produced by nonlinear chemical reactions, blown around by wind, and driven by emissions that are only partly known. Traditional chemical transport models encode the physics but rely on simplified assumptions and expensive solvers, while purely data-driven models are accurate short-term but drift badly when rolled forward for many steps. The paper's goal is to make long-horizon forecasts stay physically consistent and stable by letting physics and data correct each other at every forecasting step, rather than only merging them at the final output.

Key Contributions

1. A problem reformulation. Long-term chemically reactive air quality forecasting is cast as closed-loop latent-state inference under imperfect physical dynamics — the correction happens in latent state space before the next simulation step, not on the final prediction.

2. The AeroChem architecture. A two-branch framework: a Universal Differential Equation (UDE) process branch acting as a differentiable, structured-but-imperfect physical prior, and an S-Mamba-based virtual measurement branch producing probabilistic virtual observations, coupled by a Neural Kalman Fusion module with a learnable affine correction term.

3. A meteorology-gated neural reaction mechanism. Secondary pollutant formation is modeled with self-attention over concatenated latent state and emissions, element-wise gated by scalar meteorological variables (e.g., radiation, temperature) — a data-driven analog of reaction-rate modulation, rather than fixed stoichiometry or a generic neural residual.

4. The HanoiAir dataset. A new real-world hourly dataset for northern Vietnam with 27 monitoring stations spanning 2018–2025, integrating pollutant concentrations, meteorological variables, and downscaled emission inputs.

5. Empirical validation. Experiments on HanoiAir and the public KnowAir benchmark (355 stations, 2016–2023) showing improved long-horizon stability and sudden-change forecasting, plus ablations isolating each physical mechanism and the fusion module.

Main Findings

• Long-horizon advantage on KnowAir: AeroChem reaches MAE 23.00 at 6 days and 24.08 at 9 days, versus the strongest baselines at those horizons (TimeMoE 23.15 and 25.32; DLinear 23.85 and 24.50). Its gains widen as the horizon lengthens, which is the paper's central claim.

• Sudden-change events are where it separates most: On KnowAir, AeroChem scores MAE 34.50 / RMSE 49.50 / SMAPE 0.52 on sudden-change instances, against 36.54 / 49.12 / 0.53 for TimeMoE and 37.92 / 50.81 / 0.54 for ModernTCN. On HanoiAir the gap is larger: 68.80 / 91.16 / 0.75 versus 71.54 / 93.12 / 0.78 for TimeMoE.

• Short horizons are not a clean win: On KnowAir at 3 days, several purely data-driven baselines have lower MAE than AeroChem's 21.95 (Fredformer 21.45, TimeMoE 21.68). On HanoiAir at 3 days, Fredformer (41.60), TimeMoE (41.85), Crossformer (42.00), ModernTCN (42.33) and DLinear (42.57) all beat AeroChem's 42.74 MAE — though AeroChem still has the best 3-day RMSE there (63.11 vs 63.80 for TimeMoE).

• Removing the closed-loop fusion hurts more than removing the physics: In the ablation, dropping Neural Kalman Fusion (AeroChem-K) is the largest degradation by far — HanoiAir 3-day MAE rises to 56.40 and 9-day to 63.50, versus 42.74 and 46.80 for the full model. That is a larger drop than removing physics entirely (AeroChem-P: 47.85 and 53.20).

• All three physical mechanisms contribute, with reactions mattering most: On HanoiAir 3-day MAE, removing diffusion gives 44.31, advection 42.81, and the reaction operator 45.15, against 42.74 for the full model. Removing all physics gives 47.85.

• Element-wise Kalman fusion beats an ensemble-style alternative on MAE: AeroChem-EnKF records HanoiAir 3-day MAE 43.62 and 9-day 47.85 versus 42.74 and 46.80 for AeroChem, and KnowAir 3-day MAE 24.08 versus 21.95. (The remaining visible cells for this variant are cut off in the provided text.)

• Theory supports the design: The appendix sketches Neural Kalman Fusion as an approximate variational correction of an imperfect physical prior — reducing to a diagonal Kalman update when residuals are Gaussian, and motivating the learnable affine term when they are not — and argues that closed-loop posterior error stays uniformly bounded under Lipschitz and bounded-noise assumptions when the correction factor is contractive, whereas the open-loop variant can amplify initial errors.

• Not reported in the provided text: Results for CRPS, CRPS_sum, and PICP are defined as evaluation metrics but their tables are absent from the excerpt; hyperparameter values and full preprocessing details are deferred to appendices that were truncated.

Methodology in Plain English

The system runs two models side by side, step by step, for the whole forecast.

Branch one — a differentiable simulator. The latent atmospheric state is evolved as a continuous-time stochastic differential equation with three named pieces. Diffusion spreads pollution from high to low concentration along the sensor graph. Advection moves it with the wind: edge weights are computed from wind velocity projected onto the direction between two stations, so transport is directional and changes hourly. Reactions handle secondary pollutant formation: the latent state is concatenated with emission inputs, self-attention learns which species interact, and the result is multiplied element-wise by a gate computed from scalar meteorology, so chemistry switches on or off depending on conditions like sunlight and temperature. A stochastic term (Wiener process) absorbs what the deterministic physics cannot see — unresolved variability, boundary leakage, unobserved emissions.

Branch two — a virtual sensor. An S-Mamba encoder (a selective state space model) reads the historical window and produces an embedding. Two projection heads turn that into a base virtual observation plus a base uncertainty. During rollout, the virtual observation is nudged by a bounded function of the previous corrected state, and the uncertainty is inflated over the horizon by a learnable factor — modeling growing epistemic uncertainty without ever looking at future pollutant readings.

The handshake — Neural Kalman Fusion. The simulator's prediction becomes a prior mean, with a process uncertainty produced by a learned linear layer plus Softplus. A Kalman-style gain divides prior uncertainty by total uncertainty, and the innovation — virtual observation minus prior — is corrected by an extra learnable affine term that handles non-Gaussian residuals. The resulting posterior state is fed back as the starting point of the next simulation step. That feedback loop is the whole point: the data branch pulls the physics back on track, and the physics branch smooths and regularizes noisy data-driven evidence.

Output and training. Two shallow MLPs decode the posterior mean and covariance into a predictive mean and variance, giving a Gaussian predictive distribution. The loss combines MAE, a heteroscedastic Gaussian negative log-likelihood, and a Frobenius-norm penalty on drift magnitude. Diagonal covariance approximation keeps the fusion at O(N) instead of O(N³). Scheduled sampling is used so the model starts training with the virtual observation stop-gradient-ed and gradually switches to full closed-loop operation, mitigating exposure bias.

Evaluation. Four scenarios — 3 days (1–72 h), 6 days (1–144 h), 9 days (1–216 h), and sudden change, where ground-truth PM2.5 exceeds 75 µg/m³ and shifts by more than ±20 µg/m³ within three hours. Metrics: MAE, RMSE, SMAPE, CRPS, CRPS_sum, PICP, averaged over 5 runs after inverse transformation to physical scale. Chronological splits prevent leakage: KnowAir trains 2016–2019, validates 2020–2021, tests 2022–2023; HanoiAir trains 2018–2021, validates 2022–2023, tests 2024–2025.

Why This Matters

For research: The paper reframes the PIML debate. Most hybrid air quality models fuse physics and neural components only at the output level; AeroChem argues that the correction must happen in the latent state before propagation, and backs it with an ablation where the fusion module matters more than any individual physical term. That is a transferable design principle for any domain with imperfect simulators — climate, hydrology, epidemiology, battery modeling.

Real-world applications:

• Public health early warning — multi-day PM2.5 and ozone alerts give schools, hospitals, and vulnerable populations time to act, and the sudden-change scenario targets exactly the spikes that matter.
• Emission control and policy — the model consumes emission inventories as explicit forcing, so regulators can probe how NOx, VOCs, SO2, and primary PM2.5 reductions propagate into secondary pollution.
• Regional air quality management in data-sparse settings — HanoiAir demonstrates the framework on a 27-station network with downscaled emission inputs, a far cry from the dense monitoring of KnowAir's 355 stations.
• Operational forecasting with uncertainty — the predictive variance supports probabilistic decision-making (PICP, CRPS) rather than single-point guesses.

Industry relevance: Air quality forecasting is a commercial service for utilities, insurers, logistics, outdoor-event planning, and smart-city platforms. A model that degrades gracefully over multi-day horizons and carries calibrated uncertainty is more deployable than one that is only accurate for the next few hours. The emission-inventory interface also makes it attractive for environmental consultancies doing scenario analysis.

Future Directions

• Closing the short-horizon gap. AeroChem does not dominate at 3 days, where Fredformer and TimeMoE are competitive or better on MAE. Understanding when the physics prior helps versus when it adds bias — and possibly blending regimes — is an open engineering question.

• Replacing the diagonal covariance approximation. The O(N) diagonal simplification enables scalability but discards cross-station and cross-species correlations that matter for chemically coupled pollutants. Structured or low-rank covariances are the obvious next step.

• Learning the reaction structure explicitly. The attention-based reaction operator learns "conditional interaction affinities" rather than stoichiometric coefficients. Recovering interpretable, physically checkable reaction terms would strengthen both trust and scientific utility.

• Broader geographic and chemical validation. Results rest on two datasets, one of them newly introduced by the authors. Independent replication across more regions, more pollutants, and different emission-inventory quality levels is needed — as is a closer look at whether the gains persist when emission inventories are highly uncertain, which the paper's own motivation suggests is the common case.

Target Audience

Researchers and graduate students in physics-informed machine learning, spatio-temporal forecasting, and atmospheric/environmental modeling; practitioners building operational air quality or climate forecasting systems who need multi-day stability; and data assimilation specialists interested in how learned, rather than hand-specified, uncertainty models can replace classical filtering assumptions. Readers should be comfortable with differential equations and deep sequence models — the paper provides a plain-language on-ramp via its architecture figure and theoretical appendix, but the density is high.

Authors’ abstract

Accurate long-term air quality forecasting is challenging due to long-range transport, meteorological variability, and nonlinear chemical reactions underlying secondary pollutant formation. Purely data-driven models often suffer from long-horizon drift and lack physical consistency, while existing physics-informed approaches commonly follow open-loop designs that fuse neural and physical components only at the output level. We propose AeroChem, a closed-loop physics-informed state space framework for chemically reactive air quality forecasting. AeroChem adopts a two-branch co-correction design: a UDE-based process branch provides a structured but imperfect physical prior, while a Mamba-based virtual measurement branch captures long-range temporal context and produces probabilistic virtual observations. Neural Kalman Fusion couples the two branches by correcting the latent physical state at each forecasting step and feeding the posterior state back into the simulator, reducing long-horizon drift. We further introduce HanoiAir, a real-world dataset integrating pollutant concentrations, meteorology, and emission inventories. Experiments on HanoiAir and a public benchmark show that AeroChem improves long-horizon stability and sudden-change forecasting.