Forecasting India's air quality with a neural operator, and the validation bug that unlocked it

Summary

Framed in the Google XYZ format (accomplished X, as measured by Y, by doing Z):

Achieved a private-leaderboard score of 0.8833 (2nd of all teams) at the ANRF AISEHack Phase 2 pollution-forecasting competition, as measured by the official composite SMAPE and correlation metric, by rebuilding an AFNO neural operator on leak-free validation and layering domain physics, a two-seed ensemble, and per-horizon calibration on top.
Broke a long-standing accuracy plateau of roughly 0.86, as measured by the official private metric, by identifying and removing a validation data leak that had made every prior model comparison unreliable.
Delivered a reproducible, documented model (about 7M parameters), as measured by published weights and code on Hugging Face and Kaggle, by saving EMA checkpoints, normalization statistics, and calibration factors from the training run.

The problem

The ANRF AISEHack Phase 2, Theme 2 track - run by IIT Delhi, IIIT Hyderabad, ANRF, and IBM - asks for a spatiotemporal forecast: given the last 10 hours of gridded meteorology and pollution data over India, predict PM2.5 for the next 16 hours on a 140 x 124 grid at roughly 25 km resolution.

The difficulty is not the average. It is the spikes. Pollution episodes, the sharp peaks that matter most for public health, are rare and hard to predict, and the competition metric weights them explicitly.

The attempts

The path to the final model was a sequence of attempts, most of which plateaued in the same place. That plateau turned out to be the important clue.

Attempt	Approach	Private score	Outcome
Persistence	repeat the last observed frame	0.851 (NormGlobalSMAPE)	baseline to beat
ConvLSTM	recurrent convolutional model	~0.85 - 0.86	plateaued
Early AFNO	first neural-operator attempt	~0.856	dropped mid-hackathon
AFNO v3	leak-free validation, official STL episode masks	0.8609	first honest gain
AFNO v4	v3 plus physics features, multiscale AFNO, ensemble, calibration	0.8833	2nd on final leaderboard

Several architectures, all landing near 0.86, is not a modelling problem. It is a signal that the measurement itself is wrong.

The turning point: a validation leak

The models were being evaluated on a random split across overlapping time windows. Because consecutive windows share almost all of their hours, nearly every validation window had a near-duplicate in the training set. The model was being scored on data it had effectively already seen, so every measured improvement was partly measuring leakage rather than skill.

The fix was unglamorous: hold out a contiguous block of time at the end of each month, separated by a gap so no training window overlaps a validation window. Once the evaluation was honest, the scores dropped, became real, and for the first time tracked the actual leaderboard. This single change was worth more than any architectural idea that followed.

What worked

With trustworthy validation in place, it became possible to tell which ideas genuinely helped.

Physics as input features. Rather than raw weather alone, the model receives quantities that atmospheric science already associates with pollution:

a ventilation coefficient (boundary-layer height times wind speed), which captures how well the lower atmosphere disperses pollutants
an advection tendency, the first-order transport term from the chemistry-transport equation
a diurnal climatology anomaly, the deviation of the current value from the normal value for that hour of day, mirroring how the competition defines an episode

Multiscale spectral mixing. The Fourier mixing runs at two spatial scales so the model captures both broad regional structure and finer transport, paired with a SimVP translator that maps the 10 input hours to the 16 output hours.

Reliable gains on top. Two independently trained models are averaged, and a per-forecast-hour calibration is fitted on the honest validation set to correct a slight bias at longer lead times.

Architecture

The network is a per-frame 2D U-Net encoder feeding a neural-operator core and a temporal translator, decoded back to a 16-hour forecast expressed as a change on top of the last observed frame.

flowchart TD
    A["Input: 10 hours x 20 channels x 140 x 124"] --> B["2D U-Net encoder<br/>20 to 64 to 128 to 256, MaxPool x2"]
    B --> C["Multiscale AFNO at 35 x 31<br/>1 spectral block"]
    C --> D["AvgPool to 18 x 16 bottleneck"]
    D --> E["Bottleneck AFNO at 18 x 16<br/>6 spectral blocks, 8 heads, global 2D FFT"]
    E --> F["SimVP temporal translator<br/>Incep1D x4, 10 to 16 frames"]
    F --> G["Lead-time embedding (16)"]
    G --> H["Post-translator AFNO<br/>2 spectral blocks"]
    H --> I["2D U-Net decoder<br/>256 to 128 to 64 to 32, last and mean skip fusion"]
    I --> J["Output head<br/>Conv 32 to 32 to 1"]
    J --> K["Persistence residual<br/>output = last observed frame + delta"]
    K --> L["Forecast: 16 hours x 140 x 124"]

The 20 input channels group as follows:

flowchart LR
    subgraph Observed [Observed - 9]
        O["PM2.5, q2, t2, u10, v10,<br/>pblh, psfc, swdown, rain"]
    end
    subgraph Engineered [Engineered - 5]
        E2["wind speed, ventilation coefficient,<br/>advection, PM tendency,<br/>diurnal climatology anomaly"]
    end
    subgraph Time [Time - 2]
        T["hour sin, hour cos"]
    end
    subgraph Static [Static - 4]
        S["lat, lon, PM2.5 climatology,<br/>episode sigma"]
    end
    Observed --> M["20 input channels"]
    Engineered --> M
    Time --> M
    Static --> M

At about 7M parameters, the model predicts a delta on top of persistence, so it only has to learn how the pollution field moves rather than reconstruct it from scratch.

Results

Metric	Value
Private score	0.8833
Public score	0.8959
Final leaderboard rank	2nd (1st place 0.8893)
Leak-free validation (seeds 42 / 1337)	0.9110 / 0.9114
Prior v3 baseline	0.8609
Persistence baseline (NormGlobalSMAPE)	0.851

The metric is w1 * NormGlobalSMAPE + w2 * NormEpisodeCorr + w3 * NormEpisodeSMAPE, with weights undisclosed by the organizers.

Lessons

The move from roughly 0.86 to 0.8833 was not bought with a more complex model. It was bought by making the evaluation honest first, and only then adding domain knowledge. That ordering is the entire lesson: if the scoreboard lies, nothing downstream of it can be trusted.

Scope and limitations

This model is trained on 2016 atmospheric-simulation data with competition-specific inputs. It is a reproducible research and benchmark artifact, not a live air-quality service - it cannot ingest today's observations and warn a real city as it stands. An operational version would require retraining on reanalysis inputs such as ERA5, MERRA-2, or CAMS, validated against ground-truth measurements from monitoring networks such as OpenAQ or CPCB. That is the natural next step for turning this from a strong result into a useful tool.

References

Adaptive Fourier Neural Operator (Guibas et al., ICLR 2022)
FourCastNet (Pathak et al., 2022)
SimVP (Gao et al., CVPR 2022)
STL decomposition (Cleveland et al., 1990)

Archives

Loading