pith. machine review for the scientific record. sign in

arxiv: 2602.01486 · v2 · submitted 2026-02-01 · 💻 cs.LG

Recognition: 2 theorem links

· Lean Theorem

Multi-Scale Wavelet Transformers for Operator Learning of Dynamical Systems

Xuesong Wang , Michael Groom , Rafael Oliveira , He Zhao , Terence O'Kane , Edwin V. Bonilla
Authors on Pith no claims yet

Pith reviewed 2026-05-16 08:18 UTC · model grok-4.3

classification 💻 cs.LG
keywords wavelet transformersneural operatorsdynamical systemsspectral biasmulti-scale learningchaotic systemsclimate forecastingoperator learning
0
0 comments X

The pith

Wavelet transformers learn dynamical system operators by preserving high-frequency content across scales.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper establishes that standard neural operators suffer from spectral bias that suppresses high-frequency details essential for accurate long-term evolution of chaotic systems. MSWTs address this by mapping inputs to a tokenized wavelet domain that explicitly separates frequency bands, then applying wavelet-preserving downsampling to keep fine-scale features and wavelet-based attention to model cross-scale interactions. Experiments demonstrate lower errors and better spectral fidelity on chaotic benchmarks plus reduced climatological bias on ERA5 reanalysis data. A sympathetic reader would care because reliable fast surrogates for high-dimensional dynamics directly enable longer, more stable forecasts in weather and climate applications.

Core claim

Multi-scale wavelet transformers learn system dynamics directly in a tokenized wavelet domain. The wavelet transform decomposes the state into low- and high-frequency components across scales; a wavelet-preserving downsampling scheme retains the high-frequency coefficients without loss; and wavelet-based attention layers capture dependencies both within and across frequency bands. This architecture yields substantial reductions in prediction error and improved long-horizon spectral fidelity on chaotic dynamical systems, together with lower climatological bias when applied to ERA5 climate reanalysis.

What carries the argument

Multi-scale wavelet transformer that tokenizes the wavelet transform of the input field and performs wavelet-based attention across scales and frequency bands.

If this is right

  • Prediction error decreases substantially on standard chaotic test systems.
  • Long-horizon spectral fidelity improves, preserving energy at small scales.
  • Climatological bias is reduced when the model is trained on ERA5 reanalysis.
  • The same architecture can serve as a drop-in surrogate that runs orders of magnitude faster than traditional numerical solvers.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The wavelet-domain approach may generalize to other operator-learning tasks that require faithful representation of fine-scale structure, such as turbulence closure modeling.
  • Hybrid models that combine MSWT layers with physics-informed constraints could further improve stability without sacrificing speed.
  • Similar frequency-aware tokenization might mitigate spectral bias in other sequence or grid-based transformers used for time-series forecasting.

Load-bearing premise

The wavelet-preserving downsampling and wavelet attention retain high-frequency features without introducing artifacts that would destabilize long-horizon dynamics.

What would settle it

A controlled experiment on the Lorenz-96 system showing that MSWT long-horizon predictions have the same or higher error and the same or worse spectral energy decay as a standard Fourier neural operator baseline.

Figures

Figures reproduced from arXiv: 2602.01486 by Edwin V. Bonilla, He Zhao, Michael Groom, Rafael Oliveira, Terence O'Kane, Xuesong Wang.

Figure 1
Figure 1. Figure 1: Overview framework of Multi-Scale Wavelet Transformer (MSWT) [PITH_FULL_IMAGE:figures/full_fig_p004_1.png] view at source ↗
Figure 2
Figure 2. 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 B.2.2 [PITH_FULL_IMAGE:figures/full_fig_p006_2.png] view at source ↗
Figure 3
Figure 3. 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 B.3.1 [PITH_FULL_IMAGE:figures/full_fig_p008_3.png] view at source ↗
Figure 4
Figure 4. 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. 5.4. Parameter Sensitivity Analysis We examine how the patch size used in patch tokenization affects model performance. Specifically, we evaluate patch sizes in 2, 4, 8 with a fixed wavelet downsampling depth of L = 3 on the CKF benchmark. The results… view at source ↗
Figure 5
Figure 5. Figure 5: Chaotic Kolmogorov Flow, prediction and error comparison, rollout step =1 10 0 10 1 Wavenumber k 10 6 10 5 10 4 10 3 10 2 10 1 10 0 Energy Spectrum E(k) Ground Truth FNO Unet SAOT HFS MSWT 10 0 10 1 Wavenumber k 10 2 10 1 10 0 10 1 Enstrophy Spectrum Z(k) Ground Truth FNO Unet SAOT HFS MSWT [PITH_FULL_IMAGE:figures/full_fig_p015_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: Chaotic Kolmogorov Flow, kinetic energy spectrum and enstrophy spectrum, rollout step =1 [PITH_FULL_IMAGE:figures/full_fig_p015_6.png] view at source ↗
Figure 7
Figure 7. Figure 7: Chaotic Kolmogorov Flow, prediction and error comparison, rollout step =30 15 [PITH_FULL_IMAGE:figures/full_fig_p015_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Chaotic Kolmogorov Flow, kinetic energy spectrum and enstrophy spectrum, rollout step =30 [PITH_FULL_IMAGE:figures/full_fig_p016_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: Chaotic Kolmogorov Flow, prediction and error comparison, rollout step =64 10 0 10 1 Wavenumber k 10 6 10 5 10 4 10 3 10 2 10 1 10 0 Energy Spectrum E(k) Ground Truth FNO Unet SAOT HFS MSWT 10 0 10 1 Wavenumber k 10 2 10 1 10 0 10 1 Enstrophy Spectrum Z(k) Ground Truth FNO Unet SAOT HFS MSWT [PITH_FULL_IMAGE:figures/full_fig_p016_9.png] view at source ↗
Figure 10
Figure 10. Figure 10: Chaotic Kolmogorov Flow, kinetic energy spectrum and enstrophy spectrum, rollout step =64 [PITH_FULL_IMAGE:figures/full_fig_p016_10.png] view at source ↗
Figure 11
Figure 11. Figure 11: Shallow water equation, prediction and error comparison, rollout step =1 [PITH_FULL_IMAGE:figures/full_fig_p017_11.png] view at source ↗
Figure 12
Figure 12. Figure 12: Shallow water equation, kinetic energy spectrum and enstrophy spectrum, rollout step =1 [PITH_FULL_IMAGE:figures/full_fig_p017_12.png] view at source ↗
Figure 13
Figure 13. Figure 13: Shallow water equation, prediction and error comparison, rollout step =41 B.3. Shallow Water Equation We use the PDEArena ShallowWater-2D dataset (Gupta & Brandstetter, 2022) (https://huggingface.co/ datasets/pdearena/ShallowWater-2D ), defined on a 96 × 192 spatial grid. For our setup, we use pressure and vorticity as input variables and evaluate long-horizon stability over trajectories of 87 autoregress… view at source ↗
Figure 14
Figure 14. Figure 14: Shallow water equation, kinetic energy spectrum and enstrophy spectrum, rollout step =41 [PITH_FULL_IMAGE:figures/full_fig_p018_14.png] view at source ↗
Figure 15
Figure 15. Figure 15: Shallow water equation, prediction and error comparison, rollout step =81 [PITH_FULL_IMAGE:figures/full_fig_p018_15.png] view at source ↗
Figure 16
Figure 16. Figure 16: Shallow water equation, kinetic energy spectrum and enstrophy spectrum, rollout step =81 18 [PITH_FULL_IMAGE:figures/full_fig_p018_16.png] view at source ↗
Figure 17
Figure 17. Figure 17: Ensemble-mean annual climatology bias of LUCIE and MSWT (relative to ERA5, 2000–2010). The climatology is averaged over 10 years of simulation and 5 ensemble members. 20 [PITH_FULL_IMAGE:figures/full_fig_p020_17.png] view at source ↗
Figure 18
Figure 18. Figure 18: Ensemble-mean annual climatology bias of LUCIE and MSWT (relative to ERA5, 2000–2010). The climatology is averaged over 10 years of simulation and 5 ensemble members. 21 [PITH_FULL_IMAGE:figures/full_fig_p021_18.png] view at source ↗
Figure 19
Figure 19. Figure 19: Zonal-mean climatology of the ensemble-mean rollouts (LUCIE, MSWT) compared with ERA5 over 2000–2010. 22 [PITH_FULL_IMAGE:figures/full_fig_p022_19.png] view at source ↗
Figure 20
Figure 20. Figure 20: Power (energy) spectra of ensemble-mean rollouts (LUCIE, MSWT) compared with ERA5 over 2000–2010. B.5. Ablation study To investigate the effectiveness of the three main modules separately in our framework: (1) patch tokenizer; (2) wavelet attention block; (3) wavelet-preserving down-/up-sampling, we propose three variants of the model for ablation studies: • MSWT-V1 (no tokenizer) uses patch size of 1 for… view at source ↗
read the original 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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

3 major / 2 minor

Summary. The manuscript proposes multi-scale wavelet transformers (MSWTs) for operator learning on dynamical systems. The architecture tokenizes inputs in the wavelet domain, applies a wavelet-preserving downsampling scheme to retain high-frequency coefficients, and uses wavelet-based attention to model cross-scale dependencies. The central claim is that this construction mitigates the spectral bias of standard neural operators, yielding lower rollout errors and improved long-horizon spectral fidelity on chaotic systems together with reduced climatological bias on ERA5 reanalysis data.

Significance. If the empirical gains are reproducible and attributable to the wavelet construction rather than capacity or training details, the work would offer a concrete architectural route to stable multi-scale forecasting. The emphasis on explicit frequency separation is well-motivated for chaotic and climate applications, and successful validation would strengthen the case for wavelet-domain operators in scientific machine learning.

major comments (3)
  1. [§3] §3 (Methods): the claim that the wavelet-preserving downsampling retains high-frequency content without phase or amplitude distortion is load-bearing for the long-horizon stability argument, yet no explicit bound, preservation lemma, or controlled ablation isolating the downsampling rule from standard pooling is supplied; without this, the reported spectral-fidelity gains cannot be confidently attributed to the wavelet mechanism rather than model capacity.
  2. [§4] §4 (Experiments): the abstract asserts 'substantial error reductions' and 'improved long horizon spectral fidelity' on chaotic systems, but the manuscript must include quantitative tables with baseline comparisons (e.g., FNO, DeepONet), error bars, and ablation rows that remove either the wavelet attention or the preserving downsampling; absent these, the central empirical claim remains unverifiable.
  3. [§4.2] §4.2 (ERA5 results): the reported reduction in climatological bias is presented without a corresponding spectral decomposition or high-frequency error metric; if the bias reduction occurs only in low-frequency bands, it would weaken the paper's claim that high-frequency retention is the operative mechanism.
minor comments (2)
  1. [§3.1] Clarify the precise wavelet family, decomposition level, and padding strategy in the first paragraph of §3.1; notation for the tokenized coefficients is introduced without an explicit equation reference.
  2. [Figure 2] Figure 2 caption should state the exact time horizon and frequency band used for the spectral error curves so readers can reproduce the fidelity comparison.

Simulated Author's Rebuttal

3 responses · 0 unresolved

We thank the referee for the constructive and detailed comments. We address each major point below and describe the revisions that will be incorporated into the next version of the manuscript.

read point-by-point responses

  1. Referee: [§3] §3 (Methods): the claim that the wavelet-preserving downsampling retains high-frequency content without phase or amplitude distortion is load-bearing for the long-horizon stability argument, yet no explicit bound, preservation lemma, or controlled ablation isolating the downsampling rule from standard pooling is supplied; without this, the reported spectral-fidelity gains cannot be confidently attributed to the wavelet mechanism rather than model capacity.

    Authors: We appreciate the referee's emphasis on formal justification. The downsampling rule is constructed so that only the approximation coefficients are downsampled while all detail coefficients are retained at their native resolution; because the discrete wavelet transform is a linear isometry (up to the filter norms), this operation introduces neither phase shift nor amplitude scaling beyond the known wavelet filter bounds. We will add a short preservation lemma in the revised §3 that bounds the L2 energy of the retained high-frequency coefficients. We will also insert a controlled ablation in §4 that replaces our downsampling with standard average pooling while keeping all other components fixed, thereby isolating its contribution to spectral fidelity. revision: yes

  2. Referee: [§4] §4 (Experiments): the abstract asserts 'substantial error reductions' and 'improved long horizon spectral fidelity' on chaotic systems, but the manuscript must include quantitative tables with baseline comparisons (e.g., FNO, DeepONet), error bars, and ablation rows that remove either the wavelet attention or the preserving downsampling; absent these, the central empirical claim remains unverifiable.

    Authors: We agree that the current experimental presentation is insufficiently quantitative. In the revised manuscript we will expand §4 with tables that report rollout MSE (with standard deviations over five independent runs) for MSWT against FNO, DeepONet, and the other baselines on all chaotic-system benchmarks. We will add two ablation rows: (i) MSWT with wavelet attention replaced by standard multi-head attention, and (ii) MSWT with the preserving downsampling replaced by conventional pooling. Both ablations will include the same long-horizon spectral-fidelity metrics already used in the main results. revision: yes

  3. Referee: [§4.2] §4.2 (ERA5 results): the reported reduction in climatological bias is presented without a corresponding spectral decomposition or high-frequency error metric; if the bias reduction occurs only in low-frequency bands, it would weaken the paper's claim that high-frequency retention is the operative mechanism.

    Authors: We accept this criticism. The revised §4.2 will contain a wavelet-based spectral decomposition of the climatological bias, reporting separate errors for the approximation (low-frequency) and detail (high-frequency) coefficients. We will additionally tabulate a high-frequency-specific metric (mean squared error restricted to the detail sub-bands) to demonstrate that the bias reduction is concentrated in the high-frequency regime, thereby supporting the mechanistic role of high-frequency retention. revision: yes

Circularity Check

0 steps flagged

No significant circularity; derivation self-contained

full rationale

The provided abstract and method description introduce MSWTs via wavelet transforms, a wavelet-preserving downsampling scheme, and wavelet-based attention as independent architectural choices to mitigate spectral bias. No equations, definitions, or self-citations are shown that reduce the claimed error reductions or spectral fidelity gains to quantities defined by fitted parameters, self-referential normalizations, or prior author results by construction. The central claims rest on experimental outcomes rather than tautological reductions, making the derivation chain self-contained against external benchmarks.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 1 invented entities

The proposal rests on the domain assumption that wavelet transforms cleanly separate frequency content and that a custom downsampling scheme can retain high-frequency information without loss. No free parameters or new physical entities are named in the abstract.

axioms (1)
  • domain assumption Wavelet transforms separate low- and high-frequency content across scales without information loss when paired with appropriate downsampling
    Invoked to justify the wavelet domain representation and the claim that high-frequency features are retained.
invented entities (1)
  • Multi-scale wavelet transformers (MSWTs) no independent evidence
    purpose: Learn system dynamics in tokenized wavelet domain while preserving high-frequency content
    New model architecture introduced to address spectral bias.

pith-pipeline@v0.9.0 · 5472 in / 1190 out tokens · 25970 ms · 2026-05-16T08:18:50.104134+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Reference graph

Works this paper leans on

12 extracted references · 12 canonical work pages · 1 internal anchor

  1. [1]

    Collins, Michael S

    Bonev, B., Kurth, T., Mahesh, A., Bisson, M., Kossaifi, J., Kashinath, K., Anandkumar, A., Collins, W. D., Pritchard, M. S., and Keller, A. FourCastNet 3: A geometric ap- proach to probabilistic machine-learning weather fore- casting at scale.arXiv preprint arXiv:2507.12144,

    work page arXiv

  2. [2]

    and Hassanzadeh, P

    Chattopadhyay, A. and Hassanzadeh, P. Long-term insta- bilities of deep learning-based digital twins of the cli- mate system: The cause and a solution.arXiv preprint arXiv:2304.07029,

    work page arXiv

  3. [3]

    Gupta, G., Xiao, X., and Bogdan, P

    URL https://arxiv.org/abs/2111.13587. Gupta, G., Xiao, X., and Bogdan, P. Multiwavelet-based operator learning for differential equations.Advances in neural information processing systems, 34:24048–24062,

    work page arXiv

  4. [4]

    Gupta, J. K. and Brandstetter, J. Towards multi- spatiotemporal-scale generalized PDE modeling.arXiv preprint arXiv:2209.15616,

    work page arXiv

  5. [5]

    The ERA5 global reanalysis

    Hersbach, H., Bell, B., Berrisford, P., Hirahara, S., Hor´anyi, A., Mu ˜noz-Sabater, J., Nicolas, J., Peubey, C., Radu, R., Schepers, D., et al. The ERA5 global reanalysis. Quarterly journal of the royal meteorological society, 146(730):1999–2049,

    work page 1999

  6. [6]

    Koshizuka, T., Fujisawa, M., Tanaka, Y ., and Sato, I

    doi: 10.1038/s41467-024-49411-w. Koshizuka, T., Fujisawa, M., Tanaka, Y ., and Sato, I. Un- derstanding the expressivity and trainability of Fourier 9 Submission and Formatting Instructions for ICML 2026 neural operator: A mean-field perspective. InAdvances in Neural Information Processing Systems, volume 37, pp. 11021–11060,

    work page doi:10.1038/s41467-024-49411-w 2026

  7. [8]

    FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators

    URL https://arxiv.org/abs/2202.11214. Qin, S., Lyu, F., Peng, W., Geng, D., Wang, J., Gao, N., Liu, X., and Wang, L. L. Toward a better understanding of Fourier neural operators: Analysis and improvement from a spectral perspective.CoRR,

    work page internal anchor Pith review Pith/arXiv arXiv

  8. [9]

    A., Ross, Z

    Rahman, M. A., Ross, Z. E., and Azizzadenesheli, K. U-NO: U-shaped neural operators.arXiv preprint arXiv:2204.11127,

    work page arXiv

  9. [10]

    H., Sankaran, S., Wang, H., Pap- pas, G

    10 Submission and Formatting Instructions for ICML 2026 Wang, S., Seidman, J. H., Sankaran, S., Wang, H., Pap- pas, G. J., and Perdikaris, P. CViT: Continuous vision transformer for operator learning. InThe Thirteenth Inter- national Conference on Learning Representations,

    work page 2026

  10. [11]

    mother wavelet

    11 Submission and Formatting Instructions for ICML 2026 A. Wavelet transforms In the following, we provide a brief overview on wavelet transforms. We will focus on the real-valued case, which is more applicable to our setting, though wavelets are more generally defined in terms of complex-valued functions. More details about wavelet transforms can be foun...

    work page 2026

  11. [12]

    14 Submission and Formatting Instructions for ICML 2026 Table 4.Hyperparameter settings and parameter counts for baselines and MSWT. Model Setting # params (M) FNON layers=5, n hidden=64,truncation mode=1616.8 Unet Nlayers=4, n hidden=[16,32,64,256] 12.5 WNO Nlayers=4, n hidden=96,multiscale levels=3 14.5 SAOT Nlayers=5, n hidden=384 14.0 HFS Nlayers=5, n...

    work page arXiv 2026

  12. [13]

    ERA5 For climate-style evaluation, we follow the LUCIE setup based on coarse-resolution ERA5 fields regridded to a T30 Gaussian grid

    17 Submission and Formatting Instructions for ICML 2026 Figure 14.Shallow water equation, kinetic energy spectrum and enstrophy spectrum, rollout step =41 Figure 15.Shallow water equation, prediction and error comparison, rollout step =81 Figure 16.Shallow water equation, kinetic energy spectrum and enstrophy spectrum, rollout step =81 18 Submission and F...

    work page 2026