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

arXiv: https://arxiv.org/abs/2604.16429 · HTML: https://mlsanity.com/paper/2604.16429/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 63 references that are in the map.
- First posted on arXiv 6 Apr 2026.

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [Message Passing Neural PDE Solvers](https://mlsanity.com/paper/2202.03376/) (2022)
- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [Neural General Circulation Models for Weather and Climate](https://mlsanity.com/paper/2311.07222/) (2023)
- [WeatherBench 2: A benchmark for the next generation of data-driven global weather models](https://mlsanity.com/paper/2308.15560/) (2023)
- [Rotary Position Embedding for Vision Transformer](https://mlsanity.com/paper/2403.13298/) (2024)
- [PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers](https://mlsanity.com/paper/2308.05732/) (2023)
- [GenCast: Diffusion-based ensemble forecasting for medium-range weather](https://mlsanity.com/paper/2312.15796/) (2023)
- [Scaling transformer neural networks for skillful and reliable medium-range weather forecasting](https://mlsanity.com/paper/2312.03876/) (2023)
- [Skillful joint probabilistic weather forecasting from marginals](https://mlsanity.com/paper/2506.10772/) (2025)
- [FourCastNet 3: A geometric approach to probabilistic machine-learning weather forecasting at scale](https://mlsanity.com/paper/2507.12144/) (2025)
- [Towards Stability of Autoregressive Neural Operators](https://mlsanity.com/paper/2306.10619/) (2023)
- [Integrating Neural Operators with Diffusion Models Improves Spectral Representation in Turbulence Modeling](https://mlsanity.com/paper/2409.08477/) (2024)
- [ArchesWeather & ArchesWeatherGen: a deterministic and generative model for efficient ML weather forecasting](https://mlsanity.com/paper/2412.12971/) (2024)
- [PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations](https://mlsanity.com/paper/2505.24717/) (2025)
- [Advancing Parsimonious Deep Learning Weather Prediction using the HEALPix Mesh](https://mlsanity.com/paper/2311.06253/) (2023)
- [Fixing the Double Penalty in Data-Driven Weather Forecasting Through a Modified Spherical Harmonic Loss Function](https://mlsanity.com/paper/2501.19374/) (2025)
- [Scaling Laws in Patchification: An Image Is Worth 50,176 Tokens And More](https://mlsanity.com/paper/2502.03738/) (2025)
- [Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems](https://mlsanity.com/paper/2502.17019/) (2025)
- [Grounding Continuous Representations in Geometry: Equivariant Neural Fields](https://mlsanity.com/paper/2406.05753/) (2024)
- [BLADE: Block-Sparse Attention Meets Step Distillation for Efficient Video Generation](https://mlsanity.com/paper/2508.10774/) (2025)
- [ProxyAttn: Guided Sparse Attention via Representative Heads](https://mlsanity.com/paper/2509.24745/) (2025)

## Built on by (most cited first, 5 of 5)

- [On the sensitivity of machine-learned probabilistic weather forecast models to scale-aware scoring rules](https://mlsanity.com/paper/2607.19161/) (2026)
- [Lightweight Probabilistic Downscaling from a Deterministic Base Model](https://mlsanity.com/paper/2609.29383/) (2026)
- [Proper Scoring Rule-based Diffusion for Probabilistic Weather Forecasting](https://mlsanity.com/paper/2609.38632/) (2026)
- [SO(3)-RoPE for Spherical Transformers](https://mlsanity.com/paper/2610.06229/) (2026)
- [Cross-Domain Pretraining for Steady-State Neural CFD Surrogates](https://mlsanity.com/paper/2610.10398/) (2026)

## Abstract (arXiv)

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 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.5deg resolution with 214M parameters, Mosaic matches or outperforms models trained on 6× finer resolution on key variables and achieves state-of-the-art results among 1.5deg 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 12s on a single H100~GPU. Code is available at https://github.com/maxxxzdn/mosaic.