Scale-aware and spatially-localised optimisation of kilometre-scale MLWP models
Authors: James Chappell, Sophie Luise Arana, Jack Bowyer, Alejandro Coca-Castro, Evangeline Corcoran, Eric G. Daub, Tom Dunstan, Noushin Eftekhari, Anna-Louise Ellis, J. Scott Hosking, Manvendra Janmaijaya, David Salvador-Jasin, Ben Shipway, Nathan Simpson, Oliver Strickson, David Walters
Organizations: Met Office, Exeter, United Kingdom · The Alan Turing Institute, London, United Kingdom
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.
Figures & tables
Figure 1: Overview of the regional MLWP architecture and wavelet training objective. (a) The model produces simultaneous forecasts on a global N320 reduced Gaussian grid and a 1.5km regional grid centred on the UK; the example shows 6hour accumulated precipitation valid at 2023-09-28T00Z UTC at a lead time of 24hours . (b) The shared latent processor mesh spans the global domain and is locally refined over the regional domain. Black and orange points denote global and regional processor nodes, respectively, while curves show representative multiscale processor edges from one regional node. The edges are coloured according to their length and are represented by arcs for visual clarity. (c) During loss calculation, the prediction and target are decomposed using a J -level DT C WT. At each level, the preceding low-pass component is separated into a coarser low-pass component and six complex directional wavelet-bands. The high-pass coefficients contribute amplitude- and coherence-sensitive losses that can be weighted independently by scale, while the final low-pass residual is compared using MSE.
Figure 2: Scale-wise decomposition of the wavelet loss for two constructed prediction fields. (a) Target field and component-wise summaries for a perturbed 6hour time-offset prediction (Case 1) and a Gaussian-blurred target with σ=3 grid cells (Case 2). The scalar-loss percentages give the exact contribution of the final low-pass term and each high-pass wavelet level to the total loss. The gradient percentages give each component-gradient L2 norm as a fraction of the sum of the component-gradient L2 norms; they therefore provide a relative measure of gradient magnitude rather than an additive decomposition of the norm of the total gradient. (b,c) The two constructed predictions and the gradients of the total loss and its low-pass and level-wise components with respect to the predicted field on the original input grid. Each map shows the signed first-order sensitivity of the corresponding scalar loss component and therefore the field-space training signal supplied by that component.
Figure 3: Evolution of validation losses and relative scale-wise wavelet-loss contributions over the first 25% of Stage 3a training for (a) 500hPa geopotential, (b) 2m temperature, and (c) 6h accumulated total precipitation. All quantities are evaluated over the regional UKV domain using validation forecasts from the wavelet-trained model; the global ERA5 component is not included. Upper panels show pointwise MSE and the complete wavelet objective, both normalised by the wavelet validation loss at the first recorded training step. Lower panels show the low-pass MSE component and individual wavelet-band contributions as fractions of the total regional wavelet validation loss. Here, j=1 denotes the finest wavelet-band, with increasing j corresponding to progressively coarser scales. All decomposition levels use uniform scale weights, αj=1 .
Figure 4: Verification of the MSE- and wavelet-trained models against surface observations as a function of lead time. The top row shows HiRA CRPS for 3×3 , 7×7 and 11×11 model-grid neighbourhoods; the bottom row shows pointwise RMSE. Columns show (a,e) 2m temperature, (b,f) 10m wind speed, (c,g) 10m wind-vector error and (d,h) 6hour accumulated precipitation. Curves are aggregated over forecasts initialised between 1 January and 21 December 2023. Lower values indicate better performance.
Figure 5: Zonal power spectra of forecasts from the MSE- (blue), FFT-AMSE- (purple), and wavelet-trained (orange) models at a lead time of 24 hours, verified against UKV operational analyses (black, dashed). Columns show (a) 2m temperature, (b) 10m wind speed, and (c) 6hour accumulated precipitation. Upper panels show the absolute spectra, while lower panels show log10(Pf/Pr) , where Pf and Pr are the forecast and reference powers, respectively. Zero indicates equal forecast and reference power, while negative values indicate a forecast power deficit. Vertical dotted lines mark the internal processor-mesh spacing, Δxmesh=4\text{,}\mathrm{k}\mathrm{m}$$ , and wavelengths corresponding to 2Δxmesh , 4Δxmesh , and 8Δxmesh .
Figure 6: Precipitation forecast characteristics. (a) Mean rainfall rate over the preceding 6hours for the case valid at 0000 UTC on 14 October 2023, calculated by dividing the 6hour precipitation accumulations by six. The panels show the UKV operational NWP forecast at a lead time of 6hours and the wavelet- and MSE-trained model forecasts at a lead time of 24hours . (b) Quantile–quantile comparison of forecast and analysis 6hour accumulated precipitation over the 2023 verification period, comparing the MSE- (blue), FFT-AMSE- (purple), and wavelet-trained (orange) models at lead times of 24 (solid), 72 (dashed), and 120 (dotted) hours. (c) Instantaneous rainfall rate for the same case and valid time, with the model forecasts at a lead time of 24 hours verified against radar-derived rainfall estimates. Panels (a) and (c) use the same colour scale in mmh−1 , such that a given colour represents the same rainfall rate in both rows. (d) As in (b), but for instantaneous rainfall rate verified against radar-derived rainfall estimates. In (b) and (d), the grey dashed diagonal indicates perfect agreement.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure A.1: Comparison of first-order spatial attribution and field-space gradients for 2m temperature using MSE, Fourier-domain AMSE (FFT-AMSE), and the wavelet-based loss. Rows show the two constructed prediction fields introduced in Figure 2 : a perturbed 6hour time-offset prediction with added Gaussian noise (Case 1), and a Gaussian-blurred target with σ=3 grid cells (Case 2). For each objective, the left panel shows the attribution A(r) defined by Equation A.1 , and the right panel shows the field-space gradient ∂L/∂x(r) . Attribution is displayed on a logarithmic scale. Common colour scales are used across objectives for each diagnostic. Attribution combines the local prediction error with the sensitivity of the objective and is therefore a first-order diagnostic rather than an exact spatial decomposition of the scalar objective. Absolute loss and gradient magnitudes depend on the definition and normalisation of each objective, so the comparison emphasises the spatial structure of the resulting fields.
Stage 1
Stage 2
Stage 3
a
b
Global dataset
ERA5
ERA5
ERA5
ERA5
Global data grid
O96
N320
N320
N320
Global resolution
100 km
30 km
30 km
30 km
Global mesh refinement
5
6
6
6
Global wavelet levels, J
4
6
6
6
Appendix
Table B.1: Training configuration for each stage of the model used in this study. A dash indicates that no regional dataset or regional mesh was used in that stage. All training was performed on the Isambard-AI (Phase 2) HPC system [ 24 ] using NVIDIA GH200 superchips.
Figure D.1: Scorecard showing the percentage difference in verification scores between the wavelet- and MSE-trained models over lead times of up to 10 days. “OBS” and “ANL” denote verification against observations and operational analyses, respectively. Surface variables are 10m wind speed (10ws), 10m wind vector error (10wv), the zonal and meridional components of 10m wind (10u and 10v), 2m temperature (2t), total precipitation (tp), surface pressure (sp) and skin temperature (skt). Upper-air variables are temperature (t), specific humidity (q), geopotential (z), and wind speed (ws), with the corresponding pressure levels given in hPa . Columns show pointwise root-mean-squared error (RMSE) and high-resolution assessment (HiRA) continuous ranked probability score (CRPS), with HiRA CRPS calculated using neighbourhoods of 3×3 , 7×7 , and 11×11 model grid points. Scores are aggregated over forecasts initialised between 1 January and 21 December 2023. Blue indicates lower errors for the wavelet-trained model, whereas red indicates lower errors for the MSE-trained model.
Figure D.2: Lead-time dependence of the zonal power spectra for the MSE- (blue/green) and wavelet-trained (orange) models, verified against UKV operational analyses (black, dashed). Columns show (a) 2m temperature, (b) 10m wind speed, and (c) 6hour accumulated precipitation. Upper panels show the absolute spectra at lead times of 6, 24, and 72 hours, with colour shade indicating lead time. Lower panels show log10(Pf/Pr) , where Pf and Pr are the forecast and reference powers, respectively. Vertical dotted lines mark the internal processor-mesh spacing, Δxmesh=4\text{,}\mathrm{k}\mathrm{m}$$ , and wavelengths corresponding to 2Δxmesh , 4Δxmesh , and 8Δxmesh . The reduction in intermediate- and short-wavelength power deficits under wavelet training persists across the evaluated lead times.
Figure D.3: As Figure D.2 , but for upper-air variables: (a) 850hPa temperature, (b) 700hPa specific humidity, and (c) 500hPa geopotential.
Figure D.4: Scorecard showing the percentage difference in verification scores between the wavelet- and FFT-AMSE-trained models over lead times of up to 10 days. “OBS” and “ANL” denote verification against observations and operational analyses, respectively. Columns show pointwise RMSE and HiRA CRPS calculated using neighbourhoods of 3×3 , 7×7 , and 11×11 model grid points. Scores are aggregated over forecasts initialised between 1 January and 21 December 2023. Blue indicates lower errors for the wavelet-trained model, whereas red indicates lower errors for the FFT-AMSE-trained model.
Figure D.5: Fractions Skill Score (FSS) for 6hour accumulated precipitation, ‘tp’, at a lead time of 24hours , evaluated across precipitation thresholds and neighbourhood sizes over the regional domain. (a) Absolute FSS for the wavelet-trained model. (b,c) FSS difference between the wavelet-trained model and the MSE- and FFT-AMSE-trained models, respectively; positive values (green) indicate improved performance for the wavelet-trained model. Physical neighbourhood scales, calculated from the regional grid spacing, are shown along the upper axis. The same difference colour scale is used in panels (b) and (c).
Weather forecasting predicts future atmospheric conditions, from local storms to global patterns, using physical models, observations, and machine learning.
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.
Machine learning weather prediction (MLWP) models have achieved impressive forecasting performance at a small fraction of the computational costs required for traditional physics-based methods. However, they are primarily (1) data-driven and (2) evaluated using pixel-wide error metrics (e.g., RMSE), so there are no guarantees that their forecasts are consistent with known physical laws. We introduce PhysMetrics.Weather, an evaluation framework that assesses the physical realism of MLWP models across three types of metrics: conservation, spectral, and dynamical. By quantifying physical realism, this tool guides the development of physics-informed architectures and helps evaluate whether MLWP models are reliable for operational use. Our framework is available on Github at https://github.com/Emmakast/PhysMetrics.Weather.
Emma Kasteleyn, Timo Maier, Axel Lauer +3
University of Amsterdam, Amsterdam, The Netherlands · Deutsches Zentrum für Luft- und Raumfahrt (DLR), Institut für Physik der Atmosphäre, Oberpfaffenhofen, Germany · University of Bremen, Institute of Environmental Physics (IUP), Bremen, Germany +1
Recent machine learning weather prediction (MLWP) models have demonstrated remarkable forecasting skill on global reanalysis-based benchmarks. However, their performance remains unclear in challenging environments such as Northern Norway, where narrow fjords and rapidly changing weather result in highly variable local wind conditions. In this case study, we evaluate FourCastNet3 (FCN3), GraphCast, and ECMWF High Resolution Forecast (HRES) for wind speed forecasting using multi-year station observations from Northern Norway, focusing on their relative performance, generalization beyond the training period, and performance under high-wind conditions. Our results show that HRES slightly outperforms FCN3 and GraphCast, with an overall RMSE of 2.89 ms−1, compared to 2.96 ms−1 for FCN3 and 2.94 ms−1 for GraphCast. Notably, the MLWP models maintain comparable performance beyond their respective training periods, with no clear evidence of noticeable degradation. FCN3 performs best under high-wind conditions, although all models substantially underestimate strong winds. Our findings suggest that MLWP has become competitive with NWP for local wind, but further refinements are still needed to capture complex terrain better.
Siyan Chen, Lars Uebbing, Eirik Mikal Samuelsen +5
UiT The Arctic University of Norway · The Norwegian Meteorological Institute · Norwegian Computing Center +2