The importance of tensor handling and implementation of physics-informed DeepONets

Nahum Banks; Gustavo A. M. de Almeida; Sergio Maldonado

cs.LGOct 15, 2026

The importance of tensor handling and implementation of physics-informed DeepONets

Authors: Nahum Banks, Gustavo A. M. de Almeida, Sergio Maldonado

Organizations: Faculty of Engineering and Physical Sciences, University of Southampton, Southampton, UK

Abstract

Physics-informed neural networks (PINNs) reduce or eliminate the need for large labelled datasets by enforcing governing equations during training. Physics-informed Deep Operator Networks (PI-DeepONets) extend this approach to learning solution operators, enabling rapid prediction across families of problems. Despite these advantages, we show that their application to higher-dimensional problems remains limited by the computational cost of training. To address this issue, we propose a modified PI-DeepONet architecture that reduces this cost by removing redundancy in how the Branch network processes input functions across collocation points. The proposed formulation is mathematically equivalent to the original architecture but reduces input tensor sizes and the number of Branch network evaluations. We analyse the computational complexity of both architectures and derive theoretical scaling laws for training time with respect to the number of collocation points and input functions. The proposed architecture is evaluated on 1D, 2D, and 3D linear advection problems. In all cases, it achieves comparable predictive accuracy to the original PI-DeepONet while consistently reducing training time. In 2D experiments, a maximum speed-up of $3.75\times$ is observed, with empirical training-time scaling agreeing with the predicted linear dependence on collocation points and input functions. Furthermore, training a 3D model on an NVIDIA A100 GPU is only feasible using the proposed architecture, as the original formulation exceeds memory limits. These results demonstrate that careful consideration of how PI-DeepONets are implemented can yield substantial gains, enhancing the practical performance of PI-DeepONets and enabling their application to otherwise intractable physics-informed learning problems.

Figures & tables

Figure 1: Schematic diagram of a physics-informed Deep Operator Network (PI-DeepONet).

u

is the input function discretised at

m

sensor points, and

y

is the query location for the output function of dimensionality

d

. Derivatives of the output with respect to the input coordinates are calculated via automatic differentiation; these operations are denoted by

\partial_{x}^{(o)}

and

\partial_{t}^{(o)}

, where

o

is the order of derivative.

\mathcal{L}_{{\text{PDE}}}

\mathcal{L}_{\text{BC}}

\mathcal{L}_{\text{IC}}

, and

\mathcal{L}_{\text{data}}

denote the PDE, boundary condition, initial condition, and labelled data loss terms, respectively.

Symbol	Description
$n$	Number of input functions used during training
$p$	Number of collocation points per input function
$m$	Number of sensor (discretisation) points for input functions
$d$	Dimension of the spatiotemporal domain
$s$	Latent dimension (output size of Branch/Trunk networks)
$\mathbf{u}^{(j)}$	$j$ th input function

Table 1: Glossary of key variables used throughout this work.

Architecture	Branch input	Trunk input	Branch output	Trunk output	Combination method
Vanilla	$(np,m)$	$(np,d)$	$(np,s)$	$(np,s)$	Inner product
ATenH	$(n,m)$	$(n,p,d)$	$(n,s)$	$(n,p,s)$	Inner product (with broadcasting)

Table 2: A comparison of data processing choices in the two PI-DeepONet architectures herein presented. Here,

m

is the number of discretisation points for the input functions (the number of sensor points),

d

is the physical dimensionality of the collocation points,

n

is the number of input functions,

p

is the number of collocation points per function, and

s

is the length of the output layer of the (Branch and Trunk) subnetworks

Panel of Figure 2 — Figure 2: Training and validation loss profiles as a function of training epoch for 1D (left) and 2D (right) linear advection, trained using the ATenH PI-DeepONet. The training and validation loss curves remain closely aligned throughout training indicating no overfitting. The rate of decrease in the loss begins to plateau after approximately 50,000 epochs for 1D linear advection and 100,000 epochs for 2D linear advection

Panel of Figure 3 — Figure 3: Grid searches exploring the depth and width of the subnetworks for (left) 1D linear advection and (right) 2D linear advection. The same width and depth was used across Branch and Trunk networks to reduce the complexity of the hyperparameter space, alongside keeping the layer widths constant across hidden layers

Panel of Figure 4 — Figure 4: Representative comparison of the exact solution and the predictions of the Vanilla and ATenH PI-DeepONet architectures for a test input function in 1D linear advection. The models are trained with 10 collocation points per dimension and 500 training input functions. (a) The Vanilla PI-DeepONet’s prediction. (b) The ATenH PI-DeepONet’s prediction. (c) The exact solution. (d) The absolute error for the Vanilla architecture. (e) The absolute error for the ATenH architecture

Panel of Figure 5 — Figure 5: Comparison of the exact solution and the predictions of an ATenH PI-DeepONet for a representative combination of boundary inputs in 2D linear advection at two time instances. The Vanilla architecture yields qualitatively similar results and is therefore omitted for brevity. (a)–(c) Exact solution, prediction, and absolute error at $t=0.48$ . (d)–(f) Exact solution, prediction, and absolute error at $t=1.00$

Panel of Figure 6 — Figure 6: Mean training time of the Vanilla and ATenH architectures for 2D linear advection with (left) varying numbers of collocation points with a fixed number of input functions, and (right) varying numbers of input functions with a fixed number of collocation points. Solid lines indicate linear regression fits to the training times. Right $y$ -axis shows the computed ratio of training times for both architectures (i.e. ATenH’s speed-up factor). The two largest collocation-point data points are excluded for the Vanilla architecture due to insufficient memory (left)

Panel of Figure 7 — Figure 7: Comparison of the exact solution and the predictions of an ATenH PI-DeepONet for a representative combination of boundary inputs in 3D linear advection at $x_{3}=0.44$ and two time instances. The Vanilla architecture fails to train due to a memory error and is therefore omitted. (a)–(c) Exact solution, prediction, and absolute error at $t=0.44$ . (d)–(f) Exact solution, prediction, and absolute error at $t=1.00$

	Vanilla		ATenH		Theoretical speed-up
Spatial dimensionality	Time (s)	MRRMSE (%)	Time (s)	MRRMSE (%)	$\phi$
1D	$3.14\times 10^{2}$	$0.471\pm 0.293$	$2.76\times 10^{2}$	$0.308\pm 0.078$	1.64
2D	$1.18\times 10^{4}$	$2.094\pm 0.846$	$3.36\times 10^{3}$	$1.524\pm 0.513$	7.46
3D	—	—	$1.40\times 10^{5}$	$5.438\pm 1.108$	580

Table 3: Mean training time and MRRMSE scaling with spatial dimensionality in linear advection problems for the Vanilla and ATenH architectures. Standard deviations are not shown for the training times since they are negligible. There is no data for the Vanilla architecture in the 3D linear advection problem due to memory limitations during training. Also shown is the theoretical speed-up

\phi

for ATenH when compared to the Vanilla PI-DeepONet architecture, for each problem with their corresponding hyperparameters, as specified in Section 2.5

Physics-Informed Neural Operators

Physics-informed neural operators learn mappings between physical fields while embedding governing equations or structural constraints for faster, more reliable simulation.

See all 72 papers

An Optimisation Framework for the Well-Conditioned Training of Physics-Informed Neural Networks
Jul 2, 2026Joseph Webb, Sadok Jerad, Coralia CartisGauss-Newton Optimization PDE Solving
Mathematical Institute, University of Oxford, Oxford, United Kingdom
On the training of physics-informed neural operators for solving parametric partial differential equations
Jun 4, 2026Nanxi Chen, Chuanjie Cui, Airong Chen +2Neural Network Optimization Gradient Interference
College of Civil Engineering, Tongji University, Shanghai 200092 · Department of Engineering Science, University of Oxford, Oxford OX1 3PJ · Institute for Foundations of Data Science, Yale University, New Haven, CT 06520
Improving the accuracy of physics-informed neural networks via last-layer retraining
Mar 4, 2026Saad Qadeer, Panos StinisNeural PDE Solvers Physics-Informed ML
Advanced Computing, Mathematics and Data Division, Pacific Northwest National Laboratory, Richland, WA

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

An Optimisation Framework for the Well-Conditioned Training of Physics-Informed Neural Networks

Joseph Webb, Sadok Jerad, Coralia Cartis

Mathematical Institute, University of Oxford, Oxford, United Kingdom

Gauss-Newton Optimization PDE Solving Second-Order Optimization Physics-Informed ML

Open paper AskPDF arXiv

Physics-informed neural networks (PINNs) have emerged as a promising route to solve partial differential equations, yet they have struggled to reach the precision of classical solvers. The obstacle is increasingly understood to be one of optimisation, owing to the severely ill-conditioned loss landscape. We present

\textbf{DSGNAR}

: Doubly-Sketched Gauss-Newton with Adaptive Ratio, a scalable second-order optimisation framework that confronts this ill-conditioning and, in doing so, obtains unprecedented accuracy and speed.

\textbf{DSGNAR}

couples a doubly-sketched Gauss-Newton model with a novel strategy that carefully controls both regularisation and step length. Across a suite of problems spanning nonlinear, chaotic, multi-scale, high-dimensional, and Navier-Stokes, the framework greatly improves on the state of the art: able to attain relative

\ell_2

errors as low as

3\times10^{-16}

in double precision, improve contemporary results by five orders of magnitude on the canonical Burgers' equation, and as much as eight orders on a high-dimensional Poisson problem, while remaining markedly faster. We further show that, in single precision, solutions at the limit of round-off error can be obtained very quickly: Burgers' equation to

\ell_2^{\text{rel}} = 4.75 \times 10^{-7}

in under ten seconds. The framework is also robust to the choice of architecture, arithmetic precision, and initial hyperparameters. The code is available at https://www.github.com/wephy/physics-informed-neural-networks

Abstract

Physics-Informed Neural Operators

Explore similar work

An Optimisation Framework for the Well-Conditioned Training of Physics-Informed Neural Networks

On the training of physics-informed neural operators for solving parametric partial differential equations

Improving the accuracy of physics-informed neural networks via last-layer retraining