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

Feb 25, 2026

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

Operator learning has the potential to strongly impact scientific computing by learning solution operators for differential equations, potentially accelerating multi-query tasks such as design optimization and uncertainty quantification by orders of magnitude. Despite proven universal approximation properties, deep operator networks (DeepONets) often exhibit limited accuracy and generalization in practice, which hinders their adoption. Understanding these limitations is therefore crucial for further advancing the approach. This work analyzes performance limitations of the classical DeepONet architecture. It is shown that the approximation error is dominated by the branch network when the internal dimension is sufficiently large, and that the learned trunk basis can often be replaced by classical basis functions without a significant impact on performance. To investigate this further, a modified DeepONet is constructed in which the trunk network is replaced by the left singular vectors of the training solution matrix. This modification yields several key insights. First, for examples involving the KdV and Burgers equations, a spectral bias in the branch network is observed, with coefficients of dominant, low-frequency modes learned more effectively. Second, through the interplay of the singular-value weighting and the optimizer's neglect of small modes, the branch error is dominated by modes with large and intermediate singular values. Third, using a shared branch network for all mode coefficients, as in the standard architecture, improves generalization of small modes compared to a stacked architecture in which coefficients are computed separately. Finally, detrimental coupling between modes in parameter space is identified, which weakens with increasing network width.

Alexander Heinlein, Johannes Taraz

Delft Institute of Applied Mathematics, Delft University of Technology, Delft, 2628CD, The Netherlands

Deep Operator Networks Neural Operators Neural Network Generalization Neural Network Training Dynamics

Ask PDF

Aug 15, 2026

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

Deep operator networks can become statistically unstable when partial differential equation inputs are observed at thousands of strongly correlated sensors but only a small number of operator samples is available. We introduce FAST-DeepONet, a branch representation combining a fixed spectral path with a regularized projection of the orthogonal residual, in which the directional penalty acts on the effective residual map after each of its rows is normalized. On Navier--Stokes flow a plain DeepONet degrades from

0.0394

0.1556

mean relative

L_2

error as the branch grows from

129

8193

coordinates, while FAST-DeepONet stays near

0.04

, so the sensor grid can be refined without a statistical penalty. Across independent test sets for Navier--Stokes flow, Darcy flow, and signed terminal wavefield prediction it lowers mean relative

L_2

error by

4.7\%

37.0\%

with three to seven times fewer trainable parameters. A spectral-only branch sharing the same basis separates the two paths: the fixed spectral path carries the improvement on Navier--Stokes and Darcy, while terminal wave prediction requires the residual path together with its directional penalty. FAST-DeepONet targets coordinate-query architectures and trains on solution values alone.

Jiyong Kwon, Bongseok Kim, Guang Lin

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

PDE Operator Learning Deep Operator Networks PDE Surrogate Modeling Operator Learning

Ask PDF

Oct 3, 2026

Localized Operator Learning with Adaptive Partition-of-Unity Mixture-of-Expert Networks

Operator learning methods such as DeepONets and FNOs often struggle with PDE families featuring sharp interfaces, heterogeneous coefficients, and localized multiscale structures. We introduce a partition-of-unity (POU) mixture-of-experts framework for localized operator learning, in which geometry-aware gating networks produce smooth spatial partitions which blend the contributions of local expert networks. Our main contribution is HiRefPOU, a residual-style hierarchical POU architecture for DeepONets that organizes localized representations through nested parent-child partitions while preserving global continuity. We also show that the same POU principle can be incorporated into Fourier Neural Operators to introduce spatial adaptivity without modifying the underlying spectral layers. On heterogeneous Darcy and reaction-diffusion benchmarks, HiRefPOU achieves substantially lower error than global DeepONet and static POU-MoE baselines, while the broader operator-learning experiments show that the benefits of localization depend on the PDE structure and the chosen neural-operator backbone. The learned partitions are interpretable and align with interfaces and regions of rapid solution variation. These results show that explicit geometric localization can improve both accuracy and interpretability in neural operator learning.

Madison Cooley, Ramansh Sharma, Shandian Zhe +2

Scientific Computing Institute, University of Utah, Salt Lake City, UT · Kahlert School of Computing, University of Utah, Salt Lake City, UT

PDE Operator Learning Operator Learning Deep Operator Networks Mixture-of-Experts Models Fourier Neural Operator

Ask PDF

Abstract

PDE Operator Learning

Explore similar work

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

Localized Operator Learning with Adaptive Partition-of-Unity Mixture-of-Expert Networks