Adapting Aurora for Basin-Scale Scalar Prediction

[Jinal4502]

Really impressed by Aurora, it looks like a fantastic step forward for Earth system modeling. For my use case, I was wondering how should I adopt it for basin-scale scalar prediction tasks (eg. River Discharge)? Is it best to use the encoder as a frozen feature extractor on gridded atmospheric inputs and pool spatial tokens before a regression head, or is fine-tuning typically necessary? Also, should inputs match the full pretraining variable stack and temporal resolution?

Thank you so much for making your amazing work available to use!

[wesselb]

Hey @Jinal4502! Thanks for opening an issue and for your kind works.

Is it best to use the encoder as a frozen feature extractor on gridded atmospheric inputs and pool spatial tokens before a regression head, or is fine-tuning typically necessary?

If you also require forecasting in time, so to make forecasts at the basin scale at $t+6$ given inputs at $t$, then I think downscaling the forecasts of Aurora is a good start. The downscaling model could first be trained separately, but can eventually be trained end-to-end where the weights of Aurora are simultaneously fine-tuned. This is probably the approach that I'd try first because hooking into the internals of the model.

Also, should inputs match the full pretraining variable stack and temporal resolution?

Yes, I would recommend to use complete collection of variables and pressure levels used during pretraining as inputs.