Aurora Air Pollution value explosion and collapse

[duncanmartyn]

Hi all,

I've been using the air pollution fine-tune of Aurora and found that values in many variables explode and collapse (into NaN) between 60-80 rollout steps.

For context, I'm researching the model for use in long-term climate applications, so the lengths of rollouts I'm running far exceed the model's intended 5-day horizon. I've seen the paper caveats the air pollution fine-tune in several ways so perhaps this result isn't unexpected, but I've encountered no such issue with AuroraPretrained over rollouts exceeding one year (1,460 rollout steps).

Explosion and collapse always occurred earlier when simulate_indexing_bug was set to True relative to False when using the same data as the initial state. The location of collapse also changed, with the former setting producing increasingly extreme values over the Himalayas / Tibet and the latter over Northern China and South America. In all cases, it appears to be identical pixels that precipitate explosion every time, with gradually more extreme values propagating outwards from these over 10-20 steps to produce the blocky artefacts seen in the below images.

simulate_indexing_bug=True:

simulate_indexing_bug=False:

Thus far, I've found the only solution to be filtering of the 2t surface variable. With a Gaussian or uniform filter applied to 2t in each rollout step and the result assigned to the prediction Batch object's .surf_vars["2t"] attribute, I can run rollouts of arbitrary length.

What else I've tried:

• Using CAMS data from the train and test periods, as well as data outside of these - no effect
• Replaced CAMS climatic variables with coarsened ERA5 0.25 data - no effect
• Clamping 2t to sensible values (record global extremes) - no effect
• Per variable re-use of initial state data (i.e. replacing each rollout step's prediction for a given variable with said variable's initial state) - doing this with 2t alone resolved the issue
• Reduced window size - produced less variation in predictions between steps and accelerated collapse
• Changed timestep - 6 hour timestep accelerated collapse

Just looking for your perspectives or thoughts on this, thanks!

[wesselb]

Hey @duncanmartyn! This is super interesting. I've rolled out the model once or twice beyond the 5-day range, but have never gone to 60-80 rollout steps. I did then notice that the predictions start to degrade once you cross 5 days.

Explosion and collapse always occurred earlier when simulate_indexing_bug was set to True relative to False when using the same data as the initial state.

It could be an interesting experiment to fine-tune the model some more on CAMS data with simulate_indexing_bug=False. It sounds like this might improve stability.

but I've encountered no such issue with AuroraPretrained over rollouts exceeding one year (1,460 rollout steps).

That the pretrained version is stable for such a long time was a surprising finding to us. I'm not sure precisely why AuroraPretrained is so stable. AuroraAirPollution definitely is less well trained, which is mainly due to the quality of the CAMS data: there is much less of it, and the data significantly changes in distribution whenever CAMS undergoes an update.

What else I've tried: (...)

Everything that you've tried sounds very sensible! Probably the best way to enable long roll-outs is to try to train a better version of the air pollution model. In particular two-step roll-out fine-tuning with gradients propagated all the way back could help. The current roll-out fine-tuning procedure uses a stop-gradient trick, which doesn't generate well beyond the roll-out fine-tuning lead time range of 5 days.

The only other thing that I can think of is that perhaps another variable than 2T might be blowing up. If that's the case, then some clipping/smoothing to keep this variable in check could possibly help. In particular, some of the air pollution variables near the surface could be problematic, e.g. SO2 at 1000 hPa.

[duncanmartyn]

Hi Wessel, thanks for getting back to me! That makes sense, and fine-tuning is a possibility. Sounds like a fine-tune with data from more recent (or even the latest alone) CAMS updates given distribution changes is one to look at.

Re. your last point on 2T, I did wonder if there were a more subtle (not apparent in visual inspection of plots) influence. Applying a filter to each variable in distinct rollouts only resolved the issue with 2T, though hard to tell if it's just because 2T is truly where it starts or the filter didn't sufficiently smooth the root cause.

Potentially related are predictions for atmospheric geopotential. Minimum values across all pressure levels are always in initial state data, with Aurora predicting gradually increasing then fluctuating values, never going as low (or near) original values. I wondered if you had any thoughts on this? Below graph shows minimum (across all levels) atmospheric geopotential over 75 rollout steps.

From what I hear from researchers in the atmospheric chemistry domain it's not uncommon for physical models to exhibit similar issues over high elevation, topographically rough regions like the Himalayas and require some smoothing, either.

[wesselb]

Potentially related are predictions for atmospheric geopotential. (...) I wondered if you had any thoughts on this?

Oof, this looks problematic to me! I'm not sure what's going on. The graphs seems to show that long roll-outs are converging into an unphysical regime. That's interesting.

I'm sorry I can't be of more help! Unfortunately I don't really have any great ideas for how to alleviate this behaviour (other than simply training a better model).

[duncanmartyn]

Not a problem, really appreciate the discussion, thanks. I can leave this issue open to add any further findings and / or resolutions as I go unless you'd prefer it closed?

[wesselb]

Please feel free to leave it open and add any further findings. :) That would be cool!