WaveLiT: A Parameter-Efficient Architecture for Neural PDE Solvers

Building surrogate models for complex physical systems governed by partial differential equations (PDEs) requires models capable of accurately resolving fine-scale features without prohibitive computational costs. While transformer-based models are promising, the quadratic complexity of self-attention often restricts the input sequence length, limiting the resolution of tokenized inputs. We introduce WaveLiT, a neural PDE solver designed for high-resolution problems by utilizing an efficient linear attention mechanism. WaveLiT utilizes a wavelet transform for input tokenization, generating feature-rich tokens that are then processed by its linear attention core, enabling favorable scaling to long sequences. Crucially, to enhance performance for high-frequency details, we incorporate a wavelet-domain $L_1$ loss on the prediction error. This combination allows WaveLiT to achieve exceptional performance and parameter efficiency across multiple PDE benchmarks, with competitive training speeds. Our findings underscore the power of leveraging efficient attention mechanisms to process finer-grained inputs, complemented by targeted loss functions, offering a potent and scalable recipe for building neural PDE solvers.