[Paper Review] AtmoRep: A stochastic model of atmosphere dynamics using large scale representation learning

Christian Lessig, I. Luise+4 more|arXiv (Cornell University)|Aug 25, 2023

Meteorological Phenomena and Simulations20 citations

TL;DR

AtmoRep is a task-independent stochastic atmospheric dynamics model using large-scale representation learning, providing probabilistic forecasts and zero-shot capabilities across nowcasting, interpolation, model correction, and counterfactuals.

ABSTRACT

The atmosphere affects humans in a multitude of ways, from loss of life due to adverse weather effects to long-term social and economic impacts on societies. Computer simulations of atmospheric dynamics are, therefore, of great importance for the well-being of our and future generations. Here, we propose AtmoRep, a novel, task-independent stochastic computer model of atmospheric dynamics that can provide skillful results for a wide range of applications. AtmoRep uses large-scale representation learning from artificial intelligence to determine a general description of the highly complex, stochastic dynamics of the atmosphere from the best available estimate of the system's historical trajectory as constrained by observations. This is enabled by a novel self-supervised learning objective and a unique ensemble that samples from the stochastic model with a variability informed by the one in the historical record. The task-independent nature of AtmoRep enables skillful results for a diverse set of applications without specifically training for them and we demonstrate this for nowcasting, temporal interpolation, model correction, and counterfactuals. We also show that AtmoRep can be improved with additional data, for example radar observations, and that it can be extended to tasks such as downscaling. Our work establishes that large-scale neural networks can provide skillful, task-independent models of atmospheric dynamics. With this, they provide a novel means to make the large record of atmospheric observations accessible for applications and for scientific inquiry, complementing existing simulations based on first principles.

Motivation & Objective

• Motivate the need for accurate, versatile atmospheric models that are computationally efficient and probabilistic.
• Propose a task-agnostic stochastic model that learns from historical observation data to capture complex atmospheric dynamics.
• Demonstrate intrinsic zero-shot capabilities for multiple tasks and show how the model can be extended or bias-corrected with additional data.

Proposed method

• Introduce AtmoRep as a 3.5 billion-parameter transformer model that predicts p_theta(y|x, alpha) for atmospheric states.
• Train p_theta(y|x, alpha) with a self-supervised objective that minimizes distance to the data distribution via Monte Carlo estimates.
• Use an ensemble of prediction heads to sample from the learned distribution, enabling probabilistic forecasts.
• Operate on 4D space-time neighborhoods subdivided into tokens, with masking/distortion used for self-supervised learning.
• Adopt a Multiformer architecture with per-field transformers coupled by cross-attention, and an efficient ensemble training strategy.
• Extendibility via task-specific tails (e.g., downscaling) and bias correction using observational data.

Figure 1: The AtmoRep model provides a numerical representation $p_{\theta}(y|x,\alpha)$ of the conditional probability $p(y|x,\alpha)$ for atmospheric states $x$ , $y$ subject to external conditions $\alpha$ , e.g. the time of $x$ , $y$ or their location on the globe. It is implemented as a transfo

Experimental results

Research questions

• RQ1 Can a large-scale, task-agnostic neural model capture the stochastic dynamics of the atmosphere from observational history?
• RQ2 What are the intrinsic capabilities of AtmoRep for nowcasting, temporal interpolation, model correction, and counterfactuals without task-specific training?
• RQ3 How does AtmoRep perform relative to established forecasting systems in short-term prediction and probabilistic skill?
• RQ4 Can AtmoRep be extended to additional tasks such as downscaling and bias correction using external observational datasets?

Key findings

• AtmoRep achieves state-of-the-art or competitive nowcasting skill with no or minimal task-specific training.
• The model provides probabilistic forecasts via an ensemble whose spread reflects observational variability.
• AtmoRep enables temporal interpolation with substantially lower RMSE than linear interpolation.
• The model can correct higher-frequency content in inputs toward ERA5 through its learned representation.
• Counterfactuals demonstrate that AtmoRep can generate samples reflecting alternative external conditions under learned distributions.
• Downscaling with AtmoRep outperforms a GAN baseline in RMSE and can adjust distributions toward higher-resolution references.

Figure 2: AtmoRep can be used for a diverse set of applications without task-specific training (shaded areas depict one standard deviation). Nowcasting: Short-term forecasting can be realized by masking tokens at the future-most time step(s) (bottom right inset). Skill is compared to Pangu-Weather a