Flowers: A Warp Drive for Neural PDE Solvers

TL;DR

Flowers is a neural architecture that efficiently learns PDE solution operators using multihead warps, achieving superior performance on various 2D and 3D benchmarks without relying on Fourier, convolution, or attention mechanisms.

Contribution

The paper introduces Flowers, a novel neural PDE solver architecture based solely on multihead warps, offering a physics-inspired, efficient alternative to existing methods.

Findings

Outperforms similar-sized Fourier, convolution, and attention baselines.

A 150M-parameter model surpasses recent transformer-based models in PDE tasks.

Achieves excellent results on diverse 2D and 3D PDE benchmarks.

Abstract

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.