Read, Write, Relax: Why Neural PDE Surrogates Need Both Global and Local Processing

Anuj Kumar; Heiko Zimmermann; Josiah Bjorgaard; Jacan Chaplais; Nikolaos Bouklas; Matteo Salvador; Alexander Lavin

cs.LGAug 21, 2026

Read, Write, Relax: Why Neural PDE Surrogates Need Both Global and Local Processing

Authors: Anuj Kumar, Heiko Zimmermann, Josiah Bjorgaard, Jacan Chaplais, Nikolaos Bouklas, Matteo Salvador, Alexander Lavin

Organizations: Pasteur Labs, Brooklyn, NY, USA · Cornell University, Ithaca, NY, USA · Institute for Simulation Intelligence, New York, NY, USA

Abstract

Recent mesh-based simulation advances have, in no small part, relied on neural surrogates of two distinct families: global models that route information through a small set of latent tokens, and local models that perform message passing across mesh edges. Consistent with both classes is the inability to perform beyond low-dimensional problems and small-scale or oversimplified meshes, the simulation regimes where industrial problems reside. Our work shows this explicitly and presents a unified formulation. In global approaches, latent-token attention acts as a spatial low-pass filter, while local message passing lacks the global reach necessary to propagate information across large mesh spaces. Viewed through the error, the two operators are the halves of a multigrid cycle: one corrects errors at the lower end of the spectrum, the other at the higher end, and neither can do the other's job. We introduce Read-Write-Relax (RWR), which interleaves latent attention with message-passing relaxation under a unified formulation. The interleaved processor lowers error across the entire spectrum, making RWR the most accurate model in nearly every comparison across our industrial and public benchmarks. It is also markedly data-efficient in the scarce-data regimes, accurate on the engineering quantities of interest, and scales full-field predictions to challenging, large-scale problems.

Figures & tables

Figure 1: One RWR block. The encoder and decoder are pointwise MLPs; the processor runs

m_{1}

message-passing sweeps ( relax ),

r

read-write iterations against

L

persistent latent tokens, and

m_{2}

further sweeps;

n_{B}

such blocks are stacked. Setting

r{=}0

recovers MGN, and setting

m_{1}{=}m_{2}{=}0

recovers RWP. MP denotes message passing, CA cross-attention, and SA self-attention.

Panel of Figure 2 — Figure 2: The two simulation environments, branched pipe (left) and centrifugal pump (right), showing a representative velocity field.

	read	process	write	write-back freq.
Perceiver IO	queries	SA	out queries	once
Transolver family	slice weights	SA	tied weights	layer
LNO	learned PE	SA	inverse proj.	once
UPT	MP + pool	SA	perceiver	once
AROMA	cross-attn	DiT	query-local	once
FLARE	queries/head	—	roles swapped	layer

Table 1: The BSA family, Eq. ( 2 ). All members route the field update through

L\ll N

latent tokens and inherit the rank bound; they differ in how the read and write are parameterized and how often the write recurs (once at decode, every layer, or every iteration). RWP’s latents persist across iterations within a block, and its read may further be restricted to the boundary points. The Transolver family row covers Transolver, Transolver++, and GeoTransolver, which share the slice read and write. SA denotes self-attention over the latent tokens, DiT a diffusion transformer, PE a positional embedding, and MP message passing; in FLARE the read and write swap the query and key-value roles between the same projections.

Figure 3: Normalized error spectra over training on held-out branched pipe cases. Each panel shows $\sqrt{E_{\mathrm{err}}(k)/E_{\mathrm{err}}^{(0)}(k)}$ , the error amplitude at wavenumber $k$ normalized by the initial error of RWP, for the two velocity components (rows) under MGN, RWP, and RWR (columns) at five training epochs. Dashed lines mark the low/mid/high band cutoffs.

Table 6

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

	Branched pipe	Centrifugal pump
		(casing)	(full)
cell number	${\sim}60$ k	${\sim}270$ k	${\sim}1.4$ M
fluid	air	water	water
turbulence model	RANS	RANS	RANS
boundary conditions	velocity inlet, pressure outlet, no-slip wall
solver	FVM (Fluent)

Appendix

Table 5: Summary of simulation configurations.

Parameter	Range
Branched pipe
angle of middle baffle plate	0–60 (deg)
aperture length of inlet baffle plate	20–37 (mm)
angle between ducts	75–120 (deg)
curvature angle of duct branch	90–180 (deg)
inlet mean velocity	3–10 (m/s)

Appendix

Table 6: Design parameters used to generate the simulation datasets, with the range each parameter is sampled over.

	Branched pipe	Centrifugal pump
modeling task	steady-state	steady-state
total samples	300	300
training-set sizes	50 / 100 / 300 samples
train/val/test split	0.7/0.15/0.15
random seed	42
field normalization	per-component standardization

Appendix

Table 7: Modeling task and dataset specification, common to all models. The two pump extents share the same design of experiments and splits.

Figure 5: Band-resolved test error of RWP as a function of the number of latents $L$ on the branched pipe. Each panel shows the fRMSE over the low, mid, and high wavenumber bands for pressure and the two velocity components, with $L$ swept from 8 to 512 and all other settings fixed.

Figure 6: Band-resolved training dynamics on the branched pipe. Each panel tracks the fRMSE of the low, mid, and high wavenumber bands over training for the two velocity components (rows) under MGN, RWP, and RWR (columns), normalized by the initial band value of RWP.

		rel. $L_{2}$		fRMSE
Field	Model	field	grad.	low	mid	high
$p$	Transolver	0.090 [0.85]	0.165 [0.95]	0.118 [0.99]	0.129 [0.94]	0.133 [0.87]
	GeoTransolver	0.135 [1.28]	0.212 [1.22]	0.159 [1.33]	0.168 [1.22]	0.174 [1.14]
	RWP	0.134 [1.27]	0.199 [1.14]	0.155 [1.30]	0.177 [1.29]	0.191 [1.25]
	RWR	0.106	0.174	0.119	0.137	0.153
$u_{x}$	Transolver	0.078 [1.01]	0.174 [1.18]	0.035 [0.95]	0.070 [1.10]	0.129 [1.37]

Appendix

Table 8: Held-out test errors on the branched pipe at matched training budget for all output fields: relative

L_{2}

error of each field and of its first spatial gradients, and band-resolved fRMSE over the low, mid, and high wavenumber bands. Bracketed values give the ratio of each model’s error to RWR’s; bold marks the best value per metric and field. The pressure results are discussed in App. C .

Steps	$p$	$u_{x}$	$u_{z}$
$(8,4,0){\times}1$	5.8	5.8	4.8
$(0,4,8){\times}1$	6.0	5.3	4.7
$(2,2,2){\times}2$	5.2	5.2	4.5
$(1,1,1){\times}4$	5.1	5.1	4.5

Appendix

Table 9: Composition ablation on the branched pipe at a fixed budget of eight relaxation sweeps and four latent iterations: held-out test normalized mean absolute error (NMAE,

\times 10^{-2}

) of pressure and the two velocity components. A configuration

(m_{1},r,m_{2}){\times}n_{B}

runs

m_{1}

relaxation sweeps,

r

latent iterations, and

m_{2}

further sweeps per block, repeated over

n_{B}

blocks:

(8,4,0)

places all relaxation before the latent path,

(0,4,8)

after it, and the remaining rows interleave the two at increasing granularity.

Figure 7: Parity plots of the wall-normal derivatives $\partial u_{x}/\partial n$ and $\partial u_{z}/\partial n$ over wall cells on held-out branched pipe cases, with the coefficient of determination $R^{2}$ per model.

Figure 8: Mean absolute error of pressure and the velocity components as a function of wall distance $d_{w}$ on held-out branched pipe cases (top row), with a near-wall zoom (bottom row).

	RWP	RWR
steps $(m_{1},r,m_{2})\times n_{B}$	$(0,2,0)\times 2$	$(2,2,2)\times 2$
latents $L$	128	128
hidden width	192	128
attention heads	4	4
latent self-attention	none	none
boundary-only read	yes	yes

Appendix

Table 10: Architectural configuration of RWP and RWR for the branched pipe analysis. Hidden widths are chosen to approximately match parameter counts across MGN, RWP, and RWR.

	Transolver			GeoTransolver			RWR					RWP
#samples	#params	#slices	width	#params	#slices	width	#params	#latents	blocks	read	width	#params	#latents	blocks	read	width
Branched pipe ( $\mathbf{{\sim}60}$ k cells)
50	3.58M	32	172	4.02M	32	144	3.58M	0 64	2	bdy	128	3.87M	512	2	bdy	128
100	3.58M	32	172	4.02M	16	144	2.79M	128	2	bdy	128	5.65M	128	3	bdy	128
300	7.72M	32	256	9.53M	32	256	5.30M	128	3	bdy	128	2.52M	128	3	bdy	128
Pump casing ( $\mathbf{{\sim}270}$ k cells)

Appendix

Table 11: Architectural details of best performing models per family for each mesh and data setting.

	NMAE ( $\times 10^{-2}$ )
Model	$p$	$u_{x}$	$u_{y}$	$u_{z}$	mean	#params	peak memory	average time/epoch	#epochs until avg. $R^{2}\geq 0.95$
Branched pipe ( $\mathbf{{\sim}60}$ k cells)
50 samples
Transolver	10.3	12.7		10.8	11.3	3.58M	0 8.1 GB	8s	470
GeoTransolver	9.5	11.8		10.2	10.5	4.02M	0 8.0 GB	7s	370
RWP	10.4	10.7		9.7	10.3	3.87M	0 4.2 GB	6s	290

Appendix

Table 12: Performance metrics of best models per family for each mesh and data setting. We report the NMAE per output field, across-field mean NMAE, parameter count, peak memory, average wall-clock time per epoch, and the number of epochs to reach a mean

R^{2}\geq 0.95

(n/r: not reached). The branched pipe is quasi-2D, so

u_{y}

is absent.

Panel of Figure 9 — Figure 9: Predicted velocity magnitude on the branched pipe for two held-out test cases (top and bottom groups), shown on the mid- $y$ slice. The left column shows the ground truth; each model column shows the prediction, with its $R^{2}$ and NMAE, above the signed pointwise error.

Panel of Figure 10 — Figure 10: Predicted pressure on the branched pipe for two held-out test cases (top and bottom groups), shown on the mid- $y$ slice. The left column shows the ground truth; each model column shows the prediction, with its $R^{2}$ and NMAE, above the signed pointwise error.

Panel of Figure 11 — Figure 11: Predicted velocity magnitude on the pump casing for two held-out test cases (top and bottom groups), shown on the $y=0.04$ m plane. The left column shows the ground truth; each model column shows the prediction, with its $R^{2}$ and NMAE, above the signed pointwise error.

Panel of Figure 12 — Figure 12: Predicted pressure on the pump casing for two held-out test cases (top and bottom groups), shown on the $y=0.04$ m plane. The left column shows the ground truth; each model column shows the prediction, with its $R^{2}$ and NMAE, above the signed pointwise error.

Panel of Figure 13 — Figure 13: Predicted velocity magnitude on the full pump mesh for two held-out test cases (top and bottom groups), shown on the impeller surface. The left column shows the ground truth; each model column shows the prediction, with its $R^{2}$ and NMAE, above the signed pointwise error.

Panel of Figure 14 — Figure 14: Predicted pressure on the full pump mesh for two held-out test cases (top and bottom groups), shown on the impeller surface. The left column shows the ground truth; each model column shows the prediction, with its $R^{2}$ and NMAE, above the signed pointwise error.

Explore similar work

Oct 12, 2026

Compute-optimal data scaling for neural surrogates via multi-fidelity training

Neural surrogates for Partial Differential Equations (PDEs) are usually trained on synthetic data from numerical simulators. As models grow and scaling laws demand larger datasets, the cost of data generation often outweighs the training cost and becomes the bottleneck in the scaling quest. Unlike data in other domains, however, each PDE sample has a tunable cost that depends on its numerical fidelity, with the resulting cost-accuracy trade-off governed by known convergence rates. We exploit this property by formalizing neural surrogate training as a function of both data quantity and fidelity under a fixed data generation compute budget, treating fidelity as a fine-grained discrete spectrum rather than a binary distinction. Building on this, we show that multi-fidelity training enables Pareto-optimal error scaling against data generation compute. By using simulations at lower fidelities as an additional training signal, training on the optimal fidelity mixture matches the accuracy of high-fidelity-only training, with up to

6.3\times

less data generation compute across our settings. We validate these findings across six PDE problems spanning structured grids and unstructured meshes, two distinct error sources (discretization and iterative-solver error), and four surrogate architectures, totaling over 7,000 training runs.

Paul Setinek, Fabian Paischer, Nils Thuerey +2

Institute for Machine Learning, JKU Linz · Mistral AI · Technical University of Munich +1

Neural Surrogate Modeling PDE Surrogate Modeling Multifidelity Modeling

Ask PDF

Mar 2, 2026

Accelerating PDE Surrogates via RL-Guided Mesh Optimization

Deep surrogate models for parametric partial differential equations (PDEs) can deliver high-fidelity approximations but remain prohibitively data-hungry: training often requires thousands of fine-grid simulations, each incurring substantial computational cost. To address this challenge, we introduce RLMesh, an end-to-end framework for efficient surrogate training under limited simulation budget. The key idea is to use reinforcement learning (RL) to adaptively allocate mesh grid points non-uniformly within each simulation domain, focusing numerical resolution in regions most critical for accurate PDE solutions. A lightweight proxy model further accelerates RL training by providing efficient reward estimates without full surrogate retraining. Experiments on PDE benchmarks demonstrate that RLMesh achieves competitive accuracy to baselines but with substantially fewer simulation queries. These results show that solver-level spatial adaptivity can dramatically improve the efficiency of surrogate training pipelines, enabling practical deployment of learning-based PDE surrogates across a wide range of problems.

Yang Meng, Ruoxi Jiang, Zhuokai Zhao +3

UChicago · Fudan University · UAlbany

PDE Surrogate Modeling Data-Efficient Learning Adaptive Mesh Refinement

Ask PDF

May 2, 2026

Mesh Based Simulations with Spatial and Temporal awareness

Machine Learning surrogates for Computational Fluid Dynamics (CFD), particularly Graph Neural Networks (GNNs) and Transformers, have become a new important approach for accelerating physics simulations. However, we identify a critical bottleneck in the field: while architectures have advanced significantly, the common underlying training paradigms remain bound to naive assumptions, such as node-wise supervision and explicit Euler time-stepping. These legacy choices ignore the stiff dynamics and local flux continuity inherent to numerous partial differential equations resolution methods, such as Finite Element, Difference, or Volume (FEM). In this work, we propose a unified framework to bridge the gap between geometric deep learning and rigorous numerical analysis. We introduce three key innovations: (1) Multi Node Prediction, a stencil-level objective that predicts field values for a node's full local topology, enforcing spatial derivative consistency; (2) Temporal Correction, replacing unstable explicit schemes with a predictor-corrector via temporal Cross-Attention; and (3) Geometric Inductive Biases, leveraging 3D Rotary Positional Embeddings (RoPE) to robustly capture rotational symmetries in unstructured meshes. We evaluate this framework across three architectures (MeshGraphNet, Transolver, and a Transformer) on diverse physics datasets. Our approach yields consistent improvements in accuracy and stability, particularly in long-horizon rollouts, while producing latent representations that generalize to unseen subtasks such as Wall Shear Stress or Pressure prediction. Code is available at https://github.com/DonsetPG/graph-physics.

Paul Garnier, Vincent Lannelongue, Elie Hachem

1CEMEF - Mines Paris PSL

Graph Neural Networks Neural PDE Solvers PDE Surrogate Modeling Computational Fluid Dynamics

Ask PDF