Efficient probabilistic weather forecasting at quarter-degree resolution with two-sided transport attention
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.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
Astra achieves quarter-degree probabilistic weather forecasting with 9.8% of GenCast's training compute and 39x less inference FLOPs by introducing two-sided transport attention (Kantorovich Attention) and a self-featurised CRPS objective, enabling high-resolution ensemble forecasting on a single consumer GPU.
Methodology
The paper introduces "Astra," a probabilistic weather forecasting model that achieves 0.25° resolution with significantly reduced computational cost compared to state-of-the-art models like GenCast. The core methodological contributions are twofold: (1) **Kantorovich Attention (KA)**, a two-sided transport attention mechanism where each grid tile selects both where it reads from (pull) and where it writes to (push). This ensures every tile is read at token resolution, avoiding the "unread" tile problem of one-sided sparse attention, while maintaining linear complexity relative to grid points. (2) **Self-featurised CRPS**, a training objective that scores ensemble members not just in physical space but also in the model's own multiscale feature space using a frozen copy of the network. This encourages the model to learn representations that are discriminative and structurally consistent, similar to how diffusion models use pretrained encoders. The architecture is a W-Net operating directly on the 0.25° grid without a latent encoder, using MeanFlow for single-step generation.
Results & evaluation
The experimental evaluation is rigorous and comprehensive. Astra is benchmarked against GenCast (the current SOTA for high-res probabilistic forecasting) and the ECMWF IFS ENS (the operational gold standard). Key results include: (1) Matching GenCast's CRPS on upper-air targets with only 9.8% of the hardware-normalised training compute. (2) Outperforming IFS ENS on all 600 variable/level/lead-time combinations in CRPS and ensemble-mean RMSE. (3) Outperforming GenCast in CRPS on 73.5% of combinations. (4) Significant efficiency gains: 39x fewer FLOPs for inference than GenCast's GPU implementation, allowing 50-member 15-day forecasts on a single consumer GPU (RTX 5090) in under 40 minutes. Spectral analysis confirms that Astra preserves fine-scale spatial structure better than GenCast and IFS ENS, particularly in the Alps and tropics. The use of WeatherBench 2 and ERA5 verification provides a standard, reproducible benchmark.
Reproducibility
The paper provides extensive methodological details, including equations for the attention mechanism, the transportation problem formulation, and the self-featurised CRPS loss. It specifies hardware requirements (8x H200 GPUs for training) and inference costs (RTX 5090). While the code is not provided in the text, the level of detail in the Methods section (e.g., specific tile sizes, number of heads, initialization schemes, and noise prior construction) is high, suggesting that a competent team could reproduce the results. The reliance on specific hardware (H200/RTX 5090) may limit immediate reproducibility for labs with older GPUs, but the algorithmic details are clear.
Limitations
The primary limitation is the hardware dependency for training; while inference is efficient, training still requires 8x H200 GPUs, which is a significant barrier for smaller labs compared to models trained on fewer resources. The model is evaluated primarily on ERA5 reanalysis data; performance on real-time operational data with missing or noisy inputs is not tested. The "self-featurised" component relies on a frozen copy of the network, which adds complexity to the training pipeline and may introduce instability if the network's features drift significantly during training. Additionally, the paper focuses on medium-range forecasting (up to 15 days); performance at longer lead times (sub-seasonal) is not evaluated.
Broader impact
This paper has high potential impact by democratizing high-resolution probabilistic weather forecasting. By reducing inference costs to the level of a single consumer GPU, it enables local, real-time probabilistic forecasting for applications such as renewable energy management, agriculture, and disaster response, where cloud-based or high-performance computing access may be limited. The two-sided transport attention mechanism is a generalizable architectural innovation that could be applied to other high-resolution spatial tasks (e.g., remote sensing, medical imaging) where global context is needed but dense attention is prohibitive. The self-featurised CRPS offers a new paradigm for training generative models to produce structurally consistent samples, potentially influencing broader generative AI research. Astra achieves quarter-degree probabilistic weather forecasting with 9.8% of GenCast's training compute and 39x less inference FLOPs by introducing two-sided transport attention (Kantorovich Attention) and a self-featurised CRPS objective, enabling high-resolution ensemble forecasting on a single consumer GPU.
Author of this paper? Send a correction.