ClimaX: A Foundation AI Model for Weather and Climate
1. Observation and Context
Today, standard weather and climate predictions rely on complex computer simulations called General Circulation Models (GCMs). These models use physics equations to simulate how the atmosphere works. However, GCMs require massive amounts of supercomputing power. This makes them too slow for fast, short-term updates (like predicting rain in a few hours) or accurate predictions far into the future (like 7+ days out).
On the other hand, artificial intelligence (AI) models can make predictions instantly. The problem is that current AI models are trained on single, specific datasets. This means they can only do one specific job and cannot handle different regions, missing data, or new climate variables.
2. Hypothesis
If we build an flexible AI model using a Vision Transformer architecture and pretrain it on many different climate simulation datasets using a randomized forecasting goal, it will act as a general-purpose “foundation model” for atmospheric science. Once trained, this single AI can be quickly adapted to perform a wide variety of weather and climate tasks better than existing single-purpose AI models and traditional computer simulations.
3. Experiment and Methodology
The researchers created an AI model named ClimaX based on the Vision Transformer (ViT) framework. To help it read messy climate data with many different variables (like pressure, temperature, and humidity), they added two main features:
- Variable Tokenization: Cuts up individual weather map variables into small image-like patches.
- Variable Aggregation: Uses an attention mechanism to combine those different variable patches at each map location, making the data easier and faster for the AI to process.
The team pretrained ClimaX using historical computer simulations from CMIP6 (spanning the years 1850 to 2015) by making it predict future weather at random time steps between 6 hours and 7 days.
They then fine-tuned and tested ClimaX on four main challenges:
- Global & Regional Forecasting: Predicting worldwide and North American weather up to 1 month ahead.
- Sub-seasonal to Seasonal Prediction: Forecasting two-week averages up to 6 weeks out.
- Climate Projection: Predicting long-term global temperature trends based on greenhouse gases like CO2 and SO2.
- Climate Model Downscaling: Taking low-resolution climate predictions and turning them into high-detail local maps.
4. Results and Data
- Global Weather Forecasting: ClimaX beat traditional numerical models (ECMWF’s IFS) for predictions 7 days out or longer. At higher map resolutions, its short-term accuracy matched operational physics models.
- Regional & Seasonal Prediction: It made fewer errors in regional North American forecasts and multi-week predictions compared to standard CNN deep learning baselines.
- Climate Projections & Downscaling: ClimaX set new top performance scores on the ClimateBench benchmark for temperature projections and successfully sharpened low-resolution climate maps while removing input bias.
- Scaling Up: The model got smarter and more accurate as the researchers gave it larger training datasets, bigger computing capacity, and higher-resolution maps.
5. Conclusion and Climate Impact
ClimaX proves that a single, pretrained “foundation” AI model can handle a wide variety of Earth science tasks. By combining fast AI processing with the flexibility to adapt to different locations, timescales, and variables, ClimaX can help communities better prepare for extreme weather events and help policymakers make smarter decisions regarding long-term climate change.
6. Full Citation
Nguyen, T., Brandstetter, J., Kapoor, A., Gupta, J. K., & Grover, A. (2023). ClimaX: A foundation model for weather and climate. arXiv preprint arXiv:2301.10343v5. https://doi.org/10.48550/arXiv.2301.10343