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

Introduction

Machine-learning weather prediction (MLWP) models can produce accurate forecasts at a fraction of the computational cost of conventional numerical weather prediction [1, 2, 3, 4, 5, 6, 7, 8]. As these models are extended to higher spatial resolutions [9, 10], a variety of regional MLWP approaches have emerged using stretched-grid [11, 12, 13], limited-area [14, 15, 16, 17], and dataset-specific encoder–decoder architectures (as used in this study). At kilometre-scale resolutions, the training data capture local weather systems and spatially varying surface influences that are poorly represented by coarser atmospheric fields, providing information important for forecasting high-impact weather and issuing local warnings [18]. However, resolving these features in the training data does not guarantee their preservation in forecasts, regardless of model architecture.

This challenge is particularly acute for deterministic models trained using pointwise mean-squared error (MSE). Kilometre-scale atmospheric fields contain intermittent and sharply localised structures, such as convective precipitation, whose precise evolution and location are difficult to predict. Because MSE strongly penalises spatial displacement errors, minimising this objective tends to favour smooth predictions over uncertain fine-scale structures, suppressing small-scale variability and reducing effective resolution [19]. While probabilistic approaches can represent uncertain fine-scale variability through ensembles of plausible realisations [20, 21, 22], deterministic models provide a computationally efficient and controlled setting for investigating how training objectives influence forecast structure.

One approach to mitigate smoothing of deterministic forecasts is incorporating spectral information into the training objective. Subich et al. [19] introduced the adjusted mean-squared error (AMSE), which decomposes forecast errors spectrally and separates contributions associated with amplitude and coherence. Originally formulated using spherical harmonics for global models, AMSE can be adapted to regional domains using Fourier transforms, while related amplitude-based spectral losses have shown promise in probabilistic regional forecasting [12]. However, their globally supported basis functions limit their ability to constrain forecast errors locally in space. Improvements in spectral fidelity may therefore come at the expense of accurately positioning individual features, potentially degrading pointwise verification metrics such as root-mean-squared error (RMSE) [8]. This motivates objectives that preserve fine-scale variability while retaining sensitivity to the location of forecast structures.

Here, we develop a deterministic loss based on the two-dimensional dual-tree complex wavelet transform (DTCWT) [23, 24], illustrated schematically in Figure 1. By decomposing fields jointly by scale, location, and orientation, the DTCWT allows loss contributions to be calculated and weighted independently across spatial scales. The approximate shift invariance [25] and directional selectivity of the DTCWT make it well suited to localised atmospheric structures, while its complex coefficients support amplitude- and coherence-sensitive terms analogous to those used in spectral AMSE. Wavelets have previously been applied to atmospheric analysis and verification [26, 27], with recent machine-learning applications focusing primarily on precipitation, including generative downscaling [28] and a composite DTCWT–Fourier loss for global prediction [29]. Our standalone objective extends wavelet-based optimisation to multivariable, kilometre-scale forecasting, aiming to preserve fine-scale spectral variability without compromising spatial accuracy. The underlying multiscale decomposition could also provide a foundation for scale-aware probabilistic training objectives, complementing recent developments in probabilistic MLWP [30, 31].

We evaluate the proposed objective using a joint global–regional configuration developed as part of ongoing work towards a regional version of the medium-range MLWP model FastNet [7, 8]. The model simultaneously produces kilometre-scale forecasts over the UK and coarser forecasts globally, providing a controlled setting in which to compare wavelet-based training with MSE and an FFT-based spectral objective. We show that the wavelet loss preserves fine-scale spectral energy while maintaining strong spatial localisation of the training gradient. It achieves pointwise performance comparable to MSE, outperforms MSE under neighbourhood-based verification, and improves spectral fidelity relative to both MSE and the purely spectral objective.

Deterministic wavelet loss

To combine scale-selective amplitude matching with greater spatial localisation, we formulate the loss in the dual-tree complex wavelet domain. Let $\hat{x}$ and $x$ denote predicted and target fields, respectively, on a two-dimensional input grid. Applying a $J$-level DTCWT to a field $x$ produces complex directional band-pass coefficient fields, $z_{j,o}[x]$, at each level $j=1,\ldots,J$ and at six approximate orientations ($\pm15^\circ$, $\pm45^\circ$, and $\pm75^\circ$), together with a final real-valued low-pass residual, $l_J[x]$:

We write $\hat{z}_{j,o}=z_{j,o}[\hat{x}]$ and $z_{j,o}=z_{j,o}[x]$ for the prediction and target coefficients, respectively, with analogous notation for the low-pass residuals.

The index $j=1$ denotes the finest level, with increasing $j$ corresponding to progressively coarser scales. Because the transform downsamples the coefficient fields at successive levels, each level has its own coefficient grid and set of spatial indices, $\mathcal{S}_j$. Thus, $r \in\mathcal{S}_j$ denotes a location on the level-$j$ grid rather than on the original input grid.

For each high-pass scale and orientation, we calculate the band-integrated wavelet powers of the prediction and target:

Following the AMSE formulation [19], the loss for each wavelet band combines an amplitude error with a coherence-dependent penalty:

where the normalised coherence is

The wavelet-band loss is summed over the six orientations and combined across scales using configurable weights, $\alpha_j$. The final low-pass residual represents the coarsest spatial structure and is compared directly using MSE. The total loss is

where $N$ is the number of points in the original input grid.

The two terms in Equation (3) measure distinct aspects of the error within each scale–orientation band. The amplitude term compares the band-integrated powers of the prediction and target, while the second term applies a power-weighted penalty for disagreement between their complex coefficient patterns. Prediction and target coefficients are paired at corresponding locations in the coherence numerator before their cross-power is summed, making the loss sensitive to their relative spatial arrangement as well as their band-integrated powers. The loss therefore distinguishes forecasts with similar band-integrated power but different spatial organisation.

Although Equation (5) is additive across levels and orientations, each band loss is computed from statistics aggregated over all locations in $S_j$ and is therefore not a sum of independent coefficient-wise losses—the corresponding locations and gradients remain coupled within each band. Moreover, each coefficient is the response of a compactly supported wavelet filter to a neighbourhood of input values, rather than a pointwise value on the original grid. Back-propagating through these filters consequently produces more spatially localised field-space gradients than transforms with globally supported basis functions, such as the Fourier transform, as shown in Appendix § 8. The following subsection demonstrates how the separate level-wise terms provide scale-selective loss contributions and field-space gradients.

Scale selectivity

Figure 2 illustrates the scale selectivity of the wavelet loss using two contrasting constructed prediction fields designed to probe different error regimes. Case 1 is the atmospheric state 6 hours before the target time with Gaussian noise added ($\sigma= 5\%$), producing errors across a broad range of spatial scales. Case 2 is obtained by applying a Gaussian blur with $\sigma= 3$ grid cells to the target, preserving its large-scale structure while selectively attenuating finer-scale variability. Uniform scale weights, $\alpha_j = 1$, are used in both cases.

The scalar loss is decomposed exactly into the wavelet-band components, $\mathcal{L}_j = N^{-1}\alpha_j\sum_{o=1}^{6}\mathcal{L}_{j,o}$, and the final low-pass component, $\mathcal{L}_{\mathrm{LP}} = N^{-1}\left\lVert\hat{\boldsymbol{l}}_J-\boldsymbol{l}_J\right\rVert_2^2$. For each component, we also calculate its gradient with respect to the predicted field on the original input grid, $\nabla_x\mathcal{L}_j$, by back-propagating through the wavelet transform. These field-space gradients represent the signed first-order sensitivity of each scalar component to changes in the predicted field, revealing the scale and spatial structure of the training signal presented to the forecasting model. The total field-space gradient is the sum of the low-pass and wavelet-band gradients, although the eventual parameter update also depends on back-propagation through the model and on the optimiser.

The two cases produce clearly different responses. For the perturbed time-offset prediction (Case 1, Figure 2b), broad structural errors dominate the scalar objective, with the low-pass component contributing 80% of the total loss. Its share of the summed component-gradient norms is much lower, at 46%, with the remaining gradient magnitude distributed across the wavelet-band levels and increasing towards the coarser levels. The wavelet components therefore provide substantial scale-dependent training signals even when the scalar loss is dominated by its low-pass component.

By contrast, the blurred prediction (Case 2, Figure 2c) retains the target’s large-scale structure, giving a negligible low-pass contribution to both the loss and gradient magnitude. Its response is concentrated at wavelet-bands $j = 2$ and $j = 3$, which together account for approximately 81% of the scalar loss and 85% of the summed component-gradient norms. This concentration reflects both the range of scales attenuated by the Gaussian filter and the distribution of variance in the target field: although the finest scales undergo strong fractional attenuation, the 2 m temperature field contains more variance within the affected wavelet-bands at $j = 2$ and $j = 3$ than at $j = 1$.

Together, these examples show that the wavelet objective produces scale-dependent scalar losses and field-space training signals. Broad structural errors activate the low-pass and coarser-scale terms, whereas selectively removing finer-scale variability concentrates the response in the affected fine-scale wavelet-band components. The spatial localisation of the complete wavelet-loss gradient is compared with that produced by MSE and FFT-AMSE in Appendix § 8.

Evolution of scale-wise loss contributions during training

Having demonstrated the scale selectivity of the loss using controlled perturbations, we next examine how its scale-wise composition evolves during model training. Figure 3 shows the magnitude and scale composition of the regional wavelet validation loss during the early adaptation of the model to the regional domain for three variables. These variables were selected to span a range of characteristic spatial structures represented within kilometre-scale atmospheric fields: 500 hPa geopotential is dominated by smooth, large-scale variability; 2 m temperature contains additional structure associated with orography, surface heterogeneity, and land–sea contrasts; and precipitation is intermittent and contains substantial variability across intermediate and fine scales. The figure covers the first 25% of the regional single-step training stage (Stage 3a; see Appendix § 9 for the complete training schedule).

For 500 hPa geopotential, the low-pass MSE contributes more than 90% of the objective throughout the displayed period, and the pointwise MSE and complete wavelet loss (Eq. (5)) follow similar trajectories. Although the absolute losses decrease at every wavelet level, the low-pass term remains dominant, consistent with the concentration of geopotential variance at large spatial scales.

For 2 m temperature and precipitation, the absolute low-pass and wavelet-band losses also decrease, but the low-pass error falls more rapidly. The wavelet-band levels therefore account for an increasing fraction of the remaining wavelet objective. For temperature, the continued separation between the wavelet validation loss and pointwise MSE indicates residual scale-dependent errors that are not distinguished by the pointwise metric. These contributions are consistent with the greater regional and fine-scale variability of near-surface temperature. Precipitation contains substantial scale-dependent error from the beginning of training. Its pointwise MSE flattens comparatively early, whereas the wavelet validation loss continues to decrease, indicating continued improvement in aspects of the precipitation field represented by the multiscale objective.

The three examples demonstrate that the distribution of the wavelet loss reflects the characteristic spatial scales of each variable. Broad-scale fields remain dominated by the low-pass component, whereas variables containing greater regional and fine-scale variability receive progressively larger contributions from the wavelet bands. The increasing relative wavelet-band contributions reflect differences in the rates at which errors are reduced, rather than growth in scale-dependent error. These diagnostics characterise the model’s early adaptation to the regional domain and should not be interpreted as the complete evolution of the loss over the full training schedule.

Forecast evaluation

We evaluate three variants of a developmental joint global–regional configuration for a regional FastNet [7, 8] model, implemented using the Anemoi [6] framework. The variants share the same architecture and differ principally in their training objectives: pointwise MSE, an FFT-based variant of the AMSE loss introduced by Subich et al. [19], and the wavelet loss introduced here. For the wavelet-trained model, all decomposition levels are weighted uniformly, with $\alpha_j = 1$. Each model produces forecasts on a global N320 reduced Gaussian grid ($\sim31$ km and a 1.5 km regional grid centred on the UK, using a separate encoder–decoder pair for each dataset. Verifying forecasts are initialised every 6 hours, with output at 6 hour intervals to a lead time of 240 hours, and are produced for the entirety of 2023. Full details of the model configuration and the training schedule are provided in Appendix § 9.

Forecasts are verified against surface observations using pointwise RMSE and the High-Resolution Assessment (HiRA) framework [32]. HiRA treats the model grid points surrounding each observation as a pseudo-ensemble, which we evaluate using the continuous ranked probability score (CRPS). This neighbourhood-based approach provides greater tolerance to small spatial displacement errors than pointwise verification and assesses whether forecasts contain locally relevant structure. We calculate HiRA CRPS using neighbourhoods of $3 \times3$, $7 \times7$, and $11 \times11$ model grid points. Details of the observation processing and forecast–observation matching are provided in Appendix § 10, with additional surface and upper-air verification presented in Appendix § 11.

Results

Verification against surface observations

The models show similar overall skill but their relative performance depends consistently on the verification metric (Figure 4). The wavelet-trained model generally achieves lower HiRA CRPS, with the clearest improvements for both 10 m wind diagnostics across all three neighbourhood sizes. Smaller improvements are present for 2 m temperature and precipitation over parts of the forecast, although the differences become less systematic at longer leads. In contrast, the MSE-trained model generally has slightly lower pointwise RMSE, particularly for precipitation and, to a lesser extent, temperature and wind diagnostics.

This contrast is consistent with the wavelet-trained model retaining more local spatial variability without uniformly improving the exact placement of forecast features. Retaining this structure can improve neighbourhood-based scores while remaining vulnerable to the double penalty imposed by pointwise RMSE. However, lower HiRA CRPS may also result from improved local bias, intensity or variability, and should not be interpreted solely as evidence of displaced but otherwise accurate forecast features.

Spectral fidelity

The spectral differences among the models are substantially larger than their differences in pointwise and neighbourhood-based verification scores, illustrating the value of training objectives that explicitly constrain multiscale structure. Figure 5 compares the three training objectives at a lead time of 24 hours; similar results are observed across lead times and are provided in Appendix § 11 for reference.

The MSE-trained model develops a systematic power deficit towards shorter wavelengths, consistent with the smoothing associated with pointwise optimisation. This deficit is modest for 2 m temperature but pronounced for 10 m wind speed and accumulated precipitation. FFT-AMSE recovers some of the variance lost under MSE training, but the wavelet-trained forecasts remain closer to the UKV reference over a broader range of intermediate and short wavelengths. The distinction is clearest for wind speed and precipitation, for which wavelet training produces a more accurate distribution of variance across the resolved scales.

The MSE-trained spectra also contain regions of excess power aligned with wavelengths associated with the internal processor mesh, most clearly for precipitation. These features are absent or strongly suppressed in the wavelet-trained spectra, indicating that the improvement is not simply an indiscriminate increase in short-wavelength variance. Instead, the wavelet objective both reduces spectral power deficits and limits non-physical mesh-related artefacts [8, 33].

Precipitation forecasts

Precipitation provides a stringent test of km-scale MLWP models because forecast quality depends on representing both the intensity and spatial organisation of highly intermittent small-scale features. The improved spectral fidelity obtained with wavelet training is reflected in both the structure and distribution of forecast precipitation (Figure 6).

The wavelet-trained model produces precipitation fields that are sharper and more intense than those generated by the MSE-trained model (Figure 6a,Figure 6c). The difference is particularly evident for instantaneous rainfall rates, where the MSE forecast is dominated by diffuse weak precipitation, while the wavelet-trained forecast contains more coherent and spatially concentrated features that more closely resemble the radar observations. Similar behaviour is evident for accumulated precipitation, with stronger local maxima and sharper gradients.

These qualitative differences are reflected in the precipitation distributions. For both accumulated precipitation and instantaneous rainfall rate, the wavelet-trained model more closely reproduces the upper tail of the observed distribution than either the MSE- or FFT-AMSE-trained models (Figure 6b,d). FFT-AMSE partially reduces the underprediction of heavy precipitation relative to MSE, while wavelet training provides a further improvement. This greater intensity and fine-scale variability is accompanied by improved spatial skill: at a lead time of 24 hours, the wavelet-trained model achieves higher Fractions Skill Scores than both alternatives across most precipitation thresholds and neighbourhood sizes, with the clearest gains for heavier precipitation and broader neighbourhoods (Appendix § 11.§ 3). These combined improvements are consistent with benefits from the spatially localised wavelet representation beyond those obtained using Fourier-based spectral constraints alone.

Despite these improvements, precipitation remains a clear weakness of all model variants. The highest quantiles are systematically underestimated, particularly for instantaneous rainfall rate, and substantial departures from one-to-one agreement persist even at short lead times. Furthermore, although wavelet training produces sharper accumulated precipitation fields, neither MLWP model fully reproduces the elongated streak-like structures evident in the operational forecast that result from precipitation features advecting during the six-hour accumulation period. The wavelet objective therefore improves the representation of heavy precipitation but does not eliminate important shortcomings in the temporal organisation, morphology and intensity of precipitation.

Conclusions

We have introduced a deterministic training objective for machine-learning weather prediction based on the two-dimensional dual-tree complex wavelet transform. The DTCWT decomposes atmospheric fields jointly by scale, location, and orientation, allowing amplitude- and coherence-sensitive loss terms to be evaluated independently at each spatial scale. Its compactly supported basis functions also produce a more spatially localised training gradient than an equivalent Fourier-domain objective, while its approximate shift invariance and directional selectivity are well suited to sharply localised atmospheric structures.

Applied to an MLWP model that simultaneously produces 1.5 km forecasts over the UK and $\sim31$ km forecasts globally, the wavelet loss provides a better balance between spatial accuracy and multiscale fidelity than either pointwise MSE or FFT-AMSE. Within the regional domain, the wavelet-trained model retains pointwise performance close to the MSE baseline, generally improves neighbourhood-based verification, and outperforms FFT-AMSE. The distinction is clearest in spectral verification: across the evaluated variables and lead times, the wavelet-trained forecasts remain closer to the reference spectra over a broader range of resolved scales than either baseline. Wavelet training substantially reduces the intermediate- and small-scale power deficits produced by MSE and suppresses non-physical excess power aligned with the internal processor-mesh scales. Compared with FFT-AMSE, it more accurately reproduces the reference distribution of variance across spatial scales without incurring the same degradation in pointwise performance.

These improvements are reflected in sharper, more spatially concentrated precipitation forecasts and more realistic intensity distributions than those produced by either baseline. Nevertheless, all three models underestimate the most intense rainfall, and wavelet training does not fully recover the morphology and temporal organisation of accumulated precipitation. The wavelet loss therefore reduces, but does not eliminate, the trade-off between preserving fine-scale variability and maintaining spatial accuracy. Future work should assess its sensitivity to scale weighting, evaluate its generality across additional regions and weather regimes, and investigate its use within probabilistic forecasting systems that represent both fine-scale structure and forecast uncertainty.

Acknowledgments

The development of FastNet, including this work and its contributors, is funded by The Alan Turing Institute and the Met Office AWAKEN project. The AWAKEN project is led by the UK Met Office and is funded through the UK Government’s International Science Partnerships Fund (ISPF), with support from the Department for Business, Innovation, Science and Trade (DBIST) and the Foreign, Commonwealth and Development Office (FCDO). The authors acknowledge the use of resources provided by the Isambard-AI National AI Research Resource (AIRR). Isambard-AI is operated by the University of Bristol and is funded by the UK Government’s Department for Business, Innovation, Science and Trade (DBIST) via UK Research and Innovation; and the Science and Technology Facilities Council [ST/R-1/1023]. The authors gratefully acknowledge all contributors to the Anemoi framework.

Statement on the use of generative AI

All results, analysis and scientific conclusions in this manuscript were produced by the authors. Microsoft 365 Copilot was used for editorial assistance to improve clarity and readability.

Spatial attribution and gradient comparison

Section 2.1 showed how the wavelet objective distributes its scalar loss and field-space training signal among decomposition levels. Here, we compare the spatial structure of the complete training signals produced by pointwise MSE, FFT-AMSE, and the wavelet-based loss for the same two constructed prediction fields.

For a scalar objective $\mathcal{L}$, prediction $\hat{x}$, and target $x$, we define the first-order spatial attribution at location $\mathbf{r}$ as

This diagnostic combines the magnitude of the local prediction error with the sensitivity of the complete objective at the same location. It represents the magnitude of the first-order change associated with moving that prediction value towards its target. For pointwise MSE, the attribution is proportional to the pointwise squared error. For transform-domain objectives such as FFT-AMSE and the wavelet-based loss, coefficients from multiple spatial locations are aggregated and coupled within the objective. The attribution is therefore not an exact additive decomposition of the finite scalar loss, but provides a common field-space diagnostic of where the different objectives are most sensitive.

Figure A.1 shows the attribution and field-space gradient for each objective and prediction regime. Attribution is shown logarithmically because it spans several orders of magnitude. The gradient is displayed separately because the two diagnostics answer different questions: the gradient shows the signed sensitivity of the objective to changes in the predicted field, whereas the attribution also weights that sensitivity by the prediction error already present. A location can therefore have a large gradient but modest attribution when its error is small, or large attribution when both error and sensitivity are appreciable.

For the perturbed time-offset prediction (Case 1), all three objectives produce broadly similar attribution patterns, highlighting many of the same broad and finer-scale errors. Their field-space gradients nevertheless differ. The MSE gradient follows the pointwise prediction error, consistent with the local definition of the objective. The FFT-AMSE gradient is more broadly distributed across the domain, consistent with the global support of the Fourier basis. The wavelet gradient retains stronger location-dependent variation around the atmospheric structures contributing to the error. This behaviour is consistent with the spatial indexing and compact support of the wavelet coefficients, despite the coupling introduced by aggregation within each scale–orientation band.

The contrast is stronger for the blurred prediction (Case 2). The MSE gradient follows the pointwise residual introduced by smoothing but does not distinguish the spatial scales contributing to that residual. FFT-AMSE responds to the missing spectral variance, but its gradient is distributed throughout much of the domain. The wavelet objective is similarly sensitive to the removed variability, while producing a gradient that is more strongly organised around the temperature features altered by smoothing. The response is concentrated primarily over land and is particularly structured over regions of complex orography, including the Scottish Highlands, with a weaker response over much of the surrounding sea. This organisation is consistent with fine-scale 2 m temperature variability associated with terrain and land–sea contrasts.

Together with the scale-wise decomposition in Figure 2, this comparison shows that the wavelet objective both identifies the affected range of spatial scales and retains location-dependent sensitivity to the corresponding atmospheric structures. Absolute loss and gradient magnitudes depend on the definition and normalisation of each objective and are not directly comparable; the comparison is therefore restricted to the spatial structure and relative localisation of the attribution and gradient fields.

Model configuration and training schedule

The model was developed and trained using the Anemoi framework [6], and follows an encode–process–decode architecture with graph-transformer GNNs used in all three components. Input fields are encoded onto a latent mesh, processed in latent space and decoded onto the output grids. Separate encoder–decoder pairs are used for two datasets: ERA5 on its native reduced Gaussian grid, representing the global atmospheric state, and Met Office UKV operational analyses on a regional grid covering the UK at approximately 1.5 km resolution. This configuration allows the datasets to use different grids and variables while sharing a common latent representation. The model therefore produces global and high-resolution regional forecasts simultaneously in a single forward pass.

The MSE-, FFT-AMSE- and wavelet-trained models use the staged training schedule summarised in Table B.1, following the approach introduced in Ref. [11]. Stage 1 establishes the global forecasting capability using ERA5 on the O96 grid and a coarse global latent mesh. In Stage 2, the resolution of the ERA5 dataset is increased to N320 and the global latent mesh is refined, allowing the model to learn a higher-resolution global representation over the same training period.

Stage 3a introduces the UKV dataset, together with a regional encoder–decoder pair and a high-resolution regional mesh embedded within the global mesh. The shared latent representation retains the global atmospheric context while adapting the model to the UK domain. Global and regional losses are assigned relative weights of 67% and 33%, respectively. Without this weighting, the regional domain would contribute less than 1% of the combined objective because of its much smaller spatial extent.

Stage 3b extends training from single-step forecasts to autoregressive rollouts of 2–12 steps, exposing the model to its own predictions and improving stability at longer lead times. The reduced learning rate in this stage allows the model to be fine-tuned over longer forecast horizons. Note. Mesh refinement denotes the subdivision frequency of the icosahedron used to construct the latent mesh with larger values corresponding to finer meshes. Refinement frequencies of 6 and 11 give average node spacings of approximately 110 km and 4 km, respectively. Wavelet-transform settings apply only to the wavelet-trained experiment; all other configuration settings are shared by the three model variants.

	Stage 1	Stage 2	Stage 3a	Stage 3b
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
Regional dataset	–	–	UKV	UKV
Regional resolution	–	–	1.5 km	1.5 km
Regional mesh refinement	–	–	11	11
Regional wavelet levels, $J$	–	–	6	6
Rollout	1	1	1	2–12
Training period	1980–2022	1980–2022	2020–2022	2020–2022
Iterations	100,000	50,000	25,000	3200
Peak learning rate	$2 \times10^{-3}$	$3.2 \times10^{-4}$	$4 \times10^{-4}$	$2.5 \times10^{-5}$
Warm-up iterations	1000	100	500	50
Batch size	128	32	32	32
Total number of GPUs	16	32	64	64
GPUs per model instance	1	1	2	2
Total GPU hours	540	720	800	400

while limiting degradation of the global and regional representations learned in the preceding stages.

The three models differ only in their training objectives. The MSE- and wavelet-loss trained models use their respective loss functions throughout all four stages and for both the global and regional components. The FFT-AMSE model used a fast Fourier Transform (FFT) variant of the AMSE loss introduced in Ref. [19] in the regional component, replacing the spherical harmonic transform with an FFT performed over the regional domain. Its global component uses MSE and is initialised from the MSE model at the end of Stage 2.

For the wavelet experiment, fields on reduced Gaussian grids are regridded to regular latitude–longitude grids before applying the two-dimensional DTCWT: O96 fields are mapped to a 1.0° grid and N320 fields to a 0.25° grid, following Ref. [8]. UKV fields are transformed directly on the native regional grid. The number of wavelet decomposition levels used for each dataset is given in Table B.1, and is calculated via

where $H$ and $W$ are the meridional and zonal dimensions of the input grid, respectively.

We compute the dual-tree complex wavelet transform using a custom Triton [35] implementation written by the authors, validating its forward transform and gradients against pytorch_wavelets [36, 37] to within numerical precision. The implementation uses the ‘near_sym_b’ and ‘qshift_b’ filter sets provided by pytorch_wavelets, corresponding to a first-level biorthogonal filter bank followed by quarter-shift filters at subsequent levels [24]. With the Triton implementation, wavelet training achieved approximately 20% lower throughput than the MSE baseline in Stages 1 and 2. No difference was observed in Stage 3, where throughput was limited by data loading rather than loss computation.

Verification methodology

Surface verification uses quality-controlled observations from 114 stations across the UK. Six-hour precipitation accumulations are derived separately from hourly Synoptic Rainfall Europe West observations, following standard operational verification procedures. The same observation datasets and processing are used to evaluate all model configurations.

Neighbourhood-based verification is performed using the High-Resolution Assessment (HiRA) framework [32]. For each observation, a square neighbourhood centred on the nearest model grid point is interpreted as a pseudo-ensemble, and the resulting distribution is evaluated using the continuous ranked probability score (CRPS). We use neighbourhoods of $3 \times3$, $7 \times7$, and $11 \times11$ model grid points. HiRA therefore provides some tolerance to small spatial displacement errors while still rewarding appropriate local values and variability.

For variables without suitable station observations, forecasts are verified against the corresponding UKV operational analyses. For each analysis field, 1,000 grid points are sampled at fixed, reproducibly selected locations within the regional domain and treated as pseudo-observations. At each location, the surrounding forecast grid points form the HiRA pseudo-ensemble, from which CRPS is calculated for neighbourhoods of $3 \times3$, $7 \times7$, and $11 \times11$ model grid points. These scores measure agreement with the gridded analysis rather than independent observational skill.

Pointwise RMSE is calculated using variable-specific collocation. Forecast 2 m temperature is bilinearly interpolated to each observation location, and observations are adjusted to the UKV model orography using a standard environmental lapse rate of $-0.0065\ \mathrm{K}\ \mathrm{m}^{-1}$ [38]. The same elevation-adjusted observations are used for the corresponding HiRA calculation. Six-hour precipitation accumulations are compared with the nearest model grid point to avoid smoothing intermittent precipitation fields. The 10 m zonal and meridional wind components are bilinearly interpolated before calculating wind speed and vector error.

Additional verification

Comparison with the MSE-trained model

Figure D.1 extends the comparison between the MSE- and wavelet-trained models across surface and upper-air variables, verification metrics, and lead times. For verification against surface observations, wavelet training generally improves HiRA CRPS, particularly for the wind diagnostics and at larger neighbourhood sizes, while the MSE-trained model generally retains a small advantage in pointwise RMSE. Differences in upper-air RMSE are also generally small and vary with variable and lead time.

The corresponding spectral differences are larger and more systematic. For the surface variables, the reduction in intermediate- and short-wavelength power deficits under wavelet training persists across lead times of 6, 24, and 72 hours (Figure D.2). The distinction is clearest for 10 m wind speed and accumulated precipitation, while the differences for 2 m temperature are smaller. This lead-time comparison supports the conclusions drawn from the 24-hour spectra in § 4.2, showing that the improved agreement with the reference spectra is not confined to a single forecast lead time.

Similar behaviour is found for the upper-air variables (Figure D.3). The MSE-trained model develops intermediate-scale power deficits for 850 hPa temperature and 700 hPa specific humidity, both of which are substantially reduced by wavelet training. For 500 hPa geopotential, the MSE-trained spectra instead contain narrow regions of excess power at wavelengths associated with the internal processor mesh; these features are strongly suppressed in the wavelet-trained model.

Together, the additional verification shows that wavelet training produces forecasts that remain closer to the reference spectra across a broader range of resolved scales, variables, and lead times, while its effect on pointwise and neighbourhood-based scores depends more strongly on the variable and verification metric.

Comparison with the FFT-AMSE-trained model

The wavelet-trained model generally outperforms the FFT-AMSE-trained model across variables, verification metrics, and lead times (Figure D.4). Improvements are particularly skeletal widespread for upper-air RMSE and are largest at short to medium lead times. Wavelet training also generally improves HiRA CRPS for surface variables, particularly at the larger neighbourhood sizes. Differences in surface RMSE are smaller and less systematic, although they predominantly favour the wavelet-trained model.

Spectral fine-tuning commonly degrades pointwise performance relative to MSE because preserving fine-scale variability limits the smoothing that minimises pointwise error. The results presented here show that wavelet training substantially reduces this trade-off. The wavelet-trained model retains pointwise performance close to that of the MSE-trained model and performs considerably better than FFT-AMSE across most pointwise and neighbourhood-based metrics. At the same time, the spectral comparisons presented in Section 4.2 show that the wavelet-trained forecasts remain closer to the reference spectra over a broader range of resolved scales than either baseline. MSE training produces the largest intermediate- and small-scale power deficits, while FFT-AMSE recovers some of this missing variance but remains less consistent with the reference than the wavelet-trained model. Wavelet training therefore reproduces the reference distribution of variance across spatial scales more accurately than FFT-AMSE, while maintaining pointwise performance much closer to that of MSE. Among the three training objectives, the wavelet loss consequently provides the most favourable combination of forecast accuracy and physically realistic multiscale structure.

Neighbourhood precipitation verification

The spectral and intensity-distribution diagnostics show that wavelet training produces sharper precipitation fields with greater fine-scale variability. Although HiRA CRPS assesses the continuous distribution of forecast values within a neighbourhood, it does not directly isolate performance for particular precipitation intensities. We therefore evaluate the Fractions Skill Score (FSS) [39] across a range of precipitation thresholds and neighbourhood sizes to determine how spatial skill varies with both event intensity and scale. FSS compares the fractions of forecast and observed grid points exceeding a given threshold within each neighbourhood. It ranges from zero to one, with one indicating perfect agreement, and provides a displacementtolerant measure of precipitation skill that avoids the double penalty associated with strict pointwise verification.

Figure D.5 shows the absolute FSS of the wavelet-trained model and its differences relative to the MSE- and FFT-AMSE-trained models at a lead time of 24 hours. Wavelet training improves FSS relative to MSE across almost all precipitation thresholds and neighbourhood sizes, with the largest gains occurring for heavier precipitation and broader neighbourhoods. Improvements over FFT-AMSE are smaller but remain consistently positive across most of the matrix and follow a similar dependence on threshold and neighbourhood size.

The improvement over FFT-AMSE indicates that the benefit of wavelet training is not explained solely by greater spectral amplitude. The spatial localisation of the wavelet objective also improves the neighbourhood-scale agreement of precipitation structures. Wavelet training therefore produces sharper and more intense precipitation features without sacrificing spatial skill, with particularly clear benefits under displacement-tolerant evaluation of heavier precipitation.