Flowers: A Warp Drive for Neural PDE Solvers

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

cs.LGFeb 17, 2026

Flowers: A Warp Drive for Neural PDE Solvers

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

Organizations: Department of Mathematics and Computer Science, University of Basel, Basel, Switzerland · Department of Mathematics and Statistics, University of Helsinki, Helsinki, Finland · Simons Chair in Computational and Applied Mathematics and Earth Science, Rice University, Houston, TX, USA

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 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.

Figures & tables

Figure 1: The multi-head SelfWarp layer. (a) Given an input field $u$ with $C$ channels, (b) at each location $x$ the layer computes (c) a local per-head value $v^{(h)}(x)\in\mathbb{R}^{C_{h}}$ via a linear (affine) projection of the input $u(x)$ and (d) a displacement field (one per head) $\varrho^{(h)}(x)$ , also pointwise, using a small MLP. Each $C_{h}$ -channel per-head value field $v^{(h)}$ (e) is warped along the same per-head displacement field $\varrho^{(h)}$ (f) to compute the warped head output (g). The outputs of all heads are concatenated at the output (h).

Figure 2: Wavefronts and rays picture: at high frequency, localized wave packets propagate along rays (see Appendix E ). The solution at $(t,x)$ is a sum of contributions arriving along rays indexed by the angular variable $\eta$ , which Flowers represent by heads: each head implements one upstream pullback at the displaced source coordinate.

The Well
	Next-step				1:20 Rollout				21:60 Rollout
Dataset	FNO	CNUnet	scOT	Flower	FNO	CNUnet	scOT	Flower	FNO	CNUnet	scOT	Flower
Acoustic Maze	0.1454	0.0129	0.0361	0.0064	0.4197	0.0874	0.1996	0.0489	0.7058	0.2199	2.0202	0.1347
Active Matter	0.1749	0.0650	0.1050	0.0249	3.2862	1.7781	4.1055	1.3905	1.8214	1.6465	2.0372	1.5251
Gray-Scott	0.0372	0.0188	0.0673	0.0102	0.9125	0.3059	0.4169	0.2074	37.044	5.1671	4.5485	4.1788
MHD 64	0.3403	0.2062	–	0.1165	1.3007	0.9534	–	0.7580	1.6366	1.4299	–	1.2827

Table 1: VRMSE for unconditioned 4

\rightarrow

1 next-step prediction, 1:20 rollout, and 21:60 rollout. Best in bold, second best colored green. All models scaled to 15M-20M parameters, see Appendix A.2 for exact counts. En-dash (–) indicates that the model does not support 3D data ( scOT ), or that the trajectory is too short to calculate 21:60 rollouts (PDEGym datasets and Viscoelastic instability). † Using a fixed Viscoelastic Instability dataset, see A.3 .

Figure 3: Comparison of model predictions on the viscoelastic instability dataset (4→1 unconditioned setting), showing the conformation tensor entry $C_{zz}$ . The top row shows the ground truth and the prediction of each model, bottom row shows the last of the four input frames and prediction errors.

Figure 4: Comparison of autoregressive rollout model predictions on the shear flow dataset (4→1 unconditioned setting), showing the tracer.

Figure 5: Comparison of model predictions on the Rayleigh-Taylor instability dataset (4→1 unconditioned setting), showing the density. Visualized using vape4d ( Koehler et al., 2024 ) .

Model	Parameters	1-Step	1:20 Rollout
Flower -Tiny	17.3M	0.0160	0.1114
Flower -Small	69.3M	0.0124	0.0850
Flower -Medium	155.8M	0.0108	0.0739
Poseidon-L	628.6M	0.0194	0.1114

Table 2: Model scaling study on a compressible Euler equations dataset, loss given as VRMSE.

Model Variant	1-Step	1:20 Rollout
w/o warping	0.4974	0.6911
single-head flow	0.1241	0.6719
non-learned warps	0.2190	0.4138
w/o coord encoding in lift	0.0638	0.4065
single conv flow (no MLP)	0.0834	0.3998
w/o IdProj	0.0712	0.3816

Table 3: Ablation study on Flower components, evaluated on the viscoelastic instability dataset (via VRMSE loss), sorted by 1:20 rollout loss. CNUnet values given for comparison.

Figure 6: Visualization of a learned SelfWarp displacement field for the Shear Flow dataset (cf. Fig. 4 ). The black arrows show the learned displacement field $\varrho^{(1)}(x)$ of the first head in Flower ’s first downsampling block, overlaid on the tracer field. Without explicit supervision, the model learns displacements that align remarkably well with the underlying fluid velocity, providing evidence that the strong performance of Flower comes from discovering physically-meaningful transport directions. For more visualizations, see Fig. 13 .

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

	FNO	CNextU-net	scOT	Flower
Optimal LR	$1\text{\times}{10}^{-3}$	$1\text{\times}{10}^{-3}$	$5\text{\times}{10}^{-4}$	$1\text{\times}{10}^{-3}$
VRMSE	0.0987	0.1425	0.3068	0.0463

Appendix

Table 4: Predictions on the 15Hz acoustic time-harmonic Helmholtz equation problem from WaveBench.

Collection	Dataset	FNO	CNUnet	scOT	Flower
The Well	acoustic_scattering_maze	$1\text{\times}{10}^{-3}$	$1\text{\times}{10}^{-3}$	$1\text{\times}{10}^{-3}$	$1\text{\times}{10}^{-3}$
	active_matter	$1\text{\times}{10}^{-3}$	$1\text{\times}{10}^{-3}$	$1\text{\times}{10}^{-3}$	$1\text{\times}{10}^{-3}$
	gray_scott_reaction_diffusion	$5\text{\times}{10}^{-4}$	$1\text{\times}{10}^{-4}$	$1\text{\times}{10}^{-3}$	$5\text{\times}{10}^{-4}$
	MHD_64	$1\text{\times}{10}^{-3}$	$1\text{\times}{10}^{-3}$	–	$5\text{\times}{10}^{-4}$
	planetswe	$1\text{\times}{10}^{-4}$	$1\text{\times}{10}^{-4}$	$1\text{\times}{10}^{-3}$	$5\text{\times}{10}^{-4}$
	post_neutron_star_merger	$5\text{\times}{10}^{-4}$	$1\text{\times}{10}^{-3}$	–	$5\text{\times}{10}^{-4}$

Appendix

Table 5: Best learning rate 4→1 next-step prediction.

Model	2D Parameters
FNO	18,929,704
CNUnet	18,570,807
scOT	17,779,868
Flower	17,329,362

Appendix

Table 6: Approximate parameter counts for the 2D 4

\to

1 benchmark configurations (the exact number depends on the input/output channels of the dataset, and thus varies slightly). 3D variants grow to

\sim

22M ( CNUnet ),

\sim

24M ( Flower ), and

\sim

300M ( FNO ); scOT does not support 3D.

Figure 7: Data efficiency on the unconditioned

4\to 1

next-step prediction task across three datasets. Models are trained using subsets of the full training data, ranging from 5% to 100%. The x-axis shows the fraction of training data used, and the y-axis reports VRMSE. Numerical values are reported in Table 7 .

Dataset	Train fraction	FNO	CNUnet	scOT	Flower
acoustic_scattering_maze	0.05	0.2966	0.0279	0.1465	0.0132
	0.10	0.2443	0.0167	0.0751	0.0085
	0.20	0.1992	0.0102	0.0539	0.0062
	0.50	0.1482	0.0125	0.0372	0.0058
	1.00	0.1454	0.0129	0.0361	0.0064
supernova_explosion_128	0.05	0.6097	0.4996	–	0.3843

Appendix

Table 7: Data efficiency in the 4

\to

1 unconditioned next-step prediction setting. Results are reported as test score when training on subsets of the full training data.

Figure 8: Comparison of model throughput (samples/sec) across resolutions for 2D and 3D problems, measured during the training runs of Table 1 . Throughput is calculated as batch size divided by time per batch and includes both forward and backward passes. Error bars indicate standard deviation across datasets.

Figure 9: Comparison of autoregressive model predictions on the viscoelastic instability dataset (4→1 unconditioned setting), showing the conformation tensor entry $C_{zz}$ . The top row shows the ground truth and the prediction of each model, bottom row shows the last of the four input frames and prediction errors.

Figure 10: Comparison of model predictions on the Rayleigh-Bénard dataset (4→1 unconditioned setting), showing the buoyancy. We use blow-outs to enhance regions of interest.

Figure 11: Comparison of autoregressive model predictions on the active matter dataset (4→1 unconditioned setting), showing the concentration. The top row shows the ground truth and the prediction of each model, bottom row shows the last of the four input frames and prediction errors.

Figure 12: Comparison of model predictions on WaveBench’s helmholtz_anisotropic 15Hz dataset. The top row shows the ground truth and the prediction of each model, bottom row shows the input (background wavespeed) and prediction errors.

Panel of Figure 13 — Figure 13: Additional visualizations of learned SelfWarp displacement fields (cf. Fig. 6 ). We depict the first four heads of the first downsampling block. Left: On the viscoelastic instability dataset, one head tracks the primary flow direction, while the remaining heads capture secondary structure. Top right: On shear flow, the four heads co-align with the laminar flow direction with slight variation. Bottom right: On Planet SWE, the displacements are more heterogeneous, reflecting the more intricate multi-physics dynamics of the rotating shallow-water system. Overall, interpretability is clearest on laminar, transport-dominated problems; on turbulent or multi-physics regimes the learned displacements remain organized but are harder to read off because there is no single coherent transport direction.

Collection	Dataset	FNO	CNextU-net	scOT	Flower
The Well	acoustic_scattering (maze)	0.1650	0.0126	0.0378	0.0077
	active_matter	0.2200	0.0918	0.1934	0.0397
	MHD_64	0.3543	0.1895	–	0.1378
	planetswe	0.0255	0.0041	0.0050	0.0018
	rayleigh_benard	0.2589	0.1395	0.2295	0.0706
	shear_flow	0.0791	0.0650	0.1383	0.0285

Appendix

Table 8: VRMSE loss for conditioned 1→1 next-step prediction. Best results marked in bold. We did not test multiple learning rates in this setting. Instead all models were trained with a learning rate of

1\text{\times}{10}^{-3}

, except scOT where

5\text{\times}{10}^{-4}

was used, as reported in the original paper ( Herde et al., 2024 ) .

	Next-step VRMSE					1:20 Rollout VRMSE
Dataset	FNO	CNUnet	scOT	Flower	Flower (single-scale)	FNO	CNUnet	scOT	Flower	Flower (single-scale)
Active Matter	0.1749	0.0650	0.1050	0.0249	0.0567	3.2862	1.7781	4.1055	1.3905	2.1708
Rayleigh-Bénard	0.2104	0.2171	0.1863	0.0807	0.1565	39.038	12.507	5.6486	2.1661	2.8323
TRL 2D	0.3250	0.2559	0.3555	0.1930	0.2658	1.2328	0.7051	0.8299	0.5491	0.6762
Viscoelastic †	0.1914	0.1623	0.2017	0.0624	0.0882	0.4284	0.3592	0.4890	0.3465	0.3921

Appendix

Table 9: Results for a single-scale Flower variant on datasets from The Well. The model embeds inputs into

16\times 16

patches and processes them at fixed resolution until de-embedding, rather than using the U-net multiscale structure. Since channel counts remain constant, IdProj is unnecessary; we reallocate this convolution to a projection after the warp and rearrange operations for best performance, giving:

u\mapsto u+\phi\circ\mathrm{Proj}\circ\textsc{SelfWarp}_{V,g}\circ\mathrm{Norm}[u]

. The results in this table should be taken as a proof of concept for a single-scale Flower : we did not significantly iterate on the design of this architecture and we expect that much better results can be achieved if one does. Other models’ results are taken from Table 1 and added for comparison. † Using a fixed Viscoelastic Instability dataset, see A.3 .

Panel of Figure 14 — Figure 14: A schematic illustration of the Flower U-Net architecture.

Neural PDE Solvers

Neural PDE solvers use trainable models to approximate solutions or solution operators for partial differential equations, often with physical constraints.

See all 269 papers

UniFluids: Unified Neural Operator Learning with Conditional Flow-matching
Mar 19, 2026Haosen Li, Qi Meng, Jiahao Li +4PDE Operator Learning PDE Surrogate Modeling
Academy of Mathematics and Systems Science, Chinese Academy of Sciences · Microsoft Research Asia · Renmin University of China +1
FlashPDE: A Drop-In Fused Triton Operator Library for Neural PDE Solvers
Jul 20, 2026Peiyu Zang, Bosen Xie, Ruoxiang Xu +1Automatic Differentiation Neural PDE Solvers
School of Mathematical Sciences, Beijing Normal University · 2Beijing University of Posts and Telecommunications
Small Models, Strong Priors: Architectural Inductive Bias for Parameter-Efficient Neural PDE Solvers
May 25, 2026Shyam Sankaran, Hanwen Wang, Paris PerdikarisNeural PDE Solvers PDE Operator Learning
Department of Mechanical Engineering and Applied Mechanics, University of Pennsylvania · Graduate Program in Applied Mathematics and Computational Science, University of Pennsylvania

Mar 19, 2026J/K move · Enter open · S save

UniFluids: Unified Neural Operator Learning with Conditional Flow-matching

Haosen Li, Qi Meng, Jiahao Li, Rui Zhang+3

Academy of Mathematics and Systems Science, Chinese Academy of Sciences · Microsoft Research Asia · Renmin University of China+1

PDE Operator Learning PDE Surrogate Modeling Conditional Flow Matching

Open paper AskPDF arXiv

Partial differential equation (PDE) simulation holds extensive significance in scientific research. Currently, the integration of deep neural networks to learn solution operators of PDEs has introduced great potential. In this paper, we present UniFluids, a conditional flow-matching framework that harnesses the scalability of diffusion Transformer to unify learning of solution operators across diverse PDEs with varying dimensionality and physical variables. Unlike the autoregressive PDE foundation models, UniFluids adopts flow-matching to achieve parallel sequence generation, making it the first such approach for unified operator learning. Specifically, the introduction of a unified four-dimensional spatiotemporal representation for the heterogeneous PDE datasets enables joint training and conditional encoding. Furthermore, we find the effective dimension of the PDE dataset is much lower than its patch dimension. We thus employ

x

-prediction in the flow-matching operator learning, which is verified to significantly improve prediction accuracy. We conduct a large-scale evaluation of UniFluids on several PDE datasets covering spatial dimensions 1D, 2D and 3D. Experimental results show that UniFluids achieves strong prediction accuracy and demonstrates good scalability and cross-scenario generalization capability. The code will be released later.

Abstract

Neural PDE Solvers

Explore similar work

UniFluids: Unified Neural Operator Learning with Conditional Flow-matching

FlashPDE: A Drop-In Fused Triton Operator Library for Neural PDE Solvers

Small Models, Strong Priors: Architectural Inductive Bias for Parameter-Efficient Neural PDE Solvers