Multi-Scale Wavelet Transformers for Operator Learning of Dynamical Systems

TL;DR

The paper addresses the spectral bias problem in neural operators for dynamical systems by introducing the Multi-Scale Wavelet Transformer (MSWT), which learns dynamics in a tokenized wavelet space. MSWT combines a patch tokenizer, wavelet-preserving down/up-sampling, and wavelet-based attention to explicitly preserve multiscale frequency content across scales. Trained with a standard relative $L^2$ loss, MSWT delivers substantial improvements in short- and long-horizon rollouts on chaotic PDE benchmarks and climate reanalysis, including reduced climatology bias on ERA5. The approach yields improved spectral fidelity and long-term stability, offering a practical, efficient surrogate for high-fidelity solvers in multi-query and forecasting tasks while maintaining discretization-invariance benefits of neural operators.

Abstract

Recent years have seen a surge in data-driven surrogates for dynamical systems that can be orders of magnitude faster than numerical solvers. However, many machine learning-based models such as neural operators exhibit spectral bias, attenuating high-frequency components that often encode small-scale structure. This limitation is particularly damaging in applications such as weather forecasting, where misrepresented high frequencies can induce long-horizon instability. To address this issue, we propose multi-scale wavelet transformers (MSWTs), which learn system dynamics in a tokenized wavelet domain. The wavelet transform explicitly separates low- and high-frequency content across scales. MSWTs leverage a wavelet-preserving downsampling scheme that retains high-frequency features and employ wavelet-based attention to capture dependencies across scales and frequency bands. Experiments on chaotic dynamical systems show substantial error reductions and improved long horizon spectral fidelity. On the ERA5 climate reanalysis, MSWTs further reduce climatological bias, demonstrating their effectiveness in a real-world forecasting setting.

Figures (20)

• Figure 1: Overview framework of Multi-Scale Wavelet Transformer (MSWT). Figure 1(a) presents the U-shaped topology. Inputs are embedded into patch tokens, processed by wavelet-attention blocks at multiple scales, and down/up-sampled via Wavelet-based sampling. Figure 1(b) illustrates the mechanism of the wavelet attention operator where the token mixing attention is achieved in the wavelet space. $\mathcal{W}$ represents the discrete wavelet transform (DWT) extracting features of four frequency bands (low-low, low-high, high-low, high-high) and $\mathcal{W}^{-1}$ represents the inverse transform (iDWT) to recover the token space. Figure 1(c) shows the training and inference of MSWT for operator learning of dynamical systems. $\mathcal{G}_\theta$ includes all the parameters $\theta$ in the operator $\mathcal{G}$.
• Figure 2: Predictions and enstrophy power spectrum on the CKF dataset for long-term rollouts. We present the 64-step rollout predictions of the baselines and our approach. More comparisons can be found in the Appendix \ref{['appd: CKF']}.
• Figure 3: Prediction and error comparisons on the SWE dataset at $t=81$. More detailed comparisons and the power spectrum results can be found in the Appendix \ref{['appd: SW2D']}.
• Figure 4: Ensemble-mean annual climatology bias of LUCIE and MSWT relative to ERA5 over the time period 2000--2010. The climatology is averaged over 5 ensemble members.
• Figure 5: Chaotic Kolmogorov Flow, prediction and error comparison, rollout step =1