PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Anna-Louise Ellis, J. Scott Hosking, Manvendra Janmaijaya, David Salvador-Jasin, Ben Shipway, Nathan Simpson, Oliver Strickson, David Walters

arXiv:2610.37325v1physics.ao-phcs.LG

TL;DR

Kilometre-scale MLWP must retain local fine-scale structure without losing spatial accuracy, a challenge for pointwise and globally supported spectral objectives. We introduce a scale-, location-, and orientation-aware DTCWT loss; on joint UK-global forecasts, it improves spectral fidelity over MSE and FFT-AMSE while retaining pointwise performance close to MSE.

  • Problem

    Kilometre-scale training data contain local weather systems and surface influences important for high-impact weather, but their preservation in forecasts is not guaranteed; pointwise MSE and globally supported spectral objectives present distinct spatial-structure challenges.

    Source: §1 · pp. 1–2

  • Method

    The paper develops a deterministic DTCWT loss that decomposes atmospheric fields jointly by scale, location, and orientation, with amplitude- and coherence-sensitive terms evaluated by scale.

    Source: §1 · p. 2, §5 · p. 10

  • Results

    The wavelet-trained model retains pointwise performance close to MSE, generally improves neighbourhood verification, and remains closer to reference spectra across more resolved scales than MSE or FFT-AMSE.

    Source: §5 · p. 10

  • Takeaways & Limitations

    Wavelet training reduces, but does not eliminate, the trade-off between preserving fine-scale variability and maintaining spatial accuracy.

    Source: §5 · p. 10

    All three models underestimate the highest rainfall quantiles, and wavelet training does not fully recover accumulated-precipitation morphology or temporal organisation.

    Source: §4.3 · p. 9, §5 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Kilometre-scale MLWP must preserve local fine-scale weather structure without sacrificing spatial accuracy. This section introduces a scale-, location-, and orientation-aware wavelet loss and reports improved spectral fidelity alongside strong spatial performance.

1 Introduction

  • Kilometre-scale forecasts must preserve local weather systems and surface influences, while MSE can smooth uncertain, sharply localised structures such as convective precipitation [19].Source: §1 · pp. 1–2
  • The joint configuration produces simultaneous forecasts on a global N320 grid and a 1.5 km UK regional grid, with a shared processor mesh refined over the regional domain.The loss compares directional wavelet-band coefficients and a final low-pass residual.Source: Figure 1 · p. 2
  • Globally supported spectral objectives can improve spectral fidelity but constrain errors less locally and may compromise the spatial positioning measured by RMSE [19].Source: §1 · p. 2
  • The DTCWT loss decomposes fields jointly by scale, location, and orientation, enabling independently weighted scale contributions and localised constraints on atmospheric structures.Its complex coefficients support amplitude- and coherence-sensitive terms analogous to spectral AMSE.Source: §1 · pp. 2–3
  • The wavelet loss improves spectral fidelity over MSE and an FFT-based spectral objective, maintains pointwise performance comparable to MSE, and improves neighbourhood-based verification.Source: §1 · p. 3

2 Deterministic wavelet loss

The dual-tree complex wavelet loss compares atmospheric fields by scale, orientation, and spatial arrangement, while back-propagation through compactly supported filters yields more localised field-space gradients than Fourier transforms. Constructed perturbations and early regional-training diagnostics show that the loss components respond differently across scales and variables.

2 Deterministic wavelet loss

  • The transform produces complex directional band-pass coefficient fields at six approximate orientations and a final real-valued low-pass residual; the low-pass residual is compared directly using MSE.Band losses combine amplitude error with a power-weighted penalty for disagreement between complex coefficient patterns, paired at corresponding locations.Source: §2 · pp. 3–4
  • Band statistics aggregate locations within each level, so locations and gradients remain coupled; compactly supported filters nevertheless yield more localised field-space gradients than globally supported Fourier bases.The coefficient grids are downsampled by level, with j = 1 finest and increasing j progressively coarser.Source: §2 · pp. 3–4

2.1 Scale selectivity

  • For the perturbed time-offset prediction, the low-pass term contributes 80% of scalar loss but 46% of summed component-gradient norms, leaving substantial scale-dependent training signals.Gradient percentages are relative component-gradient magnitudes, not an additive decomposition of the total-gradient norm.Source: §2.1 · p. 4, Figure 2 · p. 5
  • For the blurred prediction, wavelet bands j = 2 and j = 3 account for approximately 81% of scalar loss and 85% of summed component-gradient norms.The large-scale structure is retained, making the low-pass contribution to both loss and gradient magnitude negligible; the concentration also reflects greater 2 m temperature variance at j = 2–3 than j = 1.Source: §2.1 · pp. 4–5
  • Across both constructed cases, broad structural errors activate low-pass and coarser-scale terms, while selectively attenuated fine variability concentrates the response in affected wavelet bands.The component gradients are signed first-order sensitivities on the original input grid; parameter updates also depend on model back-propagation and the optimiser.Source: §2.1 · pp. 4–5

2.2 Evolution of scale-wise loss contributions during training

  • For 500 hPa geopotential, during the first 25% of Stage 3a training on the regional UKV domain, low-pass MSE contributes more than 90% of the objective, while pointwise MSE and complete wavelet loss follow similar trajectories.The figure excludes the global ERA5 component.Source: §2.2 · p. 5, Figure 3 · p. 6
  • For temperature and precipitation, low-pass error falls faster than wavelet-band losses, increasing the bands’ share of the remaining objective; precipitation’s wavelet validation loss keeps decreasing after pointwise MSE flattens.For temperature, separation between wavelet validation loss and pointwise MSE indicates residual scale-dependent errors not distinguished by the pointwise metric.Source: §2.2 · p. 6

3 Forecast evaluation

The evaluation compares three training objectives in a joint global–regional FastNet configuration and verifies forecasts against surface observations.

3 Forecast evaluation

  • The three variants share an architecture and differ principally in training objective: MSE, FFT-based AMSE, or the proposed wavelet loss.The wavelet model uses uniform weights across decomposition levels.Source: §3 · p. 6
  • Each model forecasts on a global ∼31 km grid and a UK-centred 1.5 km regional grid, with forecasts every six hours through 240 hours during 2023.Source: §3 · pp. 6–7
  • Surface verification uses pointwise RMSE and HiRA CRPS, treating surrounding model-grid points as a pseudo-ensemble across three neighbourhood sizes.The HiRA neighbourhoods are 3 × 3, 7 × 7, and 11 × 11 grid points.Source: §3 · p. 7

4 Results

Wavelet training improves spectral fidelity and precipitation structure while retaining similar overall skill to MSE training. Its verification trade-off is metric-dependent, and precipitation weaknesses remain.

4.1 Verification against surface observations

  • Wavelet training generally lowers HiRA CRPS, especially for both 10 m wind diagnostics, while MSE generally has slightly lower pointwise RMSE.The lower HiRA CRPS should not be read solely as evidence of displaced but otherwise accurate features; local bias, intensity, or variability may also contribute.Source: §4.1 · p. 7, Figure 4 · p. 7

4.2 Spectral fidelity

  • Wavelet forecasts remain closer to reference power across intermediate and short wavelengths than MSE or FFT-AMSE, most clearly for wind speed and precipitation.MSE has a short-wavelength power deficit, while its mesh-aligned excess-power features are absent or strongly suppressed with wavelet training.Source: Figure 5 · p. 8, §4.2 · p. 8

4.3 Precipitation forecasts

  • Wavelet forecasts produce sharper, more intense precipitation and reproduce the upper tail of observed accumulated and instantaneous rainfall better than MSE or FFT-AMSE.At 24 hours, they also achieve higher Fractions Skill Scores than both alternatives across most precipitation thresholds and neighbourhood sizes, with clearest gains for heavier rainfall and broader neighbourhoods.Source: §4.3 · pp. 8–9, Figure 6 · p. 9
  • Despite these gains, all model variants underestimate the highest precipitation quantiles, and neither MLWP model fully reproduces advected streak-like structures in six-hour accumulations.Source: §4.3 · p. 9

5 Conclusions

The DTCWT loss decomposes fields by scale, location, and orientation, producing spatially localised gradients. In the joint global–regional FastNet configuration, it improves multiscale fidelity over MSE and FFT-AMSE while retaining pointwise performance close to MSE.

5 Conclusions

  • The DTCWT decomposes atmospheric fields by scale, location, and orientation, enabling scale-specific amplitude- and coherence-sensitive loss terms.Compactly supported filters produce more spatially localised training gradients than an equivalent Fourier-domain objective.Source: §5 · pp. 9–10
  • Across evaluated variables and lead times, wavelet forecasts stay closer to reference spectra over more resolved scales than MSE or FFT-AMSE.They reduce MSE’s intermediate- and small-scale power deficits and suppress excess power aligned with processor-mesh scales.Source: §5 · p. 10
  • The wavelet loss provides a better balance of spatial accuracy and multiscale fidelity than pointwise MSE or FFT-AMSE in forecasts at 1.5 km over the UK and approximately 31 km globally.Regional pointwise performance remains close to MSE, neighbourhood-based verification generally improves, and performance exceeds FFT-AMSE.Source: §5 · p. 10
  • Wavelet training produces sharper, more spatially concentrated precipitation forecasts and more realistic intensity distributions, but all three models underestimate the most intense rainfall.It does not fully recover accumulated-precipitation morphology or temporal organisation; future work should test scale-weight sensitivity, other regions and weather regimes, and probabilistic forecasting.Source: §5 · p. 10

Statement on the use of generative AI

The authors produced the manuscript’s results, analysis, and scientific conclusions; Microsoft 365 Copilot assisted with editorial clarity and readability.

Statement on the use of generative AI

A Spatial attribution and gradient comparison

The comparison examines first-order attribution and field-space gradients for three objectives under two constructed prediction conditions. The wavelet objective identifies affected scales while retaining location-dependent sensitivity to atmospheric structures.

A Spatial attribution and gradient comparison

  • The wavelet objective identifies affected spatial scales while retaining location-dependent sensitivity to atmospheric structures in the two constructed prediction cases.The comparison concerns spatial structure and relative localisation; absolute loss and gradient magnitudes depend on each objective’s definition and normalisation.Source: Figure A.1 · p. 14, A Spatial attribution and gradient… · p. 15
  • Attribution combines local prediction error with objective sensitivity, whereas the field-space gradient shows signed sensitivity to changes in the predicted field.A location can have a large gradient but modest attribution if its error is small; attribution is a first-order diagnostic, not an exact additive decomposition of scalar loss.Source: A Spatial attribution and gradient… · pp. 14–15
  • For the time-offset prediction, all objectives show broadly similar attribution patterns, but FFT-AMSE gradients spread more broadly while wavelet gradients vary more around error-contributing structures.The MSE gradient follows pointwise prediction error; the wavelet pattern is consistent with spatial indexing and compact support despite within-band coupling.Source: A Spatial attribution and gradient… · p. 15
  • For the blurred prediction, the wavelet gradient is more organised around altered temperature features than the broadly distributed FFT-AMSE gradient, with a structured response over land and complex terrain.MSE follows the smoothing residual but does not distinguish the spatial scales contributing to it.Source: A Spatial attribution and gradient… · p. 15

B Model configuration and training schedule

The configuration uses a shared latent representation to produce global and UK regional forecasts from separate datasets and grids. Training progresses from global forecasting to regional adaptation and autoregressive rollouts, with three variants distinguished by their objectives.

B Model configuration and training schedule

  • Separate encoder–decoder pairs process ERA5 global fields and Met Office UKV fields on a UK regional grid, while sharing a latent representation and producing both forecasts in one forward pass.The model encodes fields onto a latent mesh, processes them there, and decodes them onto their output grids.Source: B Model configuration and training… · p. 15
  • Stage 1 establishes global forecasting on ERA5 O96 with a coarse latent mesh; Stage 2 increases ERA5 resolution to N320 and refines the global mesh.Both stages cover the same training period.Source: B Model configuration and training… · p. 15, Table B.1 · p. 16
  • Stage 3a introduces UKV, a regional encoder–decoder, and a high-resolution regional mesh embedded in the global mesh, retaining global context during UK adaptation.Global and regional losses receive weights of 67% and 33%; without this weighting, the regional domain contributes less than 1% of the combined objective.Source: B Model configuration and training… · p. 15
  • Stage 3b extends training to autoregressive rollouts of 2–12 steps, exposing the model to its own predictions to improve longer-lead stability.A reduced learning rate supports fine-tuning over longer forecast horizons while limiting degradation of the learned representations.Source: B Model configuration and training… · pp. 15–16
  • The models share an architecture but differ in training objectives: MSE and wavelet losses apply across all stages and components, while FFT-AMSE replaces MSE only in the regional component.The FFT-AMSE global component uses MSE and is initialised from the MSE model at the end of Stage 2.Source: B Model configuration and training… · p. 16
  • For wavelet training, reduced Gaussian fields are regridded to regular latitude–longitude grids, while UKV fields are transformed on their native regional grid.The wavelet transform uses levels specified by dataset; other configuration settings are shared across model variants.Source: B Model configuration and training… · p. 16
  • The authors’ custom Triton DTCWT implementation was validated against pytorch_wavelets to numerical precision; wavelet training ran about 20% slower than MSE in Stages 1 and 2.No throughput difference was observed in Stage 3, where data loading, rather than loss computation, limited throughput.Source: B Model configuration and training… · pp. 16–17

C Verification methodology

Verification combines pointwise error measures with neighbourhood-based CRPS, using station observations or UKV analyses according to variable availability. HiRA treats nearby forecast values as a pseudo-ensemble, while analysis-based scores measure agreement with gridded analyses rather than independent observational skill.

C Verification methodology

  • Surface verification uses quality-controlled observations from 114 UK stations, with six-hour precipitation accumulations drawn separately from hourly rainfall observations.The same observation datasets and processing are used for all model configurations.Source: C Verification methodology · p. 17
  • HiRA treats a square of forecast grid points around each observation as a pseudo-ensemble, scored with CRPS over 3×3, 7×7, and 11×11 neighbourhoods.This provides some tolerance to small spatial displacement while rewarding appropriate local values and variability.Source: C Verification methodology · p. 17
  • For variables without suitable station observations, 1,000 fixed regional grid points from UKV analyses serve as pseudo-observations, so scores measure agreement with analyses rather than independent observational skill.Forecast neighbourhoods around each point form the HiRA pseudo-ensemble.Source: C Verification methodology · p. 17
  • Pointwise RMSE uses variable-specific collocation, including bilinear interpolation for temperature and wind components and nearest-grid comparison for six-hour precipitation.Temperature observations are adjusted to UKV orography using a −0.0065 K m−1 lapse rate.Source: C Verification methodology · p. 17

D.1 Comparison with the MSE-trained model

Against MSE, wavelet training generally improves neighbourhood verification while retaining a small pointwise RMSE disadvantage in some comparisons. Its spectral improvements persist across lead times and upper-air variables, though pointwise and neighbourhood effects depend on the variable and metric.

D.1 Comparison with the MSE-trained model

  • Wavelet training generally improves HiRA CRPS, particularly for wind and larger neighbourhoods, while MSE generally retains a small advantage in pointwise RMSE.The scorecard compares surface and upper-air variables, observations and analyses, and lead times up to 10 days.Source: §D.1 · p. 17, Figure D.1 · p. 18
  • Reduced intermediate- and short-wavelength power deficits persist at 6-, 24-, and 72-hour leads, most clearly for 10 m wind speed and accumulated precipitation.Differences are smaller for 2 m temperature, so the improvement is not confined to a single lead time.Source: §D.1 · p. 17
  • Wavelet training substantially reduces intermediate-scale deficits in 850 hPa temperature and 700 hPa specific humidity, and suppresses mesh-scale excess power in 500 hPa geopotential.Figure D.3 covers these three upper-air variables.Source: §D.1 · pp. 17–18
  • Across resolved scales, variables, and lead times, wavelet forecasts remain closer to reference spectra, while pointwise and neighbourhood score effects vary more by variable and metric.Upper-air RMSE differences are generally small and vary with variable and lead time.Source: §D.1 · pp. 17–18

D.2 Comparison with the FFT-AMSE-trained model

Wavelet training combines more accurate variance across spatial scales than FFT-AMSE with pointwise performance closer to MSE. It generally outperforms FFT-AMSE across verification metrics, variables, and lead times.

D.2 Comparison with the FFT-AMSE-trained model

  • Wavelet training generally outperforms FFT-AMSE across variables, verification metrics, and lead times, with particularly widespread upper-air RMSE gains and improved surface HiRA CRPS at larger neighbourhoods.Surface RMSE differences are smaller and less systematic, though they predominantly favour wavelet training.Source: §D.2 · pp. 18–19, Figure D.4 · p. 20
  • The scorecard compares pointwise RMSE and HiRA CRPS for 3×3, 7×7, and 11×11 neighbourhoods over lead times up to 10 days.Verification is against observations and operational analyses, aggregated over forecasts initialised from 1 January to 21 December 2023.Source: Figure D.4 · p. 20
  • Wavelet training reproduces the reference distribution of variance across spatial scales more accurately than FFT-AMSE while maintaining pointwise performance much closer to MSE.Among the three objectives, the wavelet loss provides the most favourable combination of forecast accuracy and multiscale structure.Source: §D.2 · p. 20

D.3 Neighbourhood precipitation verification

At a 24-hour lead, wavelet training improves precipitation FSS over MSE across almost all thresholds and neighbourhood sizes, with smaller but broadly positive gains over FFT-AMSE. The gains are clearest for heavier precipitation and broader neighbourhoods.

D.3 Neighbourhood precipitation verification

  • FSS compares the fractions of forecast and observed grid points exceeding a threshold within each neighbourhood, providing displacement-tolerant precipitation skill without the pointwise double penalty.The evaluation spans precipitation thresholds and neighbourhood sizes.Source: §D.3 · pp. 20–21
  • At 24 hours, wavelet training improves precipitation FSS over MSE across almost all thresholds and neighbourhood sizes, with the largest gains for heavier precipitation and broader neighbourhoods.Improvements over FFT-AMSE are smaller but consistently positive across most of the threshold-by-neighbourhood matrix.Source: §D.3 · p. 21, Figure D.5 · p. 21
  • The FFT-AMSE comparison indicates that gains are not explained solely by greater spectral amplitude; wavelet training improves neighbourhood-scale agreement while producing sharper, more intense precipitation features.Benefits are particularly clear under displacement-tolerant evaluation of heavier precipitation.Source: §D.3 · p. 21
Loading 2610.37325v1…