Key words. Hybrid least squares gradient descent method, MIONet, Kronecker product, Khatri-Rao product, tensor permutation matrix

MSC codes. 15A69, 47-08, 65F45, 65Y10, 68T07, 68T20

1 Introduction↩︎

Thanks to the recent advances in scientific machine learning, the core architectures, including deep learning (DL) and deep neural networks (DNNs) have migrated to the field of scientific computing to enhance existing numerical methods for solving various partial differential equations (PDEs). In particular, physics-informed neural network (PINN) [1] is the most successful and widely used method, where PINN represents the solution of PDE as a DNN and finds the solution by training the DNN using a physics-informed loss (PI-loss) with the automatic differentiation method [2]. However, since PINN requires separate training for different PDE instances, the need for a mapping between components and solutions of PDEs using DL architecture has emerged, which has now been generalized to neural operator mapping between function spaces. There are various examples of neural operators, including Deep Operator Network (DeepONet) [3], Fourier Neural Operator [4], Graph Kernel Network [5], PCA-based Model Reduction [6], and Multi-Wavelet Neural Operator [7].

Among these neural operators, DeepONet is the most widely used framework for neural operators, which possesses the universal approximation property. It consists of the inner product of outputs from two neural networks, branch and trunk, where the branch network encodes input functions and the trunk network encodes coordinates of the output function domain. Based on the DeepONet architecture, many variants have been proposed, such as POD-DeepONet [8], Multifidelity DeepONet [9], NOMAD [10], Multiple-Input Operator Network (MIONet) [11], Shift-DeepONet [12], HyperDeepONet [13], and Geom-DeepONet [14].

In this paper, we focus on MIONet since it is a direct generalization of DeepONet, which maps several input functions to a single output function with the corresponding universal approximation theorem (UAT) [11]. Instead of a single branch network in DeepONet, MIONet uses multiple branch networks to encode each input function and computes the entrywise product of the outputs from each branch to perform an inner product with the output of the trunk network.

However, the training for MIONet is challenging because the entrywise product and inner product among several networks make the structure more complex, and a sufficiently large dataset is needed for meaningful training. This makes the conventional MIONet training with the Adam optimizer [15] require very high computational resources and time.

To optimize and accelerate MIONet training, we take a deeper look into the hybrid least squares/gradient descent (LSGD) method for DeepONet [16], [17]. We generalize the LSGD method into the MIONet framework with vanilla structure, where each output layer of the branch network is a fully connected linear layer. For a general type of \(L^2\) loss with regularization terms for the last layer parameters of the branch networks, we formulate a minimization problem of sums of squared multilinear functions in terms of the last layer parameters. To optimize this problem, we first fix the last layer parameters of all branches except one, then the minimization problem becomes an LS system for the unfixed last layer parameters. Now, by generating and solving LS problems for the unfixed last layer parameters, alternating branches in sequence, we can optimize the last layer parameters of all the branch networks. We call this the alternating least squares (ALS) method. Although each LS system is very large to handle, we can factor the large system matrix into smaller matrices from each branch and trunk network, where the column-wise Kronecker product (Khatri-Rao product [18]) and the usual Kronecker product are used. Additionally, we introduce tensor permutation matrices [19] to match the row order shuffled due to the Kronecker products to the lexicographic order of the dimension axes. After that, the LS system is transformed into a special type of matrix equation of the form \(AXB+\lambda X = E\), where the coefficient matrices \(A\) and \(B\) are from the component matrices of the LS system. Furthermore, we provide a theorem that helps reduce the complexity of the computation of data tensors, when the given data tensors depend only on one input function argument of MIONet. Finally, we propose the ALS plus Adam (ALS+Adam) method as a practical algorithm for LSGD for MIONet, which is a modification of the LS+Adam method [16].

This paper is organized as follows. In 2, we introduce MIONet and provide the UAT for vanilla MIONets. We also briefly summarize the LSGD method [17] and LSGD methods for DeepONets [16]. In 3, we formulate the minimization problem from the general type of squared \(L^2\) loss in terms of the last layer parameters of branch networks. After that, we will explain how this problem can be understood as an LS system of the last layer parameters and present the corresponding LSGD method for MIONets. In 4, we conduct experiments on supervised learning for a nonlinear PDE and unsupervised learning for linear PDEs to compare the training performance between conventional MIONet training with Adam and MIONet training with ALS+Adam.

2 Preliminaries↩︎

In this section, we introduce MIONet with its universal approximation property [11] and the hybrid LSGD method for DeepONets [16]. Refer to 1 for the meaning of symbols and variables used in this paper.

2.1 Multiple-Input Operator Network (MIONet)↩︎

Jin et al. [11] proved an UAT for multiple-input operators on the product of Banach spaces of functions with Schauder bases. We refer to [20]–[22] for more details of the Schauder basis and its canonical projections. The theorem states that a continuous multiple-input operator \(G\) on the product of compact subsets of Banach spaces \(X_m\) with Schauder bases can be approximated by the form \[~\label{eq:UATMION} \Big\langle \underbrace{(\mathbf{b}_1\circ \phi^1_{M_1})}_{\textrm{branch 1}} \odot \cdots \odot \underbrace{(\mathbf{b}_N \circ \phi^N_{M_N})}_{\textrm{branch N}}, \underbrace{\mathbf{t}}_{\textrm{trunk}}\Big\rangle\tag{1}\] with continuous vector functions \(\mathbf{b}_m\in C(\mathbb{R}^{M_m},\mathbb{R}^{I})\) and \(\mathbf{t}\in Y^{I}\) for sufficiently large positive integers \(I\) and \(M_m\), where \(\phi^m_{M_m}\colon X_m\to\mathbb{R}^{M_m}\) extracts the first \(M_m\) coefficients of the Schauder basis representation, \(Y\) is the target Banach space, and \(\odot\) denotes the Hadamard (entrywise) product.

Here, from the inner product structure 1 of vector functions \(\mathbf{b}_m\) and \(\mathbf{t}\), a multiple-input operator network (MIONet) can be constructed by replacing those functions with neural networks. In accordance with the UAT for DeepONet [3], we want to use discretized function values of \(u^{(m)}\) as an input for \(\mathbf{b}_m\), where \(\phi_n^m\) extracts function values at certain points. Note that for the sequence of distinct points \({\{t_i\}}_{i=1}^{\infty}\) in \([0,1]\) with \(t_1=0\) and \(t_2=1\), which is dense in \([0,1]\), there exists a Schauder basis \({\{e_i\}}_{i=1}^{\infty}\) of \(C({[0,1]})\) (called the Faber-Schauder basis) where \(e_1(t) = 1\) and \(e_{n}\) is chosen as a piecewise linear function with \(e_{n}(t_{n})=1\), where the set \(\{e_1,\dots,e_{n}\}\) forms a basis for the space of all piecewise linear functions with node points \({\{t_i\}}_{i=1}^{n}\). We refer to [20], [22] for detailed Faber-Schauder basis construction of \(C({[0,1]})\) and [23] for the multivariate extension to \(C({[0,1]}^d)\). Furthermore, without loss of generality, MIONet can have a vanilla structure similar to DeepONets in [16], [24], where the last layer of each branch network is a fully connected layer without bias and activation function. This can be achieved by adding an identity layer to the output of each branch network. See 1 for the structure of vanilla MIONet. Note that MIONet with \(N=1\) becomes DeepONet with the same structure and the universal approximation property discussed in [16].

We now address the UAT for vanilla MIONets, where the input functions are discretized function values. This theorem is a modification of [21] and [3].

The MIONet approximating \(G\) in 2 can be expressed as a nested form of summations and products: \[\label{eq:MION} \begin{alignedat}{2} G(u^{(1)},\dots,u^{(N)})(y) &\approx \left\langle\mathbf{b}_{1}(\mathbf{u}^{(1)})\odot\cdots\odot\mathbf{b}_{N} (\mathbf{u}^{(N)}),\mathbf{t}(y)\right\rangle\\ &= \sum_{i=1}^{I} \left(\prod_{m=1}^{N}b_{i}^{(m)}(\mathbf{u}^{(m)})\right) t_i(y)\\ &= \sum_{i=1}^{I}\left[\prod_{m=1}^{N}\left(\sum_{j=1}^{J_m} c_{i j}^{(m)} \tilde{b}_{j}^{(m)}(\mathbf{u}^{(m)})\right)\right]t_i(y), \end{alignedat}\tag{3}\] where \(\tilde{b}_{j}^{(m)}\) is the \(j\)-th component of the vector function \(\tilde{\mathbf{b}}_m\), and \(b_i^{(m)}\) and \(t_i\) are the \(i\)-th component of the vector function \({\mathbf{b}}_m\) and \({\mathbf{t}}_m\), respectively.

2.2 Hybrid Least Squares/Gradient Descent Method for DeepONets↩︎

Cyr et al. [17] suggested the hybrid least squares/gradient descent (LSGD) method to improve deep neural network training. This method alternates between least squares (LS) steps and gradient descent (GD) steps. The LS steps solve an LS problem for the linear last layer parameters to optimize those parameters, and the GD steps optimize the hidden layer parameters with a gradient descent type optimizer.

We have proposed the hybrid LSGD method for DeepONets in [16], where we accelerate and improve training for vanilla DeepONets in both supervised and unsupervised learning. As a starting point for the LSGD method for MIONets, we briefly explain the LSGD method for DeepONets in this section.

Consider a vanilla DeepONet, which is the vanilla MIONet in 3 with \(N=1\). We use a general type of \(l_2\) loss, which consists of the sum of squared \(l_2\) error terms and a regularization term for the last layer parameters of the branch network: \[~\label{eq:GenL2LSDON} \sum_{k=1}^{K} \epsilon_{k} {\left\|\mathcal{L}_k\left[G(\cdot)\right](u,y) - \mathcal{L}_k\left[\langle \mathbf{b}(\cdot;\theta^{B},\theta^{L}), \mathbf{t}(\cdot;\theta^{T})\rangle\right](\mathbf{u},y)\right\|}_{l_2(\chi_k)}^{2} + \lambda {\left\|\theta^{L}\right\|}_{2}^{2},\tag{4}\] where \(\chi_k={\{(\hat{u}_{d_k},\hat{y}_{d_k})\}}_{d_k=1}^{D_k}\) is a finite collection of data pairs \((u,y)\), \(\mathcal{L}_k\) is a linear operator between real-valued functions, \(f_k=\mathcal{L}_k\left[G(\cdot)\right](\hat{u}_{d_k},\hat{y}_{d_k})\in\mathbb{R}^{D_k}\) is the given data for the \(k\)-th error term, \(\epsilon_k > 0\) and \(\lambda\geq 0\) are the weights for each error term and the regularization term, and \(\theta^{B}\), \(\theta^{T}\), \(\theta^{L}\) denote the parameters for the branch network except the last layer, the trunk network, and the last layer of the branch, respectively.

We represent \(\theta^{L}\) as the row-wise vectorization of the (branch) last layer parameter matrix \(C\in\mathbb{R}^{I\times J}\): \[~\label{eq:vecCDON} \theta^{L}=\text{vec}({C^{T}})= {[c_{11},\dots,c_{1J},c_{21},\dots,c_{2J},\dots,c_{I1},\dots,c_{IJ}]}^{T} \in{\mathbb{R}}^{IJ}.\tag{5}\] Then, 4 can be reformulated as an LS problem \[~\label{eq:LLLSDON} \min_{\theta^{L}} {\sum_{k=1}^{K} \epsilon_{k} {\left\|f_k - A_k \theta^{L}\right\|}_{2}^{2} + \lambda {\left\|{\theta}^L\right\|}_{2}^{2}},\tag{6}\] where \(A_k \in{\mathbb{R}}^{D_k\times IJ}\) is the matrix whose \((d_k,J(i-1)+j)\) entry is \(\mathcal{L}_k\left[\tilde{b}_j t_i\right](\hat{\mathbf{u}}_{d_k},\hat{y}_{d_k})\).

Under additional conditions for the data \(\chi_k\) and the linear operator \(\mathcal{L}_k\) such that \[~\label{eq:CondDON} \begin{align} \chi_k = \beta \times \tau_k &= {\left\{u_p\right\}}_{p=1}^{P} \times {\left\{y_{q_k}\right\}}_{q_k=1}^{Q_k},\\ \mathcal{L}_k\left[\tilde{b}_j t_i\right](\mathbf{u},y) &= \tilde{b}_j(\mathbf{u}) \mathcal{L}_k\left[t_i\right](y), \end{align}\tag{7}\] we have a main theorem for the LS step of the LSGD method for DeepONets [16]. In the theorem, the large matrix \(A_k\) in 6 can be factored into two small component matrices, each from the branch and trunk networks, using the Kronecker product, where a commutation matrix [25], [26] is introduced to permute the row order of the product matrix.

Then, the LS problem can be reformulated to a matrix equation, which is a special case of the generalized Sylvester equation of type \(A X B + \lambda X = E\) with symmetric positive semi-definite matrices \(A\) and \(B\). Such matrix equation can be solved by using the spectral decompositions [16].

Therefore, the last layer parameters minimizing 6 can be found efficiently, which is now the LS step of the LSGD method for DeepONets. The remaining hidden parameters can be found by using a GD-type optimizer. See [16] for the workflow of the LSGD method for DeepONets and its practical modification with Adam optimizer [15], LS+Adam, respectively.

3 Hybrid Least Squares/Gradient Descent Methods for MIONets↩︎

We consider the following loss of the vanilla MIONet, which consists of the sums of squared \(l_2\) error terms indexed by \(k\) and regularization terms for the \(N\) last layer parameters: \[~\label{eq:GenL2LSMION} \begin{alignedat}{2} \sum_{k=1}^{K} \epsilon_{k} &\Big\|\mathcal{L}_k\left[G(\cdot,\dots,\cdot)\right] (u^{(1)},\dots,u^{(N)},y)\\ &-\mathcal{L}_k\left[\langle\mathbf{b}_{1}(\cdot;\theta_{1}^{B},\theta_{1}^{L}) \odot\cdots\odot\mathbf{b}_{N}(\cdot;\theta_{N}^{B},\theta_{N}^{L}), \mathbf{t}(\cdot;\theta^{T})\rangle\right](\mathbf{u}^{(1)},\dots,\mathbf{u}^{(N)},y) \Big\|_{l_2(\chi_k)}^{2}\\ &+\sum_{m=1}^{N}\lambda_m {\left\|\theta_{m}^{L}\right\|}_{2}^{2}, \end{alignedat}\tag{10}\] where \(\chi_k\) is a finite collection of data pairs \({\left\{(\hat{u}_{d_k}^{(1)},\dots,\hat{u}_{d_k}^{(N)},\hat{y}_{d_k})\right\}}_{d_k=1}^{D_k}\), \(\mathcal{L}_k\) is a linear operator between real-valued functions for the \(k\)-th error term, \(f_k\) is the data for the \(k\)-th error term given as \(f_k = \left({\mathcal{L}_k\left[G(\cdot,\dots,\cdot)\right](\hat{u}^{(1)}_{d_k},\dots, \hat{u}^{(N)}_{d_k},\hat{y}_{d_k})}\right) \in{\mathbb{R}}^{D_k}\), \(\epsilon_k > 0\) and \(\lambda_m\geq 0\) are weights for error terms and the regularization terms, and \(\theta_{m}^{B}\), \(\theta^{T}\), \(\theta_{m}^{L}\) denote the parameters for the \(m\)-th branch network except the last layer, the trunk network, and the last layer of the \(m\)-th branch, respectively, for \(m=1,\dots,N\).

For the convenience of notation, we use the big-endian lexicographic order defined below for the entries of the tensor.

Note that the above definition gives the row-wise lexicographic order of matrix entries.

Let \(C_{m}\in\mathbb{R}^{I\times J_m}\) be a parameter matrix of the \(m\)-th last layer, where \(\theta_{m}^{L}\) is the row-wise vectorization of \(C_{m}\). The entries of \(C_{m}\) are sorted lexicographically in the order 11 , i.e., \[~\label{eq:vecCMION} \theta_{m}^{L}=\text{vec}(C_{m}^{T})= {[c_{11}^{(m)},\dots,c_{1J_m}^{(m)},c_{21}^{(m)},\dots,c_{2J_m}^{(m)}, \dots,c_{I1}^{(m)},\dots,c_{IJ_m}^{(m)}]}^{T} \in{\mathbb{R}}^{IJ_m}.\tag{12}\]

Here, 10 is reformulated as the following minimization problem in terms of the \(N\) last layer parameters \(\{\theta_{1}^{L},\dots,\theta_{N}^{L}\}\): \[\label{eq:LLALSMION} \min_{\{{\theta}_{1}^{L},\dots,{\theta}_{N}^{L}\}} \sum_{k=1}^{K} \epsilon_{k} {\left\|f_k - \mathcal{A}_k(\theta_{1}^{L},\dots,\theta_{N}^{L}) \right\|}_{l_2(\chi_k)}^{2} + \sum_{m=1}^{N} \lambda_m {\left\|{\theta}_{m}^{L}\right\|}_{2}^{2},\tag{13}\] where \(\mathcal{A}_k(\theta_{1}^{L},\dots,\theta_{N}^{L}) \in {\mathbb{R}}^{D_k}\) is the long vector whose \(d_k\) entry is \[\label{eq:EntryMION} \mathcal{L}_k\left[\sum_{i=1}^{I} b_{i}^{(1)} \cdots b_{i}^{(N)}t_i\right] (\hat{\mathbf{u}}_{d_k}^{(1)},\dots,\hat{\mathbf{u}}_{d_k}^{(N)},\hat{y}_{d_k}),\tag{14}\] and it is the result of an \(N\)-linear function applied to the last layer parameters \(\theta_{1}^{L},\dots,\theta_{N}^{L}\).

Although the minimization problem 13 is not an LS problem, we can form an LS problem for \(\theta_{n}^{L}\) by fixing all last layer parameters except the \(n\)-th parameter \(\theta_n^L\). By solving the LS problems generated for each last layer parameter in turn, we can optimize the \(N\) last layer parameters. The details will be discussed after formulating the LS problem for the single last layer parameter.

Note that \(\mathcal{A}_k(\theta_{1}^{L},\dots,\theta_{N}^{L})\) is linear with respect to \(\theta_{n}^{L}\) for each \(n\), we can write \[~\label{eq:MIONA} \mathcal{A}_k(\theta_{1}^{L},\dots,\theta_{N}^{L}) = A_{k,n} \theta_{n}^{L}\tag{15}\] for each \(k\) and \(n\).

Therefore, the LS step minimizing 13 in terms of \(\theta_{n}^{L}\) is \[~\label{eq:LLLSMION} \min_{{\theta}_{n}^{L}} \sum_{k=1}^{K} \epsilon_{k} {\left\|f_{k} - A_{k,n} \theta_{n}^{L}\right\|}_{2}^{2} + \lambda_n {\left\|{\theta}_{n}^{L}\right\|}_{2}^{2}\tag{16}\] where \(A_{k,n}\in{\mathbb{R}}^{D_k\times IJ_n}\) is the matrix whose \((d_k,{[i,j]}_{I,J_n})\) entry is \[~\label{eq:MIONFactEntry} \mathcal{L}_k\left[b_{i}^{(1)}\cdots b_{i}^{(n-1)} \tilde{b}_{j}^{(n)} b_{i}^{(n+1)} \cdots b_{i}^{(N)}t_i\right](\hat{\mathbf{u}}_{d_k}^{(1)},\dots,\hat{\mathbf{u}}_{d_k}^{(N)}, \hat{y}_{d_k}).\tag{17}\]

Here, the LS system 16 cannot be handled directly if \(I\) and \(J_n\) are not small unless \(D_k\) is small, but it is not desirable to use small \(D_k\) in MIONet training. We provide the data and operator conditions for MIONet, which correspond to the conditions for DeepONet in [16].

Suppose that the data collection \(\chi_k\) can be expressed as a Cartesian product of \(N+1\) input data as follows: \[~\label{eq:CondDataMION} \chi_k = \beta_1 \times \cdots \times \beta_N \times \tau_k,\tag{18}\] where \(\beta_m = {\left\{{u}_{p}^{(m)}\right\}}_{p=1}^{P_m}\) is the set of input functions for the \(m\)-th coordinate and \(\tau_k = {\left\{{y}_{q_k}\right\}}_{q_k=1}^{Q_k}\) is the set of points of the discretized domain for \(\mathcal{L}_k [G(\cdot,\dots,\cdot)]\) with \(D_k = P_1 \cdots P_N Q_k\). Thus, the input function for each branch network is independent of the input functions for other branch networks and is used equally for all error terms. Furthermore, the same discretization over the domain of \(\mathcal{L}_k [G(\cdot,\dots,\cdot)]\) is used for all tuples of input functions \(({u}_{p_1}^{(1)},\dots,{u}_{p_N}^{(N)})\) for each error term.

We also assume that the linear operator \(\mathcal{L}_k\) satisfies \[~\label{eq:CondOpMION} \mathcal{L}_k\left[\tilde{b}_{j_1}^{(1)}\cdots\tilde{b}_{j_N}^{(N)} t_i\right] (\mathbf{u}^{(1)},\dots,\mathbf{u}^{(N)},y) = \left(\prod_{m=1}^{N}\tilde{b}_{j_m}^{(m)}(\mathbf{u}^{(m)})\right) \mathcal{L}_k\left[t_i\right](y)\tag{19}\] for all \(i,j_m,k\). This implies the linear operator \(\mathcal{L}_k\) is independent of input functions and acts on each trunk component \(t_i\).

Let \(\ast\) denote the Khatri-Rao product (column-wise Kronecker product) [18]. Note that the Khatri-Rao product of two matrices \(X\in\mathbb{R}^{R_1\times S}\) and \(Y\in\mathbb{R}^{R_2\times S}\) is an \(R_1 R_2 \times S\) matrix whose entry is represented as \[~\label{eq:KhaRaoNota} {(X \ast Y)}_{r s} = {X}_{r_1 s} {Y}_{r_2 s},\tag{20}\] where \(r = {[r_1,r_2]}_{R_1,R_2}\).

Under the conditions 18 and 19 , we show the large matrix \(A_{k,n}\) in 16 can be factored into the product of \(N+1\) smaller matrices with some permutations. For convenience, we define a sequence with the \(j\)-th term removed from the sequence \({\left\{x_i\right\}}_{i=1}^{N}\) and denote it as \((x_1,\dots,\hat{x}_j,\dots,x_N)\). Also, for the generic associative binary operator \(\circ\), we denote \(\overset{N}{\underset{\substack{i=1\\i\ne j}}{\circ}} x_i\) as the repeated application of the operation to the sequence \({\left\{x_i\right\}}_{i=1}^{N}\) with \(x_j\) omitted.

We define the tensor permutation matrix introduced in [19] in terms of the lexicographic order of the entry of a rank \(N\) tensor.

Note that \(\tilde{\mathcal{T}}\) is the tensor transpose of \(\mathcal{T}\) associated with \(\sigma\) in [27] and \({K}_{D,\sigma}\) maps the lexicographic order of the entry of \(\mathcal{T}\) to that of the entry of \(\tilde{\mathcal{T}}\).

Now, we state and prove the main theorem for a factorization of the system matrix in 16 .

Note that \(H_m\) is the final output of the \(m\)-th branch network, and since the Khatri-Rao product is associative, we generalize 20 to \[~\label{eq:KhaRaoNotaGen} {(X_1 \ast \cdots \ast X_N)}_{r s} = {(X_1)}_{r_1 s} \cdots {(X_N)}_{r_N s},\tag{24}\] where \(r = {[r_1,\dots,r_N]}_{R_1,\dots,R_N}\) for the matrices \(X_n \in\mathbb{R}^{R_n \times S}\). A similar result holds for the Kronecker product. The tensor permutation matrix \({K}_{\tilde{D}_k,\sigma_n}\) is introduced to match the mixed row order of the product \((G_n\ast T_k)\otimes B_n\) to the lexicographic row order of the data \(f_k\).

Using the result of 3, let us observe the LS problem 16 and rearrange the data vector \(f_k\in\mathbb{R}^{P_1\cdots P_N Q_k}\) into a tensor \(F_k\in\mathbb{R}^{P_1\times\cdots\times P_N\times Q_k}\). Both the \((p_1,\dots,p_N,q_k)\) entry of \(F_k\) and the \({[p_1,\dots,p_N,q_k]}_{P_1,\dots,P_N,Q_k}\) entry of \(f_k\) are \[~\label{eq:MIONDataEntry} \mathcal{L}_k\left[G(\cdot,\dots,\cdot)\right](u^{(1)}_{p_1},\dots,u^{(N)}_{p_N},y_{q_k}).\tag{26}\] Therefore, the LS problem with respect to the last layer parameter is \[~\label{eq:LLLSMIONFinal} \min_{{C}_{n}} \sum_{k=1}^{K} \epsilon_{k} {\left\|f_k - {K}_{\tilde{D}_k,\sigma_n}^{T} \left(\left(G_{n}\ast T_k\right) \otimes B_n\right) \text{vec}(C_{n}^{T}) \right\|}_{2}^{2} + \lambda_n {\left\|\text{vec}(C_{n}^{T})\right\|}_{2}^{2}.\tag{27}\]

To write the normal equation concisely, we use the fact [28] that \({(X \ast Y)}^{T} (X \ast Y) = (X^{T}X) \odot (Y^{T}Y)\) for all matrices \(X\in\mathbb{R}^{R_1\times S}\) and \(Y\in\mathbb{R}^{R_2\times S}\). Here, the normal equation of 27 is given as \[~\label{eq:NormalMION} \begin{alignedat}{2} &\left[\left(\sum_{k=1}^{K} \epsilon_{k} \left(G_n^{T} G_n\right) \odot \left(T_k^{T} T_k\right)\right) \otimes \left(B_n^{T} B_n\right)\right] \text{vec}(C_{n}^{T}) + \lambda_n \text{vec}(C_{n}^{T})\\ &=\sum_{k=1}^{K} \epsilon_{k} \left({\left(G_{n}\ast T_k\right)}^{T} \otimes B_n^{T}\right) {K}_{\tilde{D}_k,\sigma_n}^{T} f_k. \end{alignedat}\tag{28}\] Introducing a long data matrix \(\hat{F}_{k,n}\in\mathbb{R}^{P_n \times P_1\cdots \hat{P}_{n} \cdots P_N Q_k}\) whose \[\left(p_n,{[p_1,\dots,\hat{p}_n,\dots{p}_N,q]}_{P_1,\dots,\hat{P}_n,\dots,P_N,Q_k}\right)\] entry is 26 , we can write 28 as a matrix equation: \[~\label{eq:NormalMatMION} B_n^{T} B_n C_n^{T} \left(\sum_{k=1}^{K} \epsilon_{k} \left(G_n^{T} G_n\right) \odot \left(T_k^{T} T_k\right)\right) + \lambda_n C_n^{T} = B_n^{T} \left(\sum_{k=1}^{K} \epsilon_{k} \hat{F}_{k,n} \left(G_{n}\ast T_k\right)\right).\tag{29}\] Note that \((G_n^{T} G_n) \odot (T_k^{T} T_k)\) can be effectively calculated without forming huge matrices since \[~\label{eq:MIONMatReduce} (G_n^{T} G_n) \odot (T_k^{T} T_k) = \left(\overset{N}{\underset{\substack{m=1\\m\ne n}}{\odot}} (H_{m}^{T} H_{m})\right) \odot (T_k^{T} T_k),\tag{30}\] where each component is an \(I \times I\) matrix generated from each network.

In the case of unsupervised learning where the given data \(F_k\) depends only on a single input argument \(u^{(m_k)}\) of the operator \(G(u^{(1)},\dots,u^{(N)})\), we can reduce the rank \((N+1)\) data tensor \(F_k\) to the reduced matrix \(\tilde{F}_{m_k}\in\mathbb{R}^{P_{m_k}\times Q_k}\), where \(m_k\) is the corresponding index of the input function coordinate. We provide a related theorem.

This theorem enables us to efficiently compute each component of the matrix equation 29 . Finally, we can use 2 to find the minimizer \(C_n\) of the LS problem 27 for the \(n\)-th last layer parameter.

Now, we present the entire LS step for MIONet (2). Since the loss is the sum of squares of multilinear functions with last layer parameters \((\theta_{1}^{L},\dots,\theta_{N}^{L})\) of \(N\) branches, along with regularization terms, we can formulate an alternating least squares (ALS) problem. Define the optimizing order of the last layer parameters as a permutation \(\pi\in S_N\). For \(m=1,\dots,N\), fix all parameters except \(\theta_{\pi(m)}^{L}\), then solve the LS problem in terms of \(\theta_{\pi(m)}^{L}\) to update \(\theta_{\pi(m)}^{L}\). Note that a unique minimizer is guaranteed if \(\lambda_{\pi(m)}>0\). Finally, the entire process is repeated until the stopping criterion is met, such as changes in the appropriate norm of the parameters or the loss being smaller than a specific threshold. The complete LSGD algorithm for MIONet (3) is similar to the LSGD algorithm for DeepONet introduced in [16]. To speed up the ALS step, one may introduce enhanced line search (ELS) [29] or partitioned ALS (PALS) [30].

4 Experimental results↩︎

In this section, we present numerical experiments on various PDEs with 2-branch MIONet to evaluate the proposed hybrid training schemes, ALS+Adam. We report training results and compare the proposed method, ALS+Adam, with the conventional Adam training in terms of training loss convergence behavior and relative error for unseen input functions.

We propose ALS+Adam as a practical LSGD method for MIONets, which generalizes the LS+Adam method for DeepONets in [16] as follows. In the initial stage, we train all parameters using Adam for a moderate number of epochs. Then, we switch to the hybrid stage, where we use the ALS step to optimize the last layer parameters of each branch network in turn. After that, the ALS step is applied once every few Adam epochs for the hidden layer parameters. See 4.

For 2-branch MIONet training in supervised learning, we use the mean square error (MSE) loss \[~\label{eq:Loss95Sup95MION} \epsilon_{1} L_{\text{data}} + \lambda_1 {\|C_1\|}_F^2 + \lambda_2 {\|C_2\|}_F^2,\tag{38}\] and for the unsupervised learning, we use PI-loss \[~\label{eq:Loss95Unsup95MION} \epsilon_{1} L_{\text{data}} + \epsilon_{2} L_{\text{physics}} + \lambda_1 {\|C_1\|}_F^2 + \lambda_2 {\|C_2\|}_F^2,\tag{39}\] where \({\|\cdot\|}_F\) denotes the Frobenius norm, \(\epsilon_{1} = 1\), \(L_{\text{data}}\) is the \(L^2\) MSE on the data pairs \((\hat{u}_{d_1},\hat{y}_{d_1})\) where \(\hat{y}_{d_1}\) corresponds to the given data points of the governing PDE, and \(L_{\text{physics}}\) is the \(L^2\) MSE of the PDE residuals \((\hat{u}_{d_2},\hat{y}_{d_2})\) where \(\hat{y}_{d_2}\) corresponds to the residual computation points.

Supervised learning requires labeled solutions for all possible pairs of input functions \((u^{(1)},u^{(2)})\). On the other hand, unsupervised learning does not require precomputed reference solutions during the training stage. Here, we ensure that each data tensor \(F_k\) depends only on one input argument of the target operator. According to 4, we do not need to form large matrices, and efficient computation is possible in the ALS steps.

In each experiment, we use hyperparameters and network settings similar to those used for DeepONets [16]. We use the Adam optimizer with \(\text{lr} = 10^{-3}\) and \((\beta_1,\beta_2) = (0.99,0.999)\). The Adam momentums are maintained between the LS steps. He normal initialization [31] is used for parameter initialization, and the Swish function \(x/(1+e^{-x})\) is used as the activation function. For training with Adam-only, no regularization terms for the last layer parameters are used, but for training with ALS+Adam, these regularization terms are applied with the same positive weights across the last layer parameters of each branch network, i.e., \(\lambda_1 = \lambda_2 = \lambda\).

We define one work unit (WU) as one Adam epoch for training all parameters and, in the hybrid stage, as a cycle of one Adam epoch followed by one ALS step in 4. For training with ALS+Adam, we assign 50 WUs for the initial stage and then switch to the hybrid stage after applying one ALS step. We empirically found that using only one ALS step for each WU is sufficient, as it greatly reduces computational time while maintaining convergence behavior. In each experiment, the last parameters of the first branch network (branch \(f\)) is optimized first in the ALS step.

For the training data, we generate 1,000 data instances independently for each of the two input functions to generate 1,000,000 pairs of input functions, which are all possible combinations. For validation data, we generate 4,000 data pairs. In the Adam-only training stage, \(100 \times 100\) data block is used as one batch in each Adam iteration. On the other hand, training with ALS+Adam uses the batch size of \(50 \times 50\) in Adam iterations. All experiments are repeated three times using different seeds for random initialization of parameters and batching for the Adam iterations.

For each model structure and details for problem statement and data generation, refer to 2 and the corresponding subsections. All computations were performed using Google JAX jax2018? on a machine with Intel Xeon Gold 6430 processors and NVIDIA GeForce RTX 4090 with 24 GB memory.

4.1 Reaction-diffusion with variable source and diffusivity↩︎

In this section, we consider a reaction-diffusion equation with variable source and diffusivity as inputs: \[\label{eq:DiffReact} \begin{alignedat}{3} \frac{\partial u}{\partial t} = \frac{\partial}{\partial x} \left(D(x) \frac{\partial u}{\partial x}\right) &+ R(u) + f(x),\quad &&(x,t)\in(0,1)\times(0,1],\\ u(x,0) &= 0,\quad &&x\in(0,1),\\ u(0,t) = u(1,t) &= 0,\quad &&t\in(0,1], \end{alignedat}\tag{40}\] where \(R(u) = 0.01 u^2\). Here, we aim to learn a solution operator that maps the source \(f\) and diffusivity \(D\) to the solution \(u\) via MIONet. The inputs are generated from a Gaussian process (GP) in the interval \([0,1]\) with zero mean and a squared exponential covariance kernel \[\label{eq:SqExpKer} k(x_1,x_2) = \sigma^2 \text{exp}\left(-\frac{{|x_1 - x_2|}^2}{2l^2}\right),\tag{41}\] with a pair of scale and variance parameters \((l,\sigma^2) = (0.2, 1)\) for \(f\) and \((l,\sigma^2) = (0.2, 0.35)\) for \(D\). The input functions \(f\) and \(D\) are discretized at \(33\) equidistant grid points of \([0,1]\), and the output function is evaluated at \(33\times33\) equidistant grid points of \({[0,1]}^2\). The reference solutions are generated by the finite difference method (FDM) using the Crank-Nicolson scheme and the first Taylor approximation for the reaction term. The computational grid is size \(129\times257\) in the space domain and time domain. For the training data, we note that 1,000,000 reference solutions from every pair of \((f,D)\) need to be prepared, which takes about 8 GB of memory.

4.2 Constant coefficient advection with variable source and IBC↩︎

In this section, we solve a 1D constant coefficient advection equation with a source term via MIONet: \[~\label{eq:Advec} \begin{alignedat}{3} \frac{\partial u}{\partial t} + a\frac{\partial u}{\partial x} &= f(x),\quad &&(x,t)\in{(0,1]}^2,\\ u(x,0) &= P(x),\quad &&x\in[0,1],\\ u(0,t) &= Q(t),\quad &&t\in[0,1], \end{alignedat}\tag{42}\] where \(a\) is a fixed constant and \(P(0)=Q(0)\). We aim to learn a solution operator which maps the source \(f\), BC \(P\), and IC \(Q\) to the solution \(u\).

Note that the analytical solution is given as follows: \[~\label{eq:ConstAdvecSol} u^{*}(x,t) = \begin{cases} \frac{1}{a}\left(F(x)-F(x-at)\right) + P(x-at), \quad &x-at \geq 0,\\ \frac{1}{a}\left(F(x)-F(0)\right) + Q(t-\frac{x}{a}), \quad &x-at < 0, \end{cases}\tag{43}\] where \(F(x) = \int_{0}^{x} f(s)\, ds\) is the antiderivative of \(f\). Since this solution may have non-differentiable cusps along the line \(x-at=0\), it is challenging to generate such a solution by minimizing the PI-loss via automatic differentiation in unsupervised learning. Therefore, we introduce an additional condition for input functions — \(f(0)=0\) and \(P'(0) = -\frac{1}{a}Q'(0)\) — to ensure differentiability of the solution. Also, instead of having separate IC and BC as input functions, \(P\) and \(Q\), they can be naturally concatenated into one input function along the domains of IC and BC. Let \(\mathbf{p}=[P(0)\cdots P(1)]\) and \(\mathbf{q}=[Q(0)\cdots Q(1)]\) be the discretizations of \(P\) and \(Q\) along their domains, \([0,1]\times\{0\}\) and \(\{0\}\times[0,1]\), respectively. Since \(P(0)=Q(0)\), we can concatenate \(\mathbf{p}\) and \(\mathbf{q}\) by \(\mathbf{r}=[Q(1)\cdots Q(0)=P(0) \cdots P(1)]\), where \(\mathbf{q}\) is flipped and one of the duplicate values \(P(0)\) or \(Q(0)\) is removed. \(\mathbf{r}\) will be used for the second input argument of the operator.

To generate input functions \(f\), \(P\) and \(Q\), we use the GP with zero mean and a squared exponential covariance kernel in 41 . First, we sample \(g\) from the GP in the interval \([0,1]\) with scale factor \(l=0.2\) and variance \(\sigma^2 = 1\), and \(h\) from the GP in the interval \([-a,1]\) with \(l=0.2\) and \(\sigma^2 = 1\). Then, we set \[~\label{eq:AdvecInputs} \begin{alignedat}{3} f(x) &= g(x) - g(0),\quad &&x \in [0,1],\\ P(x) &= h(x),\quad &&x \in [0,1],\\ Q(t) &= h(-at),\quad &&t \in [0,1]. \end{alignedat}\tag{44}\] The generated \(f\), \(P\) and \(Q\) satisfy \(f(0)=0\), \(P(0)=Q(0)\) and \(P'(0) = -\frac{1}{a}Q'(0)\). In this problem, we choose \(a = 0.5\). The original input functions \(f\), \(P\) and \(Q\) are discretized at \(33\) equidistant grid points of \([0,1]\), and the concatenated input \(\mathbf{r}\) is a vector of length \(65\). The output function is evaluated on \(33\times33\) equidistant grid points of \({[0,1]}^2\). The antiderivative of the exact solution is computed by the composite trapezoidal rule in finer grids of size \(129\) instead of \(33\).

4.3 2D Poisson equation with variable source and Dirichlet BC↩︎

In this section, we solve a 2D Poisson equation on the unit square with Dirichlet BC via MIONet: \[~\label{eq:Poisson} \begin{alignedat}{3} -\nabla\cdot\left(\kappa\nabla u\right) &= f ,\quad &&(x,y)\in\Omega={(0,1)}^2,\\ u &= g,\quad &&(x,y)\in\partial\Omega, \end{alignedat}\tag{45}\] where \(\kappa\equiv 1\) and the MIONet takes two inputs, source \(f\) and Dirichlet BC \(g\), and generates a solution \(u\). The reference solutions are generated by the finite difference method on finer grids of size \(129 \times 129\).

The 2D input source \(f\) is generated from a GP with zero mean and a 2D squared exponential covariance kernel \[~\label{eq:SqExpKer2D} k(x_1,x_2,y_1,y_2) = \sigma^2 \text{exp}\left(-\frac{{|x_1 - x_2|}^2}{2{l_x}^2} -\frac{{|y_1 - y_2|}^2}{2{l_y}^2}\right),\tag{46}\] with scale factors \(l_x=l_y=0.2\) and variance \(\sigma^2 = 0.1\). We take \(33\times33\) equidistant grid points in \({[0,1]}^2\) for the discretization of the input function and as the evaluation points of the output function. The CNN of the branch network consists of three layers with \(3\times 3\), \(2\times 2\), and \(2\times 2\) kernels each with \(2 \times 2\) strides. Since the channel sizes are \([1,16,32,64]\), the output of the CNN is a \(4\times 4\) image of \(64\) channels.

For the BC input function \(g\), we consider a flattened 1D function \(\tilde{g}\) on the interval \([0,4]\), generated by a GP with zero mean and periodic covariance kernel \[~\label{eq:PeriodicKer} k(x_1,x_2) = \sigma^2 \text{exp}\left(-\frac{2}{{l}^2}\sin^{2} {\left(\frac{\pi|x_1 - x_2|}{p}\right)}\right),\tag{47}\] where the scale factor, the period, and the variance are \(l = 0.3\), \(p = 4\), and \(\sigma^2 = 0.1\), respectively, so that \(g(\mathbf{h}(t)) = \tilde{g}(t)\) where \(\mathbf{h}\colon[0,4]\to\partial\Omega\) is the arc length parametrization of \(\partial\Omega\) such that \[~\label{eq:PoisFlatBC} \begin{alignedat}{3} \mathbf{h}(t) &= \begin{cases} (t,0),&t\in[0,1),\\ (1,t-1),&t\in[1,2),\\ (3-t,1),&t\in[2,3),\\ (0,4-t),&t\in[3,4]. \end{cases} \end{alignedat}\tag{48}\] The input function \(\tilde{g}\) is discretized on \(129\) equidistant grid points in \([0,4]\), and the output function is evaluated on \(33\times33\) equidistant grid points in \({[0,1]}^2\).

4.4 Results↩︎

As shown in 5, training with ALS+Adam outperforms classical Adam training in terms of training loss decay and model performance (relative \(L^2\) error for unseen data) in both supervised and unsupervised learning.

6 illustrates the model errors of Adam-only’s and ALS+Adam’s for unseen pairs of test data functions at specific training time. Here, the errors for ALS+Adam training are significantly smaller than Adam-only’s.

We note that errors in both training results tend to form in specific patterns. For the reaction-diffusion example, the error grows as \(t\) increases. In the advection example, the errors tend to form lines parallel to the line \(x - 0.5 t = 0\). For the Poisson’s equation, the errors are concentrated near the boundary.

5 Conclusion↩︎

In this paper, we propose a novel method to improve the training of vanilla MIONets, which generalizes the LSGD method for vanilla DeepONets [16]. By interpreting the general \(L^2\) type of loss in terms of the last layer parameters of branch networks, we can view it as a sum of squared multilinear functions. After that, we apply the ALS step to find the optimal set of last layer parameters that minimizes the loss, where we optimize by solving the corresponding LS problem of each last layer in turn.

Since each LS system is too large to handle directly, we factor the large matrix into small matrices corresponding to each branch and trunk network. Those small matrices constitute the original large matrix using the Kronecker and Khatri-Rao products together with an appropriate tensor permutation matrix. The solution of the LS system can be found in an elementary manner without forming large matrices. Finally, the LSGD method for MIONet alternates between the ALS step optimizing the last layer parameters and the GD step for the hidden layer parameters. The numerical experiments involving a nonlinear PDE with supervised learning and linear PDEs with PI-loss show that our method (ALS+Adam) accelerates training and generates better results with unseen data compared to the conventional Adam training.

References↩︎