Efficient weather forecasting model offers high resolution and lower compute cost
Efficient probabilistic weather forecasting at quarter-degree resolution with two-sided transport attention
Machine Learning
Summary
Predicting the weather accurately usually needs a lot of computer power, especially at fine detail. The authors created Astra, a weather forecasting model that works at a quarter-degree resolution using less computing resources. Astra uses a special attention method that lets different parts of the forecast interact in both directions to make better predictions. Tests show Astra matches or beats existing models while using far fewer calculations and running on accessible hardware like a single consumer GPU.
What this means in practice
- •For weather forecasters: Run accurate high-resolution probabilistic weather forecasts faster and cheaper on standard GPUs, improving availability for local predictions.
- •For climate modelers: Use Astra’s efficient attention approach to accelerate ensemble simulations while maintaining accuracy across multiple atmospheric variables and lead times.
Authors
Junu Choi, Youngjoon Hong
Abstract
High-resolution probabilistic weather forecasting remains computationally demanding. Here we present Astra, a 0.25$^\circ$ ensemble model trained on eight H200 GPUs. It combines Kantorovich Attention, a two-sided transport attention in which each region chooses where it reads from and writes to, so every region sends and receives at token resolution, with a self-featurised continuous ranked probability score (CRPS) that also scores forecasts in the model's multiscale features. In WeatherBench 2, Astra matches GenCast's CRPS averaged over five upper-air targets and days 1-10 with 9.8% of its hardware-normalised training compute. Verified against ERA5, its CRPS and ensemble-mean root-mean-square error are lower than the ECMWF ensemble's on all 600 variable, level and lead-time combinations, and its CRPS is lower than GenCast's on 73.5% of them. A 15-day member needs 39 times fewer floating-point operations than GenCast's GPU implementation, and 50 members take under 40 min on one consumer GPU, within 21 GB.