Disentangling Attention in Deep Operator Learning: A Controlled Study of Data-Driven and Physics-Informed Architectures

Amar Alem Koric; Qibang Liu; Seid Koric

cs.LGSep 3, 2026

Disentangling Attention in Deep Operator Learning: A Controlled Study of Data-Driven and Physics-Informed Architectures

Authors: Amar Alem Koric, Qibang Liu, Seid Koric

Organizations: Department of Mechanical Engineering, Stanford University, Stanford, CA, 94305, USA · National Center for Supercomputing Applications, University of Illinois at Urbana-Champaign, 1205 W. Clark St., Urbana, IL, 61801, USA · The Grainger College of Engineering, Mechanical Science and Engineering, University of Illinois at Urbana-Champaign, 1206 W. Green St., Urbana, IL, 61801, USA

Abstract

Deep neural operators learn mappings between input functions and complete PDE solution fields, enabling forward evaluations of new problem instances orders of magnitude faster than conventional numerical solvers. Attention mechanisms have recently been introduced into neural operators, but most studies change several architectural components at once, making it difficult to identify what actually improves accuracy. This work presents a controlled and systematic study of five deep operator network (DeepONet) variants with distinct attention mechanisms, trained under both data-driven and physics-informed regimes, to isolate the effects of cross-attention, self-attention, tokenization, and attention depth. We evaluate them on a source-driven transient one-dimensional nonlinear diffusion-reaction equation, a transient one-dimensional viscous Burgers equation with variable initial conditions, and a two-dimensional Poisson heat-conduction problem with heterogeneous source fields. Per-sensor tokenization with cross-attention reduces the mean relative L_2 error of the classical DeepONet in all benchmark-training combinations by factors of 2.4-28.0, while the best attention configurations reach 3.5-32.3. Branch self-attention paired only with dot-product fusion is inconsistent, degrading the one-dimensional problems while helping the more complex two-dimensional source field; added on top of cross-attention it improves all six cases, though by less than cross-attention fusion alone. Global pre-mixing provides no consistent benefit. Increasing cross-attention depth further improves accuracy, but with diminishing returns and a substantially higher cost under physics-informed training. Overall, query-dependent cross-attention is the most reliable mechanism, whereas branch self-attention is most useful for large, spatially complex functional inputs.

Figures & tables

Figure 1: The five variants, one per blue rectangle. The red panel defines the attention block used wherever attention appears above: multi-head attention followed by a residual addition and a layer normalization. Cross-attention takes the query from the trunk and the keys and values from the branch tokens, with the residual carrying the projected query $Q$ . Self-attention forms all three from the branch tokens and takes the residual from the token stream before those projections.

Setting	Diff.-react.	Burgers	Poisson
Latent dimension $p$	128	128	128
Attention heads	4	4	4
Trunk hidden layers	$4\times 128$	$4\times 128$	$4\times 128$
Branch sensors $m$	100	101	1024
Activation	tanh	tanh	tanh
Training iterations	120,000	120,000	80,000

Table 1: Shared training and architecture settings. All five variants within a given benchmark and paradigm use identical values.

	Tr. diff.-react.	Tr. Burgers	Poisson
Variant	DD $=$ PI	DD $=$ PI	DD	PI
V1 dot product	112,384	112,512	230,656	230,656
V2 tokens $+$ CA	182,145	182,145	182,145	181,889
V3 SA $+$ dot	165,504	165,504	165,504	165,504
V4 global $+$ CA	244,737	244,865	363,009	362,753
V5 SA $+$ CA	231,553	231,553	231,553	231,297

Table 2: Trainable parameter counts.

Variant	Mean relative $L_{2}$ error ( $\times 10^{-2}$ )
	Data-driven			Physics-informed
	Tr. diff.-react.	Tr. Burgers	Poisson	Tr. diff.-react.	Tr. Burgers	Poisson
V1 dot product	$1.38\pm 0.727$	$4.06\pm 2.17$	$29.5\pm 17.6$	$0.764\pm 0.426$	$12.4\pm 7.10$	$33.4\pm 24.1$
V2 tokens $+$ CA	$0.346\pm 0.160$ $4.0\times$	$0.753\pm 0.508$ $5.4\times$	$1.05\pm 0.613$ $28.0\times$	$0.312\pm 0.156$ $2.4\times$	$1.58\pm 1.06$ $7.9\times$	$6.15\pm 2.77$ $5.4\times$
V3 SA $+$ dot	$3.98\pm 2.28$ $0.35\times$	$4.59\pm 2.31$ $0.88\times$	$3.81\pm 2.20$ $7.7\times$	$1.81\pm 1.07$ $0.42\times$	$19.2\pm 8.36$ $0.64\times$	$3.96\pm 2.39$ $8.4\times$
V4 global $+$ CA	$0.399\pm 0.170$ $3.5\times$	$0.536\pm 0.427$ $7.6\times$	$1.43\pm 0.956$ $20.6\times$	$\mathbf{0.221\pm 0.105}$ $3.5\times$	$\mathbf{0.704\pm 0.679}$ $17.6\times$	$7.04\pm 4.57$ $4.7\times$

Table 3: Mean relative

L_{2}

error on held-out inputs

\pm

standard deviation across those inputs, in units of

10^{-2}

. The smaller number beneath each entry is the improvement factor over the V1 dot-product baseline, i.e. the ratio of the V1 mean to that variant’s mean. Best mean in each column in bold. “Tr.” denotes transient, i.e. the time-dependent diffusion–reaction and Burgers benchmarks; the Poisson benchmark is steady.

Figure 2: Data-driven transient 1D diffusion–reaction. Three held-out source functions, chosen as the best, median and worst case for the V1 dot-product baseline. In each block the top row is the exact space-time solution $u(x,t)$ followed by the five variants’ predictions, and the bottom row is the input source $f(x)$ followed by each variant’s signed error (exact $-$ prediction; white is exact). The error color scale is common to all three blocks; the value beneath each error panel is that variant’s relative $L_{2}$ error on the case shown.

Figure 3: Physics-informed transient 1D diffusion–reaction, laid out as in Fig. 2 .

Figure 4: Data-driven transient 1D Burgers. Three held-out initial conditions, chosen as the best, median and worst case for the V1 dot-product baseline. In each block the top row is the reference space-time solution $u(x,t)$ followed by the five variants’ predictions, and the bottom row is the initial condition $u_{0}(x)$ followed by each variant’s normalised signed error $(\mathrm{exact}-\mathrm{prediction})/\max|\mathrm{exact}|$ , where the normalization (Norm) is by the peak magnitude of the exact field for that case; white is exact. The value beneath each error panel is that variant’s relative $L_{2}$ error on the case shown.

Figure 5: Physics-informed transient 1D Burgers, laid out as in Fig. 4 .

Figure 6: Data-driven 2D Poisson heat conduction. Three held-out source fields, chosen as the best, median and worst case for the V1 dot-product baseline. In each block the top row is the reference solution $u(x,y)$ followed by the five variants’ predictions, and the bottom row is the input source $f(x,y)$ followed by each variant’s signed error (exact $-$ prediction; white is exact). The error colour scale is common to all three blocks; the value beneath each error panel is that variant’s relative $L_{2}$ error on the case shown.

Figure 7: Physics-informed 2D Poisson heat conduction, laid out as in Fig. 6 .

	$N=1$	$N=2$	$N=4$	$N=8$
Parameters	182,401	231,809	330,625	528,257
Inference (ms)	3.05	3.97	5.82	9.62
Data-driven
Mean $L_{2}$ ( $\times 10^{-3}$ )	$3.343\pm 1.796$	$2.360\pm 1.072$	$2.287\pm 1.054$	$1.626\pm 0.698$
Training time	4m 10s	4m 02s	4m 23s	6m 52s
Physics-informed

Table 4: Effect of stacking

N

cross-attention blocks on the transient diffusion–reaction benchmark, under both training paradigms. Each error entry is the mean relative

L_{2}

error over the

100

held-out source functions

\pm

one standard deviation across those functions.

	Data-driven			Physics-informed
Variant	Tr. diff.-react.	Tr. Burgers	Poisson	Tr. diff.-react.	Tr. Burgers	Poisson
Panel A: training wall-clock time
V1 dot product	1m 49s	1m 50s	1m 11s	5m 11s	5m 59s	2m 14s
V2 tokens $+$ CA	2m 57s	2m 43s	3m 01s	31m 21s	27m 49s	16m 07s
V3 SA $+$ dot	2m 18s	2m 20s	15m 14s	6m 39s	7m 18s	16m 16s
V4 global $+$ CA	3m 18s	3m 20s	3m 08s	23m 57s	28m 51s	16m 15s

Table 5: Training wall-clock time and inference time for a full-field forward pass, on a single NVIDIA H200 GPU. Panel A: training. Panel B: inference.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Held-out error distributions for the data-driven diffusion-reaction models over 100 matched test samples. Left: all five variants. Right: the V1 dot-product baseline, mean relative $L_{2}$ error $1.38\times 10^{-2}$ , against the best variant V5 at $3.19\times 10^{-3}$ .

Figure 9: Held-out error distributions for the physics-informed diffusion-reaction models over 100 matched test samples. Left: all five variants. Right: the V1 dot-product baseline, mean relative $L_{2}$ error $7.64\times 10^{-3}$ , against the best variant V4 at $2.21\times 10^{-3}$ .

Figure 10: Held-out error distributions for the data-driven Burgers models over 400 matched test samples. Left: all five variants. Right: the V1 dot-product baseline, mean relative $L_{2}$ error $4.06\times 10^{-2}$ , against the best variant V5 at $5.27\times 10^{-3}$ .

Figure 11: Held-out error distributions for the physics-informed Burgers models over 400 matched test samples. Left: all five variants. Right: the V1 dot-product baseline, mean relative $L_{2}$ error $1.24\times 10^{-1}$ , against the best variant V4 at $7.04\times 10^{-3}$ .

Figure 12: Held-out error distributions for the data-driven Poisson models over 1000 matched test samples. Left: all five variants. Right: the V1 dot-product baseline, mean relative $L_{2}$ error $2.95\times 10^{-1}$ , against the best variant V5 at $9.13\times 10^{-3}$ .

Figure 13: Held-out error distributions for the physics-informed Poisson models over 1000 matched test samples. Left: all five variants. Right: the V1 dot-product baseline, mean relative $L_{2}$ error $3.34\times 10^{-1}$ , against the best variant V5 at $2.88\times 10^{-2}$ .

PDE Operator Learning

Partial differential equation (PDE) operator learning builds models that map inputs such as coefficients and initial conditions to solution fields.

See all 228 papers

In-Context Multi-Operator Learning with DeepOSets
Dec 18, 2025Shao-Ting Chiu, Aditya Nambiar, Ali Syed +2Operator Learning In-Context Learning
Department of Electrical and Computer Engineering, Texas A&M University · Department of Mathematics, Texas A&M University
The Error of Deep Operator Networks Is the Sum of Its Parts: Branch-Trunk and Mode Error Decompositions
Feb 25, 2026Alexander Heinlein, Johannes TarazDeep Operator Networks Neural Operators
Delft Institute of Applied Mathematics, Delft University of Technology, Delft, 2628CD, The Netherlands
FAST-DeepONet: Factor-Augmented Branch Representations for High-Dimensional PDE Inputs in the Small-Sample Regime
Aug 15, 2026Jiyong Kwon, Bongseok Kim, Guang LinPDE Operator Learning Deep Operator Networks
School of Mechanical Engineering, Purdue University, 585 Purdue Mall, West Lafayette, 47907, IN, USA · Department of Mathematics, Purdue University, 150 N. University Street, West Lafayette, 47907, IN, USA

Dec 18, 2025J/K move · Enter open · S save

In-Context Multi-Operator Learning with DeepOSets

Shao-Ting Chiu, Aditya Nambiar, Ali Syed, Jonathan W. Siegel+1

Department of Electrical and Computer Engineering, Texas A&M University · Department of Mathematics, Texas A&M University

Operator Learning In-Context Learning Deep Sets

Open paper AskPDF arXiv

An important application of neural networks in scientific computing is non-linear operator learning. In this framework, a neural network is trained to fit a non-linear map between infinite dimensional spaces, for example, the solution operator of ordinary and partial differential equations. Recently, inspired by the discovery of in-context learning for large language models, an even more ambitious paradigm has been explored, called multi-operator learning. In this approach, a neural network is trained to learn many different operators at the same time. At inference time, the network is passed a small number of example input and output pairs to disambiguate the desired operator. In this work, we provide a precise mathematical formulation of the multi-operator learning problem, propose a parameter-efficient, attention-free architecture, called DeepOSets, and prove its universality in the in-context multi-operator learning task. We perform experiments demonstrating the efficacy of DeepOSets on several benchmarks, including Burgers' equation, 1D and 2D Darcy flow, and a variety of initial-value and boundary-value differential equations. In a comparison with the transformer-based ICON model on the 1D Darcy problem, under a small training budget (10k steps), DeepOSets attained lower error at every context size with 19% of the parameters, while with an extended budget (100k steps), ICON attains 26-41% lower error, but at the cost of roughly 218

\times

more training time, 12

\times

higher inference latency, and 5.3

\times

more parameters.

Abstract

PDE Operator Learning

Explore similar work

In-Context Multi-Operator Learning with DeepOSets

The Error of Deep Operator Networks Is the Sum of Its Parts: Branch-Trunk and Mode Error Decompositions

FAST-DeepONet: Factor-Augmented Branch Representations for High-Dimensional PDE Inputs in the Small-Sample Regime