Causal Learning Meet Covariates: Empowering Lightweight and Effective Nationwide Air Quality Forecasting

1 Introduction

Air quality prediction, a fundamental task in smart cities, plays a vital role in pollution control and public health protection [Liang et al., 2023]. In recent years, spatiotemporal graph convolutional networks have become the predominant approach to this task. These models incorporate graph inductive bias by employing GCNs or Transformers as spatial modules to capture correlations among monitoring stations, while using temporal modules such as TCNs or LSTMs to model the evolution of AQI over time. Despite their success in advancing air quality forecasting, these models still face three key challenges:

∗Yang Wang and Binwu Wang are corresponding authors.

Scalability on the large-scale air quality datasets. While existing models have significantly improved accuracy, their success has been primarily demonstrated on small-scale air quality datasets. A potential reason for this limitation is the lack of large-scale open-source datasets, particularly for nationwide applications, as illustrated in Figure 1. The high computational complexity of these models hampers their practicality [Shao et al., 2022b; Liang et al., 2023]. This challenge is particularly pronounced for Transformer-based models, which currently represent the most advanced architecture. Their time complexity increases quadratically with the number of monitoring stations, complicating their scalability. Furthermore, Transformers integrate attention mechanisms and MLP sub-blocks with skip connections and normalization layers. This intricate architecture leads to high memory requirements [Wang et al., 2024c].

Imbalance pollution level distribution. Researchers may tolerate minor numerical errors in predictions but be strict for misclassifications of pollution levels, especially for severe pollution level cases. Taking PM2.5 as an example, meteorologists classify PM2.5 pollution levels into three categories based on current concentration per cubic meter: No Pollution Level (≤ 35 µg/m³), Pollution Level I (35-75 µg/m³), and Pollution Level II (≥ 75 µg/m³). In Figure 2, we reveal an imbalanced air quality pollution level distribution, with Pollution Level II being extremely rare. Deep learning models typically struggle with learning from tail data. These severe leveltwo pollution events pose serious threats to human health and should be emphasized in air quality prediction.

Causal modeling of AQI and weather conditions. Weather conditions are crucial covariates that significantly influence AQI, i.e., causal association. For example, rainy and typhoon weather typically leads to substantial decreases of PM2.5, and severe PM2.5 pollution events typically occur under stagnant weather conditions [Zhang et al., 2017]. Weather covariates should serve as valuable resources for improving prediction accuracy. Unfortunately, existing models predominantly concentrate on the spatiotemporal dynamics of AQI, treating weather covariates merely as optional auxiliary information. Most models just employ shallow neural networks to encode weather covariates, subsequently concatenating the resulting high-dimensional embeddings with the final AQI representations before passing the output to the decoder for future predictions. This simplistic approach fails to explicitly capture the causal associations between weather covariates and AQI.

To address the challenges outlined above, we first release LargeAQ, a nationwide air quality dataset collected from 1,341 monitoring stations. This dataset offers comprehensive coverage across 33 major administrative regions in China. In contrast to the commonly used and open-source air quality datasets, detailed in Table 1, LargeAQ is distinguished by its large-scale and long-term coverage. This extensive dataset is anticipated to facilitate significant advancements in air quality prediction.

Furthermore, we propose a causal learning model for nationwide air quality prediction called CauAir, which explicitly models the causal association between AQI and weather covariates to improve the model's predictive performance and its ability to forecast severe pollution levels in the tail distribution. Specifically, CauAir first uses a channel mixing module to separately model spatiotemporal dynamics within AQI and covariates. We propose using a Transformer architecture to model causal association between AQI and covariates. Considering the trade-off between performance and efficiency, we propose a Cachebased Lightweight Transformer, termed CachLormer. This model is a parallelizable adaptation of the vanilla Transformer, achieved by simplifying its complex architecture through the removal of skip connections and normalization layers. CachLormer also employs a cache-attention mechanism with learnable caches to perceive causal association between covariates and AQI. By implementing feature interactions at a coarse granularity, this approach enables robust recognition of causal patterns across heterogeneous data. We also provide a theoretical explanation for the advantages of our model. Our contributions are summarized as follows:

• We introduce LargeAQ, a nationwide large-scale air

quality dataset. We anticipate that this pioneering work will create promising avenues for the advancement of air quality prediction techniques.

• We propose CauAir, a lightweight and effective model that explicitly captures causal associations between AQI and weather covariates to achieve notable performance improvements, supported by mutual information theory.
• Experimental results on three datasets with 28 baselines demonstrate that CauAir achieves competitive performance with high efficiency and low memory usage.

2 Preliminary

2.1 Task Statement

Let Xt ∈ RN×c to denote the AQI of N air quality monitoring stations at a given time step t, where c is the number of air pollutants measurements (e.g., PM2.5, PM10, or NO2). We denote Zt ∈ RN×f as the weather covariates collected from the vicinity of each station at time step t, where f represents the number of covariates, such as temperature and wind speed. Zp and Zf refer to the covariates observed in the past and those predicted for the future, respectively.

Given observed AQI of all stations from the past T time steps Xt-T+1:t ∈ RT×N×c , observed covariates Zpt-T+1:t ∈ RT×N×f , and predicted covariates Zft+1:t+L-1 ∈ RL×N×f , we aim to learn a function F that can predict future AQI over next L time steps Yt+1:t+L-1 ∈ RL×N×c : [Xt-T +1:t, Zpt-T +1:t , Zft+1:t+L-1 ] F(·) ---→ Yt+1:t+L-1.

2.2 Transformer

Transformer has demonstrated its impressive representation capabilities in various fields. The learning process of the vanilla transformer is expressed as follows:

where Hi and Ho are the input and output of the transformer, respectively. FFN(·) uses two-layer MLP with ReLU activation function, and Norm(·) is the normalization operation. The core of Transformer is the self-attention mechanism, which can be expressed as:

where Q,K,V represent query, key, and value vectors, respectively, and they come from the mapping transformation of the input. To improve the representation ability, Transformer uses the multi-head self-attention mechanism.

3 Methodology

In this section, we provide the details of the proposed CauAir framework. CauAir initially models the dynamics of the AQI and the covariates independently. Subsequently, we design a lightweight Transformer to explicitly model the causal associations between AQI and the covariates. The overview of CauAir is shown in Figure 2.

3.1 Channel Mixing for Spatiotemporal Learning

We adopt a spatiotemporal learning module to model the spatiotemporal dynamics within both AQI and weather covariates. This module operates as follows: for any input tensor P ∈ {Xt-T+1:t,Zpt-T+1:t ,Zft+1:t+L-1 }, the module first compresses both temporal and feature dimensions to generate node-level representations H ∈ RN×di . For example, if X ∈ RT×N×c is used as input, di of the output would be equal to T × c. The output H is then fed into a feed-forward network as an encoder with positional embeddings to model spatiotemporal dynamics through mixing feature channels. The forward process of our spatiotemporal learning module is expressed as follows,

where ˜H denotes the output. E ∈ RN×dh is the learnable position embedding to capture high-level features of stations. Here FFN(·) includes two-layer MLP with SwiGLU [Shazeer, 2020] as activation function:

whereW1,W2 ∈ Rdi×e∗di andW3 ∈ Re∗di×dh are learnable parameters. e = 4 is the expanding coefficient of hidden representation. σ (·) represents the sigmoid function and ⊙ represents the Hadamard product.

We employ two parameter-independent spatiotemporal learning modules to model the dynamic of the past AQI Xt-T+1:t and historical covariates ZPt-T+1:t , and the outputs are denoted as X ∈ RN×dh and Z ∈ RN×dh , respectively.

3.2 Efficient CachLormer for Causal Learning

Weather covariates, such as precipitation and wind speed, play a crucial role in influencing AQI. To explicitly capture these causal associations, we propose employing a Transformer as the backbone of our model, leveraging its strong capability to model dependencies between variables effectively. To mitigate the computational complexity associated with traditional Transformers, we introduce a Cache-based Lightweight Transformer (CachLormer). Its efficiency improvements arise from two key modifications: a simplified architecture and an efficient cache-attention mechanism.

Simplified Architecture. In Equation 1 of a standard Transformer architecture, the structure typically consists of interleaved attention and feedforward (MLP) sub-blocks, designed with the skip connection and normalization layer. In fact, such complex designs are not strictly necessary. Inspired by [He and Hofmann, 2023], we also reformulate the complex structure into a parallelizable version by removing the skip connection and normalization layer, which is shown in Figure 2,

where δ0 and δ1 ∈ (0, 1) are parameters that control the mixing ratio of two outputs. The feed-forward network here includes a two-layer MLP with SwiGLU activation function. MH-Cachattn (·) represents our design cacheattention, which will be explained in the following section. The elimination of unnecessary components reduces the memory consumption of Transformer.

Cache-attention Mechanism. The traditional self-attention mechanism calculates the attention coefficient by taking the dot product of the key vector K ∈ RN×dh and the query vector Q ∈ RN×dh , resulting in a quadratic complexity of O (N 2dh ). This dense complexity limits the model's scalability on large-scale datasets. To address this issue, we propose cache-attention mechanism, where the underlying motivation for this approach stems from the urban hierarchy theory in the spatiotemporal data analysis field [Guo et al., 2021]: certain fine-grained nodes may exhibit similar spatiotemporal features, allowing us to coarsen these finegrained nodes into P coarse-regions. Thus, we first randomly initialize a set of caches Ep = [e1, e2, . . . , eP ] ∈ RP×dh , which is a learnable parameterized embedding vectors. The hyperparameter P is the number of caches with P (≪ N) and dh is the number of channels. And the embedding vector of k-th (k ∈ {1, . . . , P}) cache ek ∈ Rdh represents the contextual features of k-th coarse-region. Fine-grained modeling may prevent the model from perceiving robust causal association between two heterogeneous data types, covariates and AQI, while a coarse-grained perspective can alleviate this issue.

Given a query matrix Q ∈ RN×dh , where qi represents the feature embedding of node vi, we employ a soft assignment mechanism to determine the association between nodes and caches. This mechanism achieves lossless attention capacity through weighted combinations of attention scores between nodes and cache elements. Specifically, the affinity between node vi and k-th cache can be computed as:

Then we obtain two similarities by normalizing the node dimension, and the cache dimension, respectively, with the above calculated assignment weights as follows,

where s(1) ·,k records the affinity of each node to k-th cache, and s(2) i,· represents the affinity of each cache to node vi. We can indirectly calculate the similarity ai,j between nodes vi and vj as follows,

The similarity calculated by this method inherits normalized properties: for any node vi, the attention scores from all nodes to vi sum to 1, i.e.,

With the cache-attention calculation process from Equation (6) ∼ (8), the forward process of the cache-attention can be expressed as follows,

Cache-attention (Q, Ep, V)

where the query vector Q = HiWq and the value vector V = HiWv. And Wq,Wv ∈ Rdh×dh and Ep are learnable parameters. Our model also uses a multi-head attention mechanism:

where W (w) q ,W (w) v ∈ Rdh× dh Nh , E(w) p ∈ RP× dh Nh are thew-th head part of Wq,Wv and Ep, respectively. Wo ∈ Rdh×dh is a learnable parameter to integrate information fromNh heads. Computational complexity analysis. Compared with the O (N 2dh ) complexity of the self-attention mechanism in Transformer, our cache-attention mechanism is reduced to O (PNdh), where P is be much smaller than N.

Causal Learning between AQI and Covariates

We first deploy a CachLormer, which is denoted as f0, to model the causal dependencies between past AQI X and past weather covariates Z. Specifically, we adopt the addition strategy with RMSNorm normalization strategy [Zhang and Sennrich, 2019] to deeply fuse AQI and covariate information by addition, compared to independent channel concatenation strategies, the model can analyze deeper causal associations from entangled representations:

where β0 and γ0 are learnable parameters used to balance the two terms. In order to integrate future weather covariates. Specifically, we first use two parameter-independent spatiotemporal learning modules to capture the dynamics of Xt-T+1:t and Zft+1:t+L-1 , and the outputs ˜X and ˜Z are the future AQI representation and the covariate representation, respectively. Then we use a CachLormer module, denoted as f1, to learn the causal associations between ˜X and ˜Z:

Finally, we add the two components together and feed them into the decoder (e.g., FFN) to predict the future AQI:

where α0 and α1 are learnable parameters that weigh the decision strengths of the two components.

3.3 Theoretical Explanation

Theory 1. Equivalence between MSE and mutual information (MI). Minimizing MSE between the predicted values ˆY and the ground-truth values Y can be equivalent to maximize the mutual information of ˆY and Y, max I ( ˆY,Y), in the regression tasks [Jing et al., 2022]:

Next, we demonstrate the superiority of our model in modeling causal association between AQI and weather covariates from the mutual information theory.

Theory 2. MI Monotonicity. For any random variables A, B, C, and D, the following inequality of the mutual information I (·; ·) holds [Peng et al., 2020]:

According to Theory 1, our goal is to maximize I( ˆY;Y). Given the fact that weather covariates have significant causal association with air quality, we incorporate these covariates Z as auxiliary information into the learning process: the learned representations of historical weather covariates and past AQI are input into CachLormer to effectively model the causal association. The learned knowledge is then propagated to future weather covariates to make accurate predictions. Thus, according to Equation 16, we have I( ˆY;Y) ≥ I (Y;Xt-T+1:t,Z) ≥ I (Y;Xt-T+1:t), indicating that our model has a higher lower bound on the mutual information (MI) between true values Y and predicted values ˆY. Based on the equivalence between MSE and MI (Theorem 1), we can conclude that weather covariates enhance the upper bound of mutual information, enabling the model to make more accurate and stable predictions by reducing predictive uncertainty. We further demonstrate the benefits of weather covariates through ablation experiments (as shown in Section 4.5). Compared to approaches that encode weather variables through shallow networks, our model can capture causal association more precisely.

4 Experiments

4.1 Nationwide Air Quality Dataset

Data Description. Due to the lack of open-source large-scale datasets for air quality prediction, we open-source a comprehensive, long-term air quality dataset named LargeAQ. It contains AQI, including PM2.5, PM10, NO2, and so on, from 1,341 monitoring stations across 333 major Chinese cities. We concentrate on PM2.5 as the primary pollutant of interest. The dataset spans 8 years from January 1, 2016, to December 31, 2023, with hour granularity. The data are sourced from the China National Environmental Monitoring Center (CNEMC) 1 , an online platform that provides real-time AQI for prefecture-level administrative regions across China. We also collect five weather features during this period from the Chinese Weather Website as covariates: temperature, humidity, precipitation, wind shear, and wind speed. A comparison between LargeAQ and several popular air quality datasets is shown in Table 1. In the experiments, we also use two other open-source datasets.

Data Analysis. In weather science, PM2.5 pollution is categorized into three levels: No Pollution Level (≤ 35 µg/m³), Pollution Level I (35-75 µg/m³), and Pollution Level II (≥ 75 µg/m³). Figure 3 illustrates the PM2.5 distribution across these categories in LargeAQ. We observe a significant imbalance in the distribution of the three pollution levels in the LargeAQ dataset-severe Pollution Level II is rare.

Data Processing. To address missing values in the AQI and weather data, we use the last observed value method. Using the sliding window to process data, the length of both past window and future window are equal to 24. For input features - including AQI, historical covariates, and future covariates - we apply standard normalization to facilitate stable learning. Only the predicted AQI values are denormalized back to their original distribution. For the future weather covariate Zf , to avoid information leakage caused by using the actual values from our collected weather data, we simulate future weather forecasts by adding synthetic noise. Following common practices in linear regression [Murphy, 2012; Jing et al., 2022], which assume that the predicted values deviate from ground truth by additive Gaussian noise, we generate simulated forecasts by injecting Gaussian noise into the observed weather data. In addition, we also report the performance of models that do not incorporate future weather features, to provide a fair comparison.

4.2 Experimental Setup

Baselines for Comparison

We compare our CauAir with the 28 baselines that belong to the following three categories:

1 https://www.cnemc.cn/

Time-series models: CATS [Lu et al., 2024], CycleNet [Lin et al., 2024a], DLinear [Zeng et al., 2023], DSformer [Yu et al., 2023], SOFTS [Han et al., 2024b], SparseTSF [Lin et al., 2024b], TimeMixer [Wang et al., 2024d], CrossGNN [Huang et al., 2023], Umixer [Ma et al., 2024], and TimeXer [Wang et al., 2024e].

Spatiotemporal forecasting models: AGCRN [Bai et al., 2020], ASTGCN [Guo et al., 2019], BigST [Han et al., 2024a], D2STGNN [Shao et al., 2022b], GWNet [Wu et al., 2019], STAEformer [Liu et al., 2023], STGCN [Yu et al., 2017], STGODE [Fang et al., 2021], STID [Shao et al., 2022a], STNorm [Deng et al., 2021], STTN [Xu et al., 2020], and RPMixer [Yeh et al., 2024].

Air quality prediction models: AirFormer [Liang et al., 2023], AirPhyNet [Hettige et al., 2024], DeepAir [Yi et al., 2018], GAGNN [Chen et al., 2023], MGSFformer [Yu et al., 2025], and PM2.5GNN [Wang et al., 2020].

Implementation Detail

The models are implemented in PyTorch 2.2.0 running on an NVIDIA A100 GPU with 40 GB memory. We adopt Adam optimizer accompanied MultiStepLR learning rate adjustment strategy with a learning rate 0.02 and weight decay 0.004. The cache number P is 32 for LargeAQ, 108 for CCAQ and 10 for KnowAir. The channel dimension of CachLormer is set to 128. We evaluate all models on numerical regression prediction tasks and pollution level category prediction tasks, using three commonly employed metrics: MAE, RMSE, and MAPE for regression prediction and F1 score for pollution level prediction task. All baselines incorporate weather covariates (if available).

4.3 Future PM2.5 Forecasting Comparison

We report the performance metrics over 24 horizon forecasts for LargeAQ dataset in Table 2.

Time-series forecasting models typically underperform compared to spatiotemporal learning models, primarily because they focus solely on temporal dependencies while neglecting inter-station correlations. Within the spatiotemporal learning family, Transformer-based models demonstrate superior performance over GCN-based approaches. For instance, D2STGNN effectively captures PM2.5 spatiotemporal dynamics through its decoupling mechanism. Among specific air quality prediction models, DeepAir achieves surprisingly good performance despite its relatively simple architecture. This success can be attributed to its specialized strategy for encoding weather variable influences, highlighting the crucial role of weather factors in prediction accuracy.

Our CauAir explicitly models the causal associations between covariates and AQI to enhance prediction performance. As a result, our model achieves state-of-the-art performance across all metrics, showing up to 8.27% improvement in PM2.5 value prediction compared to existing methods.

4.4 Efficiency Study

Figure 4 illustrates the efficiency comparison between strong baselines and CauAir on the LargeAQ dataset. Each solid circle represents a model, with the model name, training time per epoch, and MAE reported below. The default batch size is 64; for some complex models (marked with ∗ and their batch size), we have to reduce their batch sizes to enable operation on the large-scale dataset.

We observe that advanced Transformer-based models, such as STAEformer and D2STGNN, suffer from a time complexity that scales quadratically with the number of nodes, due to their reliance on standard self-attention mechanisms. Moreover, their complex architectural designs further increase memory consumption. AirFormer simplifies the Transformer structure, thereby achieving improved efficiency. Our model not only achieves superior predictive performance but also attains the highest computational efficiency, outperforming the SOTA D2STGNN in time and space complexity.

4.5 Ablation study

We conduct ablation experiments on the LargeAQ dataset to verify the effectiveness of each component.

Weather Covariates. We develop the following variants: (1). w/o Zp : Remove the past weather covariate Zp term in input; (2). w/o Zf : We do not take the future weather covariate Zf as input; (3). w/o Zp & Zf : We do not use any weather covariates. The experimental results are shown in Table 3, which shows that both past and future weather covariates are beneficial to improving prediction performance.

Transformer Variant. To evaluate the effectiveness of the proposed CachLormer, we develop the following variants:

(1). -Sattn: We use standard self-attentionmechanism instead of our cache-attention mechanism, using our lightweight architecture; (2). -VanArc: We use standard Transformer architecture in Equation 1 with our cache-attention mechanism; (3). +VanTran: We use standard Transformer architecture to replace CachLormer. We find that CachLormer achieves comparable performance to Transformer while significantly reducing memory consumption and training time (s/epoch).

4.6 Hyperparameter Sensitive Annalysis

We evaluate the sensitivity of the hyperparameter the number of caches P on LargeAQ dataset.As shown in Figure 5, we observe that optimal values of P is equal to 32. A smaller P is insufficient to capture adequate contextual features shared among nodes. When P exceeds these optimal values, performance does not improve significantly, as the model struggles to focus on extracting shared contextual features.

4.7 Performance Analysis on Small-scale Datasets

We evaluate our model against several state-of-the-art approaches using two open-source small-scale datasets: KnowAir and CCAQ. The 24-step average results are presented in Table 4. Our model demonstrates significant performance advantages, with an average improvement of 5.15% and a maximum improvement of 20.55%. These enhancements can be primarily attributed to our explicit modeling of the causal relationships between AQI and weather conditions.

4.8 RelatedWork

Time series analysis is a fundamental task [Liu et al., 2025a; Liu et al., 2025a; Miao et al., 2024; Huang et al., 2023], which is a specific area of air quality forecasting. Early time series analysis approaches modeled air pollutant emissions and dispersion as dynamic systems using numerical simulations, with notable examples such as the Community Multiscale Air Quality [Byun and Schere, 2006]. However, these methods require extensive theoretical knowledge, carefully selected features and region-specific parameters, making them impractical for real-time air quality monitoring systems. Data-driven models emerged with the expansion of urban-scale air quality sensor networks and advances in machine learning algorithms. Recent advances in machine learning, particularly deep learning, have significantly improved prediction accuracy [Liang et al., 2023; Wang et al., 2020]. Researchers have designed cutting-edge models to capture the internal spatiotemporal dynamics of AQI and extract high-dimensional representations. The dominant architecture employs spatiotemporal graph convolutional networks [Miao et al., 2025; Wang et al., 2023; Wang et al., 2024a; Wang et al., 2024b; Ma et al., 2025], utilizing spatiotemporal graphs to represent spatiotemporal data, RNNs or Transformers for temporal dependency modeling, and GNNs or Transformers for spatial dependency extraction [Liu et al., 2024; Shao et al., 2022a; Liu et al., 2025b; Zhang et al., 2025]. These works, however, present some difficulties (e.g., inefficiency) on a nationwide air quality dataset. Moreover, our model differs by explicitly modeling the influence of covariates on air quality, leveraging these additional variables to enhance prediction performance.

5 Conclusion

In this paper, we introduce a nationwide and long-term air quality data from 1,341 monitoring stations. We then propose a lightweight yet effective spatiotemporal causal learning model for air quality prediction. The model employs a novel CachLormer, an efficient Transformer architecture, to explicitly model the causal association between AQI and weather covariates. Compared across three datasets against 28 baselines, our model achieves dominant performance while maintaining high time and memory efficiency.

Contribution Statement

Jiaming Ma and Zhiqing Cui have equal contribution.

Acknowledgements

Thank Xiaolei Wang (homepage: https://quotsoft.net/) for the valuable open-source contribution. And this paper is partially supported by the National Natural Science Foundation of China (No.12227901). The AI-driven experiments, simulations and model training were performed on the robotic AI-Scientist platform of Chinese Academy of Sciences.