ICML2026

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

Maksim Zhdanov, Ana Lucic, Max Welling, Jan-Willem van de Meent

被引用 2 次

摘要

We introduce MOSAIC, a probabilistic weather forecasting model that addresses three failure modes of spectral degradation in ML-based weather prediction: spectral damping (statistical), high-frequency aliasing (architectural), and residual high-frequency leakage (parametric). MOSAIC generates ensemble members through learned functional perturbations and operates on native-resolution grids via mesh-aligned blocksparse 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°resolution with 214M parameters, MOSAIC matches or outperforms models trained on 6× finer resolution on key variables and achieves stateof-the-art results among 1.5°models, producing well-calibrated ensembles whose individual members exhibit near-perfect spectral alignment across all resolved frequencies. A 24member, 10-day forecast takes under 12 s on a single H100 GPU. Code is available at github.com/maxxxzdn/mosaic.