Xaurora: Probabilistic Weather Forecasting with Foundation Models and Stochastic Interpolants

In recent years, deep learning-based weather prediction (DLWP) systems have gained significant traction, rivalling their physics-based counterparts for a fraction of the computational cost. Nevertheless, the potential of current DLWPs is hindered by important limitations. For instance, they primarily rely on numerical weather prediction (NWP) model outputs as sources of training data and usually offer a very limited set of prognostic variables compared to standard NWPs. Furthermore, they cannot easily be coupled with other models and modalities. Perhaps most importantly, the majority of the DLWPs are deterministic. However, given the inherently chaotic nature of atmospheric dynamics, DLWP outputs are often unusable as they fail to capture the range of plausible futures. Classical probabilistic methods, such as input perturbation, are useful to a certain degree, but may fail to produce enough spread. Conversely, training natively stochastic DLWP models from scratch is often prohibitively expensive. In this work, we explore the idea of converting a deterministic DLWP into a probabilistic model. We present Xaurora, an operational stochastic DLWP based on Aurora, a 1.3 billion-parameter, transformer-based Earth system foundation model. Using the stochastic interpolant framework, we fine-tune Aurora as a drift prediction model using low-rank adaptation (LoRA). We compare our model with state-of-the-art NWPs and DLWPs on a variety of weather and climate modelling tasks and achieve competitive performance. Our main contribution is to demonstrate that large-scale DLWPs can be effectively converted into stochastic models at a fraction of the cost of training from scratch. Xaurora also paves the way for further research into the intersection of weather and climate modelling and generative models, which could help to overcome the current limitations of DLWPs and move beyond the forecasting-only paradigm.