Aurora air pollution model predicts NaNs / exhausts after around 13 days of autoregressive rollout

[yashdiggikar]

Hi,

I am running the Aurora air pollution model autoregressively and I am seeing that the model becomes unstable after about 13 days of forecasting. After that point, the predictions start to degrade badly and eventually produce NaN values.

What I observed
The rollout is stable initially.
Around day 13, the forecast starts to drift strongly.
After that, the model appears to exhaust / blow up and begins predicting NaNs.
In my tests, the surface variables degrade much faster, while some pressure-level variables remain relatively reasonable for longer.
My question

Is this expected behavior for the Aurora air pollution model during long autoregressive rollouts?

Since the model is mainly designed for shorter forecasts, I would like to understand:
Is there a known reason the model exhausts during longer rollouts?
What are the recommended ways to extend the forecast horizon more safely?

Specifically, I would appreciate guidance on:
Whether this is a limitation of the released checkpoint
Whether any normalization or preprocessing issues could cause this
Whether clamping is recommended
Best practices for extending inference beyond the standard forecast range

[wesselb]

Hey @yashdiggikar! Thanks for opening an issue. :)

AuroraAirPollution was indeed trained for roll-outs up to 5 days. Rolling out the model beyond this time period is uncharted territory in the sense that we don't really know what will happen. To be honest, I'm pleasantly surprised that the model holds up up to 13 days!

Is there a known reason the model exhausts during longer rollouts?

Fundamentally, there is no way to guarantee that a model can be rolled out for as long as you like. It might work, or the model might blow up after a certain number of days. The only thing that we can with reasonable confidence say is that roll-outs up to 5 days should work, because the model was trained in this way.

Specifically, I would appreciate guidance on:

I would say that this is a limitation of the released checkpoint. The air pollution data is the most difficult data that we trained on. It's incredibly finicky, which impacts the quality of training.

If you want to extend inference beyond the standard range, to really solve this, unfortunately you'll have to train a better version of AuroraAirPollution. There is a lot more data available than the data AuroraAirPollution was trained on, so this is definitely something that it possible.

Clearly that's a big undertaking. The only other thing I can imagine which might help is to correct the predictions of the model in some way, e.g. by clamping or filtering out clearly unrealistic components, which might help it to remain stable for longer. This could help a little, but is unlikely to fundamentally solve the issue.