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

Compute-optimal data scaling for neural surrogates via multi-fidelity training
Oct 12, 2026Paul Setinek, Fabian Paischer, Nils Thuerey +2Neural Surrogate Modeling PDE Surrogate Modeling
Institute for Machine Learning, JKU Linz · Mistral AI · Technical University of Munich +1
Accelerating PDE Surrogates via RL-Guided Mesh Optimization
Mar 2, 2026Yang Meng, Ruoxi Jiang, Zhuokai Zhao +3PDE Surrogate Modeling Data-Efficient Learning
UChicago · Fudan University · UAlbany
Mesh Based Simulations with Spatial and Temporal awareness
May 2, 2026Paul Garnier, Vincent Lannelongue, Elie HachemGraph Neural Networks Neural PDE Solvers
1CEMEF - Mines Paris PSL

Oct 12, 2026J/K move · Enter open · S save

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

Paul Setinek, Fabian Paischer, Nils Thuerey, Johannes Brandstetter+1

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

Neural Surrogate Modeling PDE Surrogate Modeling Multifidelity Modeling

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Accelerating PDE Surrogates via RL-Guided Mesh Optimization

Mesh Based Simulations with Spatial and Temporal awareness