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 only through sparse sampling at source coordinates, 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 |