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
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)∈RCh via a linear (affine) projection of the input u(x) and (d) a displacement field (one per head) ϱ(h)(x) , also pointwise, using a small MLP. Each Ch -channel per-head value field v(h) (e) is warped along the same per-head displacement field ϱ(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 η , which Flowers represent by heads: each head implements one upstream pullback at the displaced source coordinate.
Next-step
1:20 Rollout
21:60 Rollout
Dataset
FNO
CNUnet
scOT
Flower
FNO
CNUnet
scOT
Flower
FNO
CNUnet
scOT
Flower
The Well
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 → 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 Czz . 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 ϱ(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\times10−3
1\times10−3
5\times10−4
1\times10−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\times10−3
1\times10−3
1\times10−3
1\times10−3
active_matter
1\times10−3
1\times10−3
1\times10−3
1\times10−3
gray_scott_reaction_diffusion
5\times10−4
1\times10−4
1\times10−3
5\times10−4
MHD_64
1\times10−3
1\times10−3
–
5\times10−4
planetswe
1\times10−4
1\times10−4
1\times10−3
5\times10−4
post_neutron_star_merger
5\times10−4
1\times10−3
–
5\times10−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 → 1 benchmark configurations (the exact number depends on the input/output channels of the dataset, and thus varies slightly). 3D variants grow to ∼ 22M ( CNUnet ), ∼ 24M ( Flower ), and ∼ 300M ( FNO ); scOT does not support 3D.
Figure 7: Data efficiency on the unconditioned 4→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 → 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 Czz . 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.
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\times10−3 , except scOT where 5\times10−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×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↦u+ϕ∘Proj∘\textscSelfWarpV,g∘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 .
Figure 14: A schematic illustration of the Flower U-Net architecture.
Neural PDE solvers use trainable models to approximate solutions or solution operators for partial differential equations, often with physical constraints.
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.
Haosen Li, Qi Meng, Jiahao Li +4
Academy of Mathematics and Systems Science, Chinese Academy of Sciences · Microsoft Research Asia · Renmin University of China +1
Physics-Informed Neural Networks (PINNs) solve PDEs by incorporating physical constraints into neural-network training, but large-scale problems are limited by automatic-differentiation memory overhead and inefficient execution of grid-based PDE operators. We present FlashPDE, a drop-in fused operator library for grid-based scientific machine learning. FlashPDE replaces fragmented PyTorch finite-difference execution with differentiable Triton kernels. Each operator integrates fused stencil evaluation, an analytic discrete-adjoint backward pass, and boundary-gradient correction within a unified PyTorch autograd Function interface. The library provides 14 differentiable PDE operators covering 17 configurations across 1D--3D elliptic, parabolic, and Navier--Stokes systems, while remaining independent of neural architectures and training strategies. Experiments on an NVIDIA A100 GPU show that FlashPDE reduces peak memory usage by up to 37.0x compared with coordinate-based automatic differentiation and reduces CUDA kernel launches by up to 3.5x compared with eager PyTorch finite-difference implementations. Across six representative PDE benchmarks, FlashPDE achieves up to 2.30x end-to-end time-to-solution speedup and up to 19.2x kernel-level acceleration while maintaining numerical agreement with PyTorch finite-difference references. FlashPDE provides a hardware-efficient execution layer that bridges differentiable PDE solvers and GPU-optimized numerical computation within the PyTorch ecosystem.
Peiyu Zang, Bosen Xie, Ruoxiang Xu +1
School of Mathematical Sciences, Beijing Normal University · 2Beijing University of Posts and Telecommunications
Neural PDE solvers have followed the scaling trajectory of vision and language, with recent foundation models reaching billions of parameters. We argue that scale is a poor substitute for architectural inductive bias in this domain: structured priors deliver outsized parameter efficiency, and the pattern of where they succeed and fail is itself informative about what they capture. We instantiate this argument in WaveLiT, an architecture combining a discrete wavelet transform for lossless multi-resolution tokenization, an augmented linear attention block, a shared-weight multiscale feature pyramid, and a wavelet-domain auxiliary loss. Bespoke 1-10M-parameter WaveLiT models compete with foundation models of 100-1000× their size across eight TheWell benchmarks, with the largest gains on wave and acoustic-dominated benchmarks where the wavelet-multiscale prior fits the dominant dynamical structure and small per-step errors do not compound geometrically under rollout. Trained jointly across all eight benchmarks, a 10M-parameter foundation variant exhibits a structured, physically interpretable transfer pattern -- strongest where the wavelet-multiscale prior matches the dynamics, weakest on chaotic advection-dominated flows. The entire pipeline trains on a single GPU. The results suggest that small-model PDE performance is shaped by architectural inductive bias rather than scale, and that the structure of a prior's failures is a useful empirical signal about its content.
Shyam Sankaran, Hanwen Wang, Paris Perdikaris
Department of Mechanical Engineering and Applied Mechanics, University of Pennsylvania · Graduate Program in Applied Mathematics and Computational Science, University of Pennsylvania