Subseasonal-to-Seasonal strategies for the Earth System Foundation Model (ESFM)

Foundation models for the Earth system have gained popularity, as they are starting to surpass numerical solvers in the accuracy of predicting Earth’s condition while requiring fewer computational resources. The Earth System Foundation Model (ESFM) contributes to this research direction by further extending the foundation models' flexibility.

The forecasting capabilities of ESFM are achieved in an autoregressive manner, using data from the t₀ - Δt and t₀ timesteps to produce a prediction for t₀ + Δt. This approach is effective on weather timescales. Moreover, we find that it also delivers encouraging results for long-term forecasts, showing reasonable zero-shot subseasonal-to-seasonal (S2S) predictions (15–40 days). 

However, S2S predictions can be further improved while preserving weather skills. This work investigates strategies for this purpose. On such timescales, it is crucial to produce probabilistic predictions to better represent inherent uncertainty. Probabilistic predictions are realised with the introduction of multiple decoder heads (tails) for each variable. Each tail is intended to simulate a different possible trajectory, which, when combined, provides an estimate of the most probable outcome together with the spread of feasible values. To better estimate the distribution of possible values on the S2S timescale, additional trajectories are generated by running multiple predictive rollouts with different initial conditions.

Another strategy to improve S2S rollouts is to fine-tune the model to produce outputs for more distant steps. To this end, we leverage LoRA adapters (Hu et al., 2022), which are trained for each subsequent rollout step. This approach effectively improves predictive performance on long horizons, without significantly affecting training complexity or inference cost.

We also observe that some predictive variables of the model, such as climate forcings, are slowly evolving and can benefit from incorporating inputs from a more distant past than the t₀ - Δt and t₀ timesteps commonly used. To investigate this, we introduce an Attention Temporal Aggregator in the encoder, which leverages learned patch embeddings from an arbitrary number of previous timesteps and attends to those that are most informative for a given variable. In this way, for rapidly changing variables such as wind speed, the model focuses on the most recent data, whereas for slowly evolving variables such as sea surface temperature, it can utilise a broader range of inputs.

Overall, our experiments provide new insights into the development of foundation models for the Earth system, enabling improved predictions on S2S timescales, while conserving performance for weather forecasts.

References:
E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, W. Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022