(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
Abstract
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 82d59161-9b0b-430d-8167-7d0ef6a90551Cited by top-tier papers1
Ask how each one uses itBuilds on14
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra et al.NeurIPS 2022 · 5,493 citations
- Message Passing Neural PDE SolversJohannes Brandstetter, Daniel E. Worrall, Max WellingICLR 2022 · 410 citations
- Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse AttentionJingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo et al.ACL 2025 · 334 citations
- PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE SolversPhillip Lippe, Bas Veeling, Paris Perdikaris, Richard E. Turner et al.NeurIPS 2023 · 280 citations
- Scaling transformer neural networks for skillful and reliable medium-range weather forecastingTung Nguyen, Rohan Shah, Hritik Bansal, Troy Arcomano et al.NeurIPS 2024 · 165 citations
Related papers
- PESD-TSF: A Period-Aware and Explicit Structured Decomposition Framework for Long-Term Time Series ForecastingHua Wang, Xianhao Jiao, Fan ZhangICML 2026
- Robust Adverse Weather Removal via Spectral-based Spatial GroupingYuhwan Jeong, Yunseo Yang, Youngho Yoon, Kuk-Jin YoonICCV 2025 · 1 citation
- Multi-Scale Wavelet Transformers for Operator Learning of Dynamical SystemsXuesong Wang, Michael Groom, Rafael Oliveira, He Zhao et al.ICML 2026
- Probabilistic Weather Forecasting with Hierarchical Graph Neural NetworksJoel Oskarsson, Tomas Landelius, Marc Peter Deisenroth, Fredrik LindstenNeurIPS 2024 · 44 citations
- Multi-scale Physics-informed Transformer With Spatio-temporal Feature Adapter For Extreme Precipitation NowcastingJingyuan Zheng, Xin Zhang, Zhilin Qi, Ruiang Qiu et al.KDD 2025
