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

James Chappell; Sophie Luise Arana; Jack Bowyer; Alejandro Coca-Castro; Evangeline Corcoran; Eric G. Daub; Tom Dunstan; Noushin Eftekhari

physics.ao-phOct 27, 2026

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,

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

Abstract

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.

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.5\text{\,}\mathrm{k}\mathrm{m}$ regional grid centred on the UK; the example shows $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ accumulated precipitation valid at 2023-09-28T00Z UTC at a lead time of $24\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}\mathrm{s}$ . (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 $\mathbb{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 $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ time-offset prediction (Case 1) and a Gaussian-blurred target with $\sigma=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 $L_{2}$ norm as a fraction of the sum of the component-gradient $L_{2}$ 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) $500\text{\,}\mathrm{h}\mathrm{P}\mathrm{a}$ geopotential, (b) $2\text{\,}\mathrm{m}$ 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, $\alpha_{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\times 3$ , $7\times 7$ and $11\times 11$ model-grid neighbourhoods; the bottom row shows pointwise RMSE. Columns show (a,e) $2\text{\,}\mathrm{m}$ temperature, (b,f) $10\text{\,}\mathrm{m}$ wind speed, (c,g) $10\text{\,}\mathrm{m}$ wind-vector error and (d,h) $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ 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) $2\text{\,}\mathrm{m}$ temperature, (b) $10\text{\,}\mathrm{m}$ wind speed, and (c) $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ accumulated precipitation. Upper panels show the absolute spectra, while lower panels show $\log_{10}(P_{f}/P_{r})$ , where $P_{f}$ and $P_{r}$ 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, $\Delta x_{\mathrm{mesh}}=$4\text{\,}\mathrm{k}\mathrm{m}$$ , and wavelengths corresponding to $2\Delta x_{\mathrm{mesh}}$ , $4\Delta x_{\mathrm{mesh}}$ , and $8\Delta x_{\mathrm{mesh}}$ . — 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) $2\text{\,}\mathrm{m}$ temperature, (b) $10\text{\,}\mathrm{m}$ wind speed, and (c) $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ accumulated precipitation. Upper panels show the absolute spectra, while lower panels show $\log_{10}(P_{f}/P_{r})$ , where $P_{f}$ and $P_{r}$ 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, $\Delta x_{\mathrm{mesh}}=$ 4\text{,}\mathrm{k}\mathrm{m}$$ , and wavelengths corresponding to $2\Delta x_{\mathrm{mesh}}$ , $4\Delta x_{\mathrm{mesh}}$ , and $8\Delta x_{\mathrm{mesh}}$ .

Figure 6: Precipitation forecast characteristics. (a) Mean rainfall rate over the preceding $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}\mathrm{s}$ for the case valid at 0000 UTC on 14 October 2023, calculated by dividing the $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ precipitation accumulations by six. The panels show the UKV operational NWP forecast at a lead time of $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}\mathrm{s}$ and the wavelet- and MSE-trained model forecasts at a lead time of $24\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}\mathrm{s}$ . (b) Quantile–quantile comparison of forecast and analysis $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ 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 $\mathrm{m}\mathrm{m}\,\mathrm{h}^{-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 $2\text{\,}\mathrm{m}$ 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 $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ time-offset prediction with added Gaussian noise (Case 1), and a Gaussian-blurred target with $\sigma=3$ grid cells (Case 2). For each objective, the left panel shows the attribution $A(\boldsymbol{r})$ defined by Equation A.1 , and the right panel shows the field-space gradient $\partial\mathcal{L}/\partial\widehat{x}(\boldsymbol{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
	Stage 1	Stage 2	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 $10\text{\,}\mathrm{m}$ wind speed (10ws), $10\text{\,}\mathrm{m}$ wind vector error (10wv), the zonal and meridional components of $10\text{\,}\mathrm{m}$ wind (10u and 10v), $2\text{\,}\mathrm{m}$ 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 $\mathrm{h}\mathrm{P}\mathrm{a}$ . 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\times 3$ , $7\times 7$ , and $11\times 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) $2\text{\,}\mathrm{m}$ temperature, (b) $10\text{\,}\mathrm{m}$ wind speed, and (c) $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ 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 $\log_{10}(P_{f}/P_{r})$ , where $P_{f}$ and $P_{r}$ are the forecast and reference powers, respectively. Vertical dotted lines mark the internal processor-mesh spacing, $\Delta x_{\mathrm{mesh}}=$4\text{\,}\mathrm{k}\mathrm{m}$$ , and wavelengths corresponding to $2\Delta x_{\mathrm{mesh}}$ , $4\Delta x_{\mathrm{mesh}}$ , and $8\Delta x_{\mathrm{mesh}}$ . The reduction in intermediate- and short-wavelength power deficits under wavelet training persists across the evaluated lead times. — 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) $2\text{\,}\mathrm{m}$ temperature, (b) $10\text{\,}\mathrm{m}$ wind speed, and (c) $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ 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 $\log_{10}(P_{f}/P_{r})$ , where $P_{f}$ and $P_{r}$ are the forecast and reference powers, respectively. Vertical dotted lines mark the internal processor-mesh spacing, $\Delta x_{\mathrm{mesh}}=$ 4\text{,}\mathrm{k}\mathrm{m}$$ , and wavelengths corresponding to $2\Delta x_{\mathrm{mesh}}$ , $4\Delta x_{\mathrm{mesh}}$ , and $8\Delta x_{\mathrm{mesh}}$ . 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) $850\text{\,}\mathrm{h}\mathrm{P}\mathrm{a}$ temperature, (b) $700\text{\,}\mathrm{h}\mathrm{P}\mathrm{a}$ specific humidity, and (c) $500\text{\,}\mathrm{h}\mathrm{P}\mathrm{a}$ 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\times 3$ , $7\times 7$ , and $11\times 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 $6\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}$ accumulated precipitation, ‘tp’, at a lead time of $24\text{\,}\mathrm{h}\mathrm{o}\mathrm{u}\mathrm{r}\mathrm{s}$ , 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

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

See all 176 papers

(Sparse) Attention to the Details: Preserving Spectral Fidelity in ML-based Weather Forecasting Models
Apr 6, 2026Maksim Zhdanov, Ana Lucic, Max Welling +1Weather Forecasting Probabilistic Forecasting
AMLab · University of Amsterdam · CuspAI
PhysMetrics.Weather: An Evaluation Framework for Physical Consistency in ML Weather Models
Jun 9, 2026Emma Kasteleyn, Timo Maier, Axel Lauer +3Forecasting Benchmarks Physics-Informed ML
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
A Station-Based Evaluation of Machine Learning-based Weather Forecasting Models in Northern Norway
Aug 25, 2026Siyan Chen, Lars Uebbing, Eirik Mikal Samuelsen +5Weather Forecasting
UiT The Arctic University of Norway · The Norwegian Meteorological Institute · Norwegian Computing Center +2

Apr 6, 2026J/K move · Enter open · S save

(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

AMLab · University of Amsterdam · CuspAI

Weather Forecasting Probabilistic Forecasting Spectral Bias Block-Sparse Attention

Open paper AskPDF 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.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.

Abstract

Weather Forecasting

Explore similar work

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

PhysMetrics.Weather: An Evaluation Framework for Physical Consistency in ML Weather Models

A Station-Based Evaluation of Machine Learning-based Weather Forecasting Models in Northern Norway