Flowers: A Warp Drive for Neural PDE Solvers
About
We introduce Flowers, a neural architecture for learning PDE solution operators built entirely from multihead warps. Aside from pointwise channel mixing and a multiscale scaffold, Flowers use no Fourier multipliers, no dot-product attention, and no convolutional mixing. Each head predicts a displacement field and warps the mixed input features. Motivated by physics and computational efficiency, displacements are predicted pointwise, without any spatial aggregation, and nonlocality enters \emph{only} through sparse sampling at source coordinates, \emph{one} per head. Stacking warps in multiscale residual blocks yields Flowers, which implement adaptive, global interactions at linear cost. We theoretically motivate this design through three complementary lenses: flow maps for conservation laws, waves in inhomogeneous media, and a kinetic-theoretic continuum limit. Flowers achieve excellent performance on a broad suite of 2D and 3D time-dependent PDE benchmarks, particularly flows and waves. A compact 17M-parameter model consistently outperforms Fourier, convolution, and attention-based baselines of similar size, while a 150M-parameter variant improves over recent transformer-based foundation models with much more parameters, data, and training compute.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| 1:20 Rollout | The Well | Acoustic Scattering Maze Score0.0489 | 4 | |
| 1:20 Rollout | PDEBench | Diffusion-Reaction Error0.0241 | 4 | |
| 1→1 next-step prediction | The Well acoustic_scattering (maze) (test) | VRMSE0.0077 | 4 | |
| 1→1 next-step prediction | The Well active_matter (test) | VRMSE0.0397 | 4 | |
| 1→1 next-step prediction | The Well planetswe (test) | VRMSE0.0018 | 4 | |
| 1→1 next-step prediction | The Well rayleigh_benard (test) | VRMSE0.0706 | 4 | |
| 1→1 next-step prediction | The Well shear_flow (test) | VRMSE0.0285 | 4 | |
| 1→1 next-step prediction | The Well turbulent_radiative_layer_2D (test) | VRMSE0.1907 | 4 | |
| Neural PDE Solving | WaveBench 15Hz acoustic time-harmonic Helmholtz equation | VRMSE0.0463 | 4 | |
| Next step prediction | PDEBench | Diffusion-Reaction Error0.15 | 4 |