(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

cs.LGApr 6, 2026

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

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

Organizations: AMLab · University of Amsterdam · CuspAI

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 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° 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° 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.

Figures & tables

Figure 1: Block-sparse attention for weather forecasting. Spatially close query tokens (red block over Tampa Bay) collectively attend to both local key-value pairs (blue block over Florida) and dynamically selected, spatially distributed ones (green blocks). Sparse attention enables capturing long-range dependencies in high-resolution weather data, critical for extreme events such as hurricane formation (note the eye visible in the inset).

Panel of Figure 2 — Figure 2: Spectral analysis and efficiency of MLWP models. (a) Spectral power ratios (model / ERA5) of 10-meter wind speed at $1.5^{\circ}$ resolution, aggregated over 720 initial conditions throughout the 2020 test year (00:00 and 12:00 UTC) at 24 h lead time, with 16 ensemble members for probabilistic models. (b) Same quantity at $0.25^{\circ}$ (HRES-fc0) from a single 6 h forecast. Probabilistic models ( Mosaic , ArchesGen , GenCast ) track the reference closely, while deterministic models ( Stormer , ArchesWeather-Mx4 , GraphCast , Aurora , Pangu-Weather ) suppress fine-scale energy. Mosaic-C (a compressed variant entering at a coarser latent grid) and Mosaic-R (a residual variant predicting $x_{t+1}-x_{t}$ rather than $x_{t+1}$ directly) instead exhibit the opposite signature, with energy rising back toward $1.0$ at the highest wavenumbers, driven by two distinct failure modes: high-frequency aliasing for Mosaic-C , and residual high-frequency leakage for Mosaic-R . (c) Forecast skill vs. inference speed on a single H100 GPU; marker size indicates peak GPU memory. nRMSE is computed at 240 h lead time at $1.5^{\circ}$ resolution over a set of key variables (see Appendix C.8 for details).

Figure 3: HEALPix mesh refinement. Each pixel (left) is subdivided into four children (right), whose indices follow a Z-order curve that keeps spatially close pixels contiguous in memory.

Panel of Figure 4 — Figure 4: Block-sparse attention for weather forecasting. (a) Weather data is interpolated from a latitude-longitude grid (red) to the HEALPix mesh (purple) via cross-attention. (b–d) The three branches of block-sparse attention, illustrated for a single query block (red): (b) Compression computes attention between coarse-grained block representations (squares; color indicates attention score). (c) Selection attends at full resolution to the top-n key blocks. (d) Local attention captures fine-grained interactions within local blocks.

Panel of Figure 5 — Figure 5: Forward pass runtime vs. sequence length for Block-Sparse Attention, NSA, and full FlashAttention; measured on NVIDIA RTX A4500. See Appendix C.3 for details.

	Res.	Step	s/step †	Z500	T850	Q700	U850	V850	SP
IFS HRES	0.1º	1h	—	802.75	3.654	1848.0	6.460	6.530	748.68
Keisler (2022)	1º	6h	—	787.01	3.559	1659.5	6.106	6.121	N/A
Stormer ENS (mean)	1.4º	24h	0.06	665.88	3.001	1445.4	5.198	5.137	619.89
ArchesWeather-Mx4	1.5º	24h	0.06	693.56	3.117	1541.1	5.413	5.431	641.08
NeuralGCM ENS (50)	1.4º	12h	7.42	606.84	2.756	1374.2	4.830	4.852	—
ArchesWeatherGen (mean)	1.5º	24h	0.99	610.36	2.755	1373.7	4.830	4.848	567.15

Table 1: Comparison of

1.5^{\circ}

resolution models on RMSE scores for key weather variables with 10-day lead-time against ERA5 data, 2020 test year. Best scores in bold , second best underlined . Results for additional variables (T2m, U10m, V10m) are provided in Table 3 . † Per-member-per-step inference time (seconds) on a single NVIDIA H100. See Table 13 for full benchmarking details.

Panel of Figure 6 — Figure 6: Forecast skill evaluated against IFS HRES Analysis following the WeatherBench 2 protocol ( Rasp et al., 2023 ) on 2022 test year. Top row: RMSE; middle row: CRPS; bottom row: SSR.

Panel of Figure 7 — Figure 7: Hurricane Ian ensemble track forecasts from three initialization times: Sep 23 (7-day lead), Sep 25 (5-day lead), and Sep 27 (3-day lead), all in 2022. Each panel shows 48 ensemble member tracks (light blue), the ensemble mean (dark blue), and the observed track from HRES-fc0 (black). Markers indicate 24 h intervals. Storm centers are tracked via minimum MSLP guided by IBTrACS best-track positions ( Knapp et al., 2010 ) . Wind field evolution is shown in Fig. 9 .

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Variant	Params	nRMSE
Tiny	59.8M	1.0415
Small	117.9M	0.1746
Medium	195.0M	0.1720
Baseline	214.3M	0.1705

Appendix

Table 2: Ablation studies on 6 h forecast quality (mean nRMSE across all 82 output variables). All variants are trained under a reduced-scale protocol (ERA5 2007–2018, 100k steps,

4{\times}

A6000, single-step rollout, no HRES-fc0 finetuning). Each ablation modifies exactly one aspect relative to the baseline.

	Res.	Step	T2m	U10m	V10m
IFS HRES	0.1º	1h	–	–	–
Keisler (2022)	1º	6h	N/A	N/A	N/A
Stormer	1.4º	24h	2.239	3.535	3.735
NeuralGCM ENS (50)	1.4º	12h	N/A	N/A	N/A
ArchesWeather-Mx4	1.5º	24h	2.323	3.716	3.924
ArchesWeatherGen (mean)	1.5º	24h	2.049	3.306	3.485

Appendix

Table 3: Comparison of coarse-resolution MLWPs on RMSE scores for key weather variables with 240 h lead time against ERA5 data, 2020 test year (continuation of Table 1 ).

	Res.	Z500	T850	Q700	U850	V850	T2m	SP	U10m	V10m
IFS HRES	0.1°	41.23	0.634	537.3	1.126	1.149	0.541	58.70	0.802	0.832
IFS ENS	0.2°	40.91	0.630	503.1	1.102	1.121	0.589	59.44	0.784	0.806
Pangu	0.25°	45.01	0.681	537.2	1.180	1.218	0.624	59.54	0.796	0.825
GraphCast	0.25°	38.32	0.570	475.3	1.021	1.048	0.476	49.36	0.692	0.719
GenCast	0.25°	39.20	0.570	480.7	1.040	1.070	0.462	50.04	0.701	0.732
FGN	0.25°	32.27	0.510	443.2	0.928	0.953	0.418	42.22	0.627	0.655

Appendix

Table 4: RMSE scores for key weather variables at 24 h lead time. All metrics are evaluated at 1.5° resolution.

	Res.	Z500	T850	Q700	U850	V850	T2m	SP	U10m	V10m
IFS HRES	0.1°	809.0	3.626		6.417	6.500	2.576	753.3	4.513	4.767
IFS ENS	0.2°	621.9	2.761		4.854	4.918	1.991	574.7	3.394	3.595
Pangu	0.25°	799.8	3.573		6.267	6.334	2.648	742.4	4.378	4.633
GraphCast	0.25°	747.6	3.364		5.953	6.032	2.415	698.9	4.195	4.436
GenCast	0.25°	607.5	2.693		4.758	4.826	1.932	563.2	3.325	3.521
FGN	0.25°	583.1	2.595		4.636	4.710	1.864	542.7	3.247	3.440

Appendix

Table 5: RMSE scores for key weather variables at 240 h (10-day) lead time. All metrics are evaluated at 1.5° resolution.

Panel of Figure 8 — Figure 8: RMSE (rows 1–2), CRPS (rows 3–4), and spread-to-skill ratio (rows 5–6) as a function of lead time for additional variables not shown in the main text. Within each metric, the top row shows surface variables (10-meter U-wind, 10-meter V-wind, mean sea level pressure) and the bottom row shows pressure-level variables (U850, V850, Q700). Values close to 1.0 (dashed line) in the SSR rows indicate well-calibrated ensembles.

Figure 9: Hurricane Ian wind field evolution. Init: Sep 23, 2022 12Z. Columns show lead times (+24 h, +72 h, +120 h, +168 h). Rows: HRES-fc0 ground truth, three individual ensemble members, and ensemble mean (48 members). Variable: 10-meter wind speed (m/s). Region: Gulf of Mexico / Caribbean / southeastern US.

Lead time	Mean drift (hPa)	Relative	Max member
24 h	$-0.008\pm 0.045$	0.001%	$<0.2$ hPa
120 h	$-0.042\pm 0.112$	0.004%	$<0.5$ hPa
240 h	$-0.086\pm 0.188$	0.009%	$<1.0$ hPa

Appendix

Table 6: GMSP drift over 10-day rollouts (2020 test year, 712 init dates

\times

48 members).

Figure 10: GMSP drift (mean

\pm

\sigma

) over lead time across 34,176 ten-day rollouts. The drift remains below 0.1 hPa throughout, indicating stable mass conservation.

Model	Params	Precision	Resolution	Training	Duration
GraphCast	36M	fp32	0.25°	32 TPUv4	4 weeks
GenCast	57.5M	fp32	0.25°	32 TPUv5	5 days
FGN	57.5M	fp32	0.25°	490 TPUv5p/v6e	3 days*
Pangu	277M	fp32	0.25°	192 V100	15 days
Aurora	1,259M	bf16	0.25°	32 A100	2.5 weeks
Mosaic	214M	fp16	1.5°	8 H100	2 days

Appendix

Table 7: Comparison of Mosaic with baseline MLWPs. Mosaic achieves competitive performance despite operating at coarser resolution with significantly reduced computational resources.

Variable	Type	Unit	Channels
Surface-level variables (dynamic)
2-meter temperature	Surface	K	1
10-meter U-wind	Surface	m/s	1
10-meter V-wind	Surface	m/s	1
Mean sea level pressure	Surface	Pa	1
Pressure-level variables (dynamic)

Appendix

Table 8: Complete list of variables used by Mosaic . Pressure-level variables are provided at 13 levels: {50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000} hPa.

Stage	nside	Dim	Heads	Enc. Depth	Dec. Depth	MLP Ratio
Stage 1	64	768	12	4	2	4.0
Stage 2	32	1024	16	4	2	4.0
Bottleneck	16	1280	20	2		4.0
Global parameters
Total parameters: 214M
GQA ratio: 4, QKV compression ratio: 1

Appendix

Table 9: Model architecture details.

Stage	Steps	AR Length	Learning Rate	Weight Decay	Schedule
Pretrain	250k	1	1e-3	1e-2	Cosine (1e-3 $\to$ 1e-6), no warmup
Finetune 1	30k	1	1e-4	1e-2	Cosine decay by 1e-2
Finetune 2	10k	2	1e-5	1e-2	Cosine decay by 1e-2
Finetune 4	5k	4	5e-6	1e-2	Cosine decay by 1e-2
Finetune 8	2.5k	8	1e-6	1e-2	Cosine decay by 1e-2
Finetune 12	2.5k	12	1e-6	1e-2	Cosine decay by 1e-2

Appendix

Table 10: Training schedule details.

Variable / Level	$\boldsymbol{\alpha_{i}}$
Surface variables
2-meter temperature	1.0
10-meter U-wind	0.1
10-meter V-wind	0.1
Mean sea level pressure	0.1
Pressure levels (shared across all 6 pressure-level variables)

Appendix

Table 11: Variable-level loss weights

\alpha_{i}

. All six pressure-level variables share the same per-level weight.

Seq. len.	BSA	NSA	Dense	BSA/Dense
8k	2.97	7.52	2.68	0.9 $\times$
16k	6.31	17.95	11.20	1.8 $\times$
32k	13.13	36.49	44.75	3.4 $\times$
64k	27.84	86.04	176.85	6.4 $\times$
128k	57.15	234.12	752.74	13.2 $\times$
256k	114.72	700.32	3,383.69	29.5 $\times$

Appendix

Table 12: Forward pass runtime (ms) for BSA, NSA, and dense attention across sequence lengths. Measured on NVIDIA RTX A4500 (batch=1, 16 heads, head_dim=32). Dense uses compiled flex_attention .

Panel of Figure 11 — Figure 11: Runtime comparison of native sparse attention (NSA) vs block-sparse attention (BSA, ours) on NVIDIA A4500. NSA implementation from Yang & Zhang (2024) . BSA achieves consistent speedups across all sequence lengths.

Model	Res.	Params	Prec.	s/mem./step	Peak (GB)
Mosaic (Ours)	1.5°	214M	fp16	0.048	3.26
ArchesWeatherGen	1.5°	384M	fp16	0.986	2.40
ArchesWeather-Mx4	1.5°	339M	fp16	0.058	2.08
Stormer	1.4°	469M	fp16	0.062	11.24
NeuralGCM †	1.4°	12M	fp32	2.182	2.18
GraphCast †	0.25°	36M	bf16	42.57	36.55

Appendix

Table 13: Inference benchmarking on a single NVIDIA H100 GPU. s/mem./step is the wall-clock time per ensemble member per autoregressive step, computed from the largest ensemble run that fits in memory (24 members for Mosaic , ArchesWeatherGen, and NeuralGCM; 4 for ArchesWeather-Mx4; 1 for the rest). torch.compile is enabled for PyTorch models where applicable.

Figure 12: Forecast rollout trajectories showing 10-day evolution of wind speed fields at 850 hPa.

Figure 13: Forecast rollout trajectories showing 10-day evolution of surface temperature fields.

Panel of Figure 14 — Figure 14: Examples of global temperature spectra at 2-meter height (top row) and kinetic energy spectra at 10-meter height (bottom row) for 10-day Mosaic forecasts compared to HRES-fc0 0.25 ∘ ground truth, shown for multiple initial conditions (all start at 00:00 UTC).

Weather Forecasting

Weather forecasting predicts future atmospheric conditions, from local storms to global patterns, using physical models, observations, and machine learning.

See all 176 papers

Scale-aware and spatially-localised optimisation of kilometre-scale MLWP models
Oct 27, 2026James Chappell, Sophie Luise Arana, Jack Bowyer +13Weather Forecasting Spatiotemporal Forecasting
Met Office, Exeter, United Kingdom · The Alan Turing Institute, London, United Kingdom
Super-Resolving Coarse-Resolution Weather Forecasts With Flow Matching
Apr 1, 2026Aymeric Delefosse, Anastase Charantonis, Dominique BéréziatWeather Forecasting Probabilistic Forecasting
Inria, ARCHES, Paris, France · Sorbonne Université, CNRS, LIP6, Paris, France
Marchuk: Efficient Global Weather Forecasting from Mid-Range to Sub-Seasonal Scales via Flow Matching
Mar 25, 2026Arsen Kuzhamuratov, Mikhail Zhirnov, Andrey Kuznetsov +2Weather Forecasting Spatiotemporal Forecasting
AXXX · FusionBrain Lab

Oct 27, 2026J/K move · Enter open · S save

Scale-aware and spatially-localised optimisation of kilometre-scale MLWP models

James Chappell, Sophie Luise Arana, Jack Bowyer, Alejandro Coca-Castro+12

Met Office, Exeter, United Kingdom · The Alan Turing Institute, London, United Kingdom

Weather Forecasting Spatiotemporal Forecasting Wavelet Methods

Open paper AskPDF arXiv

Preserving fine-scale structure without degrading spatial accuracy is a central challenge for kilometre-scale machine-learning weather prediction (MLWP). We introduce a loss based on the dual-tree complex wavelet transform that constrains atmospheric fields jointly by scale, location, and orientation. We evaluate its performance using an experimental joint global--regional FastNet configuration, which produces \SI{1.5}{km} forecasts over the UK alongside

\sim

\SI{31}{km} global forecasts. Relative to both mean-squared error (MSE) and an FFT-based spectral objective, the wavelet loss substantially improves spectral fidelity. Compared with MSE training, it maintains similar pointwise performance, improves neighbourhood-based verification, and produces sharper precipitation fields with more realistic intensity distributions. These results demonstrate that spatially localised wavelet constraints can preserve multiscale atmospheric structure while largely avoiding the loss of pointwise accuracy associated with globally supported spectral objectives.

Abstract

Weather Forecasting

Explore similar work

Scale-aware and spatially-localised optimisation of kilometre-scale MLWP models

Super-Resolving Coarse-Resolution Weather Forecasts With Flow Matching

Marchuk: Efficient Global Weather Forecasting from Mid-Range to Sub-Seasonal Scales via Flow Matching