[Paper Review] A Foundation Model for the Earth System

Cristian Bodnar, Wessel P. Bruinsma+16 more|arXiv (Cornell University)|May 20, 2024

Solar and Space Plasma Dynamics35 citations

TL;DR

Aurora is a 1.3B-parameter foundation model of the atmosphere trained on over a million hours of diverse weather and climate data. It delivering fast, high-resolution forecasts that outperform state-of-the-art classical tools and other deep learning models in multiple scenarios, including 5-day global air pollution and 10-day high-resolution weather predictions.

ABSTRACT

Reliable forecasts of the Earth system are crucial for human progress and safety from natural disasters. Artificial intelligence offers substantial potential to improve prediction accuracy and computational efficiency in this field, however this remains underexplored in many domains. Here we introduce Aurora, a large-scale foundation model for the Earth system trained on over a million hours of diverse data. Aurora outperforms operational forecasts for air quality, ocean waves, tropical cyclone tracks, and high-resolution weather forecasting at orders of magnitude smaller computational expense than dedicated existing systems. With the ability to fine-tune Aurora to diverse application domains at only modest computational cost, Aurora represents significant progress in making actionable Earth system predictions accessible to anyone.

Motivation & Objective

• Motivate the use of foundation models to learn generalizable atmospheric representations from diverse, large-scale data.
• Develop Aurora, a flexible 3D foundation model for weather and atmospheric processes that handles heterogeneous inputs and multiple resolutions.
• Demonstrate Aurora's ability to deliver fast, accurate forecasts for tasks with limited data and extreme events.
• Show that data and model scaling improve forecasting performance across datasets and tasks.

Proposed method

• A 3D Swin Transformer-based processor paired with Perceiver-based encoder/decoder to handle heterogeneous inputs and outputs across space, pressure levels, and resolutions.
• Pretraining with a mixed dataset regime spanning ERA5, CMCC, IFS-HR, HRES Forecasts, GFS Analysis, and GFS Forecasts to minimize next-time-step MAE with a 6-hour lead time over 150k steps on 32 A100 GPUs.
• Two-stage fine-tuning consisting of short-lead-time fine-tuning followed by long-lead time rollout fine-tuning using LoRA (Low Rank Adaptation).
• Forecasts generated by iteratively feeding predictions back into the model to produce different lead times.
• Evaluation against CAMS analyses and IFS-HRES data to assess performance across atmospheric chemistry, pollution, and high-resolution weather tasks.
• Comparison against GraphCast and IFS-HRES at 0.25° resolution to benchmark AIWP models against traditional NWP.

Experimental results

Research questions

• RQ1 Can a large, heterogeneous-data-trained foundation model achieve operationally relevant forecasts across atmospheric chemistry, air pollution, and high-resolution weather tasks?
• RQ2 Does pretraining on diverse datasets improve performance and robustness, especially for data-sparse or extreme-event regimes?
• RQ3 Can a single model trained across resolutions and data fidelities outperform specialized, traditional systems in accuracy and efficiency?
• RQ4 What are the effects of model and data scaling on forecast skill and tail performance?
• RQ5 How does Aurora compare with existing AI weather prediction models and with IFS-HRES on observations and extreme-event cases?

Key findings

• Aurora produces 5-day global air pollution forecasts at 0.4° resolution and 10-day forecasts at 0.1° resolution with substantial speedups over IFS, outperforming state-of-the-art tools on most targets.
• Aurora matches or outperforms CAMS forecasts on 74% of targets and is within 20% RMSE of CAMS on 95% of targets in the CAMS analysis setting.
• At 0.1° resolution, Aurora outperforms IFS-HRES across lead times beyond 12 hours, with up to 60% RMSE reduction at longer lead times.
• Aurora at 0.25° resolution consistently outperforms GraphCast and approaches IFS-HRES performance on several variables and lead times, especially for wind and temperature.
• Data diversity and model scaling improve forecasting performance, with gains in tail performance and across variables when pretraining on multiple datasets (C2/C3/C4) vs ERA5 alone (C1).
• Storm Ciarán case study shows Aurora capturing abrupt wind-speed increases better than other AI models at 0.1° resolution.