Reproduce “Flowers: A Warp Drive for Neural PDE Solvers” on OpenResearch

About this paper

Till Muser, Alexandra Spitzer, Matti Lassas, Maarten V. de Hoop, Ivan Dokmanić

FLOWERS introduces a neural architecture for solving Partial Differential Equations (PDEs) that primarily uses learned non-linear coordinate warps, also known as pullbacks, as its core mechanism. This approach enables the model to effectively learn complex PDE solution operators across various physical phenomena, achieving state-of-the-art performance with high computational efficiency and strong physical grounding.

Key insights

• Physically-motivated multihead coordinate warps can serve as a sufficient and scalable primitive for constructing state-of-the-art neural PDE solvers, aligning with characteristic curves in conservation laws and ray propagation in wave equations.
• Learned, adaptive displacements are crucial for effective nonlocal mixing, and multiple heads are essential for maintaining stability and accuracy during long-horizon autoregressive rollouts by representing several concurrently active transport modes.
• The multiscale U-Net-like scaffold significantly enhances performance, especially for long-horizon predictions, by enabling the network to process information at different spatial scales.

Summary by alphaXiv.