Operator Learning, introduced by [1], involves learning operators that map one function to another. It extends classical function learning, which focuses on the mapping between vector spaces, and functional learning [2], which deals with functions mapping to fields. Operator Learning has become a pivotal tool in scientific machine learning [3]. Recent works on this topic include various neural operators [4][5][6][7][8] and extensive studies on the approximation and convergence aspects of these neural operators [1,[9][10][11]. Additionally, applications of neural operators span various scientific and engineering domains. For example, researchers have developed algorithms that use neural operators to address parametric Partial Differential Equation (PDE) problems [4,[12][13][14][15][16], dynamical systems with control [1,17,18], power engineering applications [19][20][21], design optimization problems [22,23], and challenges related to multifidelity and multiscale downscaling [22,24,25]. Notably, researchers have recently applied neural operators to climate predictions [26] and uncertainty quantification [27][28][29].
One of the popular approaches in this direction is the Deep Operator Neural Network (DeepONet) [4,7,12,22,[30][31][32]. DeepONet utilizes two subnetwork architectures consisting of a branch network and a trunk network to approximate operators by constructing a basis for the output function space. The trunk networks are responsible for learning the basis, while the branch network acquires the basis coefficients. This results in an approximation achieved through a linear combination of the learned basis and coefficients. Another approach for neural operators is the Fourier neural operator (FNO) [5,26,33]. FNO applies the Fourier neural transformation to convolve the input function with a learned translation invariant kernel function. It then utilizes skip connections and other mechanisms to nonlinearly generalize the approximation. A comprehensive comparison of DeepONet and FNO can be found in [30]. Additionally, there are several other noteworthy neural operators documented, as cited here [6,8,10,[34][35][36].
A significant advantage of DeepONet [1,4,7] and related structures [8,10] is the ability to discretize output functions freely [8,30,37]. The neural operator generates function values at specific domain points, which are determined by the trunk network input, resulting in enhanced adaptability. This flexibility allows the network to predict output function values at any domain point, accommodating different mesh definitions for output functions. However, a major challenge is achieving discretization invariance [8,30]. Specifically, the neural operator must remain invariant to the different discretizations applied to input functions. Several algorithmic extensions to DeepONet have been proposed, allowing different input functions to be discretized differently [30,38]. Additionally, the Basis Enhanced Learning (BelNet) extension of DeepONet, proposed by [8,10], has mathematically demonstrated discretization invariance. It is worth noting that these extensions still require the number of sensors used to sample input function values to be consistent across distinct input functions. In this paper, we introduce a novel and efficient distributed network architecture and training approach for neural operators. This approach permits individualized discretization of input functions using dedicated sensor distribution strategies.
Discretization invariance offers several advantages, particularly when working with datasets that have diverse properties [33]. For example, a dataset may consist of input functions with different levels of regularity. When using a uniform set of sensors to sample and discretize these input functions, it becomes necessary to accommodate the functions with the lowest regularity. This leads to a significant increase in the number of required sensors. However, it also results in higher costs for functions that are inherently smoother. To accommodate input functions with varying properties and regularity, we will draw inspiration from federated learning. Introduced in [39], federated learning is a machine learning framework that enables multiple clients to collaboratively train a consensus neural network model in a distributed manner, while keeping their training data local. This approach protects data privacy and reduces communication costs by coordinating local model updates with a central server, resulting in a significant reduction in data transfer volume.
A standard formulation of federated learning is a distributed optimization framework that addresses communication costs, client robustness, and data heterogeneity across different clients [40]. Communication efficiency is central to this formulation and motivates the most well-known efficient communication algorithm in federated learning: the federated averaging (FedAvg) algorithm [39]. FedAvg has been studied under realistic scenarios in [41,42]. Furthermore, several works have provided convergence proofs of the algorithm within the field of optimization [43][44][45][46][47].
Traditionally, neural operators have been trained using a centralized strategy that involves transferring the training data to a central location. However, this approach hinders our ability to leverage high-performance distributed/parallel computing platforms. To address these limitations, the authors in [48] proposed Fed-DeepONet. Fed-DeepONet enables distributed training of deep operator networks and has been empirically shown to achieve accuracy comparable to the centralized vanilla DeepONet while handling heterogeneous input functions. However, it should be noted that Fed-DeepONet requires the number of sensors for each local DeepONet to be the same, as it trains a consensus global DeepONet. In this paper, we provide a theoretical demonstration of Fed-DeepONet's ability to approximate nonlinear operators. Additionally, we design a dedicated distributed training strategy that does not require the same number of sensors, allowing the approach to handle heterogeneous input function spaces.
We summarize the contributions of this paper as follows:
1. We propose the distributed deep neural operator (D2NO) method that allows for sampling different input functions using distinct sets of sensors. This training method can reduce the total number of back-propagation steps and the number of trainable parameters by utilizing dedicated neural networks to handle the diverse input function spaces. 2. D2NO is motivated by the mathematical theory that underlies the algorithm. Since all functions share the same output basis, we construct dedicated networks to handle the diverse input functions. We then proceed to demonstrate the universal approximation of the network and algorithms we propose. 3. The proposed methodology can be implemented using different neural operators. In this study, we apply it to the novel Deep Operator Neural Network (DeepONet). By utilizing the Distributed-DeepONet (D-DeepONet), we tackle complex heterogeneous datasets that contain input functions with different regularities and geometries. D2NO enhances efficiency while preserving predictive accuracy.
The rest of the paper is organized as follows. In Section 2, we will review the DeepONet and the universal approximation theory. Next, we present the methodology in Section 3. Then, we prove the universal approximation theorem for Distributed-DeepONet in Section 4. Finally, we present several numerical experiments containing heterogeneous input function spaces in Section 5.
We start by defining an operator 𝐺 ∶ 𝑉 → 𝑈 , where 𝑉 and 𝑈 are two function spaces. Next, we examine the Deep Operator Neural Network (DeepONet) [4,7,12,22,30], as demonstrated in Figure 1. The DeepONet structure contains a branch and trunk ⨂ denotes the inner product in R 𝐾 . Specifically, 𝑡 𝑗 ∈ R is the trunk net output, 𝑏 𝑗 ∈ R is the brunch net output, and 𝑢 is a function in the input function space .
network. The branch net processes a discretized input function, while the trunk network processes an arbitrary location within a given output function space. The DeepONet structure avoids needing to discretize the output function. This allows more flexibility in both the training and prediction processes, distinguishing it from neural operators that rely on output function discretization on fixed grid mesh (see Fig. 1).
The DeepONet is theoretically based on the universal approximation theorem of nonlinear operators, which was established in [1,2,4], and extended in [4,10]. The main results are stated below.
Theorem 2.1 (Theorem 5, [1]). Suppose 𝐺 ∶ 𝑉 → 𝑈 is continuous. Specifically, let 𝑉 be compact in the continuous function space 𝐶(𝐾 1 ), where 𝐾 1 is compact in a Banach space. Additionally, suppose 𝐺(𝑉 ) ⊂ 𝐶(𝐾 2 ), where 𝐾 2 ⊂ R 𝑑 is also compact. Then, for any 𝜀, there exist sensors
for all 𝑥 ∈ 𝐾 2 and 𝑢 ∈ 𝑉 . Here û = [𝑢(𝑦 1 ), … , 𝑢(𝑦 𝑁 )] ⊤ .
The basis enhanced learning (BelNet) extension partially relaxes the restriction on the input function discretization. However, the number of sensors (input function mesh size) should still be the same for all input functions. Additionally, 𝐾 𝑦 ⊂ 𝐾 which supports the discretization-invariant property must be carefully crafted (see Remarks 3 and 4 in [10]). This limits the sensor placement and input function discretizations.
When dealing with datasets that have input functions with different properties, it is important to place sensors that can adapt to functions with the lowest regularity within the dataset. Fig. 2 illustrates various functions with different regularities. To effectively capture functions with features like black curves with sharp peaks, a large number of sensors must be used. Similarly, for the magenta curves with smooth oscillations, sensors need to be strategically placed across the entire domain. Using a uniform sensor configuration for all input functions would require a significant increase in the number of sensors. This would result in higher computational costs and additional expenses, especially for smoother functions. We propose the D2NO which can effectively address this problem.
In [48], the authors proposed Fed-DeepONet, a distributed/federated strategy that trains a consensus/universal DeepONet model for 𝐶 clients and a centralized server. This federated training strategy involves looping over the following three steps:
1. Broadcast to clients: The centralized server broadcasts the most up-to-date DeepONet consensus model to all clients. 2. Client local updates: For any 𝑐 ∈ {1, 2, … , 𝐶}, the 𝑐th client receives the most up-to-date DeepONet consensus model and performs 𝐾 ≥ 1 local stochastic gradient (or variants, e.g., Adam) DeepONet updates using the client's dedicated dataset. 3. Global synchronization: The centralized server aggregates the local DeepONet parameters into a unique consensus/universal DeepONet every 𝐾 local update.
Since Fed-DeepONet learns a consensus model, each client must use the same number of sensors. In the next section, we will expand on the previously mentioned federated training strategy and develop a dedicated distributed training strategy for DeepONets with 𝐶 clients. This strategy will efficiently handle heterogeneous input function spaces by allowing different numbers of sensors for each client 𝑐 ∈ {1, 2, … , 𝐶}.
Consider a nonlinear operator 𝐺 that maps an input function 𝑢 ∈ 𝑉 to an output function 𝐺(𝑢)(𝑥) evaluated at an arbitrary location 𝑥 ∈ 𝐾 2 . According to Theorem 2.1, we can approximate 𝐺(𝑢)(𝑥) as ∑ 𝐾 𝑘=1 𝑝 𝑘 (𝑢)𝑏 𝑘 (𝑥), where 𝑏 𝑘 (𝑥) and 𝑝 𝑘 (𝑢) are learnable functions and functionals. In the DeepONet literature, 𝑝 𝑘 (⋅) is commonly referred to as the branch net, while 𝑏 𝑘 (⋅) is known as the trunk network. To simplify our notation, we will omit the parameters of the learnable functions whenever possible. It is important to note that 𝑏 𝑘 (𝑥) can be seen as the basis of the output function space, which is a shared element among all 𝑢 and is independent of 𝑢 ∈ 𝑉 . On the other hand, 𝑝 𝑘 (𝑢) represents the coefficients in the linear combination, which depend on the specified input function 𝑢. Consequently, the approach taken is to develop a dedicated distributed approach that shares 𝑏 𝑘 (𝑥) across all samples to create a universal/consensus trunk network model. In this dedicated approach, individual clients independently manage their respective input function datasets using their branch nets (or projection in BelNet). They send their individual basis 𝑏 𝑘 (𝑥) to a centralized server to learn a universal/consensus model applicable to all samples. Please refer to Fig. 3 for an illustration and a comparison with classical training (Fig. 4).
This means that each client independently manages input functions in its dataset through its unique branch (or projection) nets. Additionally, they transmit their local trunk networks to the central server responsible for acquiring and sending the universal/consensus basis 𝑏 𝑘 (𝑥) across all samples. Since each dataset is managed by its individual branch (or projection), it becomes possible to use designated sensors tailored to the input function within each dataset.
Let 𝐷 denote the entire heterogeneous dataset, and 𝐷 1 , … ., 𝐷 𝐶 to be sub-datasets. Moreover, each 𝐷 𝑖 exclusively comprises functions belonging to a single class, sharing similar regularity assumptions. Let us denote 𝛼 𝑐 as the trainable parameters of 𝑐th client's dedicated neural network. Let 𝛽 denote the trainable parameter of the shared network. Thus in this task, 𝛼 𝑐 are the branch net parameters, and 𝛽 are the trunk net parameters. We then denote 𝛼 = ⋃ 𝐶 𝑐=1 𝛼 𝑐 and 𝜃 = 𝛼 ⋃ 𝛽. The local loss function for 𝑐th client based on its corresponding dataset 𝐷 𝑐 will be:
where 𝑁 𝑐 denotes the number of input functions in dataset 𝐷 𝑐 , 𝑢 𝑐 𝑖 and û𝑐 𝑖 are the 𝑖th function in 𝑐th sub-dataset and its discretization at all sensors. We use the local loss to update the local parameters 𝛼 𝑐 associated with its own dedicated dataset 𝐷 𝑐 , the shared parameter 𝛽 remains unchanged during this process but plays a role in formulating the loss functions and optimizing the system. Please note that we normalize the loss function with the number of samples during the implementation, and present the unnormalized version Fig. 4. Classical DeepONet training process demonstration. for a better illustration of the concept. We can then define the global loss function used to update the shared parameter 𝛼 given all the local parameters 𝛽 and using all data, 𝐿(𝛽; 𝛼) = 𝐶 ∑ 𝑐=1 𝐿 𝑐 (𝛼 𝑐 ; 𝛽),
where 𝐶 denotes the total number of clients and datasets. We now summarize the pseudo-algorithm in Algorithm 1 with an illustration in Fig. 3.
Algorithm 1: Distributed Deep Neural Network (D2NO) Algorithm 1 Initialization: weights 𝑤 𝑐 , all networks trainable parameters 𝜃, learning rate 𝜂 𝑐 and 𝜂, the total number of iterations 𝑁. for 𝑘 = 1 to 𝑁 do 2 Update the unshared parameters 𝛼 𝑐 𝑘 for 𝑐th client using its dedicated sub-dataset, 𝑘 denotes the 𝑘th optimization step. for 𝑐 = 1 to 𝐶 do 3 𝛼 𝑐 𝑘 = 𝛼 𝑐 𝑘 -𝜂 𝑐 ∇ 𝛼 𝑐 𝐿 𝑐 (𝛼 𝑐 ; 𝛽). Update the shared construction net parameter 𝛽 using all data 𝛽 𝑘 = 𝛽 𝑘 -𝜂∇ 𝛽 𝐿(𝛽; 𝛼).
This section provides a brief analysis of the efficiency of the proposed D2NO in terms of data and computational cost. To train scientific machine learning models, a significant amount of data is needed, much of which could be privacy-sensitive. Surrogates trained using this data have the potential to greatly enhance scientific and engineering discovery and simulation. However, transferring this data to a centralized location is costly, and the quantity and security of the data may make training in such a location impractical using the traditional DeepONet approach. The proposed D2NO addresses this issue by allowing the training data to remain distributed locally. Instead of sharing large datasets, the client only shares Deep Neural Operator parameters (or a partial model, e.g., the trunk network), which are aggregated using locally computed updates that reduce training costs.
To compare the computational training costs of the traditional DeepONet and D2NO, we proceed as follows. First, consider a centralized model with a vector of trainable parameters 𝜃 ∈ R 𝑝 . Without loss of generality, let us assume we select one point 𝑥 (𝑖) in each output function's domain, we then minimize the following loss function,
over the dataset of 𝑁 triplets cent = {( û(𝑖) , 𝑥 (𝑖) ), 𝐺(𝑢 (𝑖) )(𝑥 (𝑖) )} 𝑁 𝑖=1 , here û(𝑖) is the function value sampled at sensors. A standard deep learning package can compute the gradient of the loss function with respect to each trainable parameter in 𝛽. In this analysis, we will focus on the number of gradients computed for each data loss function 𝑖 , where 𝑖 = 1, … , 𝑁, with respect to the trainable parameters. This results in a total of 𝑁 ⋅ 𝑝 ⋅ 𝑛 epochs gradient computations over 𝑛 epochs of training. This analysis does not take into account batching, but extending it to the stochastic setting is direct.
For the D2NO case, we assume that the data is distributed among 𝐶 clients, so that 𝑁 = ∑ 𝐶 𝑗=1 𝑛 𝑗 , where 𝑛 𝑗 represents the number of data points available to the 𝑗th client for training a local model with a vector of 𝑝 ′ ≤ 𝑝 trainable parameters. Therefore, for the 𝑗th client, the number of gradients computed for each data loss function over 𝑛 epochs is 𝑛 𝑗 ⋅ 𝑝 ′ ⋅ 𝑛 epochs . As a result, we can establish the following inequality:
The inequality above indicates that D2NO has a training cost similar to or even less than that of the traditional DeepONet. Furthermore, this inequality is strict for certain examples presented later in this paper. For instance, when comparing the proposed D2NO (with two clients having a number of sensors 𝑚 1 = 10 and 𝑚 2 = 100) to the traditional DeepONet, which uses 𝑚 = 100 sensors, they both perform an equal number of stochastic gradient updates and use the same amount of training data samples. However, our proposed framework is more efficient because over half of these stochastic gradient updates reduce the input function size by 10x, resulting in a decrease in the computational cost of training.
We conclude this section with the following remark. Although Algorithm 1 uses multiple client networks, which increases the total number of trainable parameters, it is important to note that, as demonstrated in this section, the total number of gradient updates does not increase.
Without loss of generality, we consider two clients (𝐶 = 2) in the remainder of this paper. Let 𝐺 ∶ 𝑉 → 𝑈 be a continuous operator. We assume that 𝐺 1 ∶ 𝑉 1 → 𝑈 and 𝐺 2 ∶ 𝑉 2 → 𝑈 are the restrictions of 𝐺 on 𝑉 1 and 𝑉 2 respectively. Specifically, 𝐺(𝑢) = 𝐺 1 (𝑢) for any 𝑢 ∈ 𝑉 1 and 𝐺(𝑢) = 𝐺 2 (𝑢) for any 𝑢 ∈ 𝑉 2 . We will construct the universal approximation for both 𝐺 1 and 𝐺 2 using a shared output function basis. Let us make the following assumptions. Assumption 4.1. We make the following assumptions on the input/output function spaces and the operator.
1. 𝑉 1 ⊂ 𝐶(𝐾 1 ), 𝑉 2 ⊂ 𝐶(𝐾 2 ) are compact and are defined on compact domains 𝐾 1 ⊂ 𝐾 ′ 1 and 𝐾 2 ⊂ 𝐾 ′ 2 , where 𝐾 ′ 1 and 𝐾 ′ 2 are Banach spaces. Additionally, 𝑉 = 𝑉 1 ∪ 𝑉 2 . 2. Let 𝐺, 𝐺 1 and 𝐺 2 are continuous, and 𝐺(𝑢) = 𝐺 1 (𝑢) for any 𝑢 ∈ 𝑉 1 and 𝐺(𝑢) = 𝐺 2 (𝑢) for any 𝑢 ∈ 𝑉 2 . 3. 𝑈 = 𝐺 1 (𝑉 1 ) ∪ 𝐺 2 (𝑉 2 ) has a compact domain 𝐾 𝑥 ⊂ R 𝑑 . The following theorem establishes the universal approximation of continuous nonlinear operators. Theorem 4.3. Let 𝐺 ∶ 𝑉 → 𝑈 , 𝐺 1 ∶ 𝑉 1 → 𝑈 and 𝐺 2 ∶ 𝑉 2 → 𝑈 satisfy Assumption 4.1. For any 𝜀 > 0, there exist weights 𝑐 1 𝑖,𝑘 , 𝑐 2 𝑖,𝑘
, and shared
2 𝑖,𝑘 𝑏 2 𝑖,𝑘 ( û2 )𝑏 𝑘 (𝑥)| < 𝜀, for any 𝑢 1 ∈ 𝑉 1 and 𝑢 2 ∈ 𝑉 2 , here û1 = [𝑢(𝑦 1 ), … , 𝑢(𝑦 𝑁 1 )] ⊤ and û2 = [𝑢(𝑧 1 ), … , 𝑢(𝑧 𝑁 2 )] ⊤ . Proof. Without loss of generality, we will prove the approximation for the first client. For any 𝜀 2 > 0, by Lemma 4.1, there exist 𝐾 networks 𝑏 𝑘 (𝑥), functional 𝑎 𝑘 (⋅) on 𝑈 such that,
for any 𝑢 ∈ 𝑉 and 𝑥 ∈ 𝐾 𝑥 . We prove the case for the first client, hence assume 𝑢 ∈ 𝑉
1 , and 𝐺 1 (𝑢) = 𝐺(𝑢) for 𝑢 ∈ 𝑉 1 , it follows that |𝐺 1 (𝑢)(𝑥) -∑ 𝐾 𝑘=1 𝑎 𝑘 ( 𝐺 1 (𝑢) ) 𝑏 𝑘 (𝑥)| < 𝜀 2 . Denote 𝐿 = ∑ 𝐾 𝑘=1 sup 𝑥∈𝐾 𝑥 |𝑏 𝑘 (𝑥)|, it follows by Lemma 4.2 that, for 𝜀 2𝐿 and any 𝑘 = 1, … , 𝐾, there exist weights 𝑐 1 𝑖,𝑘 , [𝑦 1 , … , 𝑦 𝑁 1 ] ⊂ 𝐾 1 and networks 𝑏 1 𝑖,𝑘 ∶ R 𝑁 1 → R such that, |𝑎 𝑘 (𝐺(𝑢)) -
where û1 = [𝑢(𝑦 1 ), … , 𝑢(𝑦 𝑁 1 )] ⊤ . Let 𝐼 1 be large enough and set 𝑐 1 𝑖,𝑘 = 0 if necessary, substitute Eq. ( 5) into Eq. ( 4), it follows that,
for any 𝑢 ∈ 𝑉 1 and 𝑥 ∈ 𝐾 𝑥 . The approximation for 𝐺 2 is similar given 𝑢 ∈ 𝑉 2 . □
In this section, we use four numerical experiments to demonstrate the effectiveness of D2NO in approximating nonlinear operators with heterogeneous input function spaces.
Notably, in real engineering and industrial scenarios, each client possesses its distinct dataset, facilitating automatic dataset division. In our simulation experiments, we partition the dataset according to the input function generators. For instance, as depicted in Fig. 2 of the numerical experiments detailed in Section 5.1.2, we employ Gaussian distributions to generate these functions. Subsequently, the input functions are automatically segregated based on the different variances of the distribution.
We begin by studying the viscous Burgers' equation. Our main objective is to understand and approximate the operator that maps the initial condition of the PDE to the solution at the terminal time. Specifically, the viscous Burgers' equation is defined as follows:
with the initial condition 𝑢 𝑠 (𝑥, 0) = 𝑢 0 𝑠 (𝑥) and the boundary conditions 𝑢 𝑠 (0, 𝑡) = 𝑢 𝑠 (2𝜋, 𝑡). Here, 𝑢 0 𝑠 (𝑥) is the initial condition that depends on the parameter 𝑠, and 𝛼 is the viscosity of the equation. Our goal is to learn the mapping from 𝑢 0 𝑠 to 𝑢 𝑠 (𝑡) [8,10,49,50].
Z. Zhang et al.
Left: the input functions exhibit peaks, with these peaks being concentrated in two distinct regions (black and magenta). Such a dataset poses significant challenges to the conventional operator learning framework, which mandates uniform sensor placement for all input functions. This is due to the necessity for dense sampling to effectively capture the sparsely distributed peaks. However, this approach results in a substantial increase in problem dimensionality, leading to the presence of numerous redundant entries in the discretized vectors, please check the green dots in the middle image. (For interpretation of the references to color in this figure legend, the reader is referred to the web version of this article.)
The average 𝐿 2 -relative error computed from 50 independent training runs for both the DeepONet and D2NO described in Algorithm 1. As a reference, we predict by the average of the training labels, the relative error is 97.59%.
𝐿 2 -relative error 18.64% 8.10%
Firstly, we focus our attention on the challenging scenario described in Fig. 5, where input functions display heterogeneous properties. Specifically, these input functions exhibit distinct sharp peaks, which are concentrated in two separate regions (depicted by the black and magenta curves in Fig. 5).
In the traditional operator learning or DeepONet framework, uniform sensor placement is required for all input functions. Thus, only a limited number of sensors are used to sample the peak of one input function, and most of these sensors provide the same function values. This uniformity of data significantly affects the quality of training.
By using Algorithm 1, we can use two clients (𝐶 = 2) and two dedicated branch networks to capture the input functions individually in the clustered regions. As a result, we can deploy separate sensor sets for these two datasets. Fig. 5 illustrates how a function from the left region is sampled using dedicated sensors tailored to that specific region. In contrast, Fig. 5 shows the same function sampled using uniform sensors for comparison.
To evaluate the performance of D2NO (see Algorithm 1), we trained a total of 50 independent D2NO models and calculated the average 𝐿 2 -relative error. We compared D2NO with the traditional DeepONet model. Specifically, we trained 50 traditional DeepONet models using the same network structure and hyperparameters as D2NO. It is important to note that for the traditional DeepONet, we used uniform sensors for all input functions. Table 1 presents a summary of the results, which demonstrate that the proposed D2NO algorithm significantly reduces relative errors while maintaining a comparable computational cost.
In this experiment, we explore the scenario depicted in Fig. 2. This illustration reflects many other situations where using different numbers of sensors for different input functions can result in improved solutions. Previous attempts have focused on expanding the capabilities of neural operators to discretize different input functions in different ways. However, these efforts have been limited by the need for an equal number of sensors. Our new algorithm overcomes this limitation by allowing the use of varying quantities of sensors for function discretization.
In our example, we present two different types of input functions: (a) functions with humps, represented by the black curve in Fig. 2, and (b) functions with relatively smooth profiles, represented by the magenta curve in the same figure. Typically, the conventional approach involves using the same number of sensors for all functions. However, due to the presence of humps in the first function category, a larger number of sensors is required for both types of functions. The numerical examples demonstrate that DeepONet-style networks can effectively handle challenges presented by smooth input functions, even with a limited number of sensors. However, for functions with peaks, a higher number of sensors is required, which increases the computational cost when dealing with smoother functions.
Since the dedicated branch networks are not shared among clients, our algorithm enables the creation of separate branch networks for these two function classes. As a result, the sensor count and sensor placements can vary depending on the specific function being considered. An illustration of this is in Fig. 2.
We use two clients, each with its own dedicated branch network. For functions that have peak characteristics, we use a larger two-layer branch network with dimensions 75 × 100 → 100 × 1. On the other hand, for smoother functions, we employ a smaller two-layer branch network with dimensions 6 × 50 → 50 × 1. Thus, the proposed D2NO requires the use of 75 sensors and 6 sensors, respectively.
Note that our proposed distributed algorithm (D2NO) adjusts the sensor count for different input function spaces, unlike the traditional DeepONet approach which uses a uniform configuration of 75 sensors for all input functions. Table 2 presents the aggregated results after training 50 independent models using both the traditional DeepONet and the proposed strategy.
The average L2-relative error of 50 independent trainings for DeepONet and D2NO (see Algorithm 1). Two dedicated branch networks with distinct structures are used to capture hump input functions and smooth input functions, respectively. Note that this approach helps to reduce the number of trainable parameters. As a reference, if we were to predict the solution by the average of the training samples, the relative error would be 810.61%.
𝐿 2 -relative error 7.53% 5.34% # unshared parameters 77𝐾 41.5𝐾 # shared parameters 51.8𝐾 51.8𝐾 The input functions exhibit different frequencies. Such a dataset poses significant challenges to the conventional operator learning framework, which mandates uniform sensor placement for all input functions. This is due to the necessity for dense sampling to effectively capture the oscillatory behavior. However, this approach results in a substantial increase in problem dimensionality, leading to the presence of numerous redundant entries in the discretized vectors.
In this experiment, we use the proposed distributed strategy D2NO to approximate the solution operator of the nonlinear pendulum system with dynamics:
Here, 𝑠(𝑡) = (𝑠 1 (𝑡), 𝑠 2 (𝑡)) ⊤ is the state defined over the time domain 𝑡 ∈ [0, 1] and 𝑢(𝑡) is the external force. We aim to approximate the solution operator 𝐺 ∶ 𝑢(𝑡) ↦ 𝑠(𝑡) for a fixed initial condition (𝑠 1 (0), 𝑠 2 (0)) = (0.0, 0.0), 𝑘 = 1.0, and an external force 𝑢(𝑡) sampled from the mean-zero Gaussian Random Field (GRF):
where the covariance kernel
2 ) is the radial-basis function (RBF) with length-scale parameter 𝓁 > 0, which we will use to generate heterogeneous input function spaces, as described next.
Heterogeneous Input Functions with Different Frequencies. We first consider two input function spaces with different frequencies. To this end, we let the first input space 𝑉 1 be a GRF with length scale 𝓁 = 0.1 and the second input space 𝑉 2 be a GRF with length scale 𝓁 = 1.0. Fig. 6 illustrates sample functions from these two input spaces. The centralized DeepONet framework requires uniformly placing a collection of 𝑚 sensors to discretize the input function 𝑢. In general, as depicted in Fig. 6, a large number of sensors (e.g., 𝑚 = 100) is required to capture effectively the input function 𝑢, which, in turn, increases the training cost.
We note that for the input function sampled from the space 𝑉 2 , using 𝑚 = 100 sensors is excessive. One may require only 𝑚 = 10 sensors (see Fig. 6) to effectively capture input functions sampled from 𝑉 2 . However, as illustrated in Fig. 6, using less number of sensors will deteriorate our ability to describe the oscillations from the input function sampled from the input space 𝑉 1 . The proposed distributed framework D2NO can effectively tackle the above problem by using two clients (𝐶 = 2), each with dedicated branch networks, as depicted in Fig. 7. More specifically, the first client uses a branch network with 𝑚 = 100 sensors to capture and discretize input functions sampled from the input space 𝑉 1 . The second client uses a branch network with 𝑚 = 10 sensors to capture and discretize input functions sampled from 𝑉 2 . We let the two clients share and synchronize the trunk network.
Training and Testing. To ensure a fair comparison between the centralized DeepONet and D2NO, we used the same neural network architectures, with the exception of the input layer of the dedicated branch networks for D2NO. Specifically, we used a feed-forward neural network architecture with one hidden layer and leaky-ReLU activation functions for both branch and trunk networks. The hidden and output layers of these networks consist of 50 neurons. Table 3 summarizes the number of parameters for the resulting DeepONets and D2NO.
The number of parameters for the resulting DeepONets with 𝑚 = 100 and 𝑚 = 10 sensors, as well as D2NO. Note that for the DeepONets, we detail the number of parameters for both the trunk and branch networks. On the other hand, for D2NO, we specify the number of parameters for the shared trunk network and the dedicated branch networks for clients 1 and 2.
DeepONet (𝑚 = 10) D2NO
Trunk/shared trunk 2650 2650 2650 Branch/branch-1 7600 3100 7600 Branch-2 --3100
The average 𝐿 2 -relative error for 100 test trajectories sampled from the input space 𝑉 1 (GRF with length-scale 𝓁 = 0.1) and for 100 test trajectories sampled from the input space 𝑉 2 (GRF with length-scale 𝓁 = 1.0), using the centralized DeepONet (𝑚 = 100 and 𝑚 = 10 sensors) and the distributed framework D2NO detailed in Algorithm 1.
= 100) DeepONet (𝑚 = 10) D2NO 𝐿 2 -rel. error on 𝑉 1 2.90% 12.35% 4.26% 𝐿 2 -rel. error on 𝑉 2 0.62% 1.25% 1.57%
To train both DeepONets and D2NO, we sampled 𝑁 𝑉 1 = 𝑁 𝑉 2 = 1, 000 input trajectories from each input space. The traditional/centralized DeepONet was trained using stochastic gradient-based optimization for 1000 epochs. Subsequently, each model within the proposed D2NO approach was trained using 1000 synchronization rounds after one local stochastic gradient update. It should be noted that the accuracy of D2NO could be improved by increasing the number of local gradient updates, as per federated learning approaches. However, we avoided this strategy to ensure a fair comparison of computational costs. Finally, to test the proposed framework, we sampled 100 test trajectories from both 𝑉 1 and 𝑉 2 .
Performance of D2NO. Table 4 shows the performance (i.e., in terms of the 𝐿 2 -relative error) of the proposed D2NO and the centralized DeepONet (with 𝑚 = 100 and 𝑚 = 10 sensors) for the test trajectories sampled from the input spaces 𝑉 1 and 𝑉 2 . The results clearly illustrate that the proposed D2NO framework is competitive with the centralized DeepONet with 𝑚 = 100 sensors and vastly outperforms the more efficient DeepONet with 𝑚 = 10 sensors. Specifically, D2NO delivers a predictive accuracy for both input spaces 𝑉 1 and 𝑉 2 , comparable to a DeepONet with 𝑚 = 100 sensors, matches the accuracy of a DeepONet with 𝑚 = 10 sensors on the input space 𝑉 2 , and significantly outperforms a DeepONet with 𝑚 = 10 sensors on 𝑉 1 .
Computational Cost Analysis. Table 5 compares the number gradient computations for D2NO and DeepONet using 𝑚 = 10 and 𝑚 = 100 sensors after 𝑛 epochs = 1000 epochs of training. The results highlight the advantages of adapting the number of sensors with D2NO. Specifically, D2NO delivers a predictive accuracy for both input spaces 𝑉 1 and 𝑉 2 , comparable to a DeepONet with 𝑚 = 100 sensors (as seen in Table 4) but with fewer gradient computations. D2NO also matches the accuracy of a DeepONet with 𝑚 = 10 sensors on the input space 𝑉 2 . Moreover, D2NO significantly outperforms a DeepONet with 𝑚 = 10 sensors on test trajectories generated from inputs sampled from 𝑉 1 . Therefore, D2NO strikes a balance between efficiency and accuracy for DeepONets trained using heterogeneous input spaces. Finally, from these results, we anticipate that as the number of distributed clients increases (for example, due to more heterogeneous input functions), the savings from D2NO will also increase.
Here, we test the proposed D2NO framework on an imbalanced scenario where the number of training trajectories sampled from 𝑉 1 is 𝑁 𝑉 1 = 400 and the number of training trajectories sampled from 𝑉 2 is 𝑁 𝑉 2 = 1, 600.
Comparison of the number of gradient computations for D2NO, which balances efficiency and accuracy, a DeepONet that uses 𝑚 = 10 sensors for efficiency but lacks accuracy, and a costly but accurate DeepONet that uses 𝑚 = 100 sensors, after 𝑛 epochs = 1000 epochs of training for the pendulum example.
DeepONet (𝑚 = 10) D2NO grad. computations 3.28 × 10 8 1.84 × 10 8 2.56 × 10 8
The average 𝐿 2 -relative error for 100 test trajectories sampled from the input space 𝑉 1 (GRF with length-scale 𝓁 = 0.1) and for 100 test trajectories sampled from the input space 𝑉 2 (GRF with lengthscale 𝓁 = 1.0), using the centralized DeepONet (𝑚 = 100 and 𝑚 = 10 sensors) and the distributed framework D2NO detailed in Algorithm 1. The frameworks were trained using 𝑁 𝑉
1 = 400 input trajectories sampled from 𝑉 1 and 𝑁 𝑉 2 = 1600 input trajectories sampled from 𝑉 2 . DeepONet (𝑚 = 100) DeepONet (𝑚 = 10) D2NO 𝐿 2 -rel. error on 𝑉 1 5.87% 16.09% 8.09% 𝐿 2 -rel. error on 𝑉 2 0.84% 1.01% 1.48% Table 6 describes the performance of the proposed D2NO and the centralized DeepONet. The results clearly demonstrate the advantages of the proposed D2NO framework. Specifically, it outperforms the DeepONet with 𝑚 = 10 sensors on 𝑉 1 and matches its performance on 𝑉 2 , even in a scenario that favors this DeepONet architecture.
In our last experiment, we used the proposed D2NO to approximate the solution operator of a two-dimensional PDE, specifically, the 2D-elliptic equation:
and we consider the nonlinear mapping from the permeability 𝜅(𝑥) to the solution 𝑢(𝑥).
Heterogeneous Input Spaces. To approximate the solution operator, we considered heterogeneous functions sampled from two input spaces, denoted as 𝑉 1 and 𝑉 2 . We constructed samples from 𝑉 1 using the Gaussian distribution with the mean centered in the region [0, 0.5] × [0, 0.5] as shown in Fig. 8 Left. Similarly, we obtained samples from 𝑉 2 , but the mean of the Gaussian distribution is from region [0.5, 1] × [0.5, 1] (shown in Fig. 8 Right). We constructed a centralized DeepONet and a D2NO with two clients. These were used to approximate the solution operator of the 2D-elliptic equation using inputs sampled from the heterogeneous spaces 𝑉 1 and 𝑉 2 .
Fig. 8 illustrates random samples from 𝑉 1 and 𝑉 2 . We generated input samples for 𝑉 1 and 𝑉 2 using a uniform grid of size 100 × 100 over the entire input domain [0.0, 1.0] × [0.0, 1.0]. The inputs to the branch network of the centralized DeepONet were discretized using 𝑚 = 1000 subsamples from the uniform grid that covers the entire input domain. In contrast, the dedicated branch networks of D2NO used discretized input samples with 𝑚 𝑉 1 = 640 and 𝑚 𝑉 2 = 560 sensors, respectively. These sensors were located by subsampling the uniform grid over smaller regions: [0.
2, 0.8] × [0.2, 0.8] for 𝑉 1 and [0.2, 1.0] × [0.3, 1.0] for 𝑉 2 .
Training and Testing. As with the pendulum example, we ensured a fair comparison between the centralized DeepONet and D2NO by using identical neural network architectures, except for the input layer of the dedicated branch networks for D2NO. Specifically, we employed a feed-forward neural network architecture with two hidden layers and leaky-ReLU activation functions for both branch
The number of parameters for the centralized DeepONet and the proposed D2NO. Note that for the DeepONet, we detail the number of parameters for both the trunk and branch networks. On the other hand, for D2NO, we specify the number of parameters for the shared trunk network and the dedicated branch networks for clients 1 and 2.
Trunk/shared trunk 20,500 20,500 Branch/branch-1 120,300 84,300 Branch-2 -76,300 and trunk networks. Each network's hidden and output layers contain 100 neurons. Table 7 summarizes the parameter count for the resulting centralized DeepONet and D2NO.
To train both DeepONets and D2NO, we sampled 𝑁 𝑉 1 = 𝑁 𝑉 2 = 80 input functions from each input space. The traditional/centralized DeepONet was trained using stochastic gradient-based optimization for 1000 epochs. Subsequently, each model within the proposed D2NO approach was trained using 1000 synchronization rounds after one local stochastic gradient update. Finally, to test the proposed framework, we sampled 20 test input functions from both 𝑉 1 and 𝑉 2 .
Performance of D2NO. Table 8 displays the performance of the proposed D2NO and the centralized DeepONet, evaluated by the 𝐿 2 -relative error, for the test input functions sampled from the input spaces 𝑉 1 and 𝑉 2 . The results reveal that D2NO provides predictive accuracy that is comparable to or even surpasses, the centralized DeepONet for both input spaces, 𝑉 1 and 𝑉 2 .
Computational Cost Analysis. We conclude this example by illustrating the savings in gradient computations for D2NO applied to the 2D elliptic PDE in Table 9. These savings occurred after 𝑛 epochs = 1000 epochs of training. It is important to note that D2NO matches the accuracy of the centralized DeepONet, as demonstrated in Table 8, while carrying out fewer gradient computations.
In this paper, we have developed a dedicated distributed training strategy for deep neural operators (D2NO) that can efficiently handle heterogeneous input function spaces. Our strategy involves partitioning the input function space based on the properties and regularity of the input functions. Each subset of the input function space is then processed using a dedicated branch network, which is trained locally. Additionally, a consensus trunk network is trained in a centralized manner. This trunk network acts as a global basis for the local coefficients obtained from the dedicated branch networks. We have demonstrated that D2NO serves as a universal approximator of continuous nonlinear operators. Finally, we have demonstrated that D2NO not only addresses the challenges posed by input functions with heterogeneous properties but also offers improved flexibility and computational efficiency using four numerical examples, each with heterogeneous input function spaces. In the future, we will consider non-uniform sensors within each sub-dataset and apply the algorithms to discretization invariant neural operators such as BelNet. Another potential work is to study reducing the samples used in training the shared structure and study the error growth.
The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.