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× 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).
Figure 2: Spectral analysis and efficiency of MLWP models. (a) Spectral power ratios (model / ERA5) of 10-meter wind speed at 1.5∘ 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∘ (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 xt+1−xt rather than xt+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∘ 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.
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.
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∘ 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.
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.
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× 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.
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±0.045
0.001%
<0.2 hPa
120 h
−0.042±0.112
0.004%
<0.5 hPa
240 h
−0.086±0.188
0.009%
<1.0 hPa
Appendix
Table 6: GMSP drift over 10-day rollouts (2020 test year, 712 init dates × 48 members).
Figure 10: GMSP drift (mean ± 1 σ ) 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 → 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
α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 α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 ×
16k
6.31
17.95
11.20
1.8 ×
32k
13.13
36.49
44.75
3.4 ×
64k
27.84
86.04
176.85
6.4 ×
128k
57.15
234.12
752.74
13.2 ×
256k
114.72
700.32
3,383.69
29.5 ×
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 .
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.
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 predicts future atmospheric conditions, from local storms to global patterns, using physical models, observations, and machine learning.
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 ∼\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.
James Chappell, Sophie Luise Arana, Jack Bowyer +13
Met Office, Exeter, United Kingdom · The Alan Turing Institute, London, United Kingdom
Machine learning-based weather forecasting models now surpass state-of-the-art numerical weather prediction systems, but training and operating these models at high spatial resolution remains computationally expensive. We present a modular framework that decouples forecasting from spatial resolution by applying learned generative super-resolution as a post-processing step to coarse-resolution forecast trajectories. We formulate super-resolution as a stochastic inverse problem, using a residual formulation to preserve large-scale structure while reconstructing unresolved variability. The model is trained with flow matching exclusively on reanalysis data and is applied to global medium-range forecasts. We evaluate (i) design consistency by re-coarsening super-resolved forecasts and comparing them to the original coarse trajectories, and (ii) high-resolution forecast quality using standard ensemble verification metrics and spectral diagnostics. Results show that super-resolution preserves large-scale structure and variance after re-coarsening, introduces physically consistent small-scale variability, and achieves competitive probabilistic forecast skill at 0.25° resolution relative to an operational ensemble baseline, while requiring only a modest additional training cost compared with end-to-end high-resolution forecasting.
Accurate subseasonal weather forecasting remains a major challenge due to the inherently chaotic nature of the atmosphere, which limits the predictive skill of conventional models beyond the mid-range horizon (approximately 15 days). In this work, we present Marchuk, a generative latent flow-matching model for global weather forecasting spanning mid-range to subseasonal timescales, with prediction horizons of up to 30 days. Marchuk conditions on current-day weather maps and autoregressively predicts subsequent days' weather maps within the learned latent space. We replace rotary positional encodings (RoPE) with trainable positional embeddings and extend the temporal context window, which together enhance the model's ability to represent and propagate long-range temporal dependencies during latent forecasting. Marchuk offers two key advantages: high computational efficiency and strong predictive performance. Despite its compact architecture of only 276 million parameters, the model achieves performance comparable to LaDCast, a substantially larger model with 1.6 billion parameters, while operating at significantly higher inference speeds. We open-source our inference code and model at: https://v-gen-ai.github.io/Marchuk/