(Sparse) Attention to the Details: Preserving Spectral Fidelity in ML-based Weather Forecasting Models

We introduce \textsc{Mosaic}, a probabilistic weather forecasting model that addresses two sources of spectral degradation in ML-based weather prediction: training to predict the ensemble mean deterministically and compressive encoding creating an information bottleneck. \textsc{Mosaic} combines learned functional perturbations for ensemble forecasting with block-sparse attention, a hardware-aligned formulation that shares keys and values across spatially adjacent queries, enabling each block to dynamically attend to the most relevant regions. By capturing arbitrarily long-range dependencies at linear cost, \textsc{Mosaic} processes high-resolution weather data without compression. On IFS HRES data, \textsc{Mosaic} at 1.5° resolution matches or outperforms models trained on 0.25° data, with individual ensemble members exhibiting near-perfect spectral alignment across all resolved frequencies.

Mosaic is a weather forecasting model that is physically consistent with weather data. It identifies three sources of violation of the consistency and addresses them all. First, it uses efficient sparse attention to operate on large Earth grids. Then, it generates multiple possible predictions. Third, it predicts the next weather state directly. It is extremely fast and efficient.