🤖 machine learning

Multi-Scale Wavelet Transformers for Operator Learning of Dynamical Systems

This paper introduces Multi-Scale Wavelet Transformers (MSWTs), a novel architecture that learns dynamical system dynamics in a tokenized wavelet domain to overcome the spectral bias of existing neural operators, thereby significantly improving long-horizon stability and high-frequency fidelity in chaotic and real-world climate forecasting tasks.

Original authors: Xuesong Wang, Michael Groom, Rafael Oliveira, He Zhao, Terence O'Kane, Edwin V. Bonilla

Published 2026-05-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xuesong Wang, Michael Groom, Rafael Oliveira, He Zhao, Terence O'Kane, Edwin V. Bonilla

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Blurry Lens"

Imagine you are trying to predict the weather or the flow of water in a river using a computer. Scientists have two main ways to do this:

The Old Way (Numerical Solvers): Like a super-precise but incredibly slow accountant who calculates every single drop of water. It's accurate but takes forever.
The New Way (AI Models): Like a fast, intuitive artist who learns the "vibe" of the water and predicts the next frame instantly.

The problem with the "Fast Artist" (specifically a type of AI called a Neural Operator) is that they suffer from "Spectral Bias."

Think of this like a camera lens that is slightly out of focus. It captures the big, smooth shapes of the clouds (low-frequency details) perfectly, but it blurs out the tiny, sharp edges (high-frequency details) like the jagged tips of a lightning bolt or the ripples on a wave.

In the real world, these tiny details matter. If your AI blurs them out, small errors pile up. Over time, a tiny blur in a weather model can turn into a massive, wrong prediction days later. The AI becomes unstable, like a house of cards that collapses after a few seconds.

The Solution: The "Multi-Scale Wavelet Transformer" (MSWT)

The authors propose a new AI architecture called MSWT. To understand how it works, let's use a few analogies.

1. The "Magic Zoom Lens" (Wavelet Transforms)

Most AI models look at an image or a weather map as one giant, flat picture. If they try to zoom out to see the big picture, they often lose the small details.

MSWT uses a Wavelet Transform. Imagine you have a high-resolution photo of a forest.

A standard AI looks at the whole forest at once.
The MSWT uses a "Magic Zoom Lens" that splits the image into four distinct layers simultaneously:
- The Big Picture: The general shape of the forest (Low-Low).
- The Vertical Edges: The tall tree trunks (High-Low).
- The Horizontal Edges: The tree branches (Low-High).
- The Tiny Details: The leaves and twigs (High-High).

Crucially, the MSWT never throws away the "Tiny Details" layer. It keeps all four layers separate and processes them together. This ensures the AI never "blurs out" the sharp, chaotic parts of the system.

2. The "Tokenized" Inbox

The paper also introduces a Patch Tokenizer. Imagine you have a massive spreadsheet with millions of rows of data. Reading it line-by-line is slow and confusing.

The MSWT chops the spreadsheet into small, manageable "post-it notes" (patches).
It turns each post-it note into a single "token" (a summary card).
This makes the data much smaller and easier for the AI to process quickly, without losing the essential information inside each patch.

3. The "Smart Conversation" (Wavelet Attention)

Once the data is split into these different frequency layers (the big picture vs. the tiny details), the AI needs to understand how they talk to each other.

Standard AI might ignore the tiny details because they seem "noisy."
The MSWT uses a special Wavelet Attention mechanism. Think of this as a moderator in a meeting who ensures the person whispering about the tiny details (the high-frequency noise) gets heard just as loudly as the person shouting about the big trends.
This allows the model to learn how a tiny ripple in a river affects the massive current downstream, and vice versa.

How They Tested It

The authors tested their new "Magic Zoom Lens" AI on three difficult challenges:

Chaotic Water Flow (Kolmogorov Flow): A simulation of swirling, chaotic water.
- Result: Other AIs got confused and the simulation exploded (became unstable) after a short time. The MSWT kept the water swirling realistically for much longer, with significantly less error.
Shallow Water Waves: Simulating waves moving across a surface.
- Result: The MSWT was better at keeping the waves looking like real waves, rather than smoothing them out into flat lines.
Real-World Weather (ERA5): They tested it on 10 years of real global weather data.
- Result: When running a "free-running" forecast (predicting the future without correcting the data every day), the MSWT stayed stable for a decade. It didn't drift into a fake climate. It kept the "climate bias" (the average error over time) much lower than the previous best models.

The Takeaway

The paper claims that by explicitly separating and preserving both the big, smooth trends and the tiny, sharp details using a mathematical tool called wavelets, they built an AI that doesn't just guess the future; it respects the physics of the small stuff.

This prevents the "avalanche of errors" that usually happens when AI models try to predict complex systems like weather or fluid dynamics over long periods. It's like giving the AI a pair of glasses that never blurs, allowing it to see the whole picture clearly, from the mountain range down to the single leaf.

Technical Summary: Multi-Scale Wavelet Transformers for Operator Learning of Dynamical Systems

Problem Statement
Data-driven surrogates for dynamical systems, particularly neural operators, offer significant speed advantages over numerical solvers for partial differential equations (PDEs). However, these models frequently suffer from spectral bias (or oversmoothing), where neural networks learn low-frequency components more readily than high-frequency ones. In the context of operator learning, this manifests as a systematic underestimation of high-wavenumber content. This limitation is critical for multi-scale and chaotic systems, where small-scale errors can rapidly degrade large-scale predictions, leading to long-horizon instability and unphysical solutions in applications like weather forecasting. Existing mitigation strategies, such as multi-scale architectures or frequency-aware representations, often fail to preserve high-frequency features during aggressive downsampling or rely on computationally expensive attention mechanisms applied directly to high-resolution physical grids.

Methodology
The authors propose the Multi-Scale Wavelet Transformer (MSWT), an operator learning framework designed to learn system dynamics within a tokenized wavelet domain. The architecture addresses spectral bias by explicitly separating low- and high-frequency content across scales while maintaining computational efficiency. The MSWT consists of three primary components:

Patch Tokenizer: The physical field is mapped to a token space via non-overlapping patchification. This reduces the effective resolution for attention operations, smoothing the optimization landscape and improving efficiency for high-resolution PDEs.
Multi-Scale Wavelet Structure: The model employs a U-shaped topology with $L$ $L$ scales of downsampling and upsampling.
- Wavelet-Preserving Downsampling: Unlike standard pooling or strided convolutions that discard high-frequency information, the encoder uses the Discrete Wavelet Transform (DWT) (specifically Haar wavelets) to decompose features into four subbands: Low-Low (LL), Low-High (LH), High-Low (HL), and High-High (HH). All subbands are retained and passed to the next scale, ensuring high-frequency details are not lost.
- Wavelet-Preserving Upsampling: The decoder utilizes the Inverse DWT (iDWT) to reconstruct higher-resolution features from the subbands, fusing them with skip connections from the encoder.
Wavelet Attention Block (WAttn): Operating within the wavelet space, this block models dependencies across different frequency subbands. It applies a dimension compression, performs DWT to extract wavelet coefficients, and then utilizes a combination of convolution and self-attention (specifically local window attention) to mix features across subbands. This allows the model to amplify high-frequency components that typically possess low energy but are crucial for fine-scale structure. Finally, an iDWT maps the coefficients back to the token space.

The model is trained using a standard relative $\ell_2$ (MSE-type) loss on one-step prediction pairs, without requiring specialized objective functions to penalize high-frequency errors, isolating the contribution of the architectural inductive bias.

Key Contributions

Mitigation of Spectral Bias via Wavelet Attention: The paper introduces a transformer architecture that models multiband wavelet content across scales. By operating in the wavelet domain, the model explicitly preserves and mixes high-frequency subbands, resulting in improved spectral fidelity compared to Fourier-based or standard convolutional operators.
Improved Efficiency and Long-Horizon Stability: By combining tokenization with wavelet-based down/up-sampling, the method significantly reduces the computational cost of attention mechanisms (operating on reduced spatial resolutions) while maintaining stability. This enables accurate long-horizon rollouts of chaotic dynamical systems where other models typically diverge or accumulate error.
Long-Horizon Evaluation on Real-World Climate Data: The authors validate the method on the ERA5 climate reanalysis dataset. The results demonstrate reduced climatological bias in free-running forecasts, indicating improved long-term physical consistency compared to state-of-the-art baselines.

Experimental Results
The MSWT was evaluated on three benchmarks: the Chaotic Kolmogorov Flow (CKF), the Shallow Water Equation (SWE), and the ERA5 climate reanalysis.

Chaotic Kolmogorov Flow (CKF): MSWT achieved substantial error reductions over baselines (FNO, U-Net, WNO, SAOT, HFS). At step 64, MSWT reduced the Relative L2 error by over 60% compared to the second-best stable model (FNO). It also showed superior performance in spectral metrics (EMLR and EMAE), with improvements of 76.5% and 80.6% respectively, and maintained spectral fidelity closely matching the ground truth, whereas baselines exhibited energy overestimation or unphysical solutions.
Shallow Water Equation (SWE): MSWT demonstrated lower error accumulation in long-horizon rollouts (steps 41 and 81), improving upon the second-best stable baseline (FNO) by approximately 30% in Relative L2, EMLR, and EMAE. It successfully avoided boundary artifacts observed in other methods.
ERA5 Climate Reanalysis: In decade-scale autoregressive rollouts, MSWT consistently outperformed the Lightweight Uncoupled ClImate Emulator (LUCIE). It achieved lower Root Mean Square Error (RMSE) across all variables (temperature, humidity, winds, pressure, precipitation) and notably reduced climatological bias, particularly in surface pressure near polar regions and meridional wind bias over the South Pole.

Significance
The paper claims that the MSWT framework effectively isolates the role of representation and inductive bias in mitigating spectral bias. By learning dynamics in a tokenized wavelet space, the method achieves a balance between computational efficiency and the preservation of fine-scale structures. The results suggest that explicit wavelet-based sampling and attention mechanisms are superior for long-term stability in chaotic and multi-scale systems, offering a robust solution for real-world forecasting tasks where maintaining physical consistency over long horizons is critical. The authors emphasize that these improvements are achieved without modifying standard training objectives, highlighting the architectural design as the primary driver of performance.