Simple yet Effective: Low-Rank Spatial Attention for Neural Operators

Zherui Yang; Haiyang Xin; Tao Du; Ligang Liu

cs.LGApr 4, 2026

Simple yet Effective: Low-Rank Spatial Attention for Neural Operators

Authors: Zherui Yang, Haiyang Xin, Tao Du, Ligang Liu

Organizations: University of Science and Technology of China · Tsinghua University · Shanghai Qi Zhi Institute

Abstract

Neural operators have emerged as data-driven surrogates for solving partial differential equations (PDEs), and their success hinges on efficiently modeling the long-range, global coupling among spatial points induced by the underlying physics. In many PDE regimes, the induced global interaction kernels are empirically compressible, exhibiting rapid spectral decay that admits low-rank approximations. We leverage this observation to unify representative global mixing modules in neural operators under a shared low-rank template: compressing high-dimensional pointwise features into a compact latent space, processing global interactions within it, and reconstructing the global context back to spatial points. Guided by this view, we introduce Low-Rank Spatial Attention (LRSA) as a clean and direct instantiation of this template. Crucially, unlike prior approaches that often rely on non-standard aggregation or normalization modules, LRSA is built purely from standard Transformer primitives, i.e., attention, normalization, and feed-forward networks, yielding a concise block that is straightforward to implement and directly compatible with hardware-optimized kernels. In our experiments, such a simple construction is sufficient to achieve high accuracy, yielding an average error reduction of over 17% relative to second-best methods, while remaining stable and efficient in mixed-precision training.

Figures & tables

Figure 1: Compressibility of PDE interactions and the unified low-rank paradigm. Left: (a) original dense kernel, i.e., the Green function of 1D Poisson; (b–d) underlying low-rank properties, including: reconstructed kernel, fast spectral decay, and approximation error, derived from the numerical factorization illustrated in the middle. Middle: numerical low-rank approximation of global interactions via $K_{r}\approx U_{r}\Sigma_{r}V_{r}^{\top}$ . Right: diverse neural-operator global-mixing modules unified as a learnable compress–process–reconstruct template.

Figure 2: Overview of the neural operator backbone and the Low-Rank Spatial Attention (LRSA) block. LRSA routes global information through a compact latent bottleneck using only standard Transformer primitives.

Figure 3: Qualitative performance comparison across diverse discretizations. From top-left to bottom-right: Navier-Stokes (regular grid), Elasticity (point cloud), Airfoil and Plasticity (structured grid). Error maps are visualized on the same scale for each task. LRSA yields lower relative errors and preserves sharper physical patterns in high-frequency regions compared to Transolver.

Method	Structured Mesh			Regular Grid		Point Cloud
Method	Airfoil	Pipe	Plasticity	Navier–Stokes	Darcy	Elasticity
FNO ( 2020a )	–	–	–	0.1556	0.018	–
F-FNO ( 2023 )	0.0078	0.0050	0.0047	0.2231	0.0077	0.0263
LSM ( 2023 )	0.0059	0.0050	0.0025	0.1535	0.0065	0.0218
HPM ( 2025 )	0.0047	0.0030	0.0008	0.0734	0.0046	–
LNO ( 2024 )	0.0053	0.0031	0.0028	0.0734	0.0063	0.0066

Table 1: Standard benchmarks. Comparison of test relative

L_{2}

errors on six standard PDE benchmarks. Bold indicates the best result, and an underscore indicates the second-best result. “–” indicates the method is not applicable.

Method	Irregular	Pipe	Heat	Comp.
Method	Darcy	Turbulence	Transfer	Comp.
GraphSAGE	6.73e-2	2.36e-1	–	2.09e-1
DeepONet	1.36e-2	9.36e-2	7.20e-4	1.88e-2
NORM	1.05e-2	1.01e-2	1.11e-3	1.00e-2
HPM	7.36e-3	8.26e-3	1.84e-4	9.34e-3
Ours	7.56e-3	4.89e-3	1.49e-4	8.71e-3

Table 2: Irregular domain benchmarks. Relative

L_{2}

error is reported in the table.

Model	AirfRANS ( Bonnet et al., 2022 )				ShapeNet Car ( Umetani and Bickel, 2018 )
Model	Volume $\downarrow$	Surface $\downarrow$	$C_{L}\downarrow$	$\rho_{L}\uparrow$	Volume $\downarrow$	Surface $\downarrow$	$C_{D}\downarrow$	$\rho_{D}\uparrow$
GraphSAGE ( 2017 )	0.0087	0.0184	0.1476	0.9964	0.0461	0.1050	0.0270	0.9695
GraphUNet ( 2019 )	0.0076	0.0144	0.1677	0.9949	0.0471	0.1102	0.0226	0.9725
MeshGraphNet ( 2020 )	0.0214	0.0387	0.2252	0.9945	0.0354	0.0781	0.0168	0.9840
GNO ( 2020b )	0.0269	0.0405	0.2016	0.9938	0.0383	0.0815	0.0172	0.9834
GEO-FNO ( 2022b )	0.0361	0.0301	0.6161	0.9257	0.1670	0.2378	0.0664	0.8280

Table 3: Performance comparison on AirfRANS and ShapeNet Car datasets. In addition to the field-level MSE for Volume and Surface quantities, we report the relative errors for lift (

C_{L}

) and drag (

C_{D}

) coefficients, along with their Spearman’s rank correlations (

\rho_{L},\rho_{D}

). A correlation close to

1.0

indicates that the model correctly preserves design rankings, which is critical for optimization tasks. Lower values (

\downarrow

) are better for error metrics, while higher values (

\uparrow

) are better for correlation coefficients.

Model	221 $\times$ 51	111 $\times$ 51	221 $\times$ 26	111 $\times$ 26
Transolver	0.524	7.850	1.260	7.680
HPM	0.438	0.537	1.690	1.740
Ours	0.423	0.468	0.644	0.673

Table 4: Relative

L_{2}

error (

\times 10^{-2}

) comparison across different spatial resolutions. The models were trained on the

221\times 51

mesh. Bold indicates the best performance.

Prob.	Training	Transolver	LNO	HPM	Ours
Prob.	Number	Transolver	LNO	HPM	Ours
Darcy	200	1.75	5.62	1.10	0.97
	600	0.69	4.40	0.60	0.50
	1000	0.52	0.63	0.44	0.43
NS	200	37.6	40.0	18.1	26.2
	600	31.4	25.0	11.7	7.50

Table 5: Relative

L_{2}

errors (

\times 10^{-2}

) with limited training samples.

Figure 4: Training stability and efficiency. Left: relative

L_{2}

error under FP32/BF16/FP16;

\boldsymbol{\times}

denotes divergence. Right: per-step training latency (forward+backward, normalized to Transolver-FP32) and peak memory (ratio relative to Transolver-FP32; smaller is better) on three representative tasks. Memory Saving is calculated as the ratio of peak training memory consumption of the evaluated model to that of the baseline model (Transolver-FP32). A lower factor indicates better memory efficiency.

Figure 5: Rank and component ablations. Top: sensitivity to latent size

M

. Bottom: component variants of LRSA (Full, w/o latent self-attention, and enforcing symmetric compression and reconstruction).

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Test Case	Geometry	#Dim	Mesh Size	Dataset Split		Input Type
Test Case	Geometry	#Dim	Mesh Size	Train	Test	Input Type
Elasticity ( Li et al., 2020a )	Point Cloud	2	972	1000	200	Domain Shape
Darcy ( Li et al., 2020a )	Regular Grid	2	$85\times 85$	1000	200	Coefficients
Navier Stokes ( Li et al., 2020a )	Regular Grid	2+1	$64\times 64$	1000	200	Previous States
Plasticity ( Li et al., 2022b )	Structured Mesh	2+1	$101\times 31$	900	80	External Force
Airfoil ( Li et al., 2022b )	Structured Mesh	2	$221\times 51$	1000	200	Boundary Shape

Appendix

Table 6: Overview of Standard Benchmarks. Details on geometry type, spatial dimension, discretization size, dataset splits, and input features.

Test Case	Geometry	# Dim	Num. Nodes ( $N$ )	Dataset Split		Input Type
Test Case	Geometry	# Dim	Num. Nodes ( $N$ )	Train	Test	Input Type
Irregular Darcy Chen et al. (2024)	Point Cloud	2	2290	1000	200	Coefficients
Pipe Turbulence Chen et al. (2024)	Point Cloud	2	2673	300	100	Previous State
Heat Transfer Chen et al. (2024)	Point Cloud	3	7199	100	100	Boundary Shape
Composite Chen et al. (2024)	Point Cloud	3	8232	400	100	Temperature

Appendix

Table 7: Overview of Irregular Domain Benchmarks.

Test Case	Geometry	#Dim	Mesh Size	Dataset Split		Input Type
Test Case	Geometry	#Dim	Mesh Size	Train	Test	Input Type
ShapeNet Car ( Umetani and Bickel, 2018 )	Unstructured Mesh	3	32186	789	100	Structure
AirfRANS ( Bonnet et al., 2022 )	Unstructured Mesh	2	32000	800	200	Structure

Appendix

Table 8: Datasets and their split used in our experiments.

Problems	Model Configurations				Training Configurations
Problems	Depth	Width	#Heads	$M$	#Params	Max Lr	Weight Decay	Epochs	Batch Size	Loss
Elasticity	8	128	8	64	2.1M	1e-3	1e-5	500	1
Navier Stokes	8	256	8	32	12.0M	1e-3	1e-5	500	2	Relative
Plasticity	8	128	8	64	2.8M	1e-3	1e-5	500	8	L2
Airfoil	8	128	4	64	2.8M	1e-3	1e-5	500	4
Pipe	8	128	4	32	2.8M	1e-3	1e-5	500	4

Appendix

Table 9: Training and model configurations of LRSA. Training configurations are directly from previous works without extra tuning ( Bonnet et al., 2022 ; Hao et al., 2023 ; Deng et al., 2024 ) . Here

\mathcal{L}_{\mathrm{v}}

and

\mathcal{L}_{\mathrm{s}}

represent the mse loss on volume and surface fields respectively. As for Darcy, we adopt an additional spatial gradient regularization term

\mathcal{L}_{\mathrm{g}}

following ONO ( Xiao et al., 2024 ) .

Dataset	Airfoil	Pipe	Plasticity	Navier-Stokes	Darcy	Elasticity
Transolver	0.53 $\pm$ 0.01	0.33 $\pm$ 0.02	0.12 $\pm$ 0.01	9.00 $\pm$ 0.13	0.57 $\pm$ 0.01	0.64 $\pm$ 0.02
Ours	0.42 $\pm$ 0.03	0.23 $\pm$ 0.01	0.046 $\pm$ 0.001	4.84 $\pm$ 0.09	0.43 $\pm$ 0.01	0.33 $\pm$ 0.01

Appendix

Table 10: Statistical Mean and Std : Relative

L_{2}

error is reported in the table. (

\times 10^{-2}

)

Method	Navier–Stokes			Airfoil			ShapeNet Car
Method	Mem	Fwd	Bwd	Mem	Fwd	Bwd	Mem	Fwd	Bwd
Configurations	$N=4096,M=64,B=20$			$N=11271,M=64,B=20$			$N=32186,M=64,B=8$
Transolver-FP32	8.7	70.2	129.3	18.4	127.9	253.0	28.2	163.3	273.3
Transolver-BF16 †	6.9	43.2	85.3	17.1	84.7	214.3	20.5	88.2	206.4
Transolver-FP16 †	6.9	43.4	85.2	17.1	82.2	214.3	20.5	88.1	206.5
Ours-FP32	10.5	48.0	98.1	12.1	53.2	178.4	26.2	103.3	238.4

Appendix

Table 11: Detailed Efficiency Breakdown. Comparison of Peak Memory (Mem, in GB), Forward Latency (Fwd, in ms), and Backward Latency (Bwd, in ms) across different precisions. Entries marked with

\dagger

indicate configurations that resulted in training divergence (NaN loss) or severe accuracy degradation.

Variants \ Dataset	Airfoil	Pipe	Plasticity	Navier-Stokes	Darcy	Elasticity
Ablation: Components
w/o Intra Attention	4.30e-3	2.64e-3	7.09e-4	5.44e-2	4.46e-3	3.60e-3
Enforce Symmetric	4.60e-3	2.50e-3	4.71e-4	5.93e-2	4.63e-3	4.70e-3
Ablation: Latent Rank $M$
8	4.32e-3	3.01e-3	4.87e-4	1.03e-2	4.85e-3	3.71e-3
16	4.47e-3	2.32e-3	4.74e-4	9.10e-2	4.75e-3	3.47e-3

Appendix

Table 12: Full Ablation Results. Relative

L_{2}

errors across standard benchmarks under various ablation settings.

Neural Operators

Neural operators learn mappings between functions, enabling fast surrogates for PDEs and other continuous physical systems.

See all 304 papers

LLT: Local Linear Transformer for PDE Operator Learning
Jul 4, 2026Oded Ovadia, Eli TurkelPDE Surrogate Modeling Transformer Attention
School of Mathematical Sciences, Tel Aviv University, Tel Aviv, Israel
SFO: Learning PDE Operators via Spectral Filtering
Jan 23, 2026Noam Koren, Rafael Moschopoulos, Kira Radinsky +1PDE Surrogate Modeling PDE Operator Learning
Department of Computer Science Technion - Israel Institute of Technology · Department of Computer Science Princeton University
MoNo: Multiscale Optimal Transport Neural Operator for Solving PDEs on General Geometries
Aug 10, 2026Zijiang Yang, Xiaomeng Wu, Dongmei FuNeural PDE Solvers PDE Operator Learning
School of Automation and Electrical Engineering, University of Science and Technology Beijing · 2Beijing Engineering Research Center of Industrial Spectrum Imaging

Jul 4, 2026J/K move · Enter open · S save

LLT: Local Linear Transformer for PDE Operator Learning

Oded Ovadia, Eli Turkel

School of Mathematical Sciences, Tel Aviv University, Tel Aviv, Israel

PDE Surrogate Modeling Transformer Attention Linear Attention PDE Operator Learning Neural Operators Transformer Neural Operator

Open paper AskPDF arXiv

Neural operators have become a common approach for learning PDE solution maps and accelerating numerical simulations. Transformer-based neural operators are of particular interest, since attention can learn long-range dependencies in the computational domain. However, standard attention has two major limitations when applied to PDEs: it scales quadratically with the number of computational nodes, and it lacks an explicit bias toward local interactions. To address these issues, we introduce Local Linear Transformer (LLT) for PDE operator learning. The architecture combines linear global attention with local spatial mixing, and incorporates coordinate and geometry information. We evaluate LLT on several PDE problems, including elasticity, plasticity, airfoil flow, pipe flow, and Darcy flow. The reference data for these problems span finite-element, finite-volume, and finite-difference discretizations on structured and unstructured meshes. Compared with other neural-operator and transformer baselines from prior studies, LLT achieves competitive or lower relative

L_2

error across these problems. On matched structured discretizations, wall-clock time per training iteration is reduced by factors of 1.8 to 2.5 relative to Transolver. We also scale the approach and apply it to a three-dimensional car aerodynamics dataset with 32,186 unstructured mesh points per sample. Together, these results indicate that LLT provides an accurate and computationally efficient operator for PDE problems across discretizations, mesh types, and problem settings.

Abstract

Neural Operators

Explore similar work

LLT: Local Linear Transformer for PDE Operator Learning

SFO: Learning PDE Operators via Spectral Filtering

MoNo: Multiscale Optimal Transport Neural Operator for Solving PDEs on General Geometries