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

arXiv:2604.16429v1 Announce Type: cross Abstract: We introduce Mosaic, a probabilistic weather forecasting model that addresses two principal sources of spectral degradation in ML-based weather prediction: (1) deterministic training against ensemble means and (2) compressive encoding creating an information bottleneck. Mosaic generates ensemble members through learned functional perturbations and operates on native-resolution grids via 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.5eg resolution with 214M parameters, Mosaic matches or outperforms models trained on 6 times finer data on headline upper-air variables and achieves state-of-the-art results among 1.5eg 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 seconds on a single H100 GPU.

Related

• Post-Hoc Guidance for Consistency Models by Joint Flow Distribution Learning
• MCLR: Improving Conditional Modeling via Inter-Class Likelihood-Ratio Maximization and Unifying Classifier-Free Guidance with Alignment Objectives
• Diagnosing and Improving Diffusion Models by Estimating the Optimal Loss Value
• SHIFT: Steering Hidden Intermediates in Flow Transformers

Source: arXiv cs.CV | 2026-04-21