Separable Operator Networks
Authors: Xinling Yu
Sean Hooten
Ziyue Liu
Yequan Zhao
Marco Fiorentino
Thomas Van Vaerenbergh
Zheng Zhang
Venue: Accepted by TMLR ()

Abstract: Operator learning has become a powerful tool in machine learning for modeling complex physical systems governed by partial differential equations (PDEs). Although Deep Operator Networks (DeepONet) show promise, they require extensive data acquisition. Physics-informed DeepONets (PI-DeepONet) mitigate data scarcity but suffer from inefficient training processes. We introduce Separable Operator Networks (SepONet), a novel framework that significantly enhances the efficiency of physics-informed operator learning. SepONet uses independent trunk networks to learn basis functions separately for different coordinate axes, enabling faster and more memory-efficient training via forward-mode automatic differentiation. We provide a universal approximation theorem for SepONet proving the existence of a separable approximation to any nonlinear continuous operator. Then, we comprehensively benchmark its representational capacity and computational performance against PI-DeepONet. Our results demonstrate SepONet's superior performance across various nonlinear and inseparable PDEs, with SepONet's advantages increasing with problem complexity, dimension, and scale. For 1D time-dependent PDEs, SepONet achieves up to 112× faster training and 82× reduction in GPU memory usage compared to PI-DeepONet, while maintaining comparable accuracy. For the 2D time-dependent nonlinear diffusion equation, SepONet efficiently handles the complexity, achieving a 6.44% mean relative $\ell_{2}$ test error, while PI-DeepONet fails due to memory constraints. This work paves the way for extreme-scale learning of continuous mappings between infinite-dimensional function spaces. Open source code is available at https://github.com/HewlettPackard/separable-operator-networks.

Published in Transactions on Machine Learning Research (12/2024)
Separable Operator Networks
Xinling Yu'
24y644 Cucsb. edu
University of California, Santo. Barbara
Sean Hooten"
seon.hooten@@hpe.cmm
Hewlett Packard Labs, Hewlett Packord Enterprise
Ziyue Liu
ziyuelia edu
University of California, Santa Barbara
Yequan Zhao
yequan wwwtom__hho@uusb..oomm edu
University of California, Santa Barbara
Marco Fiorentino
mureofforentiino@@pp..ccmm
Hewlett Packard Labs, Hewlett Packord Enterprise
Thomas Van Vaerenbergh
thomas, hwwnosswwun.eerrrebeerggh......
Hewlett Packard Labs, HPE Belgium
Zheng Zhang University of California, Santa Barbara
zhengzhang Sece, wesb edu
*Contributed equally
Reviewed OpenReview: https://openreviiee net/ forum? id= RYtJmFDAvv
Abstract
Operator learning has become powerful tool in machine learning for modeling com- plex physical systems governed by partial differential equations (PDEs). Although Deep Operator Networks (DeepONet) show promise, they require extensive data acquisition. Physics-informed DeepONets (PI-DeepONet) miligate data scarcily but suffer from ineffi- cient training processes. We introduce Separable Operator Networks (SepONet), a novel framework that significantly enhances the efficiency of physics-informed operator learn- ing. SepONet uses independent trunk networks to learn basis functions separately for different coordinate axes, enabling faster and more memory-efficient training via forward- mode automatic differentiation. We provide universal approximation theorem for Se- PONet proving the existence of separable approximation to any nonlinear continuous operator. Then, we comprehensively benchmark its representational capacity and com- putational performance against PI-DeepONet. Our results demonstrate SepONet's supe- rior performance across various nonlinear and inseparable PDEs, with SepONet's advan- tages increasing with problem complexity, dimension, and scale. For 1D time-dependent PDEs, SepONet achieves up to 112x faster training and 82x reduction in GPU mem- ory usage compared Pl-DeepONet, while maintaining comparable accuracy. For the 2D time-dependent. noulinear diffusion equation, SepONet efficiently handles the com- plexity, achieving 6.44% mean relative l test error, while P1-DeepONet fails due to memory constraints. This work paves the way for extreme-scale learning of continuous mappings between infinite-dimmnnionnal function spaces. Open source code is available at wwt..t....oooee.eeeeeeeeeeeeeeee http://hh.acccttteeeeeeeoeeeeeep

---

Published in Transactions on Machine Learning Research (12/2024)
O
(0) y.k
$(x0,, v
PDE Residuals:
y (K,))
Trunk Net
at (u) = N (u, $(x(0, y(
EaBer x.
N Nonlinear differential operator
Ey (100)
Z
Forward-Mode AD for Derivatives: as(x)), σy
y0
Σ
Br G är 0
Prediction
dy
(210) y00
y0%)
Coordinate Grid
Branch Net
Batch Input Functions: u (u
b(E(w))
Trunk Not
:
bx (x(1)))
Nr (((o))
Loss N,N
y(y
1 1
x()
Ny + Ny + Ny Inputs
Input Function
Evaluated at Ny Outputs
x()
(Na)
x11)
(N+)
Figure 1: Separable operator network (SepONet) architecture for 2D problem instance. A coordinate grid of collocation points (x((0, y))) can be evaluated efficiently by separating the coordinate axes, feeding them through independent trunk networks, and combining the outputs by outer product to obtain multiple basis function maps. Mcanwhile, the branch network processes input functions and outputs coefficients, which are then used to scale and combine the trunk network basis functions by product and sum. Spatiotemporal derivatives of the output predictions are obtained efficiently by forward-mode automatic differentiation due to the independence of trunk networks along each coordinate axis.
1 Introduction
Operator learning, which aims learn mappings between infinite-dimensional function spaces, has gained significant. attention in scientific machine learning thanks its ability to model complex dynamics in physics systems. This approach has been successfully applied to wide range of applications, including climate
modeling (Kashinath et al. 2021
Pathak et al. 2022). multiphysics simulation (Liu et al. 2023 Cai ct al.
2021
Mao et al.
2021 Lin et al
2021
Kontolati et al. 2024), ; inverse design (Lu et al.
2022b Gu et al.
2022 and more Shukkl et al.
2023
Gupta & Brandstetter 2022). Various operator learning a algorithms
(Lu et al. 2021 Li et al.) 2020b a Ovadia et al. 20233 Wen et al. 2022 Ashiqur Rahman et al. 2022
have
been developed to address these applications, with Deep Operator Networks (DeepONets) {Lu al.
2021
being particularly notable due to their universal approximation guarantee for operators (Chen & Chen
1995
Lanthaler et al.) 2022 Gopalani et al.) and robustness (Lu et al. 2022a
To approximate the function operator G : U + S, DeepONets are trained in a supervised manner using a datasct of N, pair functions (())),()))))) ) where in the context of parametric partial differential equations (PDEs), each (1()) represents PDE configuration function and each (()) represents the corresponding solu- tion. Unlike traditional numerical methods which require repeated simulations for each different PDE config- uration, once well trained, a DeepONet allows for efficient parallel inference in abstract difnite-dimensionall function spaces. Given new PDE configurations, can immediately provide the corresponding solutions. However, this advantage comes with a significant challenge: achieve satisfactory generalization error, the number of required input-output function pairs grows quadratically (Lu al. 2021 Lanthaler et al. 2022; Liu et al. 2022a Gopalani et al.). Generating enough function pairs be computationally expensive even impractical in some applications, creating bottleneck in the effective deployment of DeepONets.
Physics-informed deep operator networks (PI-DeepONet) (Wang et al., 2021bb have been introduced as a solution to the costly data acquisition problem. Inspired by physics- physics-informi informed neural networks (PINNs) [Raissi et al. 2219)), PI-DeepONet learns operators by constraining the DecpONet outputs approximately satisfy the underlying governing PDE system parameterized by the input function u. This is achieved by penalizing the physics loss (including PDE residual loss, initial loss and boundary loss), thus eliminating
2

---

Published in Transactions on Machine Learning Research (12/2024)
the need for ground-truth output functions s. However, PI-DoepONet shares the same disadvantage PINNs in that the training process is both memory-intensive and time-consuming (He et al. 2023 Cho
et al. 2024). This inefficiency, common to both P1-DeepONet and PINNs, arises from the need to compute high-order derivatives of the PDE predictions with respect to numerous collocation points during physics loss optimization. This computation typically relies on reverse-mode automatic differentiation (Baydin et al. 2018), involving the backpropagation of physics loss through the computational graph to update model parameters. The inefficiency even more pronounced for Pl-DeepONet, as it requires evaluating the physics loss across multiple PDE configuratinns. While numerous studies (He et al., 2023 Zhao et al., 2023 Hu
ct al. 2024 Liu et al. 2022b; Cho et al. 2024) have proposed methods to enhance PINN training efficiency, there has been limited research focused on improving the training efficiency of PI-DeepONet specifically.
To address the inefficiency in training Pl-DeepONet, we propose Separable Operator Networks (SepONet); inspired by the separation of variables technique in solving PDEs and recent work on separable PINN (Cho al., (024).. Suppose we want to approximate a nonlinear continuous operator of d variables (d- dimensional collocation points). When optimizing the physics loss al M collocation points, Pl-DeepONet always requires M inputs to evaluate PDE predictions and their derivatives, regardless of whether the points are on a regular grid or randomly sampled. This approach becomes inefficient when M is large. By contrast, SepONet uses d independent trunk networks to learn univariate basis functions separately for each variable, which can then be combined to obtain predictions in d dimensions. In particular, if M can be decomposed as N1 x N2 x x Na, where Nn is the number of points sampled along the n-th coordinate axis (e.g., 1024 16 8 x for d 3), SepONet requires only + + Na imputs to evaluate all M collocation points regular grid, resulting in more efficient solution. The SepONet architecture for d 2 problem instance is shown in Figure I This simple yct effective modification enables fast training and memory-efficient implementation of SepONet. by leveraging forward-mode automatic differentiation (Khan & Bartonn 2015) to compute high-order derivatives of all Nı x N2 x-. x .N collocation points. It's important to note that factorizing functions in high-dimensional domains into multiple sub-functions defined over one- dimensional domains has been explored in supervised operator learning methods (Tran et al. 2021 Li et al., 2024 Kossaifi et al., 2023). Working independently and concurrently, our work and that of Mandl et al.) (2025) extended these ideas to physics-informed operator learning frameworks. While both approaches leverage separable structures to enhance computational efficiency, our method uniquely draws from classical separation of variables techniques in PDEs, leading to a distinct low-rank functional decomposition of the trunk network. Our key contributions are:
We introduce SepONet, a physics- operator learning framework that significantly enhances training time and GPU memory usage relative to PI-DeepONet, enabling extremeseeale learning of continuous mappings between infinite-dimensional function spaces.
2. We provide a theoretical foundation for SepONet through the universal approximation theorem, proving its capability to approximate any nonlinear continuous operator with arbitrary accuracy.
3. We provide extensive benchmarks validating SepONet's representational capacity and computational performance relative to PI-DeepONet on a range of ID and 2D time-dependent nonlinear PDEs. Our findings reveal that scaling up SepONet in both number of functions and collocation points con- sistently improves its accuracy, while typically outperforming PI-DeepONet. On 1D time-dependent PDEs, we achieve up to 112x training speed-up with minimal memory increase. Notably, at mod- erately large scales where training Pl-DeepONet exhausts 80GB of GPU memory, SepONet trains and operates efficiently with less than 1GB. Furthermore, we observe efficient scaling of SepONet with problem dimension, enabling accurate prediction of 2D time-dependent PDEs at scales where Pl-DeepONet fails.
2 Preliminaries
2.1 Operator Learning for Solving Parametric Partial Differential Equations
Let X' and y be Banach spaces, with K S X and K2 S y being compact sets. Consider a nonlinear continuous operator G U - S, mapping functions from one infirite-dimensional space to another, where
3

---

Published in Transactions on Machine Learning Research (12/2024)
u S C(K) and S S C(K)). The goal of operator learning is to approximate the operator G using model parameterized by 0, denoted as Gg. Here, U and S represent spaces of functions where the input and output functions have dimensions dy and d,, respectively. We focus on the scalar case where dy = d. = throughout most of this paper; however, it should be noted that the results apply to arbitrary du and dg-
In the context of solving parametric partial differential equations (PDFs); consider PDEs of the form:
Nuu, s) = 0, I(u, s) - 0, B(u, s) = O,
(1)
where N is a nonlinear differential operator, I and B represent the initial and boundary conditions, u E u denotes the PDE configurations (source terms, coefficients, initial conditions, and etc.); and 8 E S denotes the corresponding PDE solution. Assuming that for any u E U there exists unique solution E S, we can define the solution operator G: :: as G(u).
A widely used framework for approximating such an operator G involves constructing Go through three maps [Lanhaler et al. 2022:
Go N G Do
(2)
First, the encoder & u + RT maps an input function u E U to a dimte-diia feature representation. Next, the approximator R to another finite-dimest dintn-siiansional space IR. Finally, the decoder D R' + S produces the output function s(y) Gv))(v for y Ki.
R + Re transforms this encoded data within the finite-dimensional space
2.2 Deep Operator Networks (DeepONet)
The original DeepONet formulation (Lu et al. 2021 can be analyzed through the 3-step approximation framework (2). The encoder & U - Rm maps the input function u to its point-wise evaluations at m fixed sensors 21, .2....... Im E K, e.g., (u(=1),..., u(zm)) = E()). Two separate neural networks (usually multilayer perceptrons), the branch net and the trunk net, serve as the approximator and decoder, respectively. The branch net by. 1 - IR'' parameterized by & processes (u(zz),...,uu u(zm)) to produce a feature embedding (3i, B2, 3r). The trunk net 4. Rd + R" with parameters 9, takes a continuous coordinate = (y),...,, Wa) as imput and outputs a feature embedding (T1, .2,..., Ty). The final DeepONet prediction of a function u for a query y is:
Go(u)(v) = Σ 3AT = (E()) Lo(2), k=1
(3)
where is the vector dot product and 0 = (v, Φ) represents all the trainable parameters in the branch and trunk nets.
Despite DeepONet's remarkable success across range of applications in multiphysics simulation (Cai et al. Mao et al. 2021 Lin et al. 2021; inverse design (Lu et al.) 2022b), and carbon storage (Jiang et al, 2023 its supervised training process is highly dependent on the availability of training dała, which can be costly, Indeed, the generalization error of DeepONets scales quadratically with the number of training input- output function pairs (Lu et al.) 2021 Lanthalor et al. 2022 Liu et al. 2022a Gopalani et al.)). Generating a large number of high-quality training data is expensive or even impractical in some applications. For example, in simulating high Reynolds number (Re) turbulent flow (Pope
2021
2001), ) accurate numerical simulations require
fine mesh, leading to a computational cost scaling with Re Kochkov al. 2021), making the generation of sufficiently large and diverse training datasets prohibitively expensive.
To address the need for costly data acquisition, physics-informed deep operator networks (PI-DeepONet) (Wang et al.) (0211), inspired by physics-informed neural networks (PINNs) (Raissi et al. 2019). have been proposed to learn operators without relying on observed input-output function pairs. Given a dataset of N2 input training functions, N, residual points, N, initial points, and N2 boundary points: Nr D {((() i=1 (g()) N. Ji=1 (1(()) N, 6 No ) PT-DeepONets are trained by minimizing an unsuper- 1= J==
N.
vised physics loss: Ipphysies (0)D) Tresiduul (f/D) + (iiitiaa(((D)) + Ao Lhoundary (
(4)
4

---

Published in Transactions on Machine Learning Research (12/2024)
where
N₁
Tepesiduat (O
1
ΣΣ
Mn (o(
N,
Linitial (O N1
(5)
N, M L.boundary ( Ny ΣΣ B(ul), (o(
Here, A and As denote the weight coefficients for different loss terms. However, noted in the original PI-DeepONet paper (Wang et al. 2021b). the training process be both memory-intensive and time- consuming, Similar to PINNs (Raissi et al., 2019), this inefficiency arises because optimizing the physics loss requires calculating high-order derivatives of the PDE solution with respect to numerous collocation points, typically achieved via reverse-mode automatic differentiation (Baydin et al. 2018). This process involves backpropagating the physics loss through the unrolled computational graph to update the model parameters. For Pl-DeepONeet, the inefficiency is even more pronounced, as the physics loss terms (equation (5) must be ovaluated across multiple PDE configurations. Although various works (Chiu et al. 2022 He et al. 2023 Cho el al. 2024) have proposed different methods to improve the training elliciency of PINNs, little research has focused on enhancing the training efficiency of Pl-DeepONet. We propose to address this inefficiency through a separation of input variables.
2.3 Separation of Variables
The method of separation of variables seeks solutions to PDEs of the form s(y) = (vi)- Ya(ya) for an input point y (t, y1..... ya) and univariate functions T, Yı. ,,,...,YY. Suppose we have linear PDE system
Mit]s(y) = Lilvil(()) + +... Lalvals(v).
(6)
where M[t] dh + futt is first order differential operator of t, and L1 Cull) are linear second order ordinary differential operators of their respective variables yı.--: Va only. Furthermore, assume are provided Robin boundary conditions in each variable and separable initial condition s(t 0. V1..... . ya) (n & for functions on (yn) that satisfy the boundary conditions. Then, leveraging Sturm-Liouville theory and some massaging, the solution to this problem can be written
8 8(1 .ya) Σ A k
A (1)
II (1)),
where l is a lumped index that counts over infinite eigenfunctions of each Ci operator (potentially with repcats). For example, given {l,..., d), L, (yn) > for cigenvalue X E R. (t))) depends all the eigenvalues A/= corresponding to index k. A IR is coefficient determined by the initial condition. More details can be found in Appendix C The method of separation of variables applied to linear heat equation example can be found in Appendix D.2.
One may notice the resemblance between the form of the DeepONet prediction (3) with 7 provided B A& and 7 = 1 (t) I_ (²²),,, with appropriately ordered k. We leverage this similarity explicitly in the construction of SepONet. below.
3 Separable Operator Networks (SepONet)
Inspired by the method of separation of variables and recent work on separable PINN (Coo et al., 2024), we propose using separable operator networks (SepONet) to learn basis functions separately for different coordinate axes. SepONet approximates the solution operator of PDE system by, for given point
5

---

Published in Transactions on Machine Learning Research (12/2024)
y (f.....;,d))
d 3k
(ol Mal = Σ k=1 n=1
Tn.k
(8)
(E((()) (fh, (v1) t ( o t (v6))
where O is the Hadamard (element-wise) vector product and . is the vector dot product. Here, 3. - by (E (u))a is the k-th output of the branch net, as in DeepONet. However, unlike DeepONet, which employs a single trunk net that processes each collocation point y individually, SepONet uses d independent trunk nets, ta : R' for n ......d. d. In particular, Tn. = ". (Un)s denotes the k-th output of the n-th trunk net. Importanll, the parameters of the n-th trunk net on are independent of all other trunk net parameter sets. Viewed through the 3-step approximation framework 2 SepONet and DeepONet have identical encoder and approximator but different decoders.
Equation can be understood as low-rank approxination of the solution operator by truncating the basis function series (represented by the output shape of the trunk nets) at a maximal number of ranks r. SepONet not only enjoys the advantage that basis functions for different variables with potentially different scales can be learned more easily and efficiently, but also allows for fast and efficient training by leveraging forward- mode automatic differentiation, which we will discuss in Section 3.1 and Section 3.2 Moreover, despite the resemblance between (8) and the separation of variables method for linear PDEs 7 (discussed below in Section 3.1.3,, we find that SepONet can effectively approximate the solutions to nonlinear parametric PDEs. Indeed, we provide a universal approximation theorem for separable operator networks in Section 3.3 and extensive accuracy and performance scaling numerical experiments for nonlinear and inseparable PDEs in Section
Finally, it is worth noting that if one is only interested in solving deterministic PDEs under a certain configuration (i.c., is fixed), then the coefficients B& also fixed and be absorbed by the basis functions. In this case, SepONet will reduce to separable PINN Cho et al. 20224 which has been proven to be efficient and accurate in solving single-scenario PDE systems (E kin et al. 2024 Oh et al. 2024).
3.1 SepONet Architecture and Implementation Details
Suppose are provided a computation domain K₁ 1j of dimension d and an input function u. To evaluate the residual loss term in equation (5) on N, random collocation points, PI-DeepONet samples all N, points directly from Ki. However, if N₂ can be approxinately factorized as N, x N2 x ... x Na (e.g.. 1024 16x8 x for d = 3), and we relax the Monte Carlo sampling requirement for d-dimensional collocation points, SepONet only ncods to randomly sample N, points along the n-th coordinate axis, resulting in a total of Ni + N2 + +...+ Na samples.
It is important to note that SepONet's mapping from Ni + +2 inputs to N1 x N2 x... x Na outputs is most efficient when using regular grid sampling. However, we have empirically demonstrated (in Section that this per-axis grid-based sampling strategy does not degrade SepONet's accuracy compared to PI- DeepONet's Monte Carlo random sampling over the entire domain K₁. Irregular domains may be sampled by (a) dividing the irregular domain into subdomains each approximated by a regular grid, or (b) applying coordinate transformation to map the irregular domain onto regular one (o.g., converting Cartesian polar coordinates for a circular domain).
For shorthand and generality, we will denote the dataset of input points for SepONet as D D {vi O (c)
Each yn = (2 Fi-1 represents an array of N, samples along the n-th coordinate axis for total of Ni + N₂ +... Na samples. The initial and boundary points may be separately sampled from 0K1; the number of samples (Ny and N₂) and sampling strategy are equivalent for SepONet and Pl-DeepONet.
(6) (,,N.
3.1.1 Forward Pass
The forward pass of SepONet, illustrated for d 2 in Figure 1
follows the formulation S except generalizod to the computationally advantageous setting where predictions along grid of collocation points are processed

---

Published in Transactions on Machine Learning Research (12/2024)
in parallel. The formula can be expressed:
Co(u)(vi,..... ....., 1
BR
n,k
k=
n=l
(9)
by (E(U)) (to (vf o 2 o ( 2 (l 0..
where @ is the (outer) tensor product, which produces an output predictive array along a meshgrid of N1 N2 x... Na collocation points. Notably, 7, f3 (u!"), represents a vector of N, values produced by the n-th trunk net along the k-th output mode after feeding all yh points. After taking the outer product along each of = .....,d dimensions for all modes, the modes sum-reduced with the predictions of the branch net 3& (E(U))) While not shown here, implementation also batches over input functions ful Ny for Ny functions. Thus, for only N₂ + N1 +... Na inputs, SepONet produces a predictive array with shape N2 N1 X ... x Na.
3.1.2 Model Update
In evaluation of the physics loss 4 SepONet enables more efficient computation of high-order derivatives in terms of both time and memory use compared to PI-DeepONet by leveraging forward-mode. automatic differentiation (AD) (Khan & Barton, 2015). This fairly evident by the form of (9. For example, to compute derivatives of all SepONet outputs with respect the m-th variable ym:
(Go(u)(y1 ..... ym ..... Vdl (
Q m.kk
BR
(:) Tn
(10)
8ym
Oym
m
0 is a vector of derivatives of the m-th trunk net's k-th basis function evaluated along all inputs 0r the m-th coordinate axis. One may notice that 8ym can be written as a Jacobian-vector product (JVP) of the Jacobian of the m-th trunk net's Nm outputs with respect all Nm inputs, and a length Nm 1 tangent vector of l's:
where Oym
O+(1) dym
= J mm
(
(11)
where (k) selects the k-th output mode from the resulting Nm JVP output. This is equivalent to forward-mode AD. Consequently, the derivatives along the m-th coordinate axis across the entire grid of predictions can be obtained by pushing forward derivatives of the m-th trunk net, and then reusing the outputs of all other n * m trunk nets via outer product. By contrast, Pl-DeepOnet must compute derivatives @Go(vi) .....''' "...""""."""""""""""" (...,(())/8)) //8m for each input-output pair individually, ( (vi") (3) Vd for i 1,..,,,, M, where there is such computational advantage and it is more prudent to use reverse-mode AD. Fundamentally, the advantage of SepONet for using forward-mode AD can be attributed to the significantly smaller input-output relationship when evaluating along coordinate grids N - IR Nix...N,, Na compared to PI-DeopONet Mxd + R Mx when we choose M = N.. Na. The time and space complexity analysis below in 3.2 provides a more descriptive breakdown of computational scaling behavior. For more detailed explanation of forward- and reverse-mode AD, we refer readers to Cho et al. (2024); Margossian (2019). Once the physics loss is computed, often involving multiple evaluations of [10 reverse-mode AD is employed to update the model parameters 0 (1), $1....,..
Section
3.1.3 Inference
Once trained, SepONet can efficiently map input functions to accurate output function predictions Go(u) along large spatiotemporal grids. This is achieved by combining the learned coefficients from the
7

---

Published in Transactions on Machine Learning Research (12/2024)
Table 1: Time and space complexities of first-order derivatives of SepONet and Pl-DeepONet with respect to Nd collocation points using forward-mode and reversc-mode AD, respectively.
Method
Time Complexity
Space Complexity
SepONet (Forward AD)
O(N d. Lr² + NNd... rd) O(N Lr
O(N rd + N9) O(N4 L
PI-DeepONet (Reverse AD)
branch net with the basis functions learned by the trunk nets. Indeed, intriguing comparisons can be made between the results of the method of separation of variables 7 and trained SepONet predictions 88. For an initial value problem with linear, separable PDE operator treated in Section 2.3 we might expect SepONet learn initial value function dependent coefficients (E(E)))) basis functions 13., (yn)s ((,)) (provided supply one additional trunk net 12 (4)x (1) for the temporal dimension) for appropriately ordered modes i and sufficiently large r. Examples of the learned basis functions for separable 1D lime-dependent heat equation initial value problem example are provided in Appendix D.2 as a function of the number of the trunk net output shape For a small number of modes r = or r = 2, SepONet Iearns nearly the exact spatiotemporal basis functions obtained by separation of variables. For larger r, the SepONet basis functions do not converge to the analytically expected basis functions. Nevertheless, approximation error observed to improve with r, and near perfect accuracy is obtained large by comparison numerical estimates of the analytic solution.
AA and spatiotemporal
While the form of SepONet predictions resemble separation of variables, we note that the method of separation of variables typically only applies to linear PDFs with restricted properties. In spite of this, SepONet is capable of accurately approximating arbitrary operator learning problems (including nonlinear PDEs) as guaranteed by a universal approximation property, provided below in Section 3.3
3.2 Complexity Analysis
Suppose we are provided a computational domain K = [0, 1jd of dimension d. For Pl-DeepONet, collocation points are sampled randomly from the entire d-dimensional domain, with a total of M points. For SepONet, as described previously in Section 3.1 we randomly sample N1 + N2 + + Na inputs for Nn points along the n-th axis, and construct a Cartesian product grid in K₁ via 9 The resulting output of PI- DeepONet has shape M x 1, and the output of SepONet has shape Ni x N2 x.. x Na. For simplicity, we assume all trunk nets are L-layer fully connected networks with hidden and output dimensions т, and that N1 = N2 = ... N₂ = N, and M = Nd, The resulting time and space complexity to compute first-order derivatives of all SepONet and PI-DeepONet outputs is provided in Table [
The first term in SepONet's 's time and space complexity is due to the forward-mode AD computation of each of the d trunk networks derivatives with respect to N inputs per axis. The second term, containing Ad, is from computing and storing the tensor product. On the other hand, PI-DeepONet individually backpropagates all M = Nd outputs, resulting in complexity scaling with Nd.
From this analysis, in the limiting case of Ad-1 >> Lr, we observe that both SepONet and PI-DeepONer have time and space complexities that include N' due to evaluations over all points in a d-dimensional space. However, since typically d & Lr, SepONet is more efficient in practice due smaller coefficients the Ny term. Morcover, the tensor product in SepONet can be greatly accelerated by GPU parallelizatio, which is not taken into account in this analysis. In the limiting case Lr >> Nd-1, the first term in SepONet's time complexity dominates O(N - d Lr²), in other words, it scales linearly with dimension and sub-linean linearly with the total number of collocation points Nd, This situation is not uncommon in many practical 2D and 3D operator learning problems.
Note that Table ! only considered first-order derivatives. Higher-order derivatives are typically needed to evaluate physics loss functions (5). Fortunately, higher-order derivatives for SepONet are computed with similar complexity to Table D since they amount to sequentially repeating the Jacobian-vector products (JVP) from 10 and [1 Lastly, did not consider the parameter update for training with a physics loss
8

---

Published in Transactions on Machine Learning Research (12/2024)
in Table 1 since it requires further assumptions about the composition of the branch network, and it is not typically the limiting computation for either PI-DeepONet or SepONet.
3.3 Universal Approximation Property SepONet
The universal approximation property of DeepONet has been discussed in Chen Chen 1995]: Lu et al. (2021). Here we present the universal approximation theorem show that proposed separable operator networks can also approximate any nonlinear continuous operators that map infinite-dimensional function spaces to others.
Thcorem (Universal Approximation Theorem for Separable Operator Networks). Suppose that a is Twuber- Wiener funclion, 9 is a sinusoidal function, X is Banach space, K X, Ki C .Rds and K2 C Rdz are three compact sets in N, Rd and R respectively. U is compact set in C (K), G is nonlinear continuous operator, which maps U into a compact set S S CK₃ x K22,, then for any > 0, there are positive integers n, m, constants oi, She Ch ek 0, E B. points w E IKd,, , w E IRd,, Iy E K, i -,,...... .,, k -,....,,,, j, 1,...,,, m, such that
G(u)(v)
Go
Σ S (x,) + +
g (wk y1 + (k) 1 (wh. 122 + (t
<e
(12)
trunks
trunks
branch
holds for all U, y (2), E K x K2.
Proof. The proof can be found in Appendix A.2
Remark 1.
The defivition of the Tauber- Wiener function is given in Appendir A.1 It is worth noting that many common-usee used activations, such as ReLU, GELU and Tanh, are Tauber- Wiener functions.
Remark 2. Here we show the opprorimation property of separable operator network with two trunk nets, by repeatedly applying trigonometric angle addition formula, it is trivial to separate y as (y1, ......, Bd) E K1 x K2 x x... x Ka and exlend 12 to trunk nets.
Remark 3. In our assumptions, restrict the activation function for the trunk nets to be sirusooddal. This choice is motivated by the natural suitability of sinusoidal functions for constructing basis functions [Stein & Shakarchi, 2017) and their empirical effectiveness in solving PDEs (Sitzmann et al. 2000)). However, it would interesting to explore whether Theorem | still holds when g is a more general activation function, such as a Tauber- Wiener function. We will leave this investigation for future work.
Remark 4. Here we assume two-layer branch network and d one-layer trunk networks with sinusoidal ac- tivations. For practical implementation, our theoretical results can be extended to multi-layer trunk networks by leveruging the Universal Approzimation Theorem (UAT) to approsimale sinusoidal functions.
Remark 5. Theorem 1 suggests the eristence of separable operator network approzimation for any non- linear continuous operator. This does not imply error bounds nor tractable scaling laws with respect to any specific error metric, nor does it provide a prescription for how define and update model parometers. Below we will provide experimental evidence that error scaling is comparable paasics-informed operator learning. Please note that error bounds for the supervised training of DeepONet have been previously derived by Lanthaler et al. 2022): similar error bounds for SepONet not provided in this work.
Pl-DeepONet when using
4 Numerical Results
This section presents comprehensive numerical studics demonstrating the expressive power and effectiveness of SepONet compared to PI-DeepONet on various time-dependent PDEs: diffusion-reaction, advection, Burgers', and (2+1)-dimensional nonlinear diffusion equations. Both models trained by optimizing the

---

Published in Transactions on Machine Learning Research (12/2024)
Table 2: Summary of PDE test problems.
Governing Law
Domain
Equation Form
Initial Condition Boundary Condition
ds Ss
Diffusion-reaction
[0, e [0.1]
ks?
u(z)
s(x, 0)
(0.f) a(1,f)= =U
ät
ds
ds
Advection
#E [0.11, 1 [0.1]
u(z)
sfx, 0) sin(s2))
e(0, sin
ōt
Burgers'
c [0, [0
0s
8'
a(x, 0) u(a)
Os =(0,t) a(1,t),
(0,
(1.6)
0s
2D Nonlinear diffusion
Q [0, 1|, ee [0.1]
aV- (aVa)
<(x, u(x)
a(x,t) 09
physics loss (equation (4)) on a dataset D consisting of imput functions, residual points, initial points, and boundary points. PDE definitions are summarized in Table -
We set the number of residual points to N, Na Ne, where d is the problem dimension and N is integer. Here, No refers to the total number of training points. For SepONet, the residual points are generated by randomly sampling N points along each axis and constructing a Cartesian product grid. In contrast, for PI-DeepONet, the Nd residual points are randomly sampled from the entire d-dimensional domain. The number of initial and boundary points per axis is set to N, = No = N = </Ne, and these points are also randomly sampled from the solution domain. For each PDE, the model size remains fixed as Ne or Ny varies. Specifically, both Pl-DeepONet and SepONet have branch and trunk networks of the same size; the main difference is that SepONet uses d independent trunk networks, one for cach axis.
We evaluate both models by varying the number of input functions (Nş) and training points (No) across four key perspectives: test accuracy, GPU memory usage, training time, and extreme-scale learning capabilities. The main rosults are illustrated in Figure and Figure 3 with complete test results reported in Appendix Loss functions, training details, and problem-specific parameters are provided in Appendix [.] and Appendix
We provide additional ablation studies for our test results using trunk networks with hyperbolic tangent activation functions in Appendix B.4.1 and varied number of input sensors for the branch net in Appendix
B.2.
B.4.2
4.1
Test accuracy
Both PI-DeepONet and SepONet demonstrate improved accuracy when increasing either the number of training points (Na) or the number of input functions (Ny), while fixing the ot parameter. This trend is consistent across all four equations tested.
For instance, in the case of the inffusioonrraccton equation, when fixing Ny = 100 and varying N₂ from 8 to 1282, the relative le error of Pl-DeepONet decreases from 1.39% to 0.73%, while SepONet's 's error reduces from 1.49% to 0.62%. Conversely, when fixing N. = 128 and varying Nş from to 100, PI-DeepONet's error drops from 34.54% to 0.73%, and SepONet'ss reduces from 22.40% to 0.62%.
4.2 GPU memory usage
While both models show improved accuracy with increasing Ne or N,, their memory usage patterns differ significantly. This divergence is particularly evident in the case of the advection equation.
When fixing N, = 100 and varying Ne, PI-DeepONet cxhibits steep increase in GPU memory consumption during training, rising from 0.967 GB at N. = 8 to 59.806 GB at N, = 128 In contrast, SepONet maintains a relatively constant and low memory footprint during training, ranging between 0.713 GB and 0.719 GB across the same range of N.. Similarly, when fixing N. = 128² and varying N, from to 100, PI-DeopONet's memory usage escalates from 3.021 GB 59.806 GB. SepONet, however, maintains a stable memory usage throughout this range.
10

---

Published in Transactions on Machine Learning Research (12/2024)
Pl-DeepONel
SepONel
1.50
60
Nel. 1, Tereet (5)
Mcmmory (CB3
Rel. X, cal-orrr (%0
7.5
1,25
20
10
80
15.0
1.00
Training
0.75
16'
327
64"
128'
8'
16
32'
84
128
8'
16
327
64"
128'
8'
18'
32'
647
128'
8
16'
32
64"
128'
8'
16'
527
64'
128'
My
N2
N.
N:
(a) (1+1)-dim) Diffusion-reaction
1)-dim Advection
15
Ret. errer (50
Treiming time (h)
1y errur
Mestory
20
Mensory
20
40
64"
128'
16
32'
128'
35
126
167
N.
(c) (1 + 1)m Burgers'
(d) (2+1)-dim 1)mm Nonlinear diffusion
Figure 2: Performance comparison of Pl-DeepONet and SepONet with varying number of training points (Ne) and fixed number of input functions (N, 100). Results show test accuracy, GPU memory usage, and training time for four PDEs. As Ne increases, both models demonstrate improved accuracy, but PI- DeepONet exhibits significant increases in training time and memory usage, while SepONet maintains better computational efficiency.
4.3
Training time
The training time scaling exhibits a pattern similar memory usage, demonstrated by the advection equa- tion example. As No increases from 82 to 128² with fixed N 100, P1-DeepONet's training time increases significantly from 0.0787 to 8.231 hours (2.361 to 246.93 ms per iteration). In contrast, SepONet maintains relatively stable training Limes, ranging from 0.0730 to 0.0843 hours (2.19 to 2.529 ms per iteration) over the same N₂ range. Similarly, when varying N, from 5 to 100 with fixed N: = 128 PI-DeepONet's training time increases from 0.3997 8.231 hours (11.991 to 246.93 ms per iteration). SepONet, however, keeps training times between 0.0730 and 0.0754 hours. These results demonstrate SepONet's superior scalability in terms of training time. The ability to maintain near-constant training times across wide range of problem sizes is significant advantage, particularly for large-scale applications where computational efficiency is crucial.
4.4 Extreme-scale learning
The Burgers' and noulinear diffusion equations highlight SepONet's capabilities in extreme-scale learning scenarios.
For the Burgers' equation, PI-DeepONet encounters memory limitations at larger scales. As seen in Figure cc), Pl-DeepONet can only compute up to N. 642, achieving relative lo error of 13.72%. In contrast, SepONet continues to improve, reaching 7.51% error at N. 1282. The nonlinear diffusion equation further emphasizes this difference. In Figure 3 3(d), PI-DeepONet results are entirely unavailable due to out- of-memory issues. SepONet, however, efficiently handles this complex problem, achieving a relative lz error of 6.44% with N. = 1283 and N, = 100. Table demonstrates SepONet's 's ability to tackle even larger scales for the Burgers' equation. It achieves a relative &g error as low as 4.12% with N 5122 and N, = 800, while maintaining reasonable memory usage (10.485 GB) and training time (0.478 hours). These results underscore SepONet's capability to handle extreme-scale learning problems beyond the reach of Pl-DeepONet due to computational constraints.
5 Discussion
The field of operator learning faces a critical dilemma. Deep Operator Networks (DeepONets) offer fast training but require extensive data generation, which can be prohibitively expensive impractical for
11

---

Published in Transactions on Machine Learning Research (12/2024)
Pl-DeepONet
SepONet
Rel Is ereer (%g
Memnory (G3A
30
Bel. Is error (%)
(k)
7.5
20
20
35.0
Tratning
10
2.5
100
20
50
130
(a) (1+1)-dim Diffusion Diffusion-reactionn reaction
(b) (1++1) d-mm Advection
Rel f, ervor (%)
Mamory (CIU)
Rel. f, arne (5)
-30
(CB)
Training time
Veginiog Lsso
20
20
20
10
10
0.2
10
20
50
100
10
20
50
100
10
50
100
10
20
50
100
10
20
50
100
10
20
50
100
Nr
N/
(1+1)-dim Burgers'
(2+1)-dim Nonlinear diffusion
Figure 3: Performance comparison of PI-DeepONet and SepONet with increasing number of input functions (Nr) and fixed number of training points (N, = 128d, where is the problem dimension). Both models show improved accuracy with increasing Ny, but Pl-DeepONet's computational resources scale poorly compared to SepONet's more efficient scaling. Note: PI-DeepONet results for the 1i-))diiimsiinnl diffusion equation are unavailable due to memory constraints.
Table 3: Additional SepONet results for Burgers' equation, demonstrating that larger N₂ and Ny can be used enhance accuracy with minimal cost increase.
Metrics \ Ne & ,N
128 & 400
128 & 800
256² & 400
256 & 800
5122 & 400
5122 & 800
Relative 8 error (%)
6.60
6.21
5.68
4.46
5.38
4.12
Memory (GB)
0.966
1.466
2.466
4.466
5.593
10.485
Training time (hours)
0.0771
0.0957
0.1238
0.1717
0.2751
0.478
complex systems. Physics-informee DeepONets (PI-DecpONets) relax the data requirement but at the cost of resource-intensive training processes. This creates challenging trade-off: balancing resource allocation either before training (data generation) or during training (computational resources).
Our proposed Separable Operator Networks (SepONet) effectively address both of these concerns. Inspired by the separation of variables technique typically used for linear PDEs, SepONet const own basis functions to approximate nonlinear operator. SepONet's expressive power guaranteed by the universal approximation property (Theorem 1 ensuring it can approximate any nonlincar contimuous operator with arbitrary accuracy. Our numerical results corroborate this theoretical guarantee, demonstrating SepONet's ability to handle complex, nonlinear systems efficiently.
By leveraging independent trunk networks for different variables, SepONet enables an efficient implementa- tion via forward-mode automatic differentiation (AD). This approach achieves remarkable efficiency in both data utilization and computational resources. SepONet trained solely by optimizing the physics loss, elimi- nating the need for expensive simulations to generate ground truth PDE solutions. In terms of computational resources, SepONet maintains stable GPU memory usage and training time, even with increasing training data and network size, in contrast to P1-DeepONet's dramatic resource consumption increases under similar scaling. We anticipate that SepONet's advantages will allow it to tackle more challenging physics-informed operator learning problems, such as the Navier-Slokes equations (Jin et al.) 2021), where both input and output functions are vector-valued. These are problems that Pl-DeepONet may struggle to train on due resource constraints. As an example, we have considered a. (2+1)-dimensional Navier-Stokes equation, previously investigated in the context of PINN (Cho results be found in Appendix D.1
2024
Wang 2024). Some early, preliminary
12

---

Published in Transactions on Machine Learning Research (12/2024)
However, SepONet has certain limitations that warrant further research and development. The mesh grid structure of SepONet's solution, while enabling efficient training through forward-mode AD, may limit its flexibility in handling PDEs with irregular geometries. Addressing this limitation could involve developing adaptations or hybrid approaches that accommodate more complex spatial domains (Li ct al. 2023 Serrano
et al. 2024 Fang el al. 2024), potentially expanding SepONet's 's applicability to a broader range physical
problems.
Additionally, while the linear decoder allows for efficient SepONet implementation, very large number of basis functions may be needed for accurate linear representation in some problems (Seidhman et al., 2022 Developing a nonlinear decoder version of SepONet will be useful to balance accuracy and efficient training. Moreover, implementing an adaptive weighting strategy (Wang et al. 2021a 2022a b) for different loss terms in the physics loss function, instead of using predefined fixed weights, could lead to improved accuracy and faster convergence.
Finally, empirical observations suggest that training accuracy and robustness improve with an increase in imput training functions. However, the neural scaling laws in physics ahysice-informed operator learning remain unexplored, presenting intriguing theoretical challenge for future investigation.
Acknowledgments
We would like to thank Wolfger Peelaers for the valuable discussions and insightful comments. Z. Zhang was supported by NSF Awards #2328281 and #1846476
References
Md Ashiqur Rahman, Zachary Ross, and Kamyar Avizzadenesheli. U-no: U-shaped neural operators. arXiv e-prints, pp. arXiv-2200 2204, 2022.
Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyovich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: survey. Journal of machine learning research, 18(153):1-43, 2018.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman--Milne, and Qiao Zhang. JAX: compos- able transformations of Python-NumPy programs, 2018. URL http: sttp://github com/ gom/google/jax
Shengze Cai, Zhicheng Wang, Lu Lu, Tamer Zaki, and George Em Karniadakis. Deepm&mnet: Inferring the electroconvection multiphysics fields based operator approximation by neural networks. Journal of Computational Physics, 436:110296, 2021.
Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE transactions on neural networks, 6(4):911 917, 1995.
Pao-Hsiung Chiu, Jian Cheng Wong, Chinchun Ooi, My Ha Dao, and Yew-Soon Ong. Can-pinn: A fast plnysics-informod neural network based on maaplll-anttmatii-mmmeeeee diffcrentiation method. Computer Methods in Applied Mechanics and Engineering, 395:114909, 2022.
Junwoo Cho, Scungtae Nam, Hyunmo Yang, Scok-Bac Yun, Youngjoon Hong, and Eunbyung Park. Separable plysics-informed neural net workorss. Advances in Neural Information Processing Systems, 36, 2024.
Tobin A Driscoll, Nicholas Hale, and Lloyd N Trefethen. Chebfun guide, 2014.
Vasiliy A Es' kin, Danil V Davydov, Julia V Gur'eva, Alexey O Malkhanov, and Mikhail F. Smorkalov. Separable dhysics-informed neural networks for the solution of elasticity problems. arXiv preprint arXiv:2401.13486, 2024.
Zhiwci Fang, Sifan Wang, and Paris Perdikaris. Learning only on boundaries: a physics-informed neural operator for solving parametric partial differential equations in complex geometries. Neural Computation, 36(3):475-498, 2024.
13

---

Published in Transactions on Machine Learning Research (12/2024)
Pulkit Gopalani, Sayar Karmakar, DIBYAKANTI KUMAR, and Anirbit Mukherjee. Towards size- independent generalization bounds for deep operator nets. Available al SSRN 4763746.
Jiaqi Gu, Zhengqi Gao, Chenghao Feng, Hanqing Zhu, Ray Chen, Duane Boning, and David Pan. Neurolight: A physics-agnostic neural operator enabling parametric photonic device simulation. Advances in Newral Information Processing Systems, 35:14623 14636, 2022.
Jayesh K Gupta and Johannes Brandstetter. Towards mullisppttoteemponlsssaall generalized pde modeling.
arviv preprint XrX:::2209.1000 15616, 2022.
Di Hc, Shanda Li, Wenlei Shi, Xiaotian Gao, Jia Zhang, Jiang Bian, Liwei Wang, and Tic-Yan Liu. Learn- ing physics-informed neural networks without stacked back-propagation. In International Conference on Artificial Intelligence and Statistics, pp. 3034 3047. PMLR, 2023.
Zheyuan Hu, Khemraj Shukla, George Em Karniadakis, and Kenji Kawaguchi. Tackling the curse of dimen- sionality with physics-infonl informed neural networks. Neural Networks, 176:106369, 2024.
Arieh Iserles. A first course in the numerical analysis of differential equalions. Number 44. Cambridge university press, 2009.
Zhongyi Jiang, Min Zhu, Dongzhuo Li, Qiuzi Li, Yanhua O Yuan, and Tu Lu. Fourier-mionet: Fourier- enhanced multiple-input neural operators for multiphase modeling of geological carbon sequestration. arXiv preprint Xio:23000 04 778, 2023.
Xiaowei Jin, Shengze Cai, Hui Li, and George Em Karniadakis. Nsfnets (navier-stokes flow nets): Physics- informed neural networks for the incompressible navier-stokes equations. Journal of Computational Physics, 126:109951, 2021.
Karthik Kashinath, M Mustafa, Adrian Albert, JL Wu, C Jiang, Soheil Esmacilzadeh, Kamyar Azizzade- nesheli, R Wang, Ashesh Chattopadhyay, A Singh, et al. Physics-informed machine learning: case studies for weather and climate modelling, Philosophical Transactions of the Royal Society A, 22944)))) 2021.
Kamil A Khan and Paul Barton. A vector forward mode of automatic differentiation for generalized derivative evaluation. Optimization Methods and Software, 30(6):1185-122) 12122, 2015.
Patrick Kidger and Cristian Garcia. Equinox: neural networks in JAX via callable PyTrees and filtered transformations, Differentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXio preprint arXiv:1412. 6980, 2014.
Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine learning accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118 (21):22011884118 2021.
Katiana Kontolati, Somdatta Goswami, George Em Karniadakis, and Michael D Shields. Learning nonlinear operators in latent spaces for real-time predictions of complex dynamics in physical systems. Nature Communications, 15(1):5101, 2024.
Jean Kossaifi, Nikola Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar. Multi-grid tensorized fourier neural operator for high-resolution pdes. arXiv preprint arXiv:2310.002 2023.
Samuel Lanthaler, Siddhartha Mishra, and George E Karniadakis. Error estimates for deeponets: A deep learning framework in infinite dimensions. Transactions of Mathematics and Its Applications, 6(1) 2022.
Zijie Li, Dule Shu, and Amir Barati Farimani. Scalable transformer for pde surrogate modeling. Advances in Neural Information Processing Systems, 36, 2024.
14

---

Published in Transactions on Machine Learning Research (12/2024)
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint Xrx:220010. 08895, 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadeneshelli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arXiv preprint Xte:200003 03485, 2020b.
Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with lcarned deformations for pdes on general geometries. Journal of Machine Learning Research, 24(388); 1-26, 2023.
Chensen Lin, Zhen Li, Lu Lu, Shengze Cai, Martin Maxey, and George Em Karniadakis. Operator learning for predicting multiscale bubble growth dynamics. The Journal of Chemical Physies, 154(10), 2021.
Hao Liu, Haizhao Yang, Minshuo Chen, Tuo Zhao, and Wenjing Liao. Deep nonparametric estimation of operators between infinite dimensional spaces. ar preprint arXim:2201.002) 2022a.
Ziyue Liu, Xinling Yu, and Zheng Zhang. Tt-pinn: tensor-compressed neural pdc solver for edge computing, arXiv preprinl or Xrc:22200.011 2022b.
Ziyue Liu, Yixing Li, Jing Hu, Xinling Yu, Shinyu Shiau, Xin Ai, Zhiyu Zeng, and Zheng Zhang. Deepoheat: operator taariing-bsses ultra-fast thermal simulation 3d-ic design. In 2023 60th ACM/TEEE Design Automation Conference (DAC), pp. 1 6. IEEE, 2023.
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via. deeponet based on the universal approxination theorem of operators. Nature machine intelligence, 3(3):218 2292 2021.
Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, and George Em Karniadakis. A comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data. Computer Methods in Applied Mechanics and Engineering, 393:114778, 2022a.
Lu Lu,, Raphaël Pestourie, Steven G Johnson, and Giuseppe Romano. Multifidelity deep neural operators for efficient learning of partial differential equations with application to fast inverse design of nanoscale heat transport. Physical Review Research, 4(2):023210, 2022b.
Luis Mandl, Somdatta Goswami, Lena Lambers, and Tim Ricken. Separable plysics-informed deeponet: Breaking the curse of dimensionality physics-informed machine learning. Computer Methods in Applied Mechanics and Engineering, 434:117586, 2025.
Zhiping Mao, Lu Lu, Olaf Marxen, Tamer A Zaki, and George Em Karniadakis. Deepm&mmet for hyper- sonics: Predicting the coupled flow and finite-rate chemistry behind normal shock using neural-network approximation of operators. Journal of computational physics, 447:110698, 2021.
Charles C Margossian. review of automatic differentiation and its efficient implementation. Wiley inter- disciplinary reviews: data mining and knowledge discovery, (4) 2019.
Jaemin Oh, Seung Yeon Cho, Seok-Bae Yun, Eunbyung Park, and Youngjoon Hong. Separable physies- informed neural networks for solving the bgk model of the boltzmann equation. arXiv preprint arXiv:2403.0m 06342, 2024.
Oded Ovadia, Adar Kahana, Panos Stinis, Eli Turkel, and George Em Karniadakis, Vito: Vision transformer- operator arXiv preprint arXiv:2208.0880 08891, 2023.
Ravi Shankar Palani. spinop.m (spin operator).
satpaannge/71333-spnoooop 28, 2024.
2024. MATLAB Central File Exchange. Retrieved June
15

---

Published in Transactions on Machine Learning Research (12/2024)
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzaenesheli, et al. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint 2tl::2202.111444 2022.
Stephen B Popc. Turbulent flows. Measurement Science and Technology. 12(11):2020-0022 2001.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics- neural networks: deep learn- ing framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Compulational physics, 378:686 707666667777-- 2019.
Matthias Seeger. Gaussian processes for machine learning. International journal of neural systems, 14(02): 69 106, 2004.
Jacob Seidman, Georgios Kissas, Paris Perdikaris, and George J Pappas. Nomad: Nonlinear manifold decoders for operator learning. Advances in Neural Information Processing Systems, 35:5601-5611 2022.
Louis Serrano, Lise Le Boudec, Armand Kassaï Koupaï, Thomas X Wang, Yuan Yin, Jean-Noël Vittaut, and Patrick Gallinari. Operator learning with neural fields: Tackling pdes general geometries. Advances in Newrul Information Processing Systems, 36, 2024.
Khemraj Shukla, Vivek Oommen, Ahmad Peyvan, Michael Penwarden, Luis Bravo, Anindya Ghoshal, Robert M Kirby, and George Em Karniadakis. Deep neural operators can serve as accurate surrogates for shape optimization: case study for airfoils. arXiv preprint arXiv:2302.0000 2023.
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Implicit neural representations with periodic activation functions. Advances in neural information processing systems, 33: 7462-7473, 2020.
Elias M Stein and Rami Shakarchi. Fourier analysis: an introduction, volume 1. Princeton University Press, 2011.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. arXiv preprint ar Xrxin:211 138002, 2021.
Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow pathologies in physics-informed neural networks. SIAM Journal Scientific Compuding, A3(5):A3055-A33)) A3081, 2021a.
Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with ahysics-informer deeponets. Science advances, 7(40)=eabi8605, 2021b.
Sifan Wang, Hanwen Wang, and Paris Perdikaris. Improved architectures and training algorithms for deep operator networks. Journal of Scientific Computing, 92(2):35, 2022a.
Sifan Wang, Xinling Yu, and Paris Perdikaris. When and why pinns fail train: neural tangent kernel perspective. Journal of Computational Physsis, 449:110768, 2022b.
Sifan Wang, Shyam Sankaran, and Paris Perdikaris. Respecting causality for training physics-informet neural networks. Computer Methods in Applied Mechanics and Engineering, 421:116813, 2024.
Karl Weierstrass. Über die analytische darstellbarkett sogenannter willkürlicher functionen ciner reellen veränderlichen. Silzungsberichle der Königlich Preußischen Akademie der Wissenschaflen zu Berlin, 2 (633-639)3364, 1885.
Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson. U-fno enhanced fourier neural operator- deep-lcarning model for multiphase flow. Advances in Water Resowrces, 163:104180, 2022.
Yequan Zhao, Xinling Yu, Zhixiong Chen, Ziyue Liu, Sijia Liu, and Zheng Zhang. Tensor-compressed back- propagation-fiee free training for (physics- inffrres-) neural networks. arXiv preprint arXiv:2308 09858, 2023.
16

---

Published in Transactions on Machine Learning Research (12/2024)
A Universal Approximation Theorem for Separable Operator Networks
Here we present the universal approximation theorem for the proposed separable operator networks, originally
written in Theorem 1 and repeated below in Theorem 5 We begin by reviewing established theoretical results
approximating contimuous functions and functionals. Following this review, we introduce the preliminary
lemmas and proofs necessary for understanding Theorem
We refer our readers to Chen & Chen (1955
Weierstrass (1855) for detailed proofs of Theorems 2
Main notations are listed in Table
A.1 Preliminaries and Auxiliary Results
Definition 1 (Tauber-Wiener (TW)). If function g IR - R (continuous disconlinuous) satisfies that all the linear combinations Ei 4g (Aşx + 04), X E IR, 0, R, G E ERR i 1.2,...,NN N, are dense in every Cla, b), then g is colled Tauber- Wiener (TW) function.
Remark 1 (Density in bl), A set of functions said to be dense in Cla. b] if every function in the space of continuous functions on the interval [a,b] can be opprorimated arbitrarily closely by functions from the set.
Definition 2 (Compact Set). Suppose that a Banach space, V V X is called a compact set in X, if for every sequence [= with all Fn E EV, there is subsequence (Im) J; which converges to some element E V.
Theorem (Chen & Chen 1995)). Suppose that K is a compact set in R", S compact set in C(K), g (TW), then for any 0, there exist posilive integer N, real numbers 0,, veclors we E RR", i= 1,.... N,, which are independent of f E C(K) and constants c(ff,, i i= I.....N depending f, such that
N f(x) Σ Giffg (wr. x 0.) <c
(13)
holds for all x E K and f E S. Moreover, each cf(f) is linear continuous functional defined on S.
Theorem 3 (Chen & Chen (1955 ). Suppose that a E (TW), X is Banach Space, K C X is compact set, u is a compact set in C(K), f is a continuous functional defined on U, then for any >> 0, there are
positive integers N, m points Iı. Im E K, and real constants Ci, Oi, Eig, i 1,..., N, j = 1,.... m, such that
N
S()) - Σ Go 1=1
Σ Eyg (2,) + +0,
<e
(14)
holds for all u U.
Theorem 4 (Weierstrass Approximation Theorem Weierstrass 1885
Suppose I Claab), then for every <> 0, there cxists a polynomial p such that for all x in [a, b). we have (x() - p(x) <
Corollary Trigonometric polymomials are dense in the space of continuous and periodic functions Č10 2m (f C(0,2T] f(0) f(2m)}
Proof. For any / č 2T), extend it to a 27-periodic and continuous function f defined on R. It suffices to show that there exists a trigonometric polynomial that approximates f within any e >> 0. We construct the continuous even functions of 2n period g and h as:
g f(0) + f(-0) and N(0) f(0) f(-0) sin (0).
(15)
2
2
Let d6x)) = g(arccos v) and (xx = h(arccos I). Since o, th are continuous functions on [-1, by the Weierstrass Approximation Theorem 4 for any 6 > 0, there exist polynomials p and q such that
lő(x) - p(x) < and [+(x) q(x)) < 4
(16)
17

---

Published in Transactions on Machine Learning Research (12/2024)
Table 4: Notations and Symbols
X
some Banach space with norm | llx Euclidean space of dimension d
Rd
K
some compact set in Banach space
C(K)
Banach space of all continuous functions defined on K, with norm [[/coo max [f()) the space of functions in Cla, satisfying f(a) = f(b)
b]
some compact some input function the space of input functions some continuous operator
in C(K)
((x)
u
G(a)(v) s(v)
some output function that is mapped from the corresponding input function by the operator G the space of output functions all the Tauber- Wiener functions activation function for branch net and trunk nets in Theorem
(TW)
and
(21. 22.,22.-.. mm)
sensor points for identifying imput function
rank of some deep operator network or separable operator network
n.
operator network size peperparametess in Theorem 5
holds for all [-1,1]. Let cos 0, it follows that
(g(0) p(cos0) < and (1(0) - g(cos
[)| ((o)))(cos00)
(17)
for 0 E [0, T.. Because and cosine are even and 27-periodic, of g and h, and the fact |sin fl < 1, ssin? 0 < 1, we have
holds for all 0 E R. From the definitions
f(0) f(-0) sin A p(cos f) sin 0
f(0) f(-0) sin" g(cos 0) sin
and
(18)
4'
Using the triangle inequality, obtain
( sin? 0 [p(cos 0) sin? 0 + g(cos 0) sin 0]| < 2'
(19)
Applying the same analysis to
5(0) f(0 - 7)+ f and h((0) f(0 7)) f(-0 + 2)
sin(0),
(20)
2
we can find polynomials and such that
f (0+ 2 sin² [r(cos 8) sin" + s(cos 0) sin 0
(21)
2
holds for all 0. Substituting 0 with 2 gives
(() cos 0 [r(sin f)) cos² 0 - s(sin f) cos 8] <
C
(22)
2'
By the triangle inequality, combining 22 and 19
gives
[s()) [r(sin 0) cos2 ssin 0) cos p(cos 0) sin? + g(cos 0) sin 0] <e
(23)
holds for all 0. Thus, the trigonometric polynomial
rssin 0) cos 0 s(sin 0) cos 0 + p(cos 0) sin" 0 + q(cos 0) sin 0
(24)
is e-approximation to
18

---

Published in Transactions on Machine Learning Research (12/2024)
Remark 2. If p(x) is a polynomial, it is easy to verify that p(cos is trigonometric polynomial due to the fact cos" 0 = Ei=
(t)
2 cos ((n 2k)
Prior proving Theorem 5 we need establish the following lemmas.
Lemma 1. Sine is a Tauber- Wiener function.
Proof. Assuming the interval to be [0, π| first. For every contimuous function f on [0, π] and any 6 > 0, we
can extend to continuous function F on [0, 2x) so that F(x) = f(x) on [0, x] and F(2x) = F(0). By Lemma there exists a trigonometric polynomial
p(1) Og +
Σ On cos((er + b, sin(nx)
(25)
such that
sup [f(x) p(x) S
sup [F(z) - p(z)] <c. =E(0,2m]
(26)
zE(0,x]
Let Cu aus do 0, 0u 2 C2n-1
ons A n, 82n-1 0, C2n On A2n n, 82n for
1. ..... N, p(x) is redefined as
2N
p(x) Σ
sis (A,r + 0z)
(27)
0=
Thus we have
2N
[(x)
Σ
c sin (Apz + 0i) <E
(28)
1=0
for E [0, T.. Now consider continuous function 9 on Ja, b). define f(x) E C[0, :00:=]==) == g0 ("=0x + a), then by 28
we have
2N
g(x) - c sin
Σ 1=0
Tdaa b-a b-a + 0:) )
(29)
holds for all x E (a, b). Therefore, follows that for any continuous function g on [a, b] and any <> 0, we can approximate 9 within by choosing N sufficiently large and adjusting 6s As, 0 accordingly. Hence, the set of all such linear combinations of sin(x) is dense in Cla, b), confirming that sin(x) function.
Sauber-Wiener
Remark 3. It is straightforward to conclude that all sinusooddl functions are Tauber- Wiener functions. Lemma 2. Suppose that V1 S X1, V22 S X2 are two compact sets in Banach spaces X2 and X2, respectively, then their Cartesian product V1 x V2 is also compact.
Proof. For every sequence { x2] in Vi V22 since Vi is compact, {=1] has subsequence { } that converges to some element x² V1. As well, since V2 is compact, there cxists subsoquence [m2 that converges to x² e V2. It follows that {rh, 2 converges to (x', E V1 x V2, thus V12 x V/2 is compact. ]
Lemma 3. Suppose that X is a Banach space, Ki C Xi, K2 respectively. u is a compact set in C(K1), then the range G(U) of the continuous operator G from to C(K2) compact C(K2).
X2 are two compact sets in X₁ and X2,
Proof. For every sequence {fr} in U, since 4 compact, there exists a subsequence (fny) that converges some function U. Since G is continuous, the convergence fni + f in C(Ki) implies
G(fm)) in C(K2).
(30)
Thus, for every sequence (G(fx) in G(U), there exists subsequence (G(fne) }} that converges to G(f) e G(U). Therefore, the range G(U) of the continuous operator G is compact in C(K2).
19

---

Published in Transactions on Machine Learning Research (12/2024)
A.2 Universal Approximation Theorem for SepONet
Theorem 5 (Universal Approximation Theorem for Separable Operator Networks). Suppose that a E (TW), is a sirusoidal function, X is Banach Space, K S X, K1 C Rd and K2 S TR are three compact sets in X, Rd and R respectively, U is a compact set in C (K), G is nonlinear continuous operator, which maps _ C K1 x K2), then for any > 0, there are posilive inlegers n, m, conslants oi G Le, gk 0 E R, points w E Rd,, wht &Rd,, 2z E K1, I...., n, k I.....r, L..... m, such that
U into a compact set S
m
G(u)(v) ΣΣ Go Σ k=
E (2,) + +00 9 (wt y1 (k) g (wk. 122 (2) <E
(31)
holds for all u E u, y = (y1. E K₁ x K2.
Proof. Without loss of generality, we can assumc that is sine function, by Lemma | we have g E (TW): From the assumption that Ki and K₂ are compact, by Lemma 2 K x K₂ is compact; Since G is a continuous operator that maps U into C(K x K2), it follows that the range G(U) = [G(u): u U is compact in C(K₁ K2) due Lemma 3 Thus by Theorem g for any > O, there exists a positive integer N, real numbers eg (G(u)) and Sk. vectors wk Rdi+d2, k ,,..,,,, N, such that
((w)())
Σ ca (G(u)) (ux y + (c) 2 1=1
(32)
holds for all E K₁ x K2 and u E C(K). Let (wt, wā) = wh, where w E Rdd and w E Rdd, Utilizing the trigonometric angle addition formula, we have
g (wx - y + (c) = y (wh y1 + (k) W U2 + 2)) + (wk 81 + Cx + 3 (wt v2)
(33)
Let r = 2N, cN+x(G(u)) = cu( cu(G())), w = wh: wh++ = who &t = Ss. 2 = 2 for = 1, 1,..., N, and let
C = Ç. + n, 2 for N 1, 1....,"t r, equation 32 can be expressod as:
G(u)(v) - Σ (C(u)) (w) i1 + (k) y (wt y2 + (
<
(34)
2
Since G a continuous operator, according to the last proposition of Theorem 2 we conclude that for
each 1: = -1.....2 2N, es (G()) is contimuous functional defined on U. Repeatedly applying Theorem 3 for each ht l,...,.2. 2N, ca (G(u), can find positive integers ma constants of <t, 0f E R and x E K ni. j 1,..., my, such that
m
(G(()) Σ
Σ
< u((z)) + 0f)
e 2L
(35)
holds for all I,...,m r and E,, where
I. =
sup Kn
la (ut in + (L) g (N v2 -c2)
(36)
Substituting 35 into 34/, we obtain that
G(u)(v)
G Σ Stuu (z,) + 0² of
(wh i1 + (t) (wl 32 + (2) <E
(37)
k=1 1=1
20

---

Published in Transactions on Machine Learning Research (12/2024)
holds for all U, y1 E K1 and y/ K2.
Let max; ng, max; my. For all ny <is n, let of 0.
For all m& < S m, let Ex = O. Then 37
can be rewritten as:
m + 0 1 (wi in + (k) g (wk y2 + (2) <e
(v))(v)=
(
(38)
1=1
j=1
which holds for all u E U, 31 E K and 1/2 E K2. This completes the proof of Theorem 5
B
PDE Problem Definitions, Training details, and Complete Test Results
B.1
PDE Problem Definitions
All PDE test. problems exhibited in Section 4 are described in the subsections below.
B.1.1 Diffusion-Reactinn Systems
We set the diffusion coefficient D 0.01 and the reaction rate à = 0.01. The input training source terms are sampled from mean-zero Gaussian random field (GRF) (Seger 2004) with a length scale 0.2. To generate the test dataset, we sample 100 different source terms from the same GRF and apply a second-order implicit finite difference mothod fIserles 2009) to obtain the reference solutions on a uniform 128 x 128 grid. The
specific physics loss terms equation (5 are defined as follows:
Ny Nr
OG (G)( (1))
0 (Go("()) (
Presidual
Σ
D 0(
N₂
k (G)
(o) (x()),,01))) .((0))
w "
(39)
Livitial
N1N1 (a
Σ
(co(d())))
Lboundary
NyNo
B.1.2
Advection Equation
The input training variable coefficients are strictly positive by defining u(x) v(x) - min, x(2) +1, where sampled from GRF with length scale 0.2. To create the test dataset, we generate 100 new coefficients in the same manner that not used in training and apply the Lax-Wend scheme (Iserles) 2009 to solve the advection equation on uniform 128 x 128 grid. The specific physics loss terms in equation are defined as follows:
Nr
(Golu())2 (1) (00)
a0 (2(0)) OG (Go( (00)
1
ΣΣ
Cresidual
N,N,
06
())
N₁
Linitial N,N, Σ
1
(o)
(1)
sin( (2)
(40)
Lboundary N₁
1
Σ
Σ
Go
(0, ()) (0))
sin -
v (1) x1))
12
21

---

Published in Transactions on Machine Learning Research (12/2024)
B.1.3 Burgers' Equation
The imput training initial conditions are sampled from a GRF N - (- using the Chebfun
521)
package (Driscoll al. 2014), satisfying the periodic boundary conditions. Synthetic test dataset. consists of 100 unseen initial functions and their corresponding solutions, which generated from the same GRF and are solved by spectral method on 101 x 101 uniform grid using the spinOp library PPlann 2024 respectively. The corresponding physics loss terms in equation
5 are defined as:
Ny
(Golu())) ( acl
Go(u()))aa)) (10)) (Go(ul)))a (10)
1 N
Cresidual
Ox.
Go((()) (1) (10)) 2
N
Linstial
1
Go(
(1)
N1
i=
Ny
N
aGg (u))) t(0)
(Go(u()) (x) {((1)
2
Go [aa +
Cboundory
=
Ny
dx 1
Ox z=1
j=1
(11)
B.1.4 2D Nonlinear Diffusion Equation
The input training imitial conditions are generated as a sum of Gaussian functions, parameterized as:
3
u(2, v) Σ A expl wi w.((x x2)² + (3) y1)
(42)
where A; ~ U(0.2,005 are amplitudes, wy ~ ((10, 20) are width parameters, and (2,, ~ 0(-5..0.55 are center coordinates. We also generate 100 unseen test initial conditions using this method. The nonlincar diffusion equation then solved using explicit Adams method to obtain reference solutions on a uniform 101 101 spatial grid with 101 time points. Physies loss terms in equation ( for this problem are:
No (0) .11)) -av (/
N.
Cresidual
1
ΣΣ
NGN,
015)
N,
Cinitial NGNI
1
Σ
(a) - u (30
(43)
i= 1=
Lboundary Ny
ΣΣ
Go
(u( 2000))
B.2 Training Details and Hyperparameters
Both PI-DeepONet and SepONet were trained by minimizing the physics loss (equation using gradient descent. with the Adam optimizer (Kingma & Ba 2014). The initial learning rate is 1 10 and decays by factor of 0.9 every 1,000 iterations. Additionally, resample input training functions and training points (including residual, initial, and boundary points) every 100 iterations.
Across all benchmarks and on both models (SepONet and Pl-DeepONet), we apply Tanh activation for the branch net and Sine activation for the trunk net. We note that no extensive hyperparameter tuning was performed for either PI-DecpONet or SepONet. The code this study is implemented using JAX and Equinox libraries (Bradbury al. 2018 Kidger & Garcia) 2021), and all training was performed on single NVIDIA A 100 GPU with 80 GB of memory. Training hyperparameters are provided in Table
22

---

Published in Transactions on Machine Learning Research (12/2024)
Table 5: Training hyperparameters for different PDE benchmarks
Hyperparameters PDEs Diffusion-reactinn
Advection
Burgers'
2D Nonlinear diffusion
# of sensors Network depth
128
128
101
10201 (101 x 101)
6
7
7
Network width
100
100
128
# of training iterations
50k
120k
80k
80k
Weight coefficients (Ar An)
100 / 100
20
20
B.3 Complete Test Results
We report the relative & error, root mean squared error (RMSE), GPU memory usage and total training time as metrics to assess the performance of PI-DeepONel and SepONet. Specifically, the mean and standard deviation of the relative lg error and RMSE are calculated over all functions in the test dataset. The complete test results are shown in Table i and Table
Table 6: Performance comparison of PI-DeepONet and SepONet with varying number of training points (Ne) and fixed number of imput training functions (N, 100). The '_' symbol indicates that results not available due to out-of-memory issues.
Equations
Metrics
Models
g
16
32"
64
128
Pl-DeepONet
1.39
0.71
1.11 # 0.59
0.87 0.41
0.83 : 0.35
0.73
0.34
Relative &2 error (%)
SepONet
1.49 0.82
0.79 + 0.35
0.33
0.62
0.28
0.62 0.26
Diffusion- RellutonnReeation
RMSE (x10-2)
Pl-DeepONet
0.58
0.29
0.46 + 0.22
0.36
0.20
0.32 0.18
SepONet
0.62 0.28
0.35 + 0.22
0.32
0.23
0.28 0.20
0.29 0.21
Memory (GB)
PeDeerONNet
1,227
9.175
SepONet
0.715
0.717
0.715
0.717
0.719
PI DeepONet
0.0433
0.0641
0.1497
0.7252
Training time (hours)
SepONet
0.0418
Pl-DeepONet
9.27 1.94
7.55 + 1.86
6.69
1.95
5.72
Relative & error (%)
SepONet
14.29 2.65
11.96
2.17
6.14 =
5.80 + 1.57
4.99
1.40
Advection
RMSE (x10-2)
Pl-DeepONet
SepONet
5.88 1.34
9.06
7.58 + 1.55
4.79 + 1.27
4.31 1.23
3.69 + 1.07
4.24 + 1.29
3.63
3.17
2
Pl-DeepONet
0.967
1.741
5.103
17.995
59.806
Memory (GB)
SepONet
0.713
0.715
0.715
0.715
0.719
Training time (hours)
NeepoONet
0.0787
0.1411
0.4836
2.3987
8:231
SepONet
0.0843
0.0815
0.0844
0.0726
0.0730
Pl-DeepONet
29.33 +: 3.85
20.31
+: 4.31
14.17 + 5.25
13.72
= 5.59
Relative error (%)
SepONet
29.42 3.79
31.53
3.44
28.74 + 4.11
11.85
4.06
7.51 4.04
Burgers'
RMSE (x10-2)
PI-DeepONet
SepONet
4.19 +: 2.79
4.18 + 2.74
2.82 + 1.86
4.4 + 2.81
2.23 2.10
4.11
2.76
2.20 + 2.13
1.80 + 1.60
1.23 1.44
Pl-DeepONet
1.253
2.781
5.087
18.001
Memory (GB)
SepONet
0.603
0.605
0.603
0.605
0.716
Training time (houss)
PeepoOONtt
0.1497
(123355
0.6431
3.2162
SepONet
0.0706
0.0719
0,0716
0.0718
0.0605
Pl-DeepONet
17.38 + 5.56
9.90 4: 2.91
Relative error (%)
SepONet
16.10 + 4.46
12.11 + 3.89
6.81 1.98
6.73 + 1.96
6.44 1.69
Noulinear diffusion d= 3
RMSE (x10-2)
Pl-DeepONet
1.86 + 0.62
1.04 + 0.23
SepONet
1.72 + 0.49
1.29 4 0.37
0.72 0.19
0.71 + 0.17
0.68 0.15
Pl-DeepONet
6.993
37.715
Memory (GB)
SepONet
3.471
2.897
2.899
2.897
13.139
Training time (hours)
Pe-DeeONNet
0.5836
6.6399
SepONet
0.1044
0.1069
0.1056
0.1456
0.5575
23

---

Published in Transactions on Machine Learning Research (12/2024)
Table 7: Performance comparison Pl-DeepONet and SepONet with varying number of input training functions (N₂) and fixed number of training points (Ne = 1284, d indicated by problem instance). The symbol indicates that results are not available due out-of-memory issues.
Equations
Metries
Models
20
50
100
Pe-DeepONet
34.54 27.83
1.72
1.00
0.91
0.46
Relative f2 eroor (%)
SepONet
22.40 12.30
1.19
0.74
0.73
0.32
0.62
0.26
Diffusimn-Reacioon Reaction
RMSE (x10 2)
Detp
SepONet
14.50
9.34 5.37
1.36
0.50
0.29
0.34
0.28
0.25
0.29
0.21
d=
Net
2.767
9.239
17.951
Memory (GB)
SepONet
0.719
0.719
0.717
0.717
0.719
DeepONet.
0.1268
0.2218
0.5864
1.4018
2,8025
Training time (hours)
SepONet
0.0375
0.0390
0.0370
0.0317
0.0326
Pl-DeepONet
9.64
8.77 2.23
7.57
1.98
6.69 + 1.93
5.72
Relative l2 error (%)
SepONet
7.62
5.47
1.57
5.18
1.51
Advection
RMSE (x10 3)
Pl-DeepONet
4.80
1.33
4.24
1.28
3.63
SepONet
4.83
4.18
3.47
1.06
3.29
1.02
3.17
d= 2
Net
34.511
59.806
Memory (GB)
SepONet
0.713
0.715
0.719
0.719
0.719
DeepONet.
0.3997
1.0766
1.9765
4.411
Training time (hours)
SepONet
0.0754
0.0715
0.0736
0.0720
0.0730
Relative f/ error (%)
Pl-DeepONet
28.48 4.17
28.63 63
+ 4.10
28.26 4.38
12.33
SepONet
4.24
7.51
Burgers' d= 2
RMSE (x10 2)
Pl-DeepONet
SepONet
4.09 2.77
4.11 2.78
3.30
4.07
1.96 + 1.92
1.64
1.44
Memory (GB)
Pl-DeepONet
5.085
9.695
17.913
35.433
SepONet
0.605
0.607
0.607
0.609
0.716
Training time (hours)
NetpeOONeet
0.5135
2.6804
5.923
SepONet
0.0725
0.0707
0.0703
0.0612
0.0605
Pl-DeepONet
Relative f/ errer (%)
SepONet
31.94 9.18
25.48 8.95
21.16
10:21
Pl-DeepONet
Nonlinear diffusion 3
RMSE (x10-3)
SepONet
0.32
Pl-DeepONet
Memory (GB)
SepONet
2.923
3.139
4.947
6.995
Pl-DeepONet SepONet
Training time (hous)
0.1175
0.1408
0.1849
0.3262
0.5575
24

---

Published in Transactions on Machine Learning Research (12/2024)
PI-DeepONel
SepONel
200
Trasting Uime CO
X (310)
0.50
120
327
80
507
1207
607
1207
1207
(a) (1+1)-d Diffusion-reactionn
(b) (1+1)-d Advection
101 10 ae0 700
ThiRuoe tsee th
X
W
20
120
2.2
607
1.00
1207
422
607
1200
647
1200
120
(1+1)-d Burgers'
(d) (2+1)-d Nonlinear diffusion
Figure 4: Performance comparison of PI-DeepONet and SepONet with TanH trunk network activation functions, varying number of training points (Ne) and fixed number of imput functions (Ny 100). Results show test accuracy, GPU memory usage, and training time for four PDEs.
PI- De-DeepONet
SepONet
Dat. & erter (5)
Minsyy (201
De booo
1/
Nv
11
(a) (1 (1+ Diffusion-reaction
(b) (1 + 1)-d Advection
Trantin lite ma
34
Mecram (itt
Noowa ICt
THL mrze
100
100
(c) (1+1)-d Burgers'
(d) (2+ 1)-d Nonlinear diffusion
Figure 5: Performance comparison Pl-DeepONet and SepONet with TanH trunk network activation functions, increasing number input functions (Ny) and fixed number of training points (N₂ = 128d, where d is the problem dimension). Note: PI-DcepONet results for the (2+1)-dimensional diffusion equation are unavailable due to memory constraints.
B.4
Ablation Studies
B.4.1 Trunk Networks with Hyperbolic Tangent Activations
In Figure and Figure we provide complete testing results repeating our experiments from Figure 2 and Figure 3 for all PDE examples, varying N, and Ny, except we use hyperbolic tangenet (TanH) activation functions for all hidden and output layers of the trunk networks both PI-DeepONet and SepONet. The results are very similar, indicating that alternative activation functions may be chosen for multi-layer trunk networks to maintain the universal approximation property in accordance with Theorem
25

---

Published in Transactions on Machine Learning Research (12/2024)
Pl-DeepONet
SepONel
12
8.22
Het erter PNI
1.30
Brandh Japar (eLAnEAA
Stanch Ingat INsestonc0
Brunch Tepat DinscEs
Brunch Ispel DARACRRAM
Brunch Iegut DhonceaneE
(a) (1+1)-d Dlffusion-reaction
(b) (1+1)-d Advection
Kousay (001
Nenacy (281
Bel 1o
20 50
100
20
100
20 100
238 1004 8000
256 10244020
15 250 4005
Sapert Dimvetttion
Dippc Oemaraiio
Bmoohh Trpui Dimsensies
Branca taput Droesssoo
Inper Dizzendoon
Branch Ieput Dimerakon
(c) (1+1)-d Burgers'
(d) (2+ 1)-d Nonlinear diffusion
Figure Performance comparison of Pl-DeepONet and SepONet with varied number of input function sensors (branch input dimension). Note that we fix Ny 20 for all experiments, N. 32 for (a)-(c), and No 16 for (d).
B.4.2 Varying the Input Function Discretization to the Branch Network
Given an input function PDE configuration u, recall that the branch network predicts coefficients 8g (E())) for k = I.,.,T Our studies in Section use simple encoder that measures the imput function E(u) = (u(*1),..., u(zm)) (((,)) ) at points 21, Em E K in the input function domain. High-dimensional or highly oscillatory imput functions may lead to unwieldy discretizations with large branch input dimension that affect training performance. Here, Figure G we study the sensitivity of the PDE examples from Section 4 with respect the number of input function sensors (branch input dimension). Note that fix the number of imput functions N, 20 for all experiments, while N. = 32 for diffusion-reaction, advection, and Burgers' equations, and N. = 163 for nonlinear diffusion. We find that the error curves converge to the minimum value using only a fraction of the number of input sensors that we used in the main text in Figure 2 and Figure This indicates that the input functions we considered may be identified with a small number of points. Nevertheless, we find that training performance in terms of both memory consumption and training time constant with the number of sensors. This is because training complexity is mainly data-dominated by the need to compute high-order derivatives with respect to a large number of collocation points for evaluation of the physics loss, discussed in Section 3.2
C Complete Solution to Separation of Variables Example 7
Recall the linear PDE system treated in Section 2.3
MiIss(y) = = (1(v))(1)) ..... La(ua)s(v),
(44)
where M[t] = de + h(t) is first order differential operator of t, and L1 [i[y)),..,, Calya are linear second order ordinary differential operators of their respective variables yı.... Va only. Furthermore, assume we are provided Robin boundary conditions each variable and separable initial condition s(t 0, yy.....yd) yy) = II= (x,)) for functions (n,))) that satisfy the boundary conditions.
Assuming a separable solution exists, s(y) = (y)) ... Ya(ya), the PDE can be decomposed in the following form:
MT(t) (XYi(M
CaYalua)
(45)
T(I) Yi(y))
Ya(ua)
where is apparent that each term in the sequence is constant, since they are each only functions of single variable. Consequently, we may solve each of the Ln terms independently using Sturm-Liouville
26

---

Published in Transactions on Machine Learning Research (12/2024)
theory. After have found the associated eigenfunctions (Yr")) and eigenvalues (Ar),, we may manually integrate the left-hand side. Finally, we may decompose the separable initial condition into a product of sums of the orthonormal basis functions (eigenfumctions) of each variable. The resulting solution is given by
s(v) s(t, 1............... va)
E
BAT (1)
II
Yrh (Un), Y
(
n=1
(t) T(km exp (-Lo h(r)
Σ Aen
(46)
(Y (yn): ; on (B()))
Ly Y.un (yhe)
B2 k(k
Ainn
n=1 (Y (yn), Y. (ym))
Y (vm))
n I....dd Kn = 2.,......
Here, A = (ki,...,ke ka), where kn E possible products of eigenfunctions ykn with associated eigenvalues Ade ()), is an appropriate inner product associated with the separated Hilbert space of the -, operator. To obtain equation (7) in the main manuscript, one only need break up the sum over all hn indices into single ordered index.
Vn E (1,..., d}, is a lumped index that counts over all
Additional Experiments
D.1 (imensiiaa Navier-Stokes Equation
SepONet's memorr-efficient and fast-training advantages allow it to tackle more challenging physics-informed operator learning problems, which PI-DeepONet may struggle to train on due to resource constraints. As an example, we consider (2+1)-dimensional Navier-Stokes equation, previously investigated in the context.
of PINN (Cho et al. 2024 Wang al., 2024h:
O.w + 5 Vw = 0.01Aw, æ E [0, 2m) E V 8 0, x [0, tT,,
(47)
w(x, 0) = (x)), s(x, 0) so(x), a E [0, 2
where s (Sz, sy) R is the velocity field, = (x.y) denotes 2D spatial variables, r = [0. is the time s = Ogsy 0ys, is the vorticity. We aim learn the solution operator that maps the initial velocity and vorticity field u(x:) parameterized by 0 = (t), é1, Φ2, Φ3), which represents all the trainable parameters in the branch and trunk nets. The vector-valued velocity is approximated as:
window, and
wo(x)) R3 to the solution s(x, R using SepONet,
(u)
3x
Tn.k:
k=1
(48)
2r
8y (xr
Σ
3x Tnkks
n=1
where r denotes the rank, 3x = (E(U)) the k-th output of the branch net, and Ti, = 12, (1)x, 72,1 13 (4)s denote the k-th oulputs of the three trunk nets.
Cg (z),, 73,
Loss function
The physics loss for this problem are defined as:
Lphyssics Crasiduat + 5000 10000
(49)
27

---

Published in Transactions on Machine Learning Research (12/2024)
The specific loss terms are:
N
0(v (a( (100)) +000 v(V x Go(u(")(o()) ((01))
Cresidual N,N,
0.01A(V x Go (x ((ا1))
N
(50)
1
V Go(uu())(aa)) ,100)
Laie
N₂
N
1
Σ
( Go(u')) (a)))))_)))))))))))))))))1 (a( - ca()))))))
Linitial
N1
Note that periodic boundary conditions are enforced by applying the following positional encoding to the spatial variables:
(xx) = [1, sin sin( sin (3x) sin (4x), (5x), cos( c2x( cos cos(4x)) cos()
(51)
SepONet architecture The encoder maps the initial condition u(x) to its point wise evaluations at 128 128 3 sensors on uniform 128 x 128 grid over [0, 2x The branch network by is a CNN, starting with 1 X 1 convolution that increases the channels from 3 to 16, followed by four residual blocks. Each residual block consists of two 3 x 3 convolutions with GeLU activations; the first convolution in each block uses stride of 2 to halve the spatial dimensions while doubling the number of channels. Skip connections employ 1 x convolutions with a stride of whenover there is a change in dimension. After flattening, a fully connected layer produces a vector of dimension 2r.
The trunk networks te. are 7-layer modified MLPs layer, an input size of 11, and rank/outpun size of 256/512, using TanH activations. The initial velocities are sampled from Gaussian random field with maximum velocity of 5.
[Wang al. 2021a), each with 128 neurons per hidden
Training settings We consider learning the solution operator within two time windows: T = 0.1 and T Two separate SepONet models were trained, each for 100,000 iterations using the Adam optimizer, to minimize the physics loss (19) within their respective time windows. The number of residual points was set to N, 256 x 256 32 (N₂ x Ny N,), obtained by randomly sampling 256, 256, and 32 points along x, y: and axes, respectively, and constructing a mesh grid via the tensor product. The number of initial points, N2 = 128x x 128, corresponds to uniform mesh grid. Initial velocities were sampled from a Gaussian random field with a maximum velocity of 5. Both the initial conditions and collocation points were resampled every 100 iterations. We varied the number of input functions Ny to see the scaling as data is increased.
The model was evaluated on 100 unseen initial conditions, sampled from the same Gaussian random field. Reference solutions for both time windows were obtained using the JAX-CFD solver (Kochkov 2021 on a uniform 128 x 128 10 (N, x No x No) grid.
Evaluation
et al.
The results for two time windows are shown in Figure 7 We find that for T = 0.1 we can achieve very low relative lg crror of 2.5%. For T the error only scales to 40%. We think that improving the error for the longer time scale represents an interesting application direction for future work. Our architecture and implementation choices were not optimized for this example.
Results
D.2 Linear Heat Equation Example
SepONet is motivated by the classical method of separation of variables, which is often employed solve linear partial differential equations (PDFs). To illustrate the connection between these approaches, consider
28

---

Published in Transactions on Machine Learning Research (12/2024)
50-
T-1
T-1
T-0.1
T-0.1
Rel. f, etror (%)
Memory (GB)
Training timee
30
T-1
.00
T=0.1
20
10
10-
5
10
20
50
100
5
10
20
50
100
5
10
20
50
100
Nr
N,
N;
Figure 7: Navier-Stokes equation results with SepONet. Here, we varies the number of input functions Ny and kept the number of collocation points fixed. We consider two cases, T = 0.1 and T 1, corresponding to the length of the time window.
the linear heat equation:
0s 0%
Ōt n² Ox (z,t) E (0, (0,
(52)
s(x, 0) u(x), x (0, 1).
s0,tt) s(1,t) = 0, t (0.
The goal solve this equation for various initial conditions u(x) using both the separation of variables technique and the SepONet method, allowing for an intuitive comparison between the two.
D.2.1 Separation of Variables Technique We seck a solution in the form:
sx,, () (x))((I
(53)
for functions X, T to be determined. Substituting
53 into
52
yields:
X"
X and 2T
(54)
T
for some constant X. To satisfy the boundary condition, X must solve the following eigenvalue problem:
x" ("(x) 0, (0,1),
(55)
x(0) X(1) 0,
and T' must solve the ODE problem:
T'() = π2 T(t).
(56)
The cigenvalue problem
55 has sequence of solutions:
A (kT)), Xx(x)= sin(kTæ), for k 1.2....
(57)
For any X, the ODE solution for T' is T(t) Ae # for some constant A. Thus, for each eigenfunction X with corresponding eigenvalue λ have solution Tę such that the function
(x t) (x(x)T((tt
(58)
will be a solution
54 fact, an infinite series of the form
s(x,t) = Σ (x(x)T((t)
Σ
A&e 4²t skn()
(59)
Am
29

---

Published in Transactions on Machine Learning Research (12/2024)
will also be solution satisfying the differential operator and boundary condition of the heat equation 52
subject to appropriate convergence assumptions of this series. Now let s(x, 0) u(z)), we can find coefficients:
A = 2 / sin (kTT)) u u(z)dr
(60)
such that 59 is the exact solution of the heat equation (22
D.2.2
SepONet Method
In this section, apply the SepONet framework to solve the linear heat equation 552 and compare the basis functions it learns with those derived from the classical separation of variables method. Recall that a SepONet, parameterized by 0, approximates the solution operator of (52) as follows:
Go(u)(z,1) = Σ u((2) ((
(61)
where 21, ... T128 128 equi-spaced sensors in [0, 1], 3h is the k-th output of the branch net, and the basis functions (t) and (x are the k-th outputs of two independent trunk nets.
Training settings The branch and trunk networks each have a width of 5 and a depth of 50. To determine the parameters 0, we trained SepONet for 80,000 iterations, minimizing the physics loss. Specifically, we set A = 20, As = 1, N, = 100, and N₂ 1282 in the physics loss. The training functions (initial conditions) {u()) N i=1 were generated from Gaussian random field (GRF) ~ . (0.253 (-A+ + 5)
using the Chebfun package (Driscoll al.) 2014), ensuring zero Dirichlet boundary conditions. Additional training settings are detailed in Appendix B.2 of the main text.
Evaluation We evaluated the model on 100 unseen initial conditions sampled from the same GRF, using the forward Euler method to obtain reference solutions on a 128 x 128 uniform spatio-temporal grid.
Impact of the rank r Since (t) and (x) are independent of the initial condition u, learning an expressive and rich set of basis functions crucial for SepONet to generalize to unseen initial conditions. To investigate the impact of the rank r the generalization error, we trained SepONet with ranks ranging from I to 50. The mean RMSE between SepONet's predictions and the reference solutions over 100 unseen test initial conditions was reported. For comparison, we also computed the mean RMSE of the truncated analytical solution at rank r for from 1 to 15. The results are presented in Figure
As r inereases, the truncated analytical solution quickly converges to the reference solution. The nonzero RMSE arises due to numerical errors computing the coefficients A& and the inherent inaccuracies of the forward Euler method used to generate the reference solution. For SepONet, we observed that when 1,2, the mcan RMSE aligns closely with that of the truncated solution. However, r increases beyond that point, the error decreases more gradually, stabilizing around r 50. This indicates that SepONet may not necessarily learn the exact same basis functions as those from the truncated analytical solution. Instead, a higher rank r allows SepONet to develop its own set of basis functions, achieving similar accuracy to the truncated solution.
SepONet basis functions The learned basis functions for different ranks are visualized in Figure Figure 13
At r SepONet learns basis functions that closely resemble the first term of the truncated solution. For
2, the Icarned functions are quite similar to the first two terms of the truncated scries. However, when
r = 5, the basis functions diverge from the iruncated solution series, although some spatial components still resemble simusoidal functions and the temporal components remain monotonic. As r increases further, SepONet continues improve accuracy, though the learned basis functions increasingly differ from the truncated series, confirming SepONet's ability to accurately approximate the solution using its own learned basis functions.
30

---

Published in Transactions on Machine Learning Research (12/2024)
RMSE vS r
Truncated solution
SepONet
RMSE
10-2
10-3
10
20
30
40
50
Rank (r)
Figure Comparison of RMSE between the truncated analytical solution and SepONet predictions for varying rank r. The truncated analytical solution quickly converges, while SepONet shows a slower decay in error, converging around r 50.
2,6x
-0.7
-0.0
G sinmm))
11.0
0.2
114
0.86
0.8
0.0
033
0.4
0.8
1.0
Figure 9: Learned basis functions Tx( and (ar)) for 1. SepONet learns the same basis functions as the first term of the truncated solution.
E Visualization of SepONet Predictions
In this section, we showcase the performance of trained SepONets in predicting solutions for PDEs under previously unseen configurations. The SepONets were trained using N, = 100 and N, = 128d where d
denotes the dimensionality the PDE problem. The prediction results are presented Figure 14 to Figure
31

---

Published in Transactions on Machine Learning Research (12/2024)
0.5
0.905e
-24
7231
C/00
-atrrs Citei
-0.0
+03
Figure 10: Learned basis functions (t) and (t(x)) for r = 2. SepONet learns very similar basis functions as the first two terms of the truncated solution.
Ty1B)
Ty(e)
Talt)
Ty(e)
0.0d
600
G0d
G,0d
C,0d)
Figure 11: Learned basis functions Th (t) and (x(z) for r r= 5.
Tr(8)
Ta(e)
Tx(0)
Te(()
1,1()
1y(0)
1,11)
(,00
(,0x)
(,00)
6,00)
264
2,04
Figure 12: Learned basis functions (x( and Ce(x)) for r 10.
32

---

Published in Transactions on Machine Learning Research (12/2024)
7/6
TW/
5,10
Tydd
Tufft
De10
Tyx/d
5x100
Ty010
hald
Tpd/
tal0
Tua/1
Teafi
Tu.0
Tr,01
5,10
Ted
Teu0
Trxli
Te,00)
Tyyl
To/0
5,10
6,b0
Cd90
Goldl
Cylo
Gild
0r/00
Cotal
Gotdl
Gutdl
Eudah
Cutel
Culst
Gudst
Cutei
Cutst
Gurdl
C.ld)
Satai
C,úxt
Cutxi
Sutel
Suit0
Guddt
Galoo
Gudat
Gaten
Gutat
Satet
Guld
Edda)
Cuta)
Cals)
Tedel
Zuta)
Zute)
Coid
Figure 13: Learned basis functions (x()) and (x(x) for 50.
33

---

Published in Transactions on Machine Learning Research (12/2024)
1.0
0.0
Source term
-.5
0/
1.0
-1.0
-03
-3.5
-1.0
0.0
94
06
0.
1.0
0.0
0.2
0.4
R6
1.0
0.0
04
94
0.1
1.0
0.0
02
0X
10
00
82
A6
1,0
1.0
3.0
1.0
1.0
Ground troth
%0
00000
%.0
8.0
1.0
3.0
0.4
Prediction
-0.G
0.5
00
Absolute error D.5
055
0.5
0.5
05
0.0
000
0,30
"/b
Figure 14: (1+1)-d Diffusion-reaction equation.
300
2.5
PDE coeticient
2.0
3.0
1.5
1.0
1.0
1.0
0.0
0.2
08
1.0
0.0
0.2
0.6
06
1.0
0.0
02
06
03
0.0
0.1
06
0.0
0.0
044
Ground truth
%o
1.0
0.0
9.5
1.0
1.00
0.5
1.0
0.8
0.5
1.0
0.5
3.0
1.0
1.0
1.0
1.0
1.0
1.0
Prediction
3.5
%0
00
0.0
0.60
0%
3.150
-15
-05
0.5
0.5
US
2.05
00
0.0
0.0
0
Figure 15: (1+1)-d Advection equation.
34

---

Published in Transactions on Machine Learning Research (12/2024)
0.2
0.10
0.2
Initsll condition
0.0
0.0
0.00
-0.3
-0.2
-0.10
Ground troth
0.55
-0.00
2go
0.5
1.0
08.0
0.5
1.0
0.0
0.5
1.0
0.0
0.
1.0
0.
0.5
-0.10
3.0
3.0
1.0
1.0
1.0
Prediction
%
00.0
6
0/3
000
2.000
0.5
0.5
0.5
0.3
0.5
%
0.0
6:8
000
000
Figure 16: (1+1)-d Burgers' equation.
35

---

Published in Transactions on Machine Learning Research (12/2024)
Inital condition
622
Ground troth (tm0.5)
Prediction
0.02
(t=8.5)
Ground truth (t=1)
Prediction (t=1)
6002
Absolute error ([=1)
Figure 17: (2+1)-d Nonlinear diffusion equation. Two snapshots at 0.5 and 1 are presented.
36


Reviews:
[{"review_id": "ZnBOQerRkU", "reviewer": "Reviewer_nRRF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "There definitely is some merit in the idea of modifying a DeepONet to use a separate trunk net for each variable of the PDE. It is particularly impressive that the performance of this modification has been shown for the unsupervised setting, which is the hard regime for operator learning. But prima-facie this approach endangers needing a huge blow-up in the number of parameters andits not becoming clear from the text as to why this is (probably?) not happening in the implementations.\n\nThere are a bunch of unclear things,\n\n- In page 2 there is a count given of $N^d$ collocation points being required. Where is this coming from? It seems to suggest that the only way to use a SepONet is if one gets the collocation points in the domain of the PDE sampled by taking $N$ points in every dimension and then taking the Cartesian product. But its not clear wherefrom in the theory does this constraint on sampling the points comes from. Wont the mechanism work if we randomly sample some points in the domain space of the PDE solutions? \n\n- The experiments are not very convincing given that (A) the PDEs are not very large variable systems and (B) advantage in terms of error is barely there and (C) the claim about lowering of memory footprint is hard to judge because of lack of information on the model sizes in the different experiments.\n\n- The theoretical approximation guarantee that has been presented is somewhat out of place as the proof corresponds to the supervised setting while the experiments and the description in Section 2.2 is in the unsupervised setting.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper tries to motivate a new modification to DeepONets i.e to use a separate trunk net for each variable in the domain of the solution space. The authors give experimental evidence of the efficacy of this method on some small PDEs and they show that this architecture satisfies the same universal approximation theorem as the standard DeepONet.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I suggest the following edits,\n\n- In the first paragraph the authors seem to cite (Chen \\& Chen 1995) for the universal approximation theorem for operator nets. This is a highly incomplete referencing as the any-depth operator net approximation theorem happened in works like https://doi.org/10.1093/imatrm/tnac001  and https://www.jmlr.org/papers/volume22/21-0806/21-0806.pdf\n\n- At the bottom of page one, the authors say that \"to achieve satisfactory generalization error, the number of required input-output function pairs grows exponentially\". This statement makes little sense. Generalization bounds in terms of Rademacher complexity already exist for neural operators https://arxiv.org/abs/2205.11359 - where the bounds don't scale exponentially with size at any fixed depth. \n\n    In light of the above reference, the discussion presented needs to be corrected.\n\n - I think we need much stronger evidence in terms of performance studies on large variable systems - lets say even on the Heat PDE at say dimensions upto 10 or more, lets say as much is the maximum capacity on the free version of Colab-Pro. This is quite needed given that the theoretical contribution is not really novel (this proof is essentially known) and somewhat not connected to the main paper for reasons as explained above.\n\n - The authors need to *either* (1) show a proof that there exists a net which provably arbitrarily lowers the unsupervised operator net loss as given in equations 4 and 5 or (2) they need to add experiments as described above and relegate the current proof to the Appendix as it is orthogonal to the setup in the rest of the main paper or (3) add experiments in the supervised setup with large variable PDEs which would then correspond to the proof that is given.\n\n - Also, the way the experiments are presented need to be a lot more clearer. There needs to be a list of PDEs (not just names!) written down which are being tested on. And that needs to be accompanied with the corresponding loss functions and the nets being trained on.\n\n   This clear specification as requested above is really critical to be able to judge the issue of memory footprint of training - which seems to be the main claim here.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "fF8dMxoFaQ", "reviewer": "Reviewer_yRAN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* Represent the the basis of the target function space as tensor product of multiple single-dimensional bases is technically sound. Combined with the forward-mode AD, it allows efficient computation of PDE residuals in physics-informed operator learning.\n\n* The authors extend the previous theoretical result on DeepONet to the case with multiple trunk nets assuming sinusoidal activation in the trunk net.\n\n* Experimental results demonstrate that the proposed Separable DeepONet is much more efficient and can often outperform DeepONet baseline in terms of accuracy.\n\nWeakness:\n\n* The separable learnable bases addresses the curse of dimensionality in the trunk net part but the branch net will still need to handle high-dimensional inputs.\n\n* Separable basis and forward-mode AD are proposed in SPINN, so for this part, the technical advancement is moderate.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new architecture for operator learning that uses multiple single-dimensional trunk nets to learn and approximate the solution basis function. This proposed model reduces the computational cost associated with high-dimensional grid and brings performance improvement. The proposed method can be useful for physics-informed learning and operator learning tasks. It can also be potentially applicable to some other neural field applications.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* In SPINN, the authors conduct experiments on arguably more complex PDE problems such as 2+1d Navier-Stokes and multiple 3D systems. Can Sep-DeepONet handle these systems reasonably well? I think it is not a problem even if Sep-DeepONet can't perform well on all of them but it is worth providing some justification and discussion.\n\n* Is the number of collocation point to branch net crucial to the convergence?\n\n* There are some other relevant operator learning works which decompose/factorize the computation leveraging the tensor structure of the data that were not discussed in the paper: \n\n[1] Tran, Alasdair, et al. \"Factorized Fourier Neural Operators.\" The Eleventh International Conference on Learning Representations.\n\n[2] Li, Z., Shu, D., & Barati Farimani, A. (2024). Scalable transformer for pde surrogate modeling. Advances in Neural Information Processing Systems, 36.\n\n[3] Kossaifi, J., Kovachki, N., Azizzadenesheli, K., & Anandkumar, A. (2023). Multi-grid tensorized Fourier neural operator for high-resolution PDEs. arXiv preprint arXiv:2310.00120.", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "hg4E1IfW9s", "reviewer": "Reviewer_sLu2", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n1. SepONet offers a simple yet effective approach to significantly improve PI-DeepONet's efficiency. \n\n**Weaknesses**\n1. Weakness in theoretical results: the proposed theorem assumes the trunk net is one layer with a sinusoidal activation function that does not align with the practical implementation. \n2. The efficiency of SepONet heavily relies on the assumption that the grid is regular. If it is an irregular grid, SepONet will no longer have its advantage.\n3. The paper does not compare the time and space complexities of SepONet and PI-DeepONet as a function of problem dimensions, which will help us understand the behavior of SepONet under different scenarios as a general algorithm. \n4. While the SepONet applies to arbitrary problem dimensions theoretically, the authors only consider the case where the output function is scalar-valued. Common vector-valued cases, such as the Navier-Stokes equation[1], are not studied or discussed. \n\n[1] Jin, Xiaowei, et al. \"NSFnets (Navier-Stokes flow nets): Physics-informed neural networks for the incompressible Navier-Stokes equations.\" _Journal of Computational Physics_ 426 (2021): 109951.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "1. The authors propose a new method called SepONet, which enables efficient physics-informed operator learning by extending the idea of separable PINN to PI-DeepONet. \n2. The authors establish the universal approximation capability of the proposed SepONet under proper assumptions. \n3. The authors provide an empirical comparison across various PDE problems, showing that SepONet is much more efficient than PI-DeepONet.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Section 3.1.2 needs more clarity about how forward-mode AD helps reduce computation complexity. I agree with the argument about the \"input-output relationship\" implies that forward-mode AD is more efficient. However, it is unclear from the way it is written right now.\n\t1. The authors state that computing Eq. (11) with forward-mode AD leads to more efficient computation. However, compute Eq. (11) with reverse-mode AD gives the same time complexity $O(N_m)$. \n\t2. If I'm not mistaken, the reason why the forward-mode AD is faster in this case is because $\\tau_{m}$ is a $\\mathbb{R}\\rightarrow \\mathbb{R}^r$ function and $\\partial \\tau_m / \\partial y_m$  takes one evaluation for forward-mode and $r$ evaluations for reverse-mode. \n2. In terms of GPU memory comparison, the authors should provide the full context for precise performance claim. \n\t1. PI-DeepONet's GPU memory is a linear function of batch size, and it can decrease batch size to reduce GPU memory cost. What is the batch size for PI-DeepONet in the experiments? \n\t2. Does SepONet use all the collocation points on the grid in one single batch? \n3. The authors compare runtime and GPU memory costs of SepONet and PI-DeepONet through concrete experiments, which is good to have. However, as a general operator learning algorithm, the authors should derive the space and time complexities of SepONet in terms of the problem dimensions and compare them with PI-DeepONet, which reveals its efficiency under different scenarios. \n4. It would be great to expand the discussion or experiments about vector-valued cases such as Navier-Stokes equations. \n5. I think the parenthetical citations should be like (Lu et al., 2022a). However, the authors often write it as Lu et al. (2022a). Just name a few here, but the authors should check the whole paper and correct them. \n\t1.  Line 8 in the first paragraph of the Introduction: \"Lu et al. (2021)\"\n\t2.  Lines 1, 3, and 10 in the third paragraph of the Introduction: \"Wang et al. (2021b), Raissi et al. (2019), Baydin et al. (2018)\". \n6. Minor typos: Eq. (10) needs a comma.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: See above.
Decision: Accept (poster)