Small Models, Strong Priors: Architectural Inductive Bias for Parameter-Efficient Neural PDE Solvers

2026-05-25

AI Summary

This work addresses the excessive reliance of conventional neural PDE solvers on model scale while neglecting the critical role of architectural inductive biases in achieving parameter efficiency and physical consistency. The authors propose WaveLiT, a lightweight architecture that integrates strong inductive priors through discrete wavelet transforms, enhanced linear attention, shared-weight multi-scale feature pyramids, and wavelet-domain auxiliary losses to encode physically meaningful structural knowledge. With only 1–10 million parameters, WaveLiT matches or surpasses foundation models that are hundreds to thousands of times larger across all eight benchmarks in TheWell suite, demonstrating particularly superior performance on wave- and acoustics-dominated tasks. Moreover, the model enables efficient training on a single GPU, highlighting its practicality and scalability.

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$\times$ 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 10M-parameter 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.

Problem

• Neural PDE solvers
• architectural inductive bias
• parameter efficiency
• wavelet transform
• multiscale modeling

Innovation

• architectural inductive bias
• parameter-efficient PDE solvers
• wavelet transform
• multiscale modeling
• neural PDE