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

2026-04-06

AI Summary

This work addresses spectral degradation issues—such as spectral damping, high-frequency aliasing, and leakage—in machine learning–based weather forecasting by introducing the Mosaic model. Mosaic generates ensemble forecast members through functional perturbations and incorporates a grid-aligned block-sparse attention mechanism to efficiently model long-range dependencies at native resolution, achieving high-fidelity spectral recovery with linear computational complexity. Evaluated at 1.5° resolution, Mosaic matches or exceeds the performance of models trained at six times higher resolution. Moreover, it completes a 10-day, 24-member ensemble forecast in under 12 seconds on a single H100 GPU, attaining spectral fidelity approaching ideal levels.

Abstract

We introduce Mosaic, a probabilistic weather forecasting model that addresses two distinct failure modes of spectral degradation in ML-based weather prediction: (1) spectral damping caused by deterministic training against ensemble means; and (2) aliasing artifacts caused by compressive encoding onto a coarse latent grid. Mosaic generates ensemble members through learned functional perturbations and operates on native-resolution grids via mesh-aligned block-sparse attention, a hardware-aligned mechanism that captures long-range dependencies at linear cost by sharing keys and values across spatially adjacent queries. At 1.5{\deg} resolution with 214M parameters, Mosaic matches or outperforms models trained on 6$\times$ finer resolution on key variables and achieves state-of-the-art results among 1.5{\deg} models, producing well-calibrated ensembles whose individual members exhibit near-perfect spectral alignment across all resolved frequencies. A 24-member, 10-day forecast takes under 12\,s on a single H100~GPU. Code is available at https://github.com/maxxxzdn/mosaic.

Problem

• spectral fidelity
• weather forecasting
• spectral damping
• high-frequency aliasing
• residual leakage

Innovation

• block-sparse attention
• spectral fidelity
• probabilistic forecasting
• ensemble generation
• native-resolution modeling