Small Models, Strong Priors: Architectural Inductive Bias for Parameter-Efficient Neural PDE Solvers
arXiv:2605.25949v1 [cs.LG] 25 May 2026
Shyam Sankaran Department of Mechanical Engineering and Applied Mechanics, University of Pennsylvania [email protected] Hanwen Wang Graduate Program in Applied Mathematics and Computational Science, University of Pennsylvania [email protected] Paris Perdikaris Department of Mechanical Engineering and Applied Mechanics, University of Pennsylvania [email protected]
Abstract Neural PDE solvers have followed the scaling trajectory of vision and language, with recent foundation models reaching billions of parameters. We argue that scale is a poor substitute for architectural inductive bias in this domain: structured priors deliver outsized parameter efficiency, and the pattern of where they succeed and fail is itself informative about what they capture. We instantiate this argument in WaveLiT, an architecture combining a discrete wavelet transform for lossless multi-resolution tokenization, an augmented linear attention block, a shared-weight multiscale feature pyramid, and a wavelet-domain auxiliary loss. Bespoke 1–10M-parameter WaveLiT models compete with foundation models of 100–1000× their size across eight TheWell benchmarks, with the largest gains on wave- and acoustic-dominated benchmarks where the wavelet-multiscale prior fits the dominant dynamical structure and small per-step errors do not compound geometrically under rollout. Trained jointly across all eight benchmarks, a 10Mparameter foundation variant exhibits a structured, physically interpretable transfer pattern – strongest where the wavelet-multiscale prior matches the dynamics, weakest on chaotic advection-dominated flows. The entire pipeline trains on a single GPU. The results suggest that small-model PDE performance is shaped by architectural inductive bias rather than scale, and that the structure of a prior’s failures is a useful empirical signal about its content. WaveLiT Parameter Efficiency (NS)
Introduction
Neural surrogates for partial differential equations are increasingly central to scientific computing, with applications spanning weather forecasting [5], fluid dynamics, and materials science [45]. The promise is straightforward: where classical solvers cost hours per simulation, learned surrogates cost milliseconds. Realizing this promise at the accuracy required for scientific use has, in current practice, driven the field toward ever-larger models.
WaveLiT FNO DPOT-L (Fine-tuned)
10.0
L2 Error (%) ↓
1
W-Ti
CViT-L Other Baselines
FNO
W-S CV-L
DP-L (FT)
W-B
1.0
0.1M
W-L
1M
10M
100M
1B
Model Size
Figure 1: Parameter efficiency on PDEArena Preprint.
Navier-Stokes [13]. Each point is a model; lowerleft is better. WaveLiT (red stars) outperforms models with 100× more parameters.
Following the scaling regimes established in language[23] and vision [50], neural PDE solvers have moved from thousands of parameters to billions. Recent foundation models such as PhysiX (4.5B) [35], DPOT (500M) [15], Poseidon (629M) [17], and Walrus (1.2B) [34] pursue generalpurpose simulators by combining massive parameter counts with multi-task training across diverse PDE families. Training these models requires hundreds of GPU-hours and infrastructure that is inaccessible to most research groups. At the same time, the broader ML community is increasingly recognizing that scale is not the only axis of progress. Compact language models have emerged as efficient backbones for downstream systems [1, 4], demonstrating that careful architectural and data choices can substitute for raw parameter count. A complementary direction in reasoning has shown that small networks with appropriate inductive structure – HRM at 27M parameters [43] and TRM at 7M [22] – can be competitive on tasks with exploitable combinatorial structure. Neural PDE solvers occupy an analogous setting: the underlying physics imposes locality, regularity, and multi-scale structure that should, in principle, be exploitable by smaller models, yet architectural defaults in this field have not been re-examined under this lens. Most current PDE models inherit their design from vision transformers, adopting defaults that were never optimized for physical systems. Coarse 8×8 or 16×16 patch tokenization discards sub-patch physics before the first attention layer, and quadratic self-attention forces this coarse tokenization by making fine-grained alternatives prohibitively expensive. These defaults interact in a vicious cycle: quadratic attention demands fewer tokens, which demands coarser patches, which destroys fine-scale physics, which demands more parameters to compensate. We introduce WaveLiT, an integrated architecture that addresses each of these defaults. A discrete wavelet transform replaces learned patch embedding with parameter-free, lossless multi-resolution tokenization. A linear attention block with O(N ) complexity replaces quadratic self-attention, building on Mamba-Inspired Linear Attention [14] and the test-time-regression view of attention [41, 3] adapted to the bidirectional spatial setting. A shared-weight multiscale feature pyramid processes tokens at multiple resolutions without adding parameters, and a wavelet-domain auxiliary loss encourages spectral fidelity across scales. These choices compound into a parameter-efficient frontier (Figure 1). On PDEArena Navier-Stokes [13], WaveLiT matches or exceeds FNO [27], CViT-L [47], and DPOT-L [15] at 10–100× fewer parameters; on TheWell [36], bespoke 1–10M-parameter models compete with foundation-model baselines (MPP-AViT-L, Poseidon-L, DPOT-H, Walrus) of 100–1000× their size, with the largest gains on wave- and acoustic-dominated benchmarks. To probe whether these architectural priors survive cross-family training, we use the foundation-model regime as a diagnostic instrument: a single 10M-parameter model trained jointly on all eight TheWell benchmarks exhibits a structured, physically interpretable transfer pattern – strongest where the wavelet-multiscale prior matches the dynamics, weakest on chaotic advection-dominated flows. The entire pipeline trains on a single GPU. We do not claim global SOTA against billion-parameter specialists; we claim that the recipe defines an efficient frontier, and that the structure of its successes and failures provides evidence that small-model PDE performance is shaped by inductive bias, not scale alone.
2
Related Work
Neural operators for PDEs. Neural operators approximate mappings between infinite-dimensional function spaces [26], with two main paradigms: encoder-decoder operators such as DeepONet [31] and NoMaD [37], which map functions to latent spaces and back; and integral kernel operators [25] parameterized via graph message passing [28] or Fourier transforms [27]. These methods established that operator-valued targets can be learned end-to-end, but typically at single resolution and without explicit multi-scale structure — the property our design is built to exploit. Vision transformers for operator learning. ViTs [12] have been adapted to operator learning, most prominently by CViT [47], which combines ViT encoders with coordinate embeddings and cross-attention. These approaches inherit two defaults from vision: quadratic-cost softmax attention and coarse patch tokenization. The two interact: quadratic attention forces large patches to keep sequence lengths tractable, and large patches discard the fine-scale physics that PDE solutions actually 2
(a)
(b) Input
Wavelet
Token
Features
Grid
DWT
LAttn Mixer
RoPE @ (H,W)
Input
Tokens
Linear
LAttn Mixer
RoPE @ (H/2,W/2)
Σ
Multiscale
Output
Mixing
Block IDWT
LAttn Mixer
Linear
RoPE @ (H/4,W/4)
Figure 2: The WaveLiT architecture: (a) Input fields are tokenized via a single-level 2D discrete wavelet transform (DWT), projected to embedding dimension via a linear layer, and processed by NL multiscale mixing blocks. The output is recovered by a final linear projection and inverse DWT, restoring the original spatial resolution. (b) Each multiscale mixing block applies a sharedweight linear attention mixer at L + 1 resolutions obtained by successive average pooling. Per-level RoPE frequencies are calibrated to each grid size; outputs are upsampled to the base resolution and combined via learned scalar weights wℓ depend on. WaveLiT removes both defaults simultaneously — linear attention enables fine-grained tokens, and a wavelet transform supplies them losslessly. Efficient sequence modeling. Linear attention [24] reduces the O(N 2 ) cost of softmax attention to O(N ). Recent work has reframed such mechanisms through a test-time regression lens: DeltaNet [49] frames the recurrent state update as gradient descent on a per-token MSE loss, and the test-timeregression perspective [44, 3] shows that many modern sequence models can be derived from specific regression objectives. We adopt the Mamba-Inspired Linear Attention block [14] as a starting point and develop a ridge-corrected variant motivated by this perspective (§3), specifically adapted to the bidirectional, spatially-tokenized setting of grid-structured PDE data – a regime that prior linear-attention work has largely targeted only in causal sequence settings. Wavelets in deep learning. Wavelet transforms decompose signals into components localized jointly in space and frequency, providing a principled framework for multi-resolution analysis [32, 18]. In deep learning, wavelets have been used for image compression [2] and as a tokenization scheme for vision transformers [51]. Within operator learning, Wavelet Neural Operators perform kernel integration in the wavelet domain [40], and Wavelet Diffusion Neural Operators apply diffusion-based modeling in the wavelet domain for PDE simulation and control [21]. Our use of wavelets is distinct: we use the DWT not as the substrate for a specialized operator, but as a parameter-free tokenization that preserves multi-scale structure end-to-end through a transformer backbone, with the wavelet domain reused as an auxiliary loss to enforce spectral fidelity. Foundation models for physics simulation. The field has shifted toward general-purpose “foundation models” that aim to solve a diverse range of PDE problems [35, 17, 15], with recent entries reaching billions of parameters: Walrus at 1.2B [34], PhysiX at 4.5B [35]. These models achieve strong benchmark results but require infrastructure inaccessible to most research groups, and the implicit hypothesis – that scale is what makes cross-family transfer possible – has not been tested against architecturally efficient alternatives. Our foundation-model variant is designed precisely as such a test: at 10M parameters, memorization is mechanically infeasible, so any cross-family transfer that does occur must come from shared architectural priors.
3
Architecture
WaveLiT is an integrated architecture for neural PDE solvers built around four components: discrete wavelet tokenization, a gated linear attention block, a shared-weight multiscale feature pyramid, and a wavelet-domain auxiliary loss. Each component targets a specific architectural default inherited from vision transformers that we identified in Section 1. We describe each in turn, then describe the foundation-model extension WaveLiT-FM. Figure 2 summarizes the architecture. 3
3.1
Wavelet Tokenization
We replace the standard patch embedding with a single-level 2D discrete wavelet transform (DWT). The DWT decomposes each spatial field into four half-resolution subbands (one low-frequency approximation LL and three directional detail bands LH, HL, HH), yielding the same 4× channel expansion and 2× spatial reduction as a 2 × 2 patch projection but with three key advantages: the transform is (i) lossless, admitting exact reconstruction via the inverse DWT; (ii) physically meaningful, separating large-scale structure from edges and directional textures; and (iii) parameterfree, supplying an inductive bias with no learnable weights. For inputs of shape B × T × H × W × C, an ℓ-level DWT produces tokens of shape B × T × (H/2ℓ ) × (W/2ℓ ) × (C · 4ℓ ); a linear projection then maps the concatenated subbands to the model dimension. We use ℓ = 1 throughout to preserve fine-grained features at the input layer; this choice is consistent with recent findings that smaller patch sizes outperform parameter scaling in sub-quadratic vision architectures at fixed compute [42]. We use the bior2.2 wavelet throughout (ablations in Appendix C) and apply the inverse DWT at the output to recover the original resolution. In practice, wavelet tokenization yields a modest accuracy gain over a convolutional encoder-decoder (+1.4%) while reducing training time by 11%, consistent with the finding of [20] in a diffusion setting; the inductive bias primarily provides efficiency rather than a large accuracy jump. Full tokenizer ablations are in Appendix I. 3.2
Spatial Linear Attention Mixer
On high-resolution grids, the O(N 2 ) cost of softmax self-attention quickly becomes prohibitive. We therefore replace quadratic attention with a kernelized linear attention operator whose cost scales linearly in the number of tokens. Let ϕ(·) denote a positive feature map applied to queries and keys. The output at token i can then be written as oi = S ϕ(qi ),
S=
N X
vj ϕ(kj )⊤ ,
(1)
j=1
where the state S ∈ Rdv ×dk is accumulated once and reused across all query locations. This reordering removes the quadratic dependence on the number of tokens while preserving global token mixing. Ridge-corrected state update. PWe additionally apply aP ridge-corrected variant of the state update, Sλ = C(G + λI)−1 with C = j vj ϕ(kj )⊤ and G = j ϕ(kj )ϕ(kj )⊤ , which can be motivated through the test-time regression view of linear attention [41, 3] as incorporating the covariance structure of the kernel features. Full derivation in Appendix A; ablations (Appendix I.2) indicate this provides incremental rather than dominant gains relative to the spatial inductive biases described next. Additional Spatial Inductive Biases. While linear attention addresses the computational bottleneck, it does not by itself provide a strong spatial inductive bias for dense grid-structured data. To remedy this, we adopt an MILA-style mixer design that augments linear attention with positional operators tailored to vision-like inputs. Concretely, we apply rotary positional embeddings [38] (RoPE) to the query/key feature maps, use conditional positional encoding (CPE) [7] as a depthwiseconvolutional positional term on the residual stream before the attention and MLP sub-blocks, and add locally enhanced positional encoding (LePE) [11] as a depthwise-convolutional local bias on the attention output path. This combination injects both global positional structure and local spatial bias while preserving linear-time token mixing. This design is motivated by recent MILA results showing that, in vision settings, suitable positional encodings can effectively replace forget-gate-like functionality while retaining parallelizable computation [14]. Our final mixer combines these ingredients into a single block: RoPE-equipped linear attention for efficient global token interaction, a ridge-corrected state update motivated by the regression view above, and an MILA-style block design with CPE and LePE to supply the spatial inductive bias required on high-resolution grids. In practice, we use this full configuration throughout the model. We study the effect of the individual design choices and their relative contributions in Appendix I.2. 4
3.3
Multiscale Computation
PDE solutions exhibit structure across spatial scales. Rather than only relying on network depth to capture cross-scale interactions, we process information at multiple resolutions using a featurepyramid structure [29] with shared parameters. Given input tokens X ∈ RB×N ×D on a grid of size (H, W ), we construct a resolution pyramid {X (ℓ) }L ℓ=0 by successive 2× average pooling. A shared-weight linear attention block fθ is applied at each level: Y (ℓ) = fθ (X (ℓ) , g (ℓ) ), (2) where g (ℓ) encodes the grid size at level ℓ for positional embeddings. Outputs are upsampled and aggregated via learned scalar weights: X X̂ = wℓ · Upsample(Y (ℓ) ). (3) ℓ
Weight sharing provides multiscale sensitivity without increasing parameters; coarse levels process far fewer tokens, making additional levels computationally cheap. Concurrent work applies a similar resolution-pyramid strategy with shared attention weights for PDE simulation, achieving linear complexity via a multipole analogy [8]; our design differs in using average pooling for downsampling and learned scalar weights for aggregation. A single pyramid level reduces error by 17.4% over wavelet-only tokenization at 16% additional training cost; additional levels yield diminishing returns. We use L = 1 throughout; full ablations across tokenizer and FPN depth are in Appendix I. 3.4
Wavelet-Domain Auxiliary Loss
The wavelet basis used for tokenization also defines a natural multi-scale loss. In addition to a pixelspace MSE loss on predicted fields, we apply an L1 loss in the wavelet domain: predictions and targets are both DWT-transformed, and the absolute error is summed across all subbands. This encourages the model to match high-frequency detail coefficients as well as low-frequency approximations, which the MSE loss alone tends to underweight. Appendix D reports an ablation showing that the auxiliary loss reduces error power across the entire spatial-frequency spectrum. 3.5
Foundation-Model Extension: WaveLiT-FM
WaveLiT-FM extends the bespoke architecture to joint training across all eight TheWell benchmarks. The key additions over the single-task design, illustrated in Figure 3, are a per-channel lifting embedding that gates absent fields to zero, a per-dataset task conditioning vector, and per-dataset output heads. Unified input representation. We define a canonical channel space of dimension Ctotal by taking the union of all unique physical variables (e.g., u-velocity, pressure, buoyancy, tracer) across TheWell. Each dataset maps its available fields to their designated canonical channels and zero-pads the rest. Channel identity is encoded through a per-channel lifting matrix Wlift ∈ RCtotal ×T nbands ×D : after the wavelet decomposition, each channel’s coefficients are projected independently and summed, P z = c xc Wlift [c]. Because there is no bias term, absent (zero-padded) channels contribute exactly zero to the embedding — the shared trunk receives a representation that is naturally gated by which fields are present. Channels shared across datasets (e.g., u-velocity in both shear_flow and rayleigh_benard) reuse the same rows of Wlift , encouraging a common latent representation for each physical field type. Task-specific conditioning. After the initial linear projection, we add a learnable per-dataset embedding vector to every token. This informs the shared trunk which PDE system it is solving without breaking weight sharing. Shared backbone. The wavelet embedding, MILA stack, and reconstruction layer serve as a shared trunk trained jointly on all eight datasets; per-task specialization is confined entirely to the task-conditioning vector and output head, with no dataset-specific parameters in the core computation. 5
Canonical Channel Mapping
Central Processing & Conditioning
Per-Dataset Decoder Heads
Figure 3: Foundation model design: (i) input standardization via a per-channel lifting matrix without bias, so absent (zero-padded) channels contribute exactly zero to the trunk input; (ii) problemembedding conditioning of a shared multiscale transformer trunk; and (iii) per-dataset decoder heads decoding the shared latent representation into native output channels. Specialized output heads. Each dataset has a separate lightweight linear head that decodes the final shared latent representation into its specific output channels. The head contributes negligible parameters relative to the trunk.
4
Results
4.1
Bespoke Models: Parameter Efficiency
Implementation details. We train two model sizes (1.2M and 9.5M parameters) on each TheWell benchmark separately using AdamW [30] with exponential learning rate decay and gradient clipping. Both variants use a single-level DWT tokenization, the combined MSE and wavelet loss (Section 3.4), and exponential parameter averaging at evaluation. We additionally train rollout-finetuned (FT) variants by continuing from the pretrained checkpoint with an autoregressive rollout objective to improve long-horizon stability. Our rollout-finetuning protocol and the comparison that motivates it are detailed in Appendix G. All training is performed in JAX [6]/Flax [16] on a single GPU; full hyperparameters are in Appendix E. We report each model’s median VRMSE normalized to the per-dataset best baseline. WaveLiT bespoke models match or exceed foundation-model baselines with 100–1000× more parameters (Figure 4; absolute VRMSE in Appendix J). We compare against MPP-AViT-L (409M) [33], PoseidonL (629M) [17], DPOT-H (∼1B) [15], and Walrus (1.2B) [34]. One-step prediction (panel a). Averaged across the eight TheWell benchmarks, WaveLiT-9.5M ranks as the second-best model in the one-step regime with a mean per-dataset rank of 1.94, trailing only Walrus (1.44) and outperforming all other foundation-model (FM) baselines. Remarkably, even WaveLiT-1.2M—a model ∼1000× smaller than Walrus—outranks the remaining three FM baselines, achieving a mean rank of 3.12 compared to Poseidon-L (4.12), MPP-AViT-L (4.86), and DPOT-H (5.38) (Table 18). On the two wave- and acoustic-dominated benchmarks, ASM and HS, our architecture achieves state-of-the-art results despite its size. Specifically, WaveLiT-9.5M outperforms every baseline including Walrus on ASM (0.0016 vs. 0.0099), and WaveLiT-1.2M achieves the same feat on HS (0.0003 vs. 0.0005). Across the remaining six datasets, WaveLiT consistently surpasses 6
48× 61×
(a) One-step
Relative VRMSE
20× 10× 5×
2×
(b) Short-horizon rollout T ∈ [1, 20]
41×
35×
25×
1× (best)
Relative VRMSE
20× 10× 5×
2×
25×
1× (best)
(c) Long-horizon rollout T ∈ [21, 60]
Relative VRMSE
20× 10× 5×
2×
1× (best) Acoustic Scattering
MPP-AViT-L
Helmholtz Staircase
Rayleigh–Bénard Shear Flow
Poseidon-L
DPOT-H
Walrus
Turbulent Radiative
Active Matter
WaveLiT-1.2M
WaveLiT-9.5M
Gray–Scott
1st
2nd
Viscoelastics
3rd
Figure 4: WaveLiT bespoke models vs. foundation model baselines across eight TheWell benchmarks and three evaluation windows, normalized per dataset to the best model (1× = best, log scale). Each WaveLiT bar shows the better of pretrained and rollout-finetuned variants. Gold, silver, and bronze markers denote the top three results per benchmark. Absolute VRMSE values and full PT/FT breakdown in Appendix J.
MPP-AViT-L, Poseidon-L, and DPOT-H. The performance gap to Walrus remains narrow on RB, SF, and VI, but widens on TRL2D, AM, and GS, where Walrus’s billion-parameter capacity and design likely provide a decisive advantage. Longer rollouts (panels b, c). However, this initial one-step advantage does not carry uniformly to longer rollouts. WaveLiT-9.5M’s mean rank degrades from 1.94 at one-step to 2.38 at T ∈ [1, 20] and 3.00 at T ∈ [21, 60], with Walrus’ margin widening accordingly (Table 18). When the learned map has a local Lipschitz constant LF > 1, the per-step error ϵ accumulates as ϵ(LnF − 1)/(LF − 1). Rollout finetuning does not reduce ϵ in isolation; rather, it trades ϵ against LF . This trade-off is consistent with the worse one-step accuracy of our finetuned (FT) variants relative to their pre-trained (PT) counterparts (Appendix H). On chaotic dynamics, LF is bounded below by eλ∆t , where λ is the leading Lyapunov exponent of the underlying flow and ∆t is the time step; geometric error growth therefore survives finetuning. Consequently, our long-horizon wins occur on non-chaotic systems 7
32× 39×
48× 61×
(a) One-step
Relative VRMSE
20× 10× 5×
F
2×
25×
1× (best)
(b) Long-horizon rollout T ∈ [21, 60]
Relative VRMSE
20× 10× 5×
2× 1× (best) Acoustic Scattering
Helmholtz Staircase
Rayleigh–Bénard Shear Flow
MPP-AViT-L Poseidon-L
DPOT-H Walrus
Turbulent Radiative
Best WaveLiT WaveLiT-FM (PT)
Active Matter
Gray–Scott
Viscoelastics
WaveLiT-FM (FT) FM beats Walrus
Figure 5: Foundation model diagnostic. Relative VRMSE (normalized per dataset to the best model; 1× = best, log scale) for all four baselines, the best bespoke WaveLiT, and both WaveLiT-FM variants. ⋆ marks the one cell where a multi-task model (Fnd-PT on ASM, one-step) outperforms Walrus. where LF ≤ 1, such as ASM (9.5M-FT: 0.0268 vs. Walrus 0.0560) and HS (1.2M-PT: 0.0020 vs. DPOT-H 0.0031). Conversely, evolution on RB is dominated by fine-scale structures that inherently drive LF > 1, causing both our models to diverge under rollout.1 4.2
Foundation Model Diagnostic
WaveLiT-FM is a 10M-parameter model trained jointly across all √eight TheWell benchmarks on a single GPU, using the shared-backbone design of Section 3.5 with N -proportional dataset sampling. It is not designed to maximize benchmark performance; it is a diagnostic instrument for whether the architectural priors survive cross-family training. Figure 5 compares Walrus, the best bespoke WaveLiT variant, and both FM variants (pretrained and rollout-finetuned) at one-step and long-horizon rollout; full tables are in Appendix J. Summary of Results. At one-step on ASM, both multi-task variants fall below Walrus (0.0099): Fnd-PT (0.0087, ⋆) and Fnd-FT (0.0097). Together with Fnd-FT on TRL2D at long horizon (0.8471 vs. Walrus 0.8648), these are the only three cells in this comparison where a multi-task model beats Walrus. Overall, the FM’s performance gap relative to the bespoke WaveLiT ceiling is modest on most datasets (RB and SF ∼3×, ASM and GS ∼5×), with one pronounced exception: HS, where it widens to 19×. Pattern of transfer. The FM’s degradation relative to bespoke is structured. HS shows by far the largest relative degradation under shared-trunk training (19× at one-step). It is the dataset where the bespoke specialist achieved its lowest absolute error (0.0003) by aggressively exploiting datasetspecific structure, and joint training forces a shared, generalized basis that blurs the hyper-specific 1 The 9.5M-FT RB checkpoint uses early stopping at 40k steps; the final 50k checkpoint diverges further. The 1.2M-FT
checkpoint exhibits the same instability without recovery from earlier checkpoints.
8
filters it depends on. Conversely, datasets where bespoke and Walrus performance were already comparable (ASM, RB, TRL2D) show much smaller relative gaps. Ultimately, the FM inherits the bespoke model’s strengths and weaknesses in proportion to how well the inductive bias aligns with the local dynamics. This suggests that future scaling of WaveLiT-FM should prioritize specialized perdataset sampling for complex flows rather √ than merely increasing backbone capacity. Furthermore, our current joint training relies on static N -proportional sampling, meaning the model’s exposure to each dataset is fixed by volume rather than task difficulty. A promising direction for future work is the development of an adaptive sampling strategy that dynamically adjusts batch composition based on the evolution of validation loss, potentially recovering the specialization lost in standard multi-task training. An observation on transfer. A valuable outcome of running identical architectures in both regimes is the rare opportunity to directly compare their performance dynamics. Within our experimental setup, joint training across diverse PDE families did not yield positive transfer; rather, it consistently diluted specialization. Across all eight TheWell datasets, the bespoke WaveLiT specialists outperformed the multi-task FM variant, sometimes by over an order of magnitude. At this compact scale of 10M parameters, we did not anticipate strong positive transfer, as the network likely lacks the sheer capacity required to represent multiple disjoint physical manifolds simultaneously without severe interference. This raises a fascinating open question: does unlocking synergistic benefits between fundamentally different PDE families strictly require billion-parameter backbones? Currently, these mechanisms are difficult to isolate, as many pioneering PDE foundation models [34, 17, 15, 35] are evaluated primarily as unified systems without bespoke single-task ablations of their own backbones. Our results do not settle this question, but they do raise it: the implicit assumption that multi-PDE pretraining transfers benefits to individual tasks deserves explicit experimental support.
5
Discussion
Summary of findings. Bespoke WaveLiT models (1–10M parameters) can match or exceed foundation models up to 1000× larger, particularly on smooth problems where our inductive bias perfectly aligns with the underlying physics. On chaotic datasets, this prior improves one-step accuracy but cannot bypass the instability of teacher-forced rollouts. Crucially, our 10M-parameter WaveLiT-FM diagnostic model exhibits the exact same performance patterns as the single-task specialists. At this scale, brute-force memorization is impossible, confirming that multi-task transfer behavior is driven by the architectural prior rather than sheer scale. Limitations & Future Directions. Our primary limitation is autoregressive instability on chaotic dynamics, driven by fundamental mathematical bounds rather than training budgets. Under rollout, per-step error ϵ and the local Lipschitz constant LF dictate error growth: En ≤ ϵ(LnF − 1)/(LF − 1) (Appendix H). Rollout finetuning trades ϵ against LF , but it cannot suppress LF below eλ∆t , where λ is the leading Lyapunov exponent of the system; thus, geometric growth is unavoidable. Furthermore, resolving chaotic, high-frequency subbands inherently increases LF , making our finegrained tokenization a double-edged sword. A coarser patch embedding that never represents those modes yields a smaller effective LF and slower error growth, sacrificing one-step accuracy for potentially better long-term stability. Addressing this requires migrating from static, single-step maps to architectures equipped with strong temporal mechanisms. Immediate mitigations include temporal noise injection, or “jitter,” as leveraged by Walrus [34]. More fundamentally, incorporating temporal linear attention would allow the model to propagate a continuous hidden state, naturally regularizing the rollout trajectory. Extending this architecture to complex geometries via lifting wavelets [39] is a natural next step. Conclusions. WaveLiT demonstrates that combining lossless multi-resolution tokenization with linear attention and shared-weight multiscale processing yields a parameter-efficient design competitive with foundation models 100–1000× larger on the PDE families where its prior matches the underlying physics. Beyond the architecture, the results provide a data point on the inductive-bias-versus-scale question: at scales where memorization is mechanically infeasible, cross-family transfer is shaped by the structure of the prior, and the pattern of failures is physically legible. 9
References [1] Marah Abdin, Jyoti Aneja, Harkirat Behl, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, Michael Harrison, Russell J Hewett, Mojan Javaheripi, Piero Kauffmann, et al. Phi-4 technical report. arXiv preprint arXiv:2412.08905, 2024. [2] Tinku Acharya and Ping-Sing Tsai. JPEG2000 standard for image compression: concepts, algorithms and VLSI architectures. John Wiley & Sons, 2004. [3] Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, and Vahab Mirrokni. It’s all connected: A journey through test-time memorization, attentional bias, retention, and online optimization. arXiv preprint arXiv:2504.13173, 2025. [4] Peter Belcak, Greg Heinrich, Shizhe Diao, Yonggan Fu, Xin Dong, Saurav Muralidharan, Yingyan Celine Lin, and Pavlo Molchanov. Small language models are the future of agentic ai. arXiv preprint arXiv:2506.02153, 2025. [5] Cristian Bodnar, Wessel P Bruinsma, Ana Lucic, Megan Stanley, Johannes Brandstetter, Patrick Garvan, Maik Riechert, Jonathan Weyn, Haiyu Dong, Anna Vaughan, et al. Aurora: A foundation model of the atmosphere. arXiv preprint arXiv:2405.13063, 2024. [6] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. [7] Xiangxiang Chu, Zhi Tian, Bo Zhang, Xinlong Wang, and Chunhua Shen. Conditional positional encodings for vision transformers. arXiv preprint arXiv:2102.10882, 2021. [8] Alex Colagrande, Paul Caillon, Eva Feillet, and Alexandre Allauzen. Linear attention with global context: A multipole attention mechanism for vision and physics. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 3099–3108, 2025. [9] Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov. Unsupervised cross-lingual representation learning at scale. In Proceedings of the 58th annual meeting of the association for computational linguistics, pages 8440–8451, 2020. [10] Katherine Crowson. jax-wavelets: The 2D discrete wavelet transform for JAX, 2022. [11] Xiaoyi Dong, Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, and Baining Guo. Cswin transformer: A general vision transformer backbone with cross-shaped windows. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12124–12134, 2022. [12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020. [13] Jayesh K Gupta and Johannes Brandstetter. Towards multi-spatiotemporal-scale generalized pde modeling. arXiv preprint arXiv:2209.15616, 2022. [14] Dongchen Han, Ziyi Wang, Zhuofan Xia, Yizeng Han, Yifan Pu, Chunjiang Ge, Jun Song, Shiji Song, Bo Zheng, and Gao Huang. Demystify mamba in vision: A linear attention perspective. arXiv preprint arXiv:2405.16605, 2024. [15] Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying, Hang Su, Anima Anandkumar, Jian Song, and Jun Zhu. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training. arXiv preprint arXiv:2403.03542, 2024. [16] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2024. 10
[17] Maximilian Herde, Bogdan Raonic, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Emmanuel de Bézenac, and Siddhartha Mishra. Poseidon: Efficient foundation models for pdes. Advances in Neural Information Processing Systems, 37:72525–72624, 2024. [18] Eugenio Hernández and Guido Weiss. A first course on wavelets. CRC press, 1996. [19] Geoffrey E Hinton and James A Anderson. Parallel models of associative memory: updated edition. Psychology press, 2014. [20] Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, pages 13213– 13232. PMLR, 2023. [21] Peiyan Hu, Rui Wang, Xiang Zheng, Tao Zhang, Haodong Feng, Ruiqi Feng, Long Wei, Yue Wang, Zhi-Ming Ma, and Tailin Wu. Wavelet diffusion neural operator. arXiv preprint arXiv:2412.04833, 2024. [22] Alexia Jolicoeur-Martineau. Less is more: Recursive reasoning with tiny networks. arXiv preprint arXiv:2510.04871, 2025. [23] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. [24] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR, 2020. [25] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications to pdes. Journal of Machine Learning Research, 24(89):1–97, 2023. [26] Nikola B Kovachki, Samuel Lanthaler, and Andrew M Stuart. Operator learning: Algorithms and analysis. arXiv preprint arXiv:2402.15715, 2024. [27] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895, 2020. [28] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.03485, 2020. [29] Tsung-Yi Lin, Piotr Dollár, Ross Girshick, Kaiming He, Bharath Hariharan, and Serge Belongie. Feature pyramid networks for object detection. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2117–2125, 2017. [30] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. [31] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):218–229, 2021. [32] Stephane G Mallat. A theory for multiresolution signal decomposition: the wavelet representation. IEEE transactions on pattern analysis and machine intelligence, 11(7):674–693, 1989. [33] Michael McCabe, Bruno Régaldo-Saint Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, et al. Multiple physics pretraining for physical surrogate models. arXiv preprint arXiv:2310.02994, 2023. 11
[34] Michael McCabe, Payel Mukhopadhyay, Tanya Marwah, Bruno Regaldo-Saint Blancard, Francois Rozet, Cristiana Diaconu, Lucas Meyer, Kaze WK Wong, Hadi Sotoudeh, Alberto Bietti, et al. Walrus: A cross-domain foundation model for continuum dynamics. arXiv preprint arXiv:2511.15684, 2025. [35] Tung Nguyen, Arsh Koneru, Shufan Li, and Aditya Grover. Physix: A foundation model for physics simulations. arXiv preprint arXiv:2506.17774, 2025. [36] Ruben Ohana, Michael McCabe, Lucas Meyer, Rudy Morel, Fruzsina Agocs, Miguel Beneitez, Marsha Berger, Blakesly Burkhart, Stuart Dalziel, Drummond Fielding, et al. The well: a large-scale collection of diverse physics simulations for machine learning. Advances in Neural Information Processing Systems, 37:44989–45037, 2024. [37] Jacob Seidman, Georgios Kissas, Paris Perdikaris, and George J Pappas. Nomad: Nonlinear manifold decoders for operator learning. Advances in Neural Information Processing Systems, 35:5601–5613, 2022. [38] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024. [39] Wim Sweldens. The lifting scheme: A construction of second generation wavelets. SIAM journal on mathematical analysis, 29(2):511–546, 1998. [40] Tapas Tripura and Souvik Chakraborty. Wavelet neural operator: a neural operator for parametric partial differential equations. arXiv preprint arXiv:2205.02191, 2022. [41] Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Maximilian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, et al. Mesanet: Sequence modeling by locally optimal test-time training. arXiv preprint arXiv:2506.05233, 2025. [42] Feng Wang, Yaodong Yu, Guoyizhe Wei, Wei Shao, Yuyin Zhou, Alan Yuille, and Cihang Xie. Scaling laws in patchification: An image is worth 50,176 tokens and more. arXiv preprint arXiv:2502.03738, 2025. [43] Guan Wang, Jin Li, Yuhao Sun, Xing Chen, Changling Liu, Yue Wu, Meng Lu, Sen Song, and Yasin Abbasi Yadkori. Hierarchical reasoning model. arXiv preprint arXiv:2506.21734, 2025. [44] Ke Alexander Wang, Jiaxin Shi, and Emily B Fox. Test-time regression: a unifying framework for designing sequence models with associative memory. arXiv preprint arXiv:2501.12352, 2025. [45] Sifan Wang, Tong-Rui Liu, Shyam Sankaran, and Paris Perdikaris. Micrometer: Micromechanics transformer for predicting mechanical responses of heterogeneous materials. arXiv preprint arXiv:2410.05281, 2024. [46] Sifan Wang, Shyam Sankaran, and Paris Perdikaris. Respecting causality is all you need for training physics-informed neural networks. arXiv preprint arXiv:2203.07404, 2022. [47] Sifan Wang, Jacob H Seidman, Shyam Sankaran, Hanwen Wang, George J Pappas, and Paris Perdikaris. Cvit: Continuous vision transformer for operator learning. arXiv preprint arXiv:2405.13998, 2024. [48] Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami Al-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mt5: A massively multilingual pre-trained text-to-text transformer. In Proceedings of the 2021 conference of the North American chapter of the association for computational linguistics: Human language technologies, pages 483–498, 2021. [49] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464, 2024. [50] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12104–12113, 2022. 12
[51] Zhenhai Zhu and Radu Soricut. Wavelet-based image tokenizer for vision transformers. arXiv preprint arXiv:2405.18616, 2024.
13
A
Linear Attention and the Ridge Regression View
From softmax√to linear attention. Standard attention computes Attn(Q, K, V ) = softmax(QK ⊤ / d) V , whose QK ⊤ matrix is N × N and makes both memory and compute scale quadratically with sequence length. Any positive-definite kernel κ(q, k) admits a feature-map factorization κ(q, k) = ⟨ϕ(q), ϕ(k)⟩; replacing the softmax with such a factorization gives P ⊤ j ϕ(qi ) ϕ(kj ) vj oi = P . (4) ⊤ j ϕ(qi ) ϕ(kj ) Because ϕ(qi ) appears only as a left-multiplier, we can reorder the computation: accumulate S = P ⊤ dv ×dk once, then evaluate oi = S ϕ(qi ) for each query. This reduces complexity v ϕ(k j j) ∈ R j to O(N ) [24]. We follow common practice and omit the denominator, which many architectures drop without significant performance loss, giving oi = S ϕ(qi ),
S=
N X
vj ϕ(kj )⊤ .
(5)
j=1
RNN interpretation. The state S ∈ Rdv ×dk accumulates key-value associations and is queried via a dot product — it functions as an associative memory [19]. Under causal (left-to-right) ordering, S is updated token-by-token: Si = Si−1 + vi ϕ(ki )⊤ , giving an exact equivalence with an RNN whose hidden state is matrix-valued. For our setting — bidirectional spatial attention over PDE grids — S is accumulated over all tokens simultaneously, but the O(N ) cost and fixed-size state are retained. A.1
The Ridge Regression View
The accumulated state S in (5) can be given an optimization interpretation. Consider the Frobeniusregularized reconstruction loss N X min ∥vj − S ϕ(kj )∥2 + λ∥S∥2F . (6) S
j=1
Setting the gradient to zero yields the closed-form solution X X −1 S∗ = vj ϕ(kj )⊤ (G + λI) , G= ϕ(kj ) ϕ(kj )⊤ . j
(7)
j
The three regularization regimes connect directly to known architectures [44, 3, 41]: Regime λ ≫ ∥G∥ Finite λ λ→0
Solution P λ−1 j vj ϕ(kj )⊤ P ⊤ v ϕ(k ) (G + λI)−1 j j j P ⊤ G−1 j vj ϕ(kj )
Interpretation Vanilla linear attention (5) Ridge-corrected (WaveLiT) Least-squares
In the large-λ regime, (G + λI)−1 ≈ λ−1 I and the Gram correction vanishes — all feature directions are treated as equally important regardless of how frequently they are queried. The scalar λ−1 is absorbed into the output projection during training, leaving the functional form of vanilla linear attention (5). Vanilla linear attention thus corresponds to ignoring the covariance structure of the key features entirely. WaveLiT uses a finite-λ ridge correction, incorporating the Gram matrix G to account for how frequently each direction in feature space is queried. The correction adds a matrix inverse of size dk × dk , which is cheap when dk ≪ N . Ablations comparing vanilla, ridge-corrected, and least-squares variants are in Appendix I.2.
14
B Impact of Sequence Length and Attention Mechanism on Performance and Efficiency We explore the interplay between model architecture, effective sequence length (modulated by wavelet decomposition levels), and computational cost. Scope of this comparison. The study presented here isolates the choice of token-mixing primitive (DPA versus MILA) and is intended only to motivate the move away from O(N 2 ) self-attention in long-sequence wavelet-tokenized inputs. The configurations evaluated below predate the full WaveLiT mixer and do not include ridge regularization and gating. The accuracy numbers should therefore be read as a baseline-vs-baseline comparison of attention kernels, not as the performance of WaveLiT itself; ablations of the full mixer recipe are reported separately in Appendix I. We retain this earlier comparison because the compute/accuracy trade-off it exposes between DPA and MILA is the original motivation for our linear-attention design choice and remains informative on its own. Our primary goal is to identify model configurations that balance high accuracy with manageable computational cost. We compare standard Dot-Product Attention (DPA) with an enhanced linear attention mechanism (MILA) [14]. To understand these trade-offs, we conducted evaluations on the Navier-Stokes (NS) benchmark from PDEArena [13], following the experimental setup established in DPOT [15]. We compared models using both DPA and MILA across different model sizes (“WaveLiT1.2M” vs. “WaveLiT-9.5M”) and wavelet decomposition levels (1 vs. 2 levels, where fewer levels mean longer effective sequences). The results, including relative L2 error and training time per 1000 iterations, are presented in Table 1 and visualized for key configurations in Figure 6. While detailed results are shown for NS, similar trends in the DPA vs. MILA trade-off were observed across other datasets. Table 1: Performance and efficiency comparison of Dot-Product Attention (DPA) and Enhanced Linear Attention (MILA) across varying model sizes and wavelet levels. Training time is per 1000 iterations. Attention Model Wavelet Levels Params (M) Training Time (s) Relative L2 DPA MILA
WaveLiT-9.5M WaveLiT-9.5M
1 1
DPA MILA
WaveLiT-9.5M WaveLiT-9.5M
2 2
DPA MILA
WaveLiT-1.2M WaveLiT-1.2M
1 1
DPA MILA
WaveLiT-1.2M WaveLiT-1.2M
2 2
≈ 8.5 ≈ 8.5
212.80 60.00
0.00458 0.00480
≈ 8.7 ≈ 8.7
26.68 15.00
0.00711 0.00750
≈ 1.1 ≈ 1.1
53.30 18.00
0.01848 0.01950
≈ 1.2 ≈ 1.2
7.12 5.00
0.02606 0.02700
Discussion. The data in Table 1 (visualized in Figures 6 and 7) reveals several key trends regarding model performance and efficiency. Firstly, for a given model size and attention type, decreasing the number of wavelet levels (thus increasing effective sequence length) generally improves performance (lower relative L2 error), albeit at a higher computational cost, especially for DPA. Secondly, increasing model parameters (e.g., from WaveLiT-1.2M to WaveLiT-9.5M) also tends to enhance performance. However, as detailed in Appendix E, input sequence length (determined by wavelet levels) often plays a more decisive role than raw parameter count alone, particularly for smaller models. Most critically, these results highlight the compelling advantages of MILA. While DPA occasionally achieves marginally lower error, MILA consistently delivers comparable accuracy at a substantially reduced computational cost. For instance, with WaveLiT-9.5M and 1 wavelet level, MILA trains over 3.5 times faster than DPA for a very small trade-off in error. Therefore, MILA emerges as a highly practical and scalable attention mechanism. It allows leveraging the performance benefits of longer effective sequences, crucial for resolving fine-scale features in PDE solutions, without incurring the prohibitive computational costs.
15
Performance vs. Computational Cost: DPA vs. MLLA S,2WL
Legend DPA MLLA 1 Wavelet Level 2 Wavelet Levels
S,2WL
Test Relative L2 Error
0.025
0.020
S,1WL S,1WL
0.015
0.010 B,2WL
B,2WL
0.005
B,1WL
10
B,1WL
1
10
2
Training Time (s / 1000 iterations) Figure 6: MILA models (triangles) demonstrate significantly lower compute costs for comparable performance levels, especially at 1 wavelet level (longer sequences), compared to DPA models (circles)
Test Relative L2 Error
0.020 0.015
Training Time (s / 1000 iterations)
DPA vs. MLLA Comparison for 1 Wavelet Level Configurations Test Relative L2 Error (1 Wavelet Level) 0.0195 0.0185
Attention Type DPA MLLA
0.010 0.005 0.000
0.0048
0.0046
Base
Small
Model Size
Training Time (1 Wavelet Level) 212.8
Attention Type DPA MLLA
200 150 100 60.0
53.3
50
18.0
0
Base
Small
Model Size
Figure 7: Comparative performance of Dot-Product Attention (DPA) and Enhanced Linear Attention (MILA) for 1-wavelet-level (long sequence) configurations. Left: Test Relative L2 Error. Right: Training Time (sec/1000 iterations). MILA demonstrates significant reductions in training time while maintaining competitive error rates compared to DPA across all model sizes.
16
C
Wavelet Selection Choice
For the core experiments in this work, the Biorthogonal wavelet bior2.2 was selected as the default choice for the wavelet transform layer. The bior2.2 wavelet offers a good balance of properties, including symmetry and relatively compact support (filter length of 6), making it suitable for capturing features across various scales without excessive computational overhead. To validate this choice and assess the model’s robustness to the specific wavelet employed, we conducted an ablation study. We evaluated the performance of the WaveLiT-1.2M model on the Navier-Stokes benchmark dataset provided by PDEArena [13] using three distinct wavelets: haar, bior2.2, and bior4.4. These wavelets represent a range of characteristics: haar is the simplest, discontinuous orthogonal wavelet (filter length 2); bior2.2 is our chosen symmetric biorthogonal wavelet (filter length 6); and bior4.4 is a smoother, higher-order symmetric biorthogonal wavelet (filter length 10). Our wavelet selection was guided by the jax-wavelets library [10], specifically choosing from families compatible with reflect padding, which we found to yield better performance than wrap padding in our experiments. Discussion. The performance was measured using the relative-L2 error for 1-step ahead predictions and 4-step rollouts. The results, presented in Table 2, demonstrate consistency across the different wavelet choices. This minimal variation in performance across wavelets with differing characteristics suggested that our architecture is reasonably robust to the specific choice of wavelet and we stick to the use of bior2.2 as the default choice. Table 2: Ablation study on wavelet choice for the WaveLiT-1.2M model on the Navier-Stokes benchmark. Reported values are relative-L2 error for 1-step ahead and 4-step rollout predictions. Lower is better. Wavelet 1-Step Ahead 4-Step Rollout haar bior2.2 bior4.4
0.010669 0.010618 0.010896
17
0.023265 0.023349 0.023718
D
Loss Term Ablations
Table 3 presents the results of experiments investigating the impact of different loss term weightings on model performance, specifically comparing Mean Squared Error (MSE) loss and an L1 wavelet loss (L1 wavelet). The experiments were conducted for two model sizes (Small (1.2M) and Base (9.5M)) and two different numbers of wavelet decomposition levels (1 and 2). Table 3: Effect of MSE and L1 Wavelet Loss Terms on Model Performance. Performance metrics are relative L2 error on the test set for single-step prediction (Test Rel. L2 ) and 4-step rollout (Test Rel. L2 Rollout). Best results for each ablation are highlighted in bold. Model
Wavelet
Loss Weights
Test Performance
Levels
λM SE
λL1
Rel. L2
Rel. L2 Rollout (4 steps)
WaveLiT-9.5M WaveLiT-9.5M WaveLiT-9.5M
1 1 1
0 1 1
1 1 0
0.00416 0.00453 0.00563
0.01028 0.01009 0.01295
WaveLiT-9.5M WaveLiT-9.5M WaveLiT-9.5M
2 2 2
0 1 1
1 1 0
0.00662 0.00670 0.00837
0.01634 0.01636 0.01972
WaveLiT-1.2M WaveLiT-1.2M WaveLiT-1.2M
1 1 1
0 1 1
1 1 0
0.01177 0.01167 0.01280
0.02792 0.02702 0.02945
WaveLiT-1.2M WaveLiT-1.2M WaveLiT-1.2M
2 2 2
0 1 1
1 1 0
0.01407 0.01406 0.01515
0.03283 0.03295 0.03369
Power Spectral Density
Discussion A clear trend emerges from the data: the inclusion of the L1 wavelet loss term, either exclusively (λL1 = 1, λM SE = 0) or in conjunction with the MSE loss (λL1 = 1, λM SE = 1), consistently yields superior performance compared to using only the MSE loss (λL1 = 0, λM SE = 1). This benefit of the wavelet loss is further demonRadially Averaged Power Spectral Density strated by examining the spectral characteristics With Wavelet Loss of the prediction error. Figure 8 shows the RadiWithout Wavelet Loss ally Averaged Power Spectral Density (RAPSD) 10 of the prediction error for the WaveLiT-9.5M model (1 wavelet level) when trained with the combined MSE and L1 wavelet loss versus MSE 10 loss alone. The plot indicates that the incorporation of the wavelet loss term leads to a reduction in error power across the entire frequency spec10 10 Frequency trum compared to the model trained without it. This reduction is evident at low frequencies, cor- Figure 8: Radially Averaged Power Spectral Denresponding to large-scale structures, and extends sity (RAPSD) of the prediction error for the to higher frequencies, representing finer details. WaveLiT-9.5M model (1 wavelet level), comparSuch broad spectral improvement suggests that ing training with both MSE and L1 wavelet loss the L1 wavelet loss aids the model in more accu- terms (“With Wavelet Loss") versus training with rately resolving the multi-scale features inherent MSE loss alone (“Without Wavelet Loss"). The in the solution. inclusion of the wavelet loss demonstrably reduces Given these observations, we find that setting error power across all frequencies. both λM SE = 1 and λL1 = 1 provides robust, high-quality results. Therefore, for all experiments in this paper, we adopted this balanced weighting, setting both the MSE loss weight and the L1 wavelet loss weight to 1. −4
−5
−2
18
−1
E
Training Details
All models are trained using the AdamW optimizer [30], using a weight decay rate of 1 × 10−4 . We employ an exponential decay learning rate scheduler with a linear warm-up phase: the learning rate is linearly scaled up from 1 × 10−7 to 1 × 10−3 over the first 5,000 steps, after which it is exponentially decayed with a decay rate of 0.99 every 2,000 transition steps. To ensure training stability, gradients are globally clipped to a norm of 1. The batch size is set to 4 for all benchmarks. We default to a Discrete Wavelet Transform (DWT) level of 1 for most benchmarks to maximize input sequence length and provide the model with fine-grained details. Table 4 outlines the projected total training times for 500,000 steps for various WaveLiT models across different benchmarks, without intermediate evaluation. It is important to note that these times are estimates from running on a shared NVIDIA H200 server and can vary depending on factors such as data loading efficiency, software library versions, and computational facility workload; thus, they may not perfectly and consistently reflect the raw computational cost. Table 5 shows the timings training the foundation model backbone on an NVIDIA B200. Table 4: Estimated bespoke model training time per dataset (GPU-hours, single GPU). PT times estimated as titer × 500,000; FT times as titer × 50,000, where titer is the mean per-iteration wall-clock time logged during training. Pretrain (500k iters) Finetune (50k iters) Dataset
1.2M
9.5M
1.2M
9.5M
RB ASM SF AM HS GS VI TRL2D
5.5h 5.3h 18.1h 6.9h 13.2h 2.2h 27.1h 5.3h
15.3h 12.5h 25.7h 15.0h 38.3h 4.9h 42.7h 8.9h
3.3h 3.3h 5.9h 4.4h 9.6h 1.4h 12.7h 2.7h
11.1h 11.1h 20.2h 11.6h 38.3h 4.0h 39.6h 8.9h
Total
83.5h
163.4h
43.3h
144.6h
Bespoke grand total: 434.8 GPU-hours (18.1 GPU-days)
Table 5: Foundation model training time (GPU-hours, single GPU). PT estimated as titer × 1,000,000; FT as titer × 50,000. Actual wall-clock time from wandb shown in parentheses for reference. Pretrain (1M iters) Finetune (50k iters) Foundation model
41.0h (51.1h)
Total
8.4h (8.5h)
49.4 GPU-hours (2.1 GPU-days)
Table 6: Total compute budget summary. Model family GPU-hours GPU-days Bespoke 1.2M (PT + FT) Bespoke 9.5M (PT + FT) Foundation (PT + FT)
126.8h 308.0h 49.4h
5.3 12.8 2.1
Grand total
484.2h
20.2
19
F
Dataset Sampling Strategy
Training a shared-trunk foundation model across heterogeneous PDE datasets introduces a fundamental tension: the eight datasets in our corpus differ in size by up to 66× in effective token count (Table 7). The trunk must accumulate useful gradient signal from every PDE family, yet unconstrained training will concentrate update mass on whichever datasets dominate the corpus. Table 7 reports, for each dataset, the number of training trajectories, the native spatial resolution (H × W ), the effective post-wavelet token count per example (h/2 × w/2 for one decomposition level), and the resulting total epoch token count Ni = ni · τi that serves as the basis for all sampling calculations. Table 7: Dataset statistics. ni : training trajectories; τi : tokens per example after one wavelet level (spatial dims halved); Ni : total epoch tokens. The token corpus spans nearly two orders of magnitude. ni
H×W
τi
Ni (tokens)
14,000 256×256 960,000 128×128 278,600 512×128 178,304 256×512 7,200 128×384 6,487 512×512 321,600 256×256 20,384 1024×256
16,384 4,096 16,384 32,768 12,288 65,536 16,384 65,536
229,376,000 3,932,160,000 4,564,582,400 5,842,665,472 88,473,600 425,132,032 5,269,094,400 1,335,885,824
—
21,687,369,728
Dataset active_matter gray_scott_reaction_diffusion rayleigh_benard shear_flow turbulent_radiative_layer_2D viscoelastic_instability acoustic_scattering_maze helmholtz_staircase Total
—
—
F.1
Three Candidate Sampling Schemes P Let pi = Ni / j Nj denote the proportional share of the corpus. We compare three methods for assigning the sampling probability wi (the probability of drawing the next mini-batch from dataset i). Uniform.
wi = 1/K for all i, with K = 8.
Proportional, temperature T = 0.2. p i . (8) wi ∝ exp T With T → 0 all weight concentrates on arg maxi pi ; with T → ∞ the scheme recovers uniform. T = 0.2 sharpens toward the larger datasets relative to uniform. We choose this relaxation to allow sampling from sparser datasets such as turbulent_radiative_layer_2D which would otherwise be sampled at less than 1% under naive proportional sampling. √
N -proportional.
√ Ni wi = P p . Nj j
(9)
No temperature hyperparameter; wi lies between the pure-proportional and uniform extremes by construction. Table 8 lists the resulting probabilities under all three schemes alongside the raw proportional share pi for reference. Table 9 provides the most informative view: the oversampling ratio wi /pi , i.e., how many times more frequently each dataset is sampled relative to what its token share alone would dictate. F.2
Analysis
We analyze the three sampling strategies based on their theoretical properties and empirical performance: Under uniform sampling, the smallest dataset (TRL2D) is drawn 30.6× more often than 20
Table 8: Sampling probabilities under each scheme. Prop. (pure) is the raw token share pi , shown for reference only. √ Dataset Prop. (pure) Uniform Prop. N (T =0.2) active_matter gray_scott_rd rayleigh_benard shear_flow turb. rad. layer 2D viscoelastic_inst. acoustic_scattering helmholtz_staircase
0.0106 0.1813 0.2105 0.2694 0.0041 0.0196 0.2430 0.0616
0.1250 0.1250 0.1250 0.1250 0.1250 0.1250 0.1250 0.1250
0.0617 0.1448 0.1676 0.2250 0.0597 0.0645 0.1971 0.0796
0.0420 0.1737 0.1871 0.2117 0.0261 0.0571 0.2011 0.1012
DKL (wi ∥pi ) [nats]
0.000
0.766
0.214
0.099
Table 9: Oversampling ratio wi /pi under each scheme. A ratio > 1 means the dataset is sampled more often than its raw token share warrants; ratio < 1 means it is downsampled. Pure proportional gives ratio 1.0 by definition (shown for reference). √ Dataset Prop. (pure) Uniform Prop. N (T =0.2) active_matter gray_scott_rd rayleigh_benard shear_flow turb. rad. layer 2D viscoelastic_inst. acoustic_scattering helmholtz_staircase
1.0× 1.0× 1.0× 1.0× 1.0× 1.0× 1.0× 1.0×
11.8× 0.7× 0.6× 0.5× 30.6× 6.4× 0.5× 2.0×
5.8× 0.8× 0.8× 0.8× 14.6× 3.3× 0.8× 1.3×
4.0× 1.0× 0.9× 0.8× 6.4× 2.9× 0.8× 1.6×
its token share warrants. This would cause overfitting on small datasets early in training. The temperature-scaled scheme partially corrects this, dropping the maximum oversampling ratio to 14.6×. However, this requires careful tuning of the hyper-parameter T and is dependent on√the specific corpus distribution, requiring retuning if datasets are added or removed. We adopt N proportional sampling for the foundation model. It resolves the tensions of multi-physics training without introducing sensitive hyperparameters, offering distinct advantages: • Minimal Distributional Distortion: While we must smooth the sampling probabilities to prevent underfitting small datasets, we want to minimize the distortion of the natural √ data distribution pi . Measuring the Kullback-Leibler divergence DKL (wi ∥pi ) shows that N alters the true data distribution the least. • Established Precedent: This weighting belongs to the exponential smoothing family (Niα ) established in large-scale multilingual pretraining. While models like XLM-R [9] √ and mT5 [48] tuned this exponent to α = 0.3 to aggressively upsample the long tail, N (α = 0.5) serves as a balanced, parameter-free middle ground. Empirical Validation. As shown in √ Table 10, we evaluated all three strategies at the 1.2M parameter scale across the full benchmark. N -proportional sampling achieves the lowest overall median VRMSE (2.85 × 10−2 ) and outperforms the baselines on 5 out of 8 datasets. While uniform sampling narrowly minimizes error on the smallest datasets (TRL2D and VI) due to extreme early overfitting, √ N provides the robust, balanced signal required for generalized multi-physics training at scale.
21
Table 10: Ablation of dataset sampling strategies at the 1.2M parameter scale. VRMSE (median) on one-step prediction across eight datasets from The Well benchmark [36]. Best result per dataset in bold. √
Dataset
S1: Uniform
S2: Temp-Scaled (T =0.2)
AM GS ASM HS RB SF TRL2D VI
6.07e-2 8.07e-3 2.79e-2 8.48e-3 4.82e-2 1.27e-2 1.80e-1 1.79e-1
7.18e-2 6.19e-3 1.91e-2 8.42e-3 4.15e-2 8.69e-3 1.96e-1 2.13e-1
7.90e-2 5.31e-3 1.82e-2 7.23e-3 3.88e-2 8.39e-3 2.25e-1 2.28e-1
Median (overall) Dataset wins
3.81e-2 3
3.03e-2 0
2.85e-2 5
22
S3:
-Prop (ours)
G
Rollout Finetuning
Autoregressive models accumulate errors over long rollouts: small per-step inaccuracies compound, causing the predicted trajectory to drift increasingly far from the true solution. To address this, we explore several finetuning strategies that expose the model to its own predictions during training, thereby reducing the mismatch between training (teacher-forced) and inference regimes. We consider four families of methods. Scheduled Sampling combines a linear teacher-forcing curriculum with K=8 BPTT unrolling: the teacher-forcing rate is annealed from 1.0 to 0.1 over the course of training, progressively replacing ground-truth inputs with model predictions while gradients are backpropagated through the rollout. BPTT unrolls the model for K steps and backpropagates through the rollout without any teacher-forcing annealing, evaluated at K=4 and K=8. CausalBPTT [46] runs a fully autoregressive rollout of K steps and weights each step’s loss causally: P wi = exp −ϵ k<i Lk , (10) where Lk is the loss at step k (treated as a constant via stop_gradient). At ϵ=0 all steps are equally weighted, recovering plain BPTT. For ϵ>0, step i receives significant gradient only once all earlier steps are well-predicted, implementing a data-driven sequential curriculum without any teacher-forcing annealing. Pushforward unrolls K steps but propagates gradients only through the final prediction, detaching intermediate states, also evaluated at K=4 and K=8. We evaluate each method on two axes: rollout VRMSE across short (2:8), medium (9:26), and long (27:56) prediction horizons, and one-step prediction metrics to assess degradation from finetuning. The latter is important — aggressive rollout training risks degrading the model’s single-step accuracy, which is undesirable. Results are shown in Tables 11 and 12. All finetuned methods improve over the baseline at longer horizons, with gains becoming more pronounced as rollout length increases. However, the methods differ substantially in how much they degrade one-step performance. Pushforward incurs the largest one-step penalties across both K=4 and K=8, with performance worsening as K increases. Notably, BPTT (K=8) — which uses the same unroll length as Scheduled Sampling but without annealing — degrades one-step performance considerably, whereas Scheduled Sampling does not. This direct comparison isolates the effect of the teacher-forcing curriculum: gradual annealing allows the model to adapt to its own prediction errors without destabilizing single-step accuracy. Both Scheduled Sampling and CausalBPTT emerge as strong candidates: CausalBPTT achieves competitive rollout performance but at a measurable cost to one-step accuracy, while Scheduled Sampling delivers comparable rollout gains with greater robustness on single-step metrics. We therefore adopt Scheduled Sampling as our rollout finetuning strategy for all subsequent experiments. Table 11: Rollout VRMSE (median) by prediction window. All results on the TRL2D test set. Lower is better. Bold indicates the best value in each column. Method Short [2:8] Medium [9:26] Long [27:56] Baseline
0.3797
0.7873
1.0239
Sched. Sampling (K=8) BPTT (K=4) BPTT (K=8) CausalBPTT (ϵ=0) CausalBPTT (ϵ=1) CausalBPTT (ϵ=5) Pushforward (K=4) Pushforward (K=8)
0.3156 0.3188 0.3191 0.3194 0.3147 0.3169 0.3508 0.3713
0.6544 0.7107 0.6457 0.6399 0.6420 0.6588 0.6775 0.6781
0.8744 0.9310 0.8969 0.8774 0.8877 0.8893 0.9003 0.9660
23
Table 12: One-step prediction metrics. All results on the TRL2D test set. Lower is better. Bold indicates the best value in each column. Method MSE Rel. ℓ2 VRMSE (med.) Baseline
0.03710
0.1134
0.1466
Sched. Sampling (K=8) BPTT (K=4) BPTT (K=8) CausalBPTT (ϵ=0) CausalBPTT (ϵ=1) CausalBPTT (ϵ=5) Pushforward (K=4) Pushforward (K=8)
0.03515 0.03585 0.03640 0.03623 0.03621 0.03560 0.04257 0.04900
0.1123 0.1123 0.1162 0.1159 0.1151 0.1133 0.1232 0.1340
0.1505 0.1487 0.1595 0.1585 0.1577 0.1528 0.1628 0.1851
Figure 9: Median rollout VRMSE as a function of prediction step (TRL2D). All finetuned methods reduce error accumulation relative to the teacher-forced baseline, with gains widening at longer horizons. Pushforward degrades at both short and long steps as K increases, while Scheduled Sampling and CausalBPTT maintain low error throughout. Shaded bands show the interquartile range over test trajectories.
24
H
Rollout Error Analysis
All autoregressive surrogates trained with teacher forcing share a common failure mode under rollout. We formalize it here and connect it to WaveLiT’s pattern of successes and failures across datasets. Error recurrence. Let the ground-truth dynamics be governed by xn+1 = f (xn ) and the learned map by x̂n+1 = F (x̂n ). Starting from a perfect initial condition x0 = x̂0 , define En := ∥xn − x̂n ∥. Suppose F satisfies: ∥f (x) − F (x)∥ ≤ ϵ (per-step approximation error), ∥F (x) − F (y)∥ ≤ LF ∥x − y∥ (Lipschitz stability condition).
(11) (12)
Then: En+1 = ∥f (xn ) − F (x̂n )∥ ≤ ∥f (xn ) − F (xn )∥ + ∥F (xn ) − F (x̂n )∥ ≤ ϵ + LF En ,
(13)
and unrolling from E0 = 0 yields n ϵ LF − 1 , LF ̸= 1, i LF − 1 En ≤ ϵ LF = i=0 nϵ, LF = 1. n−1 X
(14)
When LF > 1, errors grow geometrically with rollout length, independent of how small ϵ is. Rollout finetuning trains on the rolled-out trajectory and so jointly optimizes ϵ and LF , in practice trading them off against each other — visible in our experiments as a degradation of one-step accuracy under FT relative to its PT counterpart, in exchange for tighter rollout error within the training horizon. The fundamental constraint is that on chaotic dynamics LF is bounded below by eλ∆t , where λ is the leading Lyapunov exponent of the underlying flow and ∆t is the time step: FT can move along the trade-off curve but cannot escape geometric error growth once the dynamics is genuinely chaotic. When LF ≤ 1 the dynamics are contractive and error grows at most linearly in n — long-horizon accuracy is then limited only by ϵ, which a well-matched prior can make small.
25
I
Architectural Ablations
This appendix collects the cumulative ablations that motivate the WaveLiT recipe. Section I.1 isolates the contributions of the wavelet tokenizer and the multiscale feature pyramid, and Section I.2 isolates the contributions of the mixer components. I.1
Tokenizer and Multiscale Feature Pyramid
To isolate the contribution of each architectural choice, we perform a cumulative ablation on the Navier-Stokes benchmark from PDEArena [13], training all configurations for 200k iterations under identical hyperparameters. Starting from a convolutional encoder-decoder baseline, we incrementally substitute each component and measure both accuracy and wall-clock training time. Results are summarized in Table 13. Table 13: Cumulative ablation on PDEArena Navier-Stokes (200k iterations). Each row adds one ingredient to the previous configuration. ∆ Error denotes relative improvement over the convolutional baseline; higher is better. ∆ Time denotes relative change in training time; negative means faster. Configuration
VRMSE ↓
Time
∆ Error
∆ Time
Conv enc/dec (baseline) Wavelet enc/dec + Multiscale (FPN-L1) + Multiscale (FPN-L2) + Multiscale (FPN-L3) + Multiscale (FPN-L4)
0.01172 0.01155 0.00968 0.00928 0.00899 0.00897
64.6 min 57.4 min 75.1 min 84.1 min 90.0 min 97.0 min
— +1.4% +17.4% +20.8% +23.3% +23.5%
— −11% +16% +30% +39% +50%
Replacing the convolutional encoder-decoder with the wavelet tokenization scheme yields only a modest accuracy improvement (+1.4%) while reducing training time by 11%. This is perhaps not surprising, since a convolutional encoder should be able to approximate wavelet-like features over the course of training; the gain we observe is consistent with the findings of [20], who leveraged wavelet tokenizers for pixel-space diffusion models. The more substantial gain comes from introducing multiscale computation via the feature pyramid: a single pyramid level (FPN-L1) reduces error by 17.4% relative to the baseline at a cost of 16% additional training time, and additional levels provide diminishing returns. To keep the training-cost–accuracy frontier balanced, we adopt FPN-L1 as the default for all subsequent experiments. I.2
Mixer Components
Table 14: Ablation of the proposed mixer design. We keep RoPE fixed in all configurations and progressively add ridge regularization, kernel gating, local positional enhancement (LePE), conditional positional encoding (CPE), and the MILA-style block design. The largest improvements arise from LePE/CPE and the MILA-style block, while ridge regularization and kernel gating provide comparatively modest gains. Comparing B3 and C1 further shows that re-enabling the kernel gate yields only a small additional improvement over the no-gate MILA configuration. Model
Ridge
A1: RoPE-LA A2: + Ridge A3: + Kernel Gate A4: + LePE A5: + CPE
✓ ✓ ✓ ✓
B1: MILA + LePE (no gate) B2: + Ridge (no gate) B3: + CPE (no gate)
✓ ✓
C1: Full model
✓
Kernel Gate
✓ ✓ ✓
LePE
✓ ✓
✓
CPE
MILA Block
✓
✓ ✓ ✓
✓
✓ ✓ ✓
✓
✓
✓
Test/Rel. ↓
Runtime ↓
0.21906 0.20256 0.19236 0.11463 0.03055
36m 37s 39m 30s 40m 38s 40m 56s 43m 07s
0.060258 0.055471 0.025324
40m 41s 42m 58s 46m 42s
0.024363
47m 25s
Table 14 presents an additive ablation of our mixer design. Starting from a RoPE-equipped linear attention baseline (A1), ridge regularization and kernel gating provide only incremental improvements 26
(A2–A3). In contrast, introducing local positional structure through LePE produces a substantial gain (A4), and adding CPE yields a further large improvement (A5), indicating that positional and locality-aware mechanisms are the dominant contributors in the standard linear-attention branch. We next evaluate an MILA-inspired branch in which the kernel gate is disabled (B1–B3) in order to isolate the contribution of the MILA-style block design from that of the additional output modulation. This branch consistently outperforms its non-MILA counterpart at comparable levels of positional modeling, and the full no-gate MILA variant (B3) already surpasses the best non-MILA configuration (A5). These results suggest that the MILA-style block design contributes independently beyond the gains obtained from LePE and CPE alone. Ultimately, each component contributes positively, making the full model (C1) the best overall configuration. The proposed recipe successfully combines two complementary perspectives: the ridge regression view of test-time regression, which motivates the regularization and gating choices, and standard positional and locality augmentations (LePE, CPE), which are critical for the bidirectional spatial setting. While the regression-inspired components provide consistent refinements, the dominant contributors remain the MILA block design and the positional/locality operators, consistent with prior findings in vision settings.
27
J
Detailed Results Tables
All tables below compare WaveLiT bespoke variants (1.2M and 9.5M parameters, pretrained and rollout-finetuned) against four foundation model baselines across eight TheWell benchmarks. Rows marked † indicate checkpoints that diverged under autoregressive rollout and are excluded from ranking. MPP-AViT-L was not evaluated on GS; VI trajectories contain only 16 rollout steps, leaving the T ∈ [21 : 60] window undefined for all models. Table 15: One-step VRMSE (median teacher-forcing next-step prediction error, lower is better). Bold = best per dataset. † = diverged checkpoint, excluded from ranking. ASM∗ : one-step aggregates all sliding-window positions including trivially predictable early steps. Baselines Dataset ∗
ASM HS RB SF TRL2D AM GS VI
Ours
MPP-AViT-L
Poseidon-L
DPOT-H
Walrus
1.2M-PT
1.2M-FT
9.5M-PT
9.5M-FT
0.0337 0.0026 0.0264 0.0071 0.1707 0.0157 — 0.1030
0.0116 0.0019 0.0215 0.0090 0.1323 0.0214 0.0048 0.0878
0.0126 0.0017 0.0288 0.0162 0.1601 0.0476 0.0061 0.1398
0.0099 0.0005 0.0059 0.0012 0.0831 0.0057 0.0001 0.0295
0.0036 0.0003 0.0139 0.0024 0.1421 0.0211 0.0007 0.1084
0.0037 0.0004 0.0159 0.0308 0.1466 0.0282 0.0008 0.1244
0.0016 0.0005 0.0065 0.0015 0.1167 0.0114 0.0006 0.0301
0.0017 0.0005 0.0077 0.0023 0.1186 0.0131 0.0008 0.0336
Table 16: Short-horizon rollout VRMSE, T ∈ [1 : 20] (median over trajectories of mean autoregressive VRMSE over steps 1–20, lower is better). Bold = best per dataset. † = diverged checkpoint. Baselines
Ours
Dataset
MPP-AViT-L
Poseidon-L
DPOT-H
Walrus
1.2M-PT
1.2M-FT
9.5M-PT
9.5M-FT
ASM HS RB SF TRL2D AM GS VI
0.0797 0.0053 0.4109 0.0377 0.4796 0.3195 — 0.1578
0.0785 0.0201 1.3819 0.0657 0.4117 0.3355 0.0411 0.1532
0.0350 0.0022 3.4868 0.0772 0.5134 0.5272 0.1354 0.1989
0.0345 0.0040 0.0992 0.0146 0.3393 0.1262 0.0278 0.0373
0.0239 0.0008 5.7239 0.1375 0.5684 0.5098 0.1086 0.4898
0.0221 0.0009 4.0865 0.1884 0.4540 0.4072 0.1002 0.3340
0.0158 0.0017 1.4417 0.0588 0.3362 0.2096 0.0762 0.1533
0.0153 0.0015 1.8252 0.0657 0.3371 0.2231 0.0824 0.1651
Table 17: Long-horizon rollout VRMSE, T ∈ [21 : 60] (median over trajectories of mean autoregressive VRMSE over steps 21–60; steps 21–46 for HS, lower is better). Bold = best per dataset. † = diverged checkpoint. VI trajectories contain only 16 rollout steps; this window is undefined. Baselines
Ours
Dataset
MPP-AViT-L
Poseidon-L
DPOT-H
Walrus
1.2M-PT
1.2M-FT
9.5M-PT
9.5M-FT
ASM HS RB SF TRL2D AM GS VI
0.1390 0.0096 0.7468 0.1814 0.8879 1.2481 — —
0.2429 0.0507 0.9586 0.2408 0.8441 1.3015 0.1295 —
0.0543 0.0031 1.2664 0.2880 0.9316 1.2867 0.6567 —
0.0560 0.0074 0.6441 0.0810 0.8648 1.2451 0.1268 —
0.1166 0.0020 2.7116 1.0732 1.0853 1.4186 0.4045 —
0.0693 0.0022 1.9080 0.8311 0.8398 1.3863 0.3996 —
0.0406 0.0047 1.1690 0.3812 0.8759 1.3016 0.3040 —
0.0268 0.0038 1.4803 0.2539 0.7986 1.3204 0.3048 —
Foundation Model: Detailed Comparison Tables 19–21 compare WaveLiT-FM against Walrus and the best bespoke WaveLiT variant per cell. The Best WaveLiT column reports the minimum across all four bespoke variants (1.2M-PT/FT, 9.5MPT/FT). ‡ marks cells where a foundation model variant beats Walrus. † marks rollout divergence. 28
Table 18: Average per-dataset rank across the eight TheWell benchmarks (lower is better). Each model’s rank is computed within each dataset; ties get the mean rank; undefined cells are excluded (MPP-AViT-L on GS for all windows; VI for long-horizon). WaveLiT rows use the better of pretrained and rollout-finetuned variants per cell, matching Figure 4. WaveLiT-9.5M is the runner-up to Walrus in every window despite being ∼130× smaller. Model
Params
One-step
T ∈ [1, 20]
T ∈ [21, 60]
1.2B 9.5M 1.2M 629M 409M ∼1B
1.44 1.94 3.12 4.12 4.86 5.38
1.75 2.38 4.25 3.62 3.86 4.88
2.14 3.00 4.14 3.86 3.50 4.00
Walrus WaveLiT-9.5M WaveLiT-1.2M Poseidon-L MPP-AViT-L DPOT-H
Table 19: One-step VRMSE — foundation model results in context. Bold = best per dataset across all four columns. ‡ = foundation model beats Walrus. Dataset
Walrus (external)
Best WaveLiT (bespoke)
Fnd-PT (multi-task)
Fnd-FT (multi-task + FT)
ASM HS RB SF TRL2D AM GS VI
0.0099 0.0005 0.0059 0.0012 0.0831 0.0057 0.0001 0.0295
0.0016 0.0003 0.0065 0.0015 0.1167 0.0114 0.0006 0.0301
0.0087‡ 0.0057 0.0171 0.0046 0.1508 0.0423 0.0032 0.0970
0.0097‡ 0.0061 0.0198 0.0069 0.1655 0.0540 0.0039 0.1100
Table 20: Short-horizon rollout VRMSE, T ∈ [1 : 20]. † = Fnd-PT diverges on RB under autoregressive rollout (value shown for completeness, excluded from ranking). Dataset
Walrus (external)
Best WaveLiT (bespoke)
Fnd-PT (multi-task)
Fnd-FT (multi-task + FT)
ASM HS RB SF TRL2D AM GS VI
0.0345 0.0040 0.0992 0.0146 0.3393 0.1262 0.0278 0.0373
0.0153 0.0008 1.4417 0.0588 0.3362 0.2096 0.0762 0.1533
0.0429 0.0179 5.7242† 0.1376 0.6729 0.7456 0.1964 0.3497
0.0494 0.0183 4.6113† 0.1091 0.5181 0.7195 0.2222 0.2852
Table 21: Long-horizon rollout VRMSE, T ∈ [21 : 60]. VI trajectories contain only 16 rollout steps; this window is undefined (− − −). ‡ = foundation model beats Walrus. Dataset
Walrus (external)
Best WaveLiT (bespoke)
Fnd-PT (multi-task)
Fnd-FT (multi-task + FT)
ASM HS RB SF TRL2D AM GS VI
0.0560 0.0074 0.6441 0.0810 0.8648 1.2451 0.1268 —
0.0268 0.0020 1.1690 0.2539 0.7986 1.3016 0.3040 —
0.1050 0.0349 2.9770 0.9487 1.1706 1.4208 0.6371 —
0.0765 0.0346 2.4407 0.5143 0.8471‡ 1.3440 0.7173 —
29
Table 22: Rank-1 finishes across 20 dataset–window pairs (among {Walrus, MPP-AViT-L, PoseidonL, DPOT-H, Fnd-PT, Fnd-FT}; GS excluded entirely since MPP-AViT-L is not evaluated on it, and VI excluded at T ∈ [21 : 60]). Model
Rank-1 count (of 20)
Walrus (1.2B, bespoke) Fnd-FT (10M, multi-task + FT) Fnd-PT (10M, multi-task) DPOT-H Poseidon-L MPP-AViT-L
30
15 0 1 3 1 0
K
Comparison with Wavelet Neural Operator
Wavelet-based methods for neural PDE solvers span three distinct families, and it is worth clarifying how WaveLiT relates to each before presenting comparisons: • WNO [40] is an integral kernel method that performs kernel integration directly in the wavelet domain — the wavelet transform defines the space in which convolution is computed, analogous to how FNO operates in the Fourier domain. • WDNO [21] is a diffusion-based generative model that learns to denoise PDE trajectories in the wavelet domain. • WaveLiT uses the DWT as a parameter-free tokenization scheme that feeds a transformer backbone, and reuses the wavelet domain as an auxiliary loss to enforce spectral fidelity. The wavelet transform is not the computational substrate of the operator; it is the input representation. A direct comparison is therefore not perfectly fair, but it is nonetheless informative. We run WNO using the official repository2 with default hyperparameters, adapting only the output dimensionality and spatial resolution to match each benchmark. This yields a WNO model of approximately 9M parameters, directly comparable in size to WaveLiT-9.5M. Table 23: One-step VRMSE comparison between WNO (∼9M parameters) and WaveLiT-9.5M on two TheWell benchmarks. Lower is better. WNO uses default repository settings; WaveLiT-9.5M uses the pretrained checkpoint. Benchmark WNO (∼9M) WaveLiT-9.5M Active Matter Turbulent Radiative Layer
0.5736 0.3939
0.0114 0.1167
WaveLiT-9.5M outperforms WNO by a large margin on both benchmarks despite identical parameter counts.
2 https://github.com/TapasTripura/WNO/tree/main/Version%202.0.0
31