Source-linked AI summary
AeroChem: Closed-loop Physics-Informed State Space Modeling for Long-term Chemically-Reactive Air Quality Forecasting--Extended Version
Kien Le Trung, Tung Kieu, Dinh Duc Nha Nguyen, Cuong Do, Bin Yang, Phi Long Nguyen
TL;DR
Long-term air-quality forecasting must handle atmospheric transport, meteorological variability, and nonlinear chemistry, while existing data-driven and physics-informed approaches can drift or fuse their components only at the output. AeroChem uses a UDE process branch, Mamba-based virtual measurements, and recurrent latent-state correction; experiments report improved long-horizon stability and sudden-change forecasting.
Problem
Long-term air-quality forecasting must model transport, meteorology, and nonlinear chemistry, while existing data-driven models can drift and many physics-informed hybrids use open-loop, output-level fusion.
Source: §1 · pp. 1–2Method
AeroChem fuses a UDE-based process prior with Mamba-based probabilistic virtual measurements and recursively feeds each corrected posterior state back into the process model.
Source: §1 · p. 2Results
Experiments on KnowAir and HanoiAir report improved long-horizon stability and sudden-change forecasting, especially under noisy observations and secondary pollutant formation.
Source: §1 · p. 2, §6 · p. 10Takeaways & Limitations
Results on KnowAir and HanoiAir support closed-loop neural state correction as a direction for long-term air-quality forecasting.
Source: §6 · p. 10The process model simplifies atmospheric chemistry and transport, and its learned reaction operator may limit extrapolation to unseen regimes.
Source: Limitations · p. 11
Abstract
from arXiv · showhide
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.
1 Introduction
Long-term air-quality forecasting must account for transport, meteorology, and nonlinear chemistry, yet existing approaches can either simplify these processes or accumulate errors. AeroChem combines a structured physical prior with data-driven virtual measurements and recurrent state correction; experiments report gains in long-horizon stability and sudden-change forecasting.
Forecasting challenges
- Traditional chemical transport models encode atmospheric processes with physical differential equations but rely on simplified assumptions, incomplete reactions, uncertain boundary conditions, and expensive solvers.Data-driven models can achieve strong short-term accuracy but may violate physical constraints, generalize poorly under distribution shifts, and accumulate errors during long autoregressive rollouts.Source: §1 · p. 2
- Existing hybrid methods often fuse physical and data-driven branches only at the output, leaving errors in the latent physical state to propagate into future steps.The introduction argues that correction should instead occur in latent state space before the next simulation step.Source: §1 · p. 2
AeroChem framework
- AeroChem corrects the latent physical state at each forecasting step by fusing a UDE process prior with Mamba-based virtual measurements, then feeds the posterior state back into the simulator.Its process branch separates transport from chemical dynamics and uses meteorology-gated neural reactions.Source: §1 · p. 2, Figure 2 · p. 3
Evaluation and contributions
- Experiments on two real-world air-quality datasets, including newly constructed HanoiAir, show improved long-horizon forecasting and performance in sudden pollution-change scenarios.HanoiAir integrates pollutant concentrations, meteorological variables, and emission inventories.Source: §1 · p. 2
2 Preliminaries
The forecasting problem represents monitoring stations as a weighted graph and combines pollutant, meteorological, and emissions data in multivariate observations. Given the graph and historical observations, the objective is a probabilistic forecast over a future horizon.
Sensor Network
- A weighted static graph represents monitoring stations as nodes and uses its adjacency matrix to encode spatial correlations, such as geographic proximity.Each station also has fixed spatial coordinates.Source: §2 · p. 3
Multivariate Observations
- Each time-step observation combines pollutant concentrations, meteorological variables, and emission rates; concentrations are forecast targets, while emissions represent external forcing.The described inputs include target pollutants and co-pollutants, environmental drivers, and estimated releases of pollutants or precursors.Source: §2 · p. 3
Problem Formulation
- Given the sensor graph and historical observations, the objective is a probabilistic distribution of future observations over the prediction horizon.This formulation supports both point prediction and uncertainty estimation.Source: §2 · p. 3
3 Methodology
AeroChem performs closed-loop forecasting by combining a physics-informed process prior with context-conditioned virtual observations and feeding corrected latent states into subsequent steps. Its method specifies the branch designs, state correction, output decoding, theoretical interpretation, and training procedure.
3.1 Framework Overview
- AeroChem couples a physics-informed process branch, a virtual measurement branch, and Neural Kalman Fusion, recursively feeding the corrected posterior state into the process branch.This closed-loop co-correction is intended to mitigate error accumulation over long horizons.Source: §3 · pp. 3–4, §3.1 · p. 4
3.2 Virtual Measurement Branch: Context-Conditioned Virtual Observations
- S-Mamba encodes historical observations into long-range context, which projection heads use to produce a base virtual observation and its uncertainty.The virtual measurements adapt to the evolving trajectory without future pollutant observations, while the uncertainty grows over the forecast horizon.Source: §3.2 · p. 4
3.3 Process Model: Universal Differential Equation
- The UDE process model separates deterministic transport and chemical drift from stochastic diffusion, using meteorology and emissions as exogenous drivers.Its drift includes graph-based diffusion, wind-driven advection, and a meteorology-gated neural reaction mechanism for nonlinear chemical interactions.Source: §3.3 · pp. 4–5
3.4 Neural Kalman Fusion
- Neural Kalman Fusion uses a diagonal covariance approximation and a learnable correction term to adapt state updates to non-Gaussian residuals.The posterior state is fed back as the next prediction's initial condition, allowing virtual measurements and the UDE prior to correct each other during rollout.Source: §3.4 · p. 5
3.5 Theoretical Interpretation
- The theoretical interpretation treats fusion as an approximate variational correction of an imperfect physical prior, with a diagonal Kalman correction for Gaussian residual potential.Under Lipschitz and bounded-noise assumptions, posterior error remains uniformly bounded when the correction factor is contractive; the open-loop variant can amplify initial errors.Source: §3.5 · p. 5
3.6 Output Generation
- A decoder maps posterior latent states and their covariance matrices into predictive means and variances in observation space.The output is a parameterized predictive distribution intended to quantify epistemic uncertainty.Source: §3.6 · p. 6
3.7 Training Objective
- Training minimizes a combined objective with an MAE term, a heteroscedastic Gaussian negative log-likelihood, and a physics regularizer on drift magnitude.Scheduled sampling begins by conditioning the process model on observed states with exponentially decaying probability before transitioning to fully closed-loop operation.Source: §3.7 · p. 6
4 Related Work
Prior air-quality forecasting spans statistical, deep-learning, and physics-informed approaches, with trade-offs in temporal modeling, physical structure, and computational assumptions. AeroChem distinguishes itself by using virtual measurements to repeatedly correct a structured process model's latent state during rollout.
Spatio-temporal Air Quality Forecasting
- Statistical and deep-learning approaches model predictive patterns, including spatiotemporal dependencies, but often leave transport and chemical dynamics unmodeled and forecast autoregressively.The cited discussion notes that such models may accumulate long-horizon errors or smooth abrupt secondary-pollution events.Source: §4 · p. 6
Physics-Informed Machine Learning
- Physics-informed approaches incorporate physical knowledge, but air-quality methods may simplify reactions or loosely couple neural and physical components, leaving latent-state errors to propagate.The passage contrasts AeroChem's repeated latent-state correction using context-conditioned virtual measurements with output-level coupling.Source: §4 · p. 6
Deep Data Assimilation and Kalman Filtering
- Classical data assimilation and Kalman filters rely on linear-Gaussian assumptions or costly ensembles, while learning-based filters typically use real observations or a standard observation model.Because future pollutant observations are unavailable during forecasting, AeroChem generates virtual observations from historical context and the previous posterior state.Source: §4 · pp. 6–7
5 Experiments
Experiments evaluate AeroChem on KnowAir and HanoiAir across forecast horizons and sudden-change events, with results emphasizing long-horizon stability and sudden-change forecasting. Ablations and sensitivity studies examine the physical mechanisms, closed-loop fusion, and hyperparameter choices.
5.1 Experimental Settings
- Evaluation uses hourly KnowAir and HanoiAir data with chronological train, validation, and test splits, and covers 3-day, 6-day, 9-day, and Sudden Change scenarios.Sudden Change instances have ground-truth PM2.5 above 75 µg/m3 and a change exceeding ±20 µg/m3 within the next three hours.Source: §5.1 · p. 7
5.2 Main Results
- AeroChem is most effective in long-horizon and sudden-change PM2.5 forecasting; its advantage is more evident on HanoiAir and its performance remains more stable than direct-mapping models as horizons lengthen.Tables 1 and 2 report the main results; the text also reports stronger performance than Transformer and SSM baselines in sudden-change settings.Source: §5.2 · pp. 7–9, Table 1 · p. 8, Table 2 · p. 8
5.3 Ablation Study
- Removing posterior feedback causes the largest degradation, especially at longer horizons; removing physical terms, the reaction module, or the learnable neural correction also degrades performance.Late Fusion performs worse than latent-state fusion, while a strictly Gaussian correction is described as insufficient for open environmental systems.Source: Table 3 · p. 8, §5.3 · p. 9
5.4 Effect of Forecast Horizons
- AeroChem’s stability advantage over physics-guided baselines widens as the prediction window extends, while those baselines show rapid error propagation.Figure 3 illustrates error evolution across forecast horizons.Source: §5.4 · p. 9, Figure 3 · p. 10
5.5 Hyperparameter Studies
- The default λ_phy = 0.1 balances physical consistency and data flexibility, d_model = 96 is the point of performance saturation, and fusion maintains accuracy across solver step sizes.Figure 4 reports the sensitivity analysis on KnowAir and marks the default hyperparameters at optimal performance regions.Source: §5.5 · p. 9, Figure 4 · p. 10
5.6 Additional Experiments
- Additional experiments examine numerical stability and internal mechanisms; sudden-event visualizations show that AeroChem can track abrupt shifts but may undershoot extreme peak amplitudes.The numerical-stability analysis and mechanism case study are in the appendix.Source: §5.6 · p. 10
6 Conclusion
AeroChem couples physical dynamics and virtual measurements through closed-loop latent-state correction, with experiments on KnowAir and HanoiAir showing improved long-horizon stability and sudden-change forecasting. The paper identifies richer chemical interaction modeling and probabilistic risk assessment as future work.
6 Conclusion
- AeroChem’s closed-loop neural state correction improves long-horizon stability and sudden-change forecasting, especially under noisy observations and secondary pollutant formation.The framework feeds posterior latent states back into the simulator during rollout.Source: §6 · p. 10
Broader Impact.
AeroChem may help agencies anticipate persistent haze, secondary pollutant formation, and sudden pollution changes, particularly in rapidly urbanizing or data-sparse regions. Its closed-loop co-correction approach may also be useful in other partially observed systems with imperfect physical simulators.
Broader Impact
- By reducing long-horizon drift, AeroChem may support anticipation of haze episodes, secondary pollutant formation, and sudden pollution changes in data-sparse or rapidly urbanizing regions.The paper also identifies climate modeling, hydrology, and energy systems as possible domains for similar approaches.Source: Broader Impact · p. 11
Limitations.
AeroChem’s limitations concern its approximate process model, restricted uncertainty representation, and dependence on auxiliary inputs.
Limitations
- The process model approximates atmospheric chemistry and transport, using effective reaction dynamics rather than explicitly resolving full chemical mechanisms.This may limit extrapolation to unseen regimes.Source: Limitations · p. 11
- Learned diagonal covariances do not capture full spatial correlations and may underestimate structured error.Source: Limitations · p. 11
- The framework relies on meteorological and emission inputs, whose inaccuracies or delays can propagate through the system.Source: Limitations · p. 11
A.1 Algorithmic Summary of AeroChem
Algorithm 1 summarizes AeroChem’s closed-loop forecasting procedure, in which posterior latent states are assimilated and fed back into the UDE solver. It takes historical observations and forecast-horizon meteorology and emissions as inputs and produces probabilistic forecasts.
A.1 Algorithmic Summary of AeroChem
- At each forecasting step, an adaptive virtual observation and UDE prediction feed Neural Kalman Fusion, whose posterior state returns to the simulator.The algorithm highlights this recursive feedback over the forecasting horizon.Source: §A.1 · p. 14
- The procedure takes a sensor graph, historical observations, forecast-horizon meteorology and emissions, and a forecast horizon as inputs, and outputs forecasted observations.Source: §A.1 · p. 14
- It initializes the posterior latent state, for example by projecting or encoding the last observed timestep.Source: §A.1 · p. 14
- The procedure decodes the corrected state into a probabilistic output.Source: §A.1 · p. 14
A.2 Variational Interpretation of Neural Kalman Fusion
The appendix recasts AeroChem as a variational objective over a latent trajectory given historical observations. It relates the Neural Kalman Fusion update to a Newton step and identifies the Gaussian special case.
A.2 Variational Interpretation of Neural Kalman Fusion
- AeroChem is recast under a single variational objective, with a per-step variational free energy defined over the latent trajectory and historical observations.Source: §A.2 · p. 14
- The objective’s terms include the UDE-predicted prior, the Mamba encoder’s virtual measurement, and a smooth residual potential for non-Gaussian discrepancies.The potential absorbs discrepancies such as boundary leakage, unobserved emissions, and model misspecification.Source: §A.2 · p. 14
- Under the diagonal-covariance regime, the equation assumes positive element-wise prior and observation covariances.Source: §A.2 · p. 14
- When the residual potential is zero, the update reduces exactly to AeroChem’s standard Kalman update.Source: §A.2 · p. 14
- With a nonzero residual potential, a first-order Taylor expansion makes the Newton step affine, and substitution yields the Neural Kalman Fusion update.The covariance update follows the standard Kalman posterior covariance identity under optimal gain.Source: §A.2 · p. 15
A.3 Closed-loop Neural State Correction
This appendix analyzes conditions for bounded posterior error under closed-loop correction and reports additional performance, physical-behavior, dataset-scope, and computational analyses. It also documents reproducibility, public access to code, HanoiAir, and the work’s ethics, impact, and safeguard status.
A.3 Closed-loop Neural State Correction
- When α < 1 under Assumption A.2, AeroChem’s closed-loop posterior error is uniformly bounded over forecast steps, whereas the open-loop variant has an exponential lower bound when LF > 0.The contraction factor is α := (1 + LF ∆t)[(1 − κmin) + κmaxLc].Source: §A.3 · pp. 15–16
A.4 Computational Complexity
- AeroChem’s diagonal-covariance Kalman correction costs O(ND), leaving ODE integration dominant, and its training curve shows stable convergence over 100 epochs with normalization supporting numerical stability.The dominant solver cost scales as O(τ · L · (N · D + |E| · D)); diagonal covariance reduces Kalman overhead from O(N^2D^2) to O(ND).Source: §A.4 · p. 16, Figure A.1 · p. 17, §A.10 · p. 19
A.5 HanoiAir Dataset Details
- HanoiAir combines hourly observations from 27 northern Vietnamese stations spanning 2018–2025 with meteorology and spatially allocated emissions, including precursor-gas measurements for modeling secondary chemistry.The listed pollutants include PM2.5, O3, PM10, NOx, SO2, CO, NO, and NO2; emissions include NOx, VOC, SO2, NH4, and primary PM2.5.Source: §A.5 · p. 17
A.9 Probabilistic forecasting performance
- Under a unified uncertainty protocol, AeroChem achieves the lowest CRPS and CRPSsum across KnowAir horizons, while its lower PICP reflects a sharpness–calibration trade-off.Point-forecasting baselines receive post-hoc Gaussian predictive layers fitted from validation residuals; VAR uses closed-form Gaussian uncertainty propagation.Source: §A.9 · pp. 18–19, Table A.3 · p. 18
A.11 Effect of Sudden Changes
- In HanoiAir sudden-change episodes, AeroChem identifies directional shifts and reduces anomaly errors versus baselines, but sometimes undershoots extreme peak magnitudes.The appendix describes full reconstruction of extreme-event amplitudes as an inherent challenge.Source: §A.11 · p. 20
A.12 Case study
- AeroChem’s learned kinetics gate tracks solar radiation, opening at high radiation and suppressing reactions at low radiation; a case study links NO2 depletion to rising O3 and shows PM2.5 plume movement along wind flow.The reported ozone episode captures a diurnal peak that baseline models often miss, while the plume shift is presented as evidence of learned advection.Source: §A.12 · pp. 20–21
A.13 Sensitivity and Data Fidelity
- Across supplementary settings, AeroChem improves long-horizon forecasting and sudden-change performance; its correction ablation beats classical data corrections, and it remains competitive without future emissions.The appendix also reports improved accuracy from temporally downscaled emissions and 9-day O3 results where AeroChem outperforms data-driven models.Source: §A.13 · p. 21, Table A.6 · p. 22, Table A.7 · p. 22
A.14 Statistical significance
- Statistical tests qualify the gains: Transformer baselines significantly outperform AeroChem on 3-day MAE in both datasets, while AeroChem’s relative advantages emerge at longer horizons and sudden events.Normalized results report nMAE improvements up to 5.58% at long horizons and sudden events, lower nRMSE in all HanoiAir settings, and 2.3–2.7% worse nMAE at 3 days.Source: §A.14 · p. 21, §A.15 · p. 22, Table A.9 · p. 23
NeurIPS Paper Checklist
- The paper reports sufficient methodological detail, experimental settings, and results as means and standard deviations across multiple runs.The checklist says the paper specifies architecture, training objectives, dataset splits, hyperparameters, forecasting horizons, and evaluation protocols.Source: §2 · p. 24
- The source code is released through an anonymized public repository with reproduction instructions, and external datasets, benchmark methods, and prior works are credited.The checklist also marks experimental setting details and access to data and code as provided.Source: §2 · pp. 24–25
- The checklist says the work complies with the NeurIPS Code of Ethics, discusses societal benefits and deployment considerations, and does not require additional safeguards.It states the work uses environmental monitoring data without human subjects or high-risk generative models.Source: §2 · p. 24
- HanoiAir is introduced as a new dataset with documentation on its data sources, preprocessing, and variable construction.The checklist marks new assets as provided.Source: §2 · p. 25