Hybrid Least Squares/Gradient Descent Methods for MIONets

# Hybrid Least Squares/Gradient Descent Methods for MIONets
Source: [https://arxiv.org/html/2607.06976](https://arxiv.org/html/2607.06976)
Chang\-Ock LeeDepartment of Mathematical Sciences, KAIST, Daejeon 34141, KOREAMinam MoonDepartment of Mathematics, Korea Military Academy, Seoul 01805, KOREA

###### Abstract

In this paper, we propose an efficient hybrid least squares/gradient descent \(LSGD\) method for MIONets to accelerate training\. This method generalizes the LSGD method for DeepONets\. Since MIONet is the sum of the entrywise product of multiple branch networks and a trunk network, it can be viewed as a multilinear function with respect to the last layer parameters of each branch network\. These sets of parameters can be optimized using the alternating least squares method, where we solve the LS system for a single branch network in turn\. To handle the large\-sized system matrix, we introduce Kronecker and Khatri\-Rao products and tensor permutation matrices to factor the large matrix into small ones\. Our method is compatible with a general type ofL2L^\{2\}loss with regularization terms for the last layer parameters of each branch, where linear operators can be applied to the MIONet output in each loss term\.

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

††footnotetext:Funding:This work was supported by Basic Science Research Program through the National Research Foundation \(NRF\) of Korea funded by the Ministry of Education \[RS2025–25397599\]\.## 1Introduction

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\)\[Raissi2019\]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\[Baydin2018\]\. 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\)\[Lu2021\], Fourier Neural Operator\[Li2020\], Graph Kernel Network\[Li2020neural\], PCA\-based Model Reduction\[Bhattacharya2021\], and Multi\-Wavelet Neural Operator\[Gupta2021\]\.

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\[Lu2022\], Multifidelity DeepONet\[Lu2022multifidelity\], NOMAD\[Seidman2022\], Multiple\-Input Operator Network \(MIONet\)\[Jin2022\], Shift\-DeepONet\[Hadorn2022\], HyperDeepONet\[Lee2023hyper\], and Geom\-DeepONet\[He2024\]\.

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\)\[Jin2022, Theorem 3\.1\]\. 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\[Kingma2017\]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\[Choi2025,Cyr2020\]\. 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 ofL2L^\{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\[Khatri1968\]\) and the usual Kronecker product are used\. Additionally, we introduce tensor permutation matrices\[Rakotonirina2005\]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 formA​X​B\+λ​X=EAXB\+\\lambda X=E, where the coefficient matricesAAandBBare 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\[Choi2025\]\.

This paper is organized as follows\. In[Section2](https://arxiv.org/html/2607.06976#S2), we introduce MIONet and provide the UAT for vanilla MIONets\. We also briefly summarize the LSGD method\[Cyr2020\]and LSGD methods for DeepONets\[Choi2025\]\. In[Section3](https://arxiv.org/html/2607.06976#S3), we formulate the minimization problem from the general type of squaredL2L^\{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[Section4](https://arxiv.org/html/2607.06976#S4), 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\.

## 2Preliminaries

In this section, we introduce MIONet with its universal approximation property\[Jin2022\]and the hybrid LSGD method for DeepONets\[Choi2025\]\. Refer to[Table1](https://arxiv.org/html/2607.06976#S2.T1)for the meaning of symbols and variables used in this paper\.

Table 1:Notation Table\.NotationSpaceDescriptionNNℕ\\mathbb\{N\}Number of branch networksIIℕ\\mathbb\{N\}Number of output nodes of branch and trunk networksJmJ\_\{m\}ℕ\\mathbb\{N\}Number of input nodes of output layer ofmm\-th branchMmM\_\{m\}ℕ\\mathbb\{N\}Number of discretization points formm\-th branch inputd0d\_\{0\}ℕ\\mathbb\{N\}Number of dimension for output function coordinate𝐮\(m\)\\mathbf\{u\}^\{\(m\)\}ℝMm\\mathbb\{R\}^\{M\_\{m\}\}Discretized input function formm\-th branchyyℝd0\\mathbb\{R\}^\{d\_\{0\}\}Coordinate for output function𝐛m​\(𝐮\(m\)\)\\mathbf\{b\}\_\{m\}\(\\mathbf\{u\}^\{\(m\)\}\)ℝI\\mathbb\{R\}^\{I\}Output ofmm\-th branch network of vanilla MIONet𝐭​\(y\)\\mathbf\{t\}\(y\)ℝI\\mathbb\{R\}^\{I\}Output of trunk network of vanilla MIONet𝐛~m​\(𝐮\(m\)\)\\tilde\{\\mathbf\{b\}\}\_\{m\}\(\\mathbf\{u\}^\{\(m\)\}\)ℝJm\\mathbb\{R\}^\{J\_\{m\}\}Output of layer before last layer ofmm\-th branchCmC\_\{m\}ℝI×Jm\\mathbb\{R\}^\{I\\times J\_\{m\}\}Last layer parameter matrix ofmm\-th branchθmB\\theta^\{B\}\_\{m\}\-Hidden layer parameters ofmm\-th branchθT\\theta^\{T\}\-Parameters of trunk networkθmL\\theta^\{L\}\_\{m\}ℝI​Jm\\mathbb\{R\}^\{IJ\_\{m\}\}Last layer parameters ofmm\-th branch,θL=vec​\(CmT\)\\theta^\{L\}=\\text\{vec\}\(C\_\{m\}^\{T\}\)\[d1,…,dN\]D1,…,DN\{\[d\_\{1\},\\dots,d\_\{N\}\]\}\_\{D\_\{1\},\\dots,D\_\{N\}\}ℕ\\mathbb\{N\}Big\-endian order for the entries ofthe rankNNtensor in[3\.2](https://arxiv.org/html/2607.06976#S3.E2)KKℕ\\mathbb\{N\}Number of loss terms except regularization termsϵk\\epsilon\_\{k\}ℝ\>0\\mathbb\{R\}^\{\>0\}Weight forkk\-th loss termλm\\lambda\_\{m\}ℝ\>0\\mathbb\{R\}^\{\>0\}Weight forL2L^\{2\}regularization term ofθmL\\theta^\{L\}\_\{m\}ℒk\\mathcal\{L\}\_\{k\}ℒ​\(C​\(ℝ\),C​\(ℝ\)\)\\mathcal\{L\}\(C\(\\mathbb\{R\}\),C\(\\mathbb\{R\}\)\)Linear operator forkk\-th loss termDkD\_\{k\}ℕ\\mathbb\{N\}Number of data pairs forkk\-th loss termχk\\chi\_\{k\}\(\(∏m=1NℝMm\)×ℝd0\)Dk\{\\left\(\(\\prod\_\{m=1\}^\{N\}\\mathbb\{R\}^\{M\_\{m\}\}\)\\times\\mathbb\{R\}^\{d\_\{0\}\}\\right\)\}^\{D\_\{k\}\}Data\(𝐮\(1\),…,𝐮\(N\),y\)\(\\mathbf\{u\}^\{\(1\)\},\\dots,\\mathbf\{u\}^\{\(N\)\},y\)collection forkk\-th loss term𝒜k\\mathcal\{A\}\_\{k\}ℝDk\\mathbb\{R\}^\{D\_\{k\}\}Long vector whosedkd\_\{k\}\-th entry is[3\.4](https://arxiv.org/html/2607.06976#S3.E4)Ak,mA\_\{k,m\}ℝDk×I​Jm\\mathbb\{R\}^\{D\_\{k\}\\times IJ\_\{m\}\}System matrix forkk\-th loss term of LS problem inθmL\\theta^\{L\}\_\{m\}fkf\_\{k\}ℝDk\\mathbb\{R\}^\{D\_\{k\}\}Data forkk\-th loss term of LS problemPmP\_\{m\}ℕ\\mathbb\{N\}Number of input functions formm\-th branchQkQ\_\{k\}ℕ\\mathbb\{N\}Number of coordinate points ofkk\-th loss termβm\\beta\_\{m\}\(ℝMm\)Pm\{\(\\mathbb\{R\}^\{M\_\{m\}\}\)\}^\{P\_\{m\}\}Set of discretized input functions formm\-th branchτk\\tau\_\{k\}\(ℝd0\)Qk\{\(\\mathbb\{R\}^\{d\_\{0\}\}\)\}^\{Q\_\{k\}\}Set of coordinate points ofkk\-th loss termKD,σK\_\{D,\\sigma\}ℝD1​⋯​Dm×D1​⋯​Dm\\mathbb\{R\}^\{\{D\_\{1\}\}\\cdots\{D\_\{m\}\}\\times\{D\_\{1\}\}\\cdots\{D\_\{m\}\}\}Tensor permutation matrix in[Definition3\.2](https://arxiv.org/html/2607.06976#S3.Thmtheorem2)BmB\_\{m\}ℝPm×Jm\\mathbb\{R\}^\{P\_\{m\}\\times J\_\{m\}\}mm\-th branch pre\-output matrix,\(b~j​\(𝐮p\(m\)\)\)\\left\(\\tilde\{b\}\_\{j\}\(\\mathbf\{u\}^\{\(m\)\}\_\{p\}\)\\right\)TkT\_\{k\}ℝQk×I\\mathbb\{R\}^\{Q\_\{k\}\\times I\}kk\-th trunk output matrix withℒk\\mathcal\{L\}\_\{k\},\(ℒk​\[ti\]​\(yq\)\)\\left\(\\mathcal\{L\}\_\{k\}\[t\_\{i\}\]\(y\_\{q\}\)\\right\)FkF\_\{k\}ℝP1×⋯×PN×Qk\\mathbb\{R\}^\{P\_\{1\}\\times\\cdots\\times P\_\{N\}\\times Q\_\{k\}\}Data forkk\-th loss term in tensor form### 2\.1Multiple\-Input Operator Network \(MIONet\)

Jin et al\.\[Jin2022\]proved an UAT for multiple\-input operators on the product of Banach spaces of functions with Schauder bases\. We refer to\[Fabian2011,Hu2025,Semadeni2006\]for more details of the Schauder basis and its canonical projections\. The theorem states that a continuous multiple\-input operatorGGon the product of compact subsets of Banach spacesXmX\_\{m\}with Schauder bases can be approximated by the form

⟨\(𝐛1∘ϕM11\)⏟branch1⊙⋯⊙\(𝐛N∘ϕMNN\)⏟branchN,𝐭⏟trunk⟩~\\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\(2\.1\)with continuous vector functions𝐛m∈C​\(ℝMm,ℝI\)\\mathbf\{b\}\_\{m\}\\in C\(\\mathbb\{R\}^\{M\_\{m\}\},\\mathbb\{R\}^\{I\}\)and𝐭∈YI\\mathbf\{t\}\\in Y^\{I\}for sufficiently large positive integersIIandMmM\_\{m\}, whereϕMmm:Xm→ℝMm\\phi^\{m\}\_\{M\_\{m\}\}\\colon X\_\{m\}\\to\\mathbb\{R\}^\{M\_\{m\}\}extracts the firstMmM\_\{m\}coefficients of the Schauder basis representation,YYis the target Banach space, and⊙\\odotdenotes the Hadamard \(entrywise\) product\.

Here, from the inner product structure \([2\.1](https://arxiv.org/html/2607.06976#S2.E1)\) of vector functions𝐛m\\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\[Lu2021\], we want to use discretized function values ofu\(m\)u^\{\(m\)\}as an input for𝐛m\\mathbf\{b\}\_\{m\}, whereϕnm\\phi\_\{n\}^\{m\}extracts function values at certain points\. Note that for the sequence of distinct points\{ti\}i=1∞\{\\\{t\_\{i\}\\\}\}\_\{i=1\}^\{\\infty\}in\[0,1\]\[0,1\]witht1=0t\_\{1\}=0andt2=1t\_\{2\}=1, which is dense in\[0,1\]\[0,1\], there exists a Schauder basis\{ei\}i=1∞\{\\\{e\_\{i\}\\\}\}\_\{i=1\}^\{\\infty\}ofC​\(\[0,1\]\)C\(\{\[0,1\]\}\)\(called the Faber\-Schauder basis\) wheree1​\(t\)=1e\_\{1\}\(t\)=1andene\_\{n\}is chosen as a piecewise linear function withen​\(tn\)=1e\_\{n\}\(t\_\{n\}\)=1, where the set\{e1,…,en\}\\\{e\_\{1\},\\dots,e\_\{n\}\\\}forms a basis for the space of all piecewise linear functions with node points\{ti\}i=1n\{\\\{t\_\{i\}\\\}\}\_\{i=1\}^\{n\}\. We refer to\[Fabian2011,Semadeni2006\]for detailed Faber\-Schauder basis construction ofC​\(\[0,1\]\)C\(\{\[0,1\]\}\)and\[Glenn2019\]for the multivariate extension toC​\(\[0,1\]d\)C\(\{\[0,1\]\}^\{d\}\)\. Furthermore, without loss of generality, MIONet can have a vanilla structure similar to DeepONets in\[Choi2025,Son2025\], 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[Figure1](https://arxiv.org/html/2607.06976#S2.F1)for the structure of vanilla MIONet\. Note that MIONet withN=1N=1becomes DeepONet with the same structure and the universal approximation property discussed in\[Choi2025\]\.

![Refer to caption](https://arxiv.org/html/2607.06976v1/pics/2/21.png)Figure 1:Structure of vanilla MIONet with fully connected layers\.We now address the UAT for vanilla MIONets, where the input functions are discretized function values\. This theorem is a modification of\[Hu2025, Theorem 2\]and\[Lu2021, Theorem 2\]\.

###### Theorem 2\.1\(Universal Approximation Theorem for vanilla MIONets\)\.

LetX~m\\tilde\{X\}\_\{m\}be Banach spaces,K~m⊂X~m\\tilde\{K\}\_\{m\}\\subset\\tilde\{X\}\_\{m\},K0⊂ℝd0K\_\{0\}\\subset\\mathbb\{R\}^\{d\_\{0\}\}be compact subsets ofX~m\\tilde\{X\}\_\{m\}andℝd0\\mathbb\{R\}^\{d\_\{0\}\}, respectively,VmV\_\{m\}be a compact set inC​\(K~m\)C\(\\tilde\{K\}\_\{m\}\), andGGbe a continuous operator which mapsV1×⋯×VNV\_\{1\}\\times\\cdots\\times V\_\{N\}intoC​\(K0\)C\(K\_\{0\}\)\. Then, for anyϵ\>0\\epsilon\>0, there exist positive integersI,MmI,M\_\{m\}, continuous vector functions𝐛m∈C​\(ℝMm,ℝI\)\\mathbf\{b\}\_\{m\}\\in C\(\\mathbb\{R\}^\{M\_\{m\}\},\\mathbb\{R\}^\{I\}\)and𝐭∈C​\(ℝd0,ℝI\)\\mathbf\{t\}\\in C\(\\mathbb\{R\}^\{d\_\{0\}\},\\mathbb\{R\}^\{I\}\), and pointsxi\(m\)∈K~mx^\{\(m\)\}\_\{i\}\\in\\tilde\{K\}\_\{m\}withm=1,…,Nm=1,\\dots,Nandi=1,…,Mmi=1,\\dots,M\_\{m\}such that

\|G​\(u\(1\),…,u\(N\)\)​\(y\)−⟨𝐛1​\(𝐮\(1\)\)⏟branch1⊙⋯⊙𝐛N​\(𝐮\(N\)\)⏟branchN,𝐭​\(y\)⏟trunk⟩\|<ϵ~\{\\Biggl\|G\(u^\{\(1\)\},\\dots,u^\{\(N\)\}\)\(y\)\-\\Big\\langle\\underbrace\{\\mathbf\{b\}\_\{1\}\(\\mathbf\{u\}^\{\(1\)\}\)\}\_\{\\textrm\{branch $1$\}\}\\odot\\cdots\\odot\\underbrace\{\\mathbf\{b\}\_\{N\}\(\\mathbf\{u\}^\{\(N\)\}\)\}\_\{\\textrm\{branch $N$\}\},\\underbrace\{\\mathbf\{t\}\(y\)\}\_\{\\textrm\{trunk\}\}\\Big\\rangle\\Biggr\|\}<\\epsilon\(2\.2\)holds for allu\(m\)∈Vmu^\{\(m\)\}\\in V\_\{m\}andy∈K0y\\in K\_\{0\}, where⟨⋅,⋅⟩\\langle\\cdot,\\cdot\\rangledenotes the inner product inℝI\\mathbb\{R\}^\{I\}and𝐮\(m\)=\[u\(m\)​\(x1\(m\)\),…,u\(m\)​\(xMm\(m\)\)\]T∈ℝMm\\mathbf\{u\}^\{\(m\)\}=\{\\left\[u^\{\(m\)\}\(x\_\{1\}^\{\(m\)\}\),\\dots,u^\{\(m\)\}\(x\_\{M\_\{m\}\}^\{\(m\)\}\)\\right\]\}^\{T\}\\in\{\\mathbb\{R\}\}^\{M\_\{m\}\}is a discretized input functionu\(m\)u^\{\(m\)\}at\{xi\(m\)\}i=1Mm\{\\\{\{x\_\{i\}^\{\(m\)\}\}\\\}\}\_\{i=1\}^\{M\_\{m\}\}\. Moreover, the functions𝐛m\\mathbf\{b\}\_\{m\}and𝐭\\mathbf\{t\}can be chosen as diverse classes of neural networks satisfying the classical UAT of functions, where the branch networks𝐛m​\(𝐮\(m\)\)=Cm∘𝐛~m​\(𝐮\(m\)\)\\mathbf\{b\}\_\{m\}\(\\mathbf\{u\}^\{\(m\)\}\)=C\_\{m\}\\circ\\tilde\{\\mathbf\{b\}\}\_\{m\}\(\\mathbf\{u\}^\{\(m\)\}\)are decomposed into a hidden layer part𝐛~m:ℝMm→ℝJm\\tilde\{\\mathbf\{b\}\}\_\{m\}\\colon\\mathbb\{R\}^\{M\_\{m\}\}\\to\\mathbb\{R\}^\{J\_\{m\}\}and a linear last layer partCm:ℝJm→ℝIC\_\{m\}\\colon\\mathbb\{R\}^\{J\_\{m\}\}\\to\\mathbb\{R\}^\{I\}with a matrix parametrizationCm=\(ci​j\(m\)\)∈ℝI×JmC\_\{m\}=\(c\_\{ij\}^\{\(m\)\}\)\\in\\mathbb\{R\}^\{I\\times J\_\{m\}\}\.

The MIONet approximatingGGin \([2\.2](https://arxiv.org/html/2607.06976#S2.E2)\) can be expressed as a nested form of summations and products:

G​\(u\(1\),…,u\(N\)\)​\(y\)\\displaystyle G\(u^\{\(1\)\},\\dots,u^\{\(N\)\}\)\(y\)≈⟨𝐛1​\(𝐮\(1\)\)⊙⋯⊙𝐛N​\(𝐮\(N\)\),𝐭​\(y\)⟩\\displaystyle\\approx\\left\\langle\\mathbf\{b\}\_\{1\}\(\\mathbf\{u\}^\{\(1\)\}\)\\odot\\cdots\\odot\\mathbf\{b\}\_\{N\}\(\\mathbf\{u\}^\{\(N\)\}\),\\mathbf\{t\}\(y\)\\right\\rangle\(2\.3\)=∑i=1I\(∏m=1Nbi\(m\)​\(𝐮\(m\)\)\)​ti​\(y\)\\displaystyle=\\sum\_\{i=1\}^\{I\}\\left\(\\prod\_\{m=1\}^\{N\}b\_\{i\}^\{\(m\)\}\(\\mathbf\{u\}^\{\(m\)\}\)\\right\)t\_\{i\}\(y\)=∑i=1I\[∏m=1N\(∑j=1Jmci​j\(m\)​b~j\(m\)​\(𝐮\(m\)\)\)\]​ti​\(y\),\\displaystyle=\\sum\_\{i=1\}^\{I\}\\left\[\\prod\_\{m=1\}^\{N\}\\left\(\\sum\_\{j=1\}^\{J\_\{m\}\}c\_\{ij\}^\{\(m\)\}\\tilde\{b\}\_\{j\}^\{\(m\)\}\(\\mathbf\{u\}^\{\(m\)\}\)\\right\)\\right\]t\_\{i\}\(y\),whereb~j\(m\)\\tilde\{b\}\_\{j\}^\{\(m\)\}is thejj\-th component of the vector function𝐛~m\\tilde\{\\mathbf\{b\}\}\_\{m\}, andbi\(m\)b\_\{i\}^\{\(m\)\}andtit\_\{i\}are theii\-th component of the vector function𝐛m\{\\mathbf\{b\}\}\_\{m\}and𝐭m\{\\mathbf\{t\}\}\_\{m\}, respectively\.

### 2\.2Hybrid Least Squares/Gradient Descent Method for DeepONets

Cyr et al\.\[Cyr2020\]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\[Choi2025\], 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 \([2\.3](https://arxiv.org/html/2607.06976#S2.E3)\) withN=1N=1\. We use a general type ofl2l\_\{2\}loss, which consists of the sum of squaredl2l\_\{2\}error terms and a regularization term for the last layer parameters of the branch network:

∑k=1Kϵk​‖ℒk​\[G​\(⋅\)\]​\(u,y\)−ℒk​\[⟨𝐛​\(⋅;θB,θL\),𝐭​\(⋅;θT\)⟩\]​\(𝐮,y\)‖l2​\(χk\)2\+λ​‖θL‖22,~\\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\},\(2\.4\)whereχk=\{\(u^dk,y^dk\)\}dk=1Dk\\chi\_\{k\}=\{\\\{\(\\hat\{u\}\_\{d\_\{k\}\},\\hat\{y\}\_\{d\_\{k\}\}\)\\\}\}\_\{d\_\{k\}=1\}^\{D\_\{k\}\}is a finite collection of data pairs\(u,y\)\(u,y\),ℒk\\mathcal\{L\}\_\{k\}is a linear operator between real\-valued functions,fk=ℒk​\[G​\(⋅\)\]​\(u^dk,y^dk\)∈ℝDkf\_\{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 thekk\-th error term,ϵk\>0\\epsilon\_\{k\}\>0andλ≥0\\lambda\\geq 0are the weights for each error term and the regularization term, andθB\\theta^\{B\},θT\\theta^\{T\},θL\\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θL\\theta^\{L\}as the row\-wise vectorization of the \(branch\) last layer parameter matrixC∈ℝI×JC\\in\\mathbb\{R\}^\{I\\times J\}:

θL=vec​\(CT\)=\[c11,…,c1​J,c21,…,c2​J,…,cI​1,…,cI​J\]T∈ℝI​J\.~\\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\}\.Then, \([2\.4](https://arxiv.org/html/2607.06976#S2.E4)\) can be reformulated as an LS problem

minθL​∑k=1Kϵk​‖fk−Ak​θL‖22\+λ​‖θL‖22,~\\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\}\},\(2\.5\)whereAk∈ℝDk×I​JA\_\{k\}\\in\{\\mathbb\{R\}\}^\{D\_\{k\}\\times IJ\}is the matrix whose\(dk,J​\(i−1\)\+j\)\(d\_\{k\},J\(i\-1\)\+j\)entry isℒk​\[b~j​ti\]​\(𝐮^dk,y^dk\)\\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χk\\chi\_\{k\}and the linear operatorℒk\\mathcal\{L\}\_\{k\}such that

χk=β×τk\\displaystyle\\chi\_\{k\}=\\beta\\times\\tau\_\{k\}=\{up\}p=1P×\{yqk\}qk=1Qk,\\displaystyle=\{\\left\\\{u\_\{p\}\\right\\\}\}\_\{p=1\}^\{P\}\\times\{\\left\\\{y\_\{q\_\{k\}\}\\right\\\}\}\_\{q\_\{k\}=1\}^\{Q\_\{k\}\},\(2\.6a\)ℒk​\[b~j​ti\]​\(𝐮,y\)\\displaystyle\\mathcal\{L\}\_\{k\}\\left\[\\tilde\{b\}\_\{j\}t\_\{i\}\\right\]\(\\mathbf\{u\},y\)=b~j​\(𝐮\)​ℒk​\[ti\]​\(y\),\\displaystyle=\\tilde\{b\}\_\{j\}\(\\mathbf\{u\}\)\\mathcal\{L\}\_\{k\}\\left\[t\_\{i\}\\right\]\(y\),\(2\.6b\)we have a main theorem for the LS step of the LSGD method for DeepONets\[Choi2025, Theorem 3\]\. In the theorem, the large matrixAkA\_\{k\}in[2\.5](https://arxiv.org/html/2607.06976#S2.E5)can be factored into two small component matrices, each from the branch and trunk networks, using the Kronecker product, where a commutation matrix\[MacRae1974,Magnus1979\]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 typeA​X​B\+λ​X=EAXB\+\\lambda X=Ewith symmetric positive semi\-definite matricesAAandBB\. Such matrix equation can be solved by using the spectral decompositions\[Choi2025, Proposition 6\]\.

###### Proposition 2\.2\.

LetA∈ℝR×RA\\in\\mathbb\{R\}^\{R\\times R\}andB∈ℝS×SB\\in\\mathbb\{R\}^\{S\\times S\}be symmetric positive semi\-definite matrices,E∈ℝR×SE\\in\\mathbb\{R\}^\{R\\times S\}be any matrix, andλ\\lambdabe a nonnegative real number\. Then, the solution of the matrix equation

A​X​B\+λ​X=E~AXB\+\\lambda X=E\(2\.7\)is given as

X=QA​\[\(𝐝A​𝐝BT\+λ​𝟏R×S\)⊙−1⊙\(QAT​E​QB\)\]​QBT,~X=Q\_\{A\}\\left\[\{\(\\mathbf\{d\}\_\{A\}\\mathbf\{d\}\_\{B\}^\{T\}\+\\lambda\\mathbf\{1\}\_\{R\\times S\}\)\}^\{\\odot\-1\}\\odot\(Q\_\{A\}^\{T\}EQ\_\{B\}\)\\right\]Q\_\{B\}^\{T\},\(2\.8\)whereA=QA​DA​QATA=Q\_\{A\}D\_\{A\}Q\_\{A\}^\{T\}andB=QB​DB​QBTB=Q\_\{B\}D\_\{B\}Q\_\{B\}^\{T\}are the spectral decompositions with orthogonal matricesQAQ\_\{A\},QBQ\_\{B\}and diagonal matricesDA=diag​\(𝐝A\)D\_\{A\}=\\text\{diag\}\(\\mathbf\{d\}\_\{A\}\),DB=diag​\(𝐝B\)D\_\{B\}=\\text\{diag\}\(\\mathbf\{d\}\_\{B\}\)when𝐝A∈ℝR\\mathbf\{d\}\_\{A\}\\in\\mathbb\{R\}^\{R\},𝐝B∈ℝS\\mathbf\{d\}\_\{B\}\\in\\mathbb\{R\}^\{S\}\. Here,⊙−1\{\\odot^\{\-1\}\}denotes the entrywise inverse, and𝟏R×S\\mathbf\{1\}\_\{R\\times S\}represents theR×SR\\times Smatrix with all entries equal to one\.

Therefore, the last layer parameters minimizing \([2\.5](https://arxiv.org/html/2607.06976#S2.E5)\) 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\[Choi2025, Algorithms 1 and 2\]for the workflow of the LSGD method for DeepONets and its practical modification with Adam optimizer\[Kingma2017\], LS\+Adam, respectively\.

## 3Hybrid Least Squares/Gradient Descent Methods for MIONets

We consider the following loss of the vanilla MIONet, which consists of the sums of squaredl2l\_\{2\}error terms indexed bykkand regularization terms for theNNlast layer parameters:

∑k=1Kϵk∥ℒk\[G\(⋅,…,⋅\)\]\(u\(1\),…,u\(N\),y\)−ℒk​\[⟨𝐛1​\(⋅;θ1B,θ1L\)⊙⋯⊙𝐛N​\(⋅;θNB,θNL\),𝐭​\(⋅;θT\)⟩\]​\(𝐮\(1\),…,𝐮\(N\),y\)∥l2​\(χk\)2\+∑m=1Nλm​‖θmL‖22,~\\begin\{aligned\} \\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\{aligned\}\(3\.1\)whereχk\\chi\_\{k\}is a finite collection of data pairs\{\(u^dk\(1\),…,u^dk\(N\),y^dk\)\}dk=1Dk\{\\left\\\{\(\\hat\{u\}\_\{d\_\{k\}\}^\{\(1\)\},\\dots,\\hat\{u\}\_\{d\_\{k\}\}^\{\(N\)\},\\hat\{y\}\_\{d\_\{k\}\}\)\\right\\\}\}\_\{d\_\{k\}=1\}^\{D\_\{k\}\},ℒk\\mathcal\{L\}\_\{k\}is a linear operator between real\-valued functions for thekk\-th error term,fkf\_\{k\}is the data for thekk\-th error term given asfk=\(ℒk​\[G​\(⋅,…,⋅\)\]​\(u^dk\(1\),…,u^dk\(N\),y^dk\)\)∈ℝDkf\_\{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\}\},ϵk\>0\\epsilon\_\{k\}\>0andλm≥0\\lambda\_\{m\}\\geq 0are weights for error terms and the regularization terms, andθmB\\theta\_\{m\}^\{B\},θT\\theta^\{T\},θmL\\theta\_\{m\}^\{L\}denote the parameters for themm\-th branch network except the last layer, the trunk network, and the last layer of themm\-th branch, respectively, form=1,…,Nm=1,\\dots,N\.

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

###### Definition 3\.1\.

Let𝒯∈ℝD1×⋯×DN\\mathcal\{T\}\\in\{\\mathbb\{R\}\}^\{D\_\{1\}\\times\\cdots\\times D\_\{N\}\}be a rankNNtensor of sizeD1×⋯×DND\_\{1\}\\times\\cdots\\times D\_\{N\}\. The lexicographic order of the entry of𝒯\\mathcal\{T\}maps the coordinate index\(d1,…,dN\)\(d\_\{1\},\\dots,d\_\{N\}\)to the corresponding big\-endian lexicographic order\[d1,…,dN\]D1,…,DN\{\[d\_\{1\},\\dots,d\_\{N\}\]\}\_\{D\_\{1\},\\dots,D\_\{N\}\}such that

\[d1,…,dN\]D1,…,DN≔∑m=1N−1\[\(∏l=m\+1NDl\)​\(dm−1\)\]\+dN\.\{\[d\_\{1\},\\dots,d\_\{N\}\]\}\_\{D\_\{1\},\\dots,D\_\{N\}\}\\coloneq\\sum\_\{m=1\}^\{N\-1\}\\left\[\\left\(\\prod\_\{l=m\+1\}^\{N\}D\_\{l\}\\right\)\(d\_\{m\}\-1\)\\right\]\+d\_\{N\}\.\(3\.2\)

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

LetCm∈ℝI×JmC\_\{m\}\\in\\mathbb\{R\}^\{I\\times J\_\{m\}\}be a parameter matrix of themm\-th last layer, whereθmL\\theta\_\{m\}^\{L\}is the row\-wise vectorization ofCmC\_\{m\}\. The entries ofCmC\_\{m\}are sorted lexicographically in the order[3\.2](https://arxiv.org/html/2607.06976#S3.E2), i\.e\.,

θmL=vec​\(CmT\)=\[c11\(m\),…,c1​Jm\(m\),c21\(m\),…,c2​Jm\(m\),…,cI​1\(m\),…,cI​Jm\(m\)\]T∈ℝI​Jm\.~\\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\}\}\.
Here, \([3\.1](https://arxiv.org/html/2607.06976#S3.E1)\) is reformulated as the following minimization problem in terms of theNNlast layer parameters\{θ1L,…,θNL\}\\\{\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\\\}:

min\{θ1L,…,θNL\}​∑k=1Kϵk​‖fk−𝒜k​\(θ1L,…,θNL\)‖l2​\(χk\)2\+∑m=1Nλm​‖θmL‖22,\\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\},\(3\.3\)where𝒜k​\(θ1L,…,θNL\)∈ℝDk\\mathcal\{A\}\_\{k\}\(\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\)\\in\{\\mathbb\{R\}\}^\{D\_\{k\}\}is the long vector whosedkd\_\{k\}entry is

ℒk​\[∑i=1Ibi\(1\)​⋯​bi\(N\)​ti\]​\(𝐮^dk\(1\),…,𝐮^dk\(N\),y^dk\),\\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\}\}\),\(3\.4\)and it is the result of anNN\-linear function applied to the last layer parametersθ1L,…,θNL\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\.

Although the minimization problem \([3\.3](https://arxiv.org/html/2607.06976#S3.E3)\) is not an LS problem, we can form an LS problem forθnL\\theta\_\{n\}^\{L\}by fixing all last layer parameters except thenn\-th parameterθnL\\theta\_\{n\}^\{L\}\. By solving the LS problems generated for each last layer parameter in turn, we can optimize theNNlast layer parameters\. The details will be discussed after formulating the LS problem for the single last layer parameter\.

Note that𝒜k​\(θ1L,…,θNL\)\\mathcal\{A\}\_\{k\}\(\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\)is linear with respect toθnL\\theta\_\{n\}^\{L\}for eachnn, we can write

𝒜k​\(θ1L,…,θNL\)=Ak,n​θnL~\\mathcal\{A\}\_\{k\}\(\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\)=A\_\{k,n\}\\theta\_\{n\}^\{L\}for eachkkandnn\.

Therefore, the LS step minimizing \([3\.3](https://arxiv.org/html/2607.06976#S3.E3)\) in terms ofθnL\\theta\_\{n\}^\{L\}is

minθnL​∑k=1Kϵk​‖fk−Ak,n​θnL‖22\+λn​‖θnL‖22~\\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\}\(3\.5\)whereAk,n∈ℝDk×I​JnA\_\{k,n\}\\in\{\\mathbb\{R\}\}^\{D\_\{k\}\\times IJ\_\{n\}\}is the matrix whose\(dk,\[i,j\]I,Jn\)\(d\_\{k\},\{\[i,j\]\}\_\{I,J\_\{n\}\}\)entry is

ℒk​\[bi\(1\)​⋯​bi\(n−1\)​b~j\(n\)​bi\(n\+1\)​⋯​bi\(N\)​ti\]​\(𝐮^dk\(1\),…,𝐮^dk\(N\),y^dk\)\.~\\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\}\}\)\.
Here, the LS system \([3\.5](https://arxiv.org/html/2607.06976#S3.E5)\) cannot be handled directly ifIIandJnJ\_\{n\}are not small unlessDkD\_\{k\}is small, but it is not desirable to use smallDkD\_\{k\}in MIONet training\. We provide the data and operator conditions for MIONet, which correspond to the conditions for DeepONet in\[Choi2025, Eqs\. \(3\.4\) and \(3\.5\)\]\.

Suppose that the data collectionχk\\chi\_\{k\}can be expressed as a Cartesian product ofN\+1N\+1input data as follows:

χk=β1×⋯×βN×τk,~\\chi\_\{k\}=\\beta\_\{1\}\\times\\cdots\\times\\beta\_\{N\}\\times\\tau\_\{k\},\(3\.6\)whereβm=\{up\(m\)\}p=1Pm\\beta\_\{m\}=\{\\left\\\{\{u\}\_\{p\}^\{\(m\)\}\\right\\\}\}\_\{p=1\}^\{P\_\{m\}\}is the set of input functions for themm\-th coordinate andτk=\{yqk\}qk=1Qk\\tau\_\{k\}=\{\\left\\\{\{y\}\_\{q\_\{k\}\}\\right\\\}\}\_\{q\_\{k\}=1\}^\{Q\_\{k\}\}is the set of points of the discretized domain forℒk​\[G​\(⋅,…,⋅\)\]\\mathcal\{L\}\_\{k\}\[G\(\\cdot,\\dots,\\cdot\)\]withDk=P1​⋯​PN​QkD\_\{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ℒk​\[G​\(⋅,…,⋅\)\]\\mathcal\{L\}\_\{k\}\[G\(\\cdot,\\dots,\\cdot\)\]is used for all tuples of input functions\(up1\(1\),…,upN\(N\)\)\(\{u\}\_\{p\_\{1\}\}^\{\(1\)\},\\dots,\{u\}\_\{p\_\{N\}\}^\{\(N\)\}\)for each error term\.

We also assume that the linear operatorℒk\\mathcal\{L\}\_\{k\}satisfies

ℒk​\[b~j1\(1\)​⋯​b~jN\(N\)​ti\]​\(𝐮\(1\),…,𝐮\(N\),y\)=\(∏m=1Nb~jm\(m\)​\(𝐮\(m\)\)\)​ℒk​\[ti\]​\(y\)~\\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\)\(3\.7\)for alli,jm,ki,j\_\{m\},k\. This implies the linear operatorℒk\\mathcal\{L\}\_\{k\}is independent of input functions and acts on each trunk componenttit\_\{i\}\.

Let∗\\astdenote the Khatri\-Rao product \(column\-wise Kronecker product\)\[Khatri1968\]\. Note that the Khatri\-Rao product of two matricesX∈ℝR1×SX\\in\\mathbb\{R\}^\{R\_\{1\}\\times S\}andY∈ℝR2×SY\\in\\mathbb\{R\}^\{R\_\{2\}\\times S\}is anR1​R2×SR\_\{1\}R\_\{2\}\\times Smatrix whose entry is represented as

\(X∗Y\)r​s=Xr1​s​Yr2​s,~\{\(X\\ast Y\)\}\_\{rs\}=\{X\}\_\{r\_\{1\}s\}\{Y\}\_\{r\_\{2\}s\},\(3\.8\)wherer=\[r1,r2\]R1,R2r=\{\[r\_\{1\},r\_\{2\}\]\}\_\{R\_\{1\},R\_\{2\}\}\.

Under the conditions \([3\.6](https://arxiv.org/html/2607.06976#S3.E6)\) and \([3\.7](https://arxiv.org/html/2607.06976#S3.E7)\), we show the large matrixAk,nA\_\{k,n\}in \([3\.5](https://arxiv.org/html/2607.06976#S3.E5)\) can be factored into the product ofN\+1N\+1smaller matrices with some permutations\. For convenience, we define a sequence with thejj\-th term removed from the sequence\{xi\}i=1N\{\\left\\\{x\_\{i\}\\right\\\}\}\_\{i=1\}^\{N\}and denote it as\(x1,…,x^j,…,xN\)\(x\_\{1\},\\dots,\\hat\{x\}\_\{j\},\\dots,x\_\{N\}\)\. Also, for the generic associative binary operator∘\\circ, we denote∘i=1i≠j𝑁​xi\\overset\{N\}\{\\underset\{\\begin\{subarray\}\{c\}i=1\\\\ i\\neq j\\end\{subarray\}\}\{\\circ\}\}x\_\{i\}as the repeated application of the operation to the sequence\{xi\}i=1N\{\\left\\\{x\_\{i\}\\right\\\}\}\_\{i=1\}^\{N\}withxjx\_\{j\}omitted\.

We define the tensor permutation matrix introduced in\[Rakotonirina2005\]in terms of the lexicographic order of the entry of a rankNNtensor\.

###### Definition 3\.2\.

Let𝒯∈ℝD1×⋯×DN\\mathcal\{T\}\\in\{\\mathbb\{R\}\}^\{D\_\{1\}\\times\\cdots\\times D\_\{N\}\}be a rankNNtensor of sizeD1×⋯×DND\_\{1\}\\times\\cdots\\times D\_\{N\}and𝐭∈ℝD1​⋯​DN\\mathbf\{t\}\\in\\mathbb\{R\}^\{D\_\{1\}\\cdots D\_\{N\}\}be the lexicographically vectorized tensor𝒯\\mathcal\{T\}whose\[d1,…,dN\]D1,…,DN\{\[d\_\{1\},\\dots,d\_\{N\}\]\}\_\{D\_\{1\},\\dots,D\_\{N\}\}entry is𝒯d1,…,dN\\mathcal\{T\}\_\{d\_\{1\},\\dots,d\_\{N\}\}\. Letσ∈SN\\sigma\\in S\_\{N\}be a permutation, whereSNS\_\{N\}is the set of all permutations of\{1,…,N\}\\\{1,\\dots,N\\\}\. Suppose𝒯~∈ℝDσ​\(1\)×⋯×Dσ​\(N\)\\tilde\{\\mathcal\{T\}\}\\in\{\\mathbb\{R\}\}^\{D\_\{\\sigma\(1\)\}\\times\\cdots\\times D\_\{\\sigma\(N\)\}\}be a transposed tensor satisfying𝒯d1,…,dN=𝒯~dσ​\(1\),…,dσ​\(N\)\\mathcal\{T\}\_\{d\_\{1\},\\dots,d\_\{N\}\}=\\tilde\{\\mathcal\{T\}\}\_\{d\_\{\\sigma\(1\)\},\\dots,d\_\{\\sigma\(N\)\}\}with the lexicographic vectorization𝐭~\\tilde\{\\mathbf\{t\}\}forDσ​\(1\)×⋯×Dσ​\(N\)D\_\{\\sigma\(1\)\}\\times\\cdots\\times D\_\{\\sigma\(N\)\}tensors\. If a permutation matrixKD,σ∈ℝD1​⋯​DN×D1​⋯​DN\{K\}\_\{D,\\sigma\}\\in\{\\mathbb\{R\}\}^\{D\_\{1\}\\cdots D\_\{N\}\\times D\_\{1\}\\cdots D\_\{N\}\}satisfies

𝐭~=KD,σ​𝐭,\{\\tilde\{\\mathbf\{t\}\}\}=\{K\}\_\{D,\\sigma\}\{\\mathbf\{t\}\},\(3\.9\)then we callKD,σ\{K\}\_\{D,\\sigma\}a tensor permutation matrix with dimensionD=\(D1,…,DN\)D=\(D\_\{1\},\\dots,D\_\{N\}\)and permutationσ\\sigma\.

Note that𝒯~\\tilde\{\\mathcal\{T\}\}is the tensor transpose of𝒯\\mathcal\{T\}associated withσ\\sigmain\[Pan2014\]andKD,σ\{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[3\.5](https://arxiv.org/html/2607.06976#S3.E5)\.

###### Theorem 3\.3\.

Under the conditions \([3\.6](https://arxiv.org/html/2607.06976#S3.E6)\) and \([3\.7](https://arxiv.org/html/2607.06976#S3.E7)\), the system matrixAk,n∈ℝP1​⋯​PN​Qk×I​JnA\_\{k,n\}\\in\\mathbb\{R\}^\{P\_\{1\}\\cdots P\_\{N\}Q\_\{k\}\\times IJ\_\{n\}\}in the least squares problem \([3\.5](https://arxiv.org/html/2607.06976#S3.E5)\) can be factored as

Ak,n=KD~k,σnT​\(\(Gn∗Tk\)⊗Bn\),~A\_\{k,n\}=\{K\}\_\{\\tilde\{D\}\_\{k\},\\sigma\_\{n\}\}^\{T\}\\left\(\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)\\otimes B\_\{n\}\\right\),with

Gn=∗m=1m≠n𝑁​Hm∈ℝP1​⋯​P^n​⋯​PN×I,Hm=Bm​CmT∈ℝPm×I,~\\begin\{aligned\} G\_\{n\}&=\\overset\{N\}\{\\underset\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}\{\\ast\}\}H\_\{m\}\\in\{\\mathbb\{R\}\}^\{P\_\{1\}\\cdots\\hat\{P\}\_\{n\}\\cdots P\_\{N\}\\times I\},\\\\ H\_\{m\}&=B\_\{m\}C\_\{m\}^\{T\}\\in\{\\mathbb\{R\}\}^\{P\_\{m\}\\times I\},\\end\{aligned\}whereBm=\(b~j\(m\)​\(𝐮p\(m\)\)\)∈ℝPm×JnB\_\{m\}=\\left\(\{\\tilde\{b\}\_\{j\}^\{\(m\)\}\(\{\\mathbf\{u\}\}\_\{p\}^\{\(m\)\}\)\}\\right\)\\in\{\\mathbb\{R\}\}^\{P\_\{m\}\\times J\_\{n\}\}andTk=\(ℒk​\[ti\]​\(yqk\)\)∈ℝQk×IT\_\{k\}=\\left\(\{\\mathcal\{L\}\_\{k\}\\left\[t\_\{i\}\\right\]\(\{y\}\_\{q\_\{k\}\}\)\}\\right\)\\in\{\\mathbb\{R\}\}^\{Q\_\{k\}\\times I\}are themm\-th branch pre\-output matrix and the trunk output data with operatorℒk\\mathcal\{L\}\_\{k\}, respectively, andKD~k,σn∈ℝP1​⋯​PN​Qk×P1​⋯​PN​Qk\{K\}\_\{\\tilde\{D\}\_\{k\},\\sigma\_\{n\}\}\\in\{\\mathbb\{R\}\}^\{P\_\{1\}\\cdots P\_\{N\}Q\_\{k\}\\times P\_\{1\}\\cdots P\_\{N\}Q\_\{k\}\}is the tensor permutation matrix with dimensionD~k=\(P1,…,PN,Qk\)\\tilde\{D\}\_\{k\}=\(P\_\{1\},\\dots,P\_\{N\},Q\_\{k\}\)and the permutation

σn=\(1⋯n−1nn\+1⋯NN\+11⋯n−1n\+1n\+2⋯N\+1n\)∈SN\+1\.\\sigma\_\{n\}=\\bigl\(\\begin\{smallmatrix\}1&\\cdots&n\-1&n&n\+1&\\cdots&N&N\+1\\\\ 1&\\cdots&n\-1&n\+1&n\+2&\\cdots&N\+1&n\\end\{smallmatrix\}\\bigr\)\\in S\_\{N\+1\}\.

Note thatHmH\_\{m\}is the final output of themm\-th branch network, and since the Khatri\-Rao product is associative, we generalize \([3\.8](https://arxiv.org/html/2607.06976#S3.E8)\) to

\(X1∗⋯∗XN\)r​s=\(X1\)r1​s​⋯​\(XN\)rN​s,~\{\(X\_\{1\}\\ast\\cdots\\ast X\_\{N\}\)\}\_\{rs\}=\{\(X\_\{1\}\)\}\_\{r\_\{1\}s\}\\cdots\{\(X\_\{N\}\)\}\_\{r\_\{N\}s\},\(3\.10\)wherer=\[r1,…,rN\]R1,…,RNr=\{\[r\_\{1\},\\dots,r\_\{N\}\]\}\_\{R\_\{1\},\\dots,R\_\{N\}\}for the matricesXn∈ℝRn×SX\_\{n\}\\in\\mathbb\{R\}^\{R\_\{n\}\\times S\}\. A similar result holds for the Kronecker product\. The tensor permutation matrixKD~k,σn\{K\}\_\{\\tilde\{D\}\_\{k\},\\sigma\_\{n\}\}is introduced to match the mixed row order of the product\(Gn∗Tk\)⊗Bn\(G\_\{n\}\\ast T\_\{k\}\)\\otimes B\_\{n\}to the lexicographic row order of the datafkf\_\{k\}\.

###### Proof\.

By \([3\.7](https://arxiv.org/html/2607.06976#S3.E7)\), for eachdkd\_\{k\}, there exist positive integersp1,…,pNp\_\{1\},\\dots,p\_\{N\}andqqsuch thatdk=\[p1,…,pN,q\]P1,…,PN,Qkd\_\{k\}=\{\[p\_\{1\},\\dots,p\_\{N\},q\]\}\_\{P\_\{1\},\\dots,P\_\{N\},Q\_\{k\}\}, whereu^dk\(m\)=upm\(m\)∈βm\\hat\{u\}\_\{d\_\{k\}\}^\{\(m\)\}=\{u\}\_\{p\_\{m\}\}^\{\(m\)\}\\in\\beta\_\{m\}andy^dk=yq∈τk\\hat\{y\}\_\{d\_\{k\}\}=\{y\}\_\{q\}\\in\\tau\_\{k\}\. Also, the tuple\(p1,…,pN,q\)\(p\_\{1\},\\dots,p\_\{N\},q\)anddkd\_\{k\}have one\-to\-one correspondence as1≤pm≤Pm1\\leq p\_\{m\}\\leq P\_\{m\}form=1,…,Nm=1,\\dots,N,1≤q≤Qk1\\leq q\\leq Q\_\{k\}, and1≤dk≤Dk=P1​⋯​PN​Qk1\\leq d\_\{k\}\\leq D\_\{k\}=P\_\{1\}\\cdots P\_\{N\}Q\_\{k\}\.

Hence, by \([3\.6](https://arxiv.org/html/2607.06976#S3.E6)\),\(\[p1,…,pN,q\]P1,…,PN,Qk,\[i,j\]I,Jn\)\(\{\[p\_\{1\},\\dots,p\_\{N\},q\]\}\_\{P\_\{1\},\\dots,P\_\{N\},Q\_\{k\}\},\{\[i,j\]\}\_\{I,J\_\{n\}\}\)entry ofAk,nA\_\{k,n\}can be expressed as

\(∏m=1m≠nNbi\(m\)​\(𝐮pm\(m\)\)\)​b~j\(n\)​\(𝐮pn\(n\)\)​ℒk​\[ti\]​\(yq\)\.~\\left\(\\prod\_\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}^\{N\}\{b\}\_\{i\}^\{\(m\)\}\(\{\\mathbf\{u\}\}\_\{p\_\{m\}\}^\{\(m\)\}\)\\right\)\\tilde\{b\}\_\{j\}^\{\(n\)\}\(\{\\mathbf\{u\}\}\_\{p\_\{n\}\}^\{\(n\)\}\)\\mathcal\{L\}\_\{k\}\\left\[t\_\{i\}\\right\]\(\{y\}\_\{q\}\)\.\(3\.11\)
On the other hand, by using similar argument in \([3\.10](https://arxiv.org/html/2607.06976#S3.E10)\),

\(\[p1,…,p^n,…,pN,q,pn\]P1,…,P^n,…,PN,Qk,Pn,\[i,j\]I,Jn\)\(\{\[p\_\{1\},\\dots,\\hat\{p\}\_\{n\},\\dots,p\_\{N\},q,p\_\{n\}\]\}\_\{P\_\{1\},\\dots,\\hat\{P\}\_\{n\},\\dots,P\_\{N\},Q\_\{k\},P\_\{n\}\},\{\[i,j\]\}\_\{I,J\_\{n\}\}\)entry of the mixed product\(Gn∗Tk\)⊗Bn\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)\\otimes B\_\{n\}is exactly \([3\.11](https://arxiv.org/html/2607.06976#S3.E11)\)\. UsingKD~k,σn\{K\}\_\{\\tilde\{D\}\_\{k\},\\sigma\_\{n\}\}to permute the row order of\(Gn∗Tk\)⊗Bn\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)\\otimes B\_\{n\}, we haveAk,n=KD~k,σnT​\(\(Gn∗Tk\)⊗Bn\)A\_\{k,n\}=\{K\}\_\{\\tilde\{D\}\_\{k\},\\sigma\_\{n\}\}^\{T\}\\left\(\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)\\otimes B\_\{n\}\\right\)as desired\. ∎

Using the result of[Theorem3\.3](https://arxiv.org/html/2607.06976#S3.Thmtheorem3), let us observe the LS problem \([3\.5](https://arxiv.org/html/2607.06976#S3.E5)\) and rearrange the data vectorfk∈ℝP1​⋯​PN​Qkf\_\{k\}\\in\\mathbb\{R\}^\{P\_\{1\}\\cdots P\_\{N\}Q\_\{k\}\}into a tensorFk∈ℝP1×⋯×PN×QkF\_\{k\}\\in\\mathbb\{R\}^\{P\_\{1\}\\times\\cdots\\times P\_\{N\}\\times Q\_\{k\}\}\. Both the\(p1,…,pN,qk\)\(p\_\{1\},\\dots,p\_\{N\},q\_\{k\}\)entry ofFkF\_\{k\}and the\[p1,…,pN,qk\]P1,…,PN,Qk\{\[p\_\{1\},\\dots,p\_\{N\},q\_\{k\}\]\}\_\{P\_\{1\},\\dots,P\_\{N\},Q\_\{k\}\}entry offkf\_\{k\}are

ℒk​\[G​\(⋅,…,⋅\)\]​\(up1\(1\),…,upN\(N\),yqk\)\.~\\mathcal\{L\}\_\{k\}\\left\[G\(\\cdot,\\dots,\\cdot\)\\right\]\(u^\{\(1\)\}\_\{p\_\{1\}\},\\dots,u^\{\(N\)\}\_\{p\_\{N\}\},y\_\{q\_\{k\}\}\)\.\(3\.12\)Therefore, the LS problem with respect to the last layer parameter is

minCn​∑k=1Kϵk​‖fk−KD~k,σnT​\(\(Gn∗Tk\)⊗Bn\)​vec​\(CnT\)‖22\+λn​‖vec​\(CnT\)‖22\.~\\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\}\.\(3\.13\)
To write the normal equation concisely, we use the fact\[Rao1970\]that\(X∗Y\)T​\(X∗Y\)=\(XT​X\)⊙\(YT​Y\)\{\(X\\ast Y\)\}^\{T\}\(X\\ast Y\)=\(X^\{T\}X\)\\odot\(Y^\{T\}Y\)for all matricesX∈ℝR1×SX\\in\\mathbb\{R\}^\{R\_\{1\}\\times S\}andY∈ℝR2×SY\\in\\mathbb\{R\}^\{R\_\{2\}\\times S\}\. Here, the normal equation of \([3\.13](https://arxiv.org/html/2607.06976#S3.E13)\) is given as

\[\(∑k=1Kϵk​\(GnT​Gn\)⊙\(TkT​Tk\)\)⊗\(BnT​Bn\)\]​vec​\(CnT\)\+λn​vec​\(CnT\)=∑k=1Kϵk​\(\(Gn∗Tk\)T⊗BnT\)​KD~k,σnT​fk\.~\\begin\{aligned\} &\\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\{aligned\}\(3\.14\)Introducing a long data matrixF^k,n∈ℝPn×P1​⋯​P^n​⋯​PN​Qk\\hat\{F\}\_\{k,n\}\\in\\mathbb\{R\}^\{P\_\{n\}\\times P\_\{1\}\\cdots\\hat\{P\}\_\{n\}\\cdots P\_\{N\}Q\_\{k\}\}whose

\(pn,\[p1,…,p^n,…​pN,q\]P1,…,P^n,…,PN,Qk\)\\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 \([3\.12](https://arxiv.org/html/2607.06976#S3.E12)\), we can write \([3\.14](https://arxiv.org/html/2607.06976#S3.E14)\) as a matrix equation:

BnT​Bn​CnT​\(∑k=1Kϵk​\(GnT​Gn\)⊙\(TkT​Tk\)\)\+λn​CnT=BnT​\(∑k=1Kϵk​F^k,n​\(Gn∗Tk\)\)\.~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\)\.\(3\.15\)Note that\(GnT​Gn\)⊙\(TkT​Tk\)\(G\_\{n\}^\{T\}G\_\{n\}\)\\odot\(T\_\{k\}^\{T\}T\_\{k\}\)can be effectively calculated without forming huge matrices since

\(GnT​Gn\)⊙\(TkT​Tk\)=\(⊙m=1m≠n𝑁​\(HmT​Hm\)\)⊙\(TkT​Tk\),~\(G\_\{n\}^\{T\}G\_\{n\}\)\\odot\(T\_\{k\}^\{T\}T\_\{k\}\)=\\left\(\\overset\{N\}\{\\underset\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}\{\\odot\}\}\(H\_\{m\}^\{T\}H\_\{m\}\)\\right\)\\odot\(T\_\{k\}^\{T\}T\_\{k\}\),where each component is anI×II\\times Imatrix generated from each network\.

In the case of unsupervised learning where the given dataFkF\_\{k\}depends only on a single input argumentu\(mk\)u^\{\(m\_\{k\}\)\}of the operatorG​\(u\(1\),…,u\(N\)\)G\(u^\{\(1\)\},\\dots,u^\{\(N\)\}\), we can reduce the rank\(N\+1\)\(N\+1\)data tensorFkF\_\{k\}to the reduced matrixF~mk∈ℝPmk×Qk\\tilde\{F\}\_\{m\_\{k\}\}\\in\\mathbb\{R\}^\{P\_\{m\_\{k\}\}\\times Q\_\{k\}\}, wheremkm\_\{k\}is the corresponding index of the input function coordinate\. We provide a related theorem\.

###### Theorem 3\.4\.

Suppose the data tensorFk∈ℝP1×⋯×PN×QkF\_\{k\}\\in\\mathbb\{R\}^\{P\_\{1\}\\times\\cdots\\times P\_\{N\}\\times Q\_\{k\}\}in \([3\.12](https://arxiv.org/html/2607.06976#S3.E12)\) is given as the case of unsupervised learning, where the data varies only along themkm\_\{k\}\-th axis and the last axis\. Then,

F^k,n​\(Gn∗Tk\)=\{𝟏Pn×1​\(⊙m=1m≠n𝑁​𝐡~mT\)forn≠mk,\[𝟏Pn×1​\(⊙m=1m≠n𝑁​𝐡~mT\)\]⊙\(F~n​Tk\)forn=mk,~\\hat\{F\}\_\{k,n\}\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)=\\begin\{cases\}\\mathbf\{1\}\_\{P\_\{n\}\\times 1\}\\left\(\\overset\{N\}\{\\underset\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}\{\\odot\}\}\\tilde\{\\mathbf\{h\}\}^\{T\}\_\{m\}\\right\)\\qquad&\\text\{for $n\\neq m\_\{k\}$\},\\\\ \\left\[\\mathbf\{1\}\_\{P\_\{n\}\\times 1\}\\left\(\\overset\{N\}\{\\underset\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}\{\\odot\}\}\\tilde\{\\mathbf\{h\}\}^\{T\}\_\{m\}\\right\)\\right\]\\odot\\left\(\\tilde\{F\}\_\{n\}T\_\{k\}\\right\)\\qquad&\\text\{for $n=m\_\{k\}$\},\\end\{cases\}where𝐡~mT∈ℝ1×I\\tilde\{\\mathbf\{h\}\}^\{T\}\_\{m\}\\in\{\\mathbb\{R\}\}^\{1\\times I\}is a row vector of lengthIIdefined by

𝐡~mT=\{𝟏1×Pm​Hmifm≠mk,𝟏1×Pmk​\(Hmk⊙\(F~mk​Tk\)\)ifm=mk\.~\\tilde\{\\mathbf\{h\}\}^\{T\}\_\{m\}=\\begin\{cases\*\}\{\\mathbf\{1\}\_\{1\\times P\_\{m\}\}\}H\_\{m\}\\qquad&\\text\{if $m\\neq m\_\{k\}$\},\\\\ \{\\mathbf\{1\}\_\{1\\times P\_\{m\_\{k\}\}\}\}\(H\_\{m\_\{k\}\}\\odot\(\\tilde\{F\}\_\{m\_\{k\}\}T\_\{k\}\)\)\\qquad&\\text\{if $m=m\_\{k\}$\}\.\\end\{cases\*\}

###### Proof\.

By expanding the summation in the matrix multiplicationF^k,n​\(Gn∗Tk\)\\hat\{F\}\_\{k,n\}\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)of intermediate dimensionP1​⋯​P^n​⋯​PN​QkP\_\{1\}\\cdots\\hat\{P\}\_\{n\}\\cdots P\_\{N\}Q\_\{k\}intoNNsummations, we have

\(F^k,n​\(Gn∗Tk\)\)pn​i=∑l\(F^k,n\)pn​l​\(Gn∗Tk\)l​i=∑p1⋯​∑pn^​⋯​∑pN∑q\(F~mk\)pmk​q​\(∏m=1m≠nN\(Hm\)pm​i\)​\(Tk\)q​i\.~\\begin\{aligned\} \{\\left\(\\hat\{F\}\_\{k,n\}\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)\\right\)\}\_\{p\_\{n\}i\}&=\\sum\_\{l\}\{\(\\hat\{F\}\_\{k,n\}\)\}\_\{p\_\{n\}l\}\{\(G\_\{n\}\\ast T\_\{k\}\)\}\_\{li\}\\\\ &=\\sum\_\{p\_\{1\}\}\\cdots\\widehat\{\\sum\_\{p\_\{n\}\}\}\\cdots\\sum\_\{p\_\{N\}\}\\sum\_\{q\}\{\(\\tilde\{F\}\_\{m\_\{k\}\}\)\}\_\{p\_\{m\_\{k\}\}q\}\\left\(\\prod\_\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}^\{N\}\{\(H\_\{m\}\)\}\_\{p\_\{m\}i\}\\right\)\{\(T\_\{k\}\)\}\_\{qi\}\.\\end\{aligned\}\(3\.16\)Ifn≠mkn\\neq m\_\{k\}, by rearranging summations with respect to their corresponding indices, \([3\.16](https://arxiv.org/html/2607.06976#S3.E16)\) becomes

\(F^k,n\\displaystyle\\Bigl\(\\hat\{F\}\_\{k,n\}\(Gn∗Tk\)\)pn​i\\displaystyle\\bigl\(G\_\{n\}\\ast T\_\{k\}\\bigr\)\\Bigr\)\_\{p\_\{n\}i\}\(3\.17\)=\(∏m=1m≠n,mkN∑pm\(Hm\)pm​i\)​\(∑pmk\(\(Hmk\)pmk​i​∑q\(F~mk\)pmk​q​\(Tk\)q​i\)\)\\displaystyle=\\left\(\\prod\_\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n,m\_\{k\}\\end\{subarray\}\}^\{N\}\\sum\_\{p\_\{m\}\}\{\(H\_\{m\}\)\}\_\{p\_\{m\}i\}\\right\)\\left\(\\sum\_\{p\_\{m\_\{k\}\}\}\\left\(\{\(H\_\{m\_\{k\}\}\)\}\_\{p\_\{m\_\{k\}\}i\}\\sum\_\{q\}\{\(\\tilde\{F\}\_\{m\_\{k\}\}\)\}\_\{p\_\{m\_\{k\}\}q\}\{\(T\_\{k\}\)\}\_\{qi\}\\right\)\\right\)=\(∏m=1m≠n,mkN\(𝟏1×Pm​Hm\)i\)​\(∑pmk\(Hmk\)pmk​i​\(F~mk​Tk\)pmk​i\)\\displaystyle=\\left\(\\prod\_\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n,m\_\{k\}\\end\{subarray\}\}^\{N\}\{\\left\(\{\\mathbf\{1\}\_\{1\\times P\_\{m\}\}\}H\_\{m\}\\right\)\}\_\{i\}\\right\)\\left\(\\sum\_\{p\_\{m\_\{k\}\}\}\{\(H\_\{m\_\{k\}\}\)\}\_\{p\_\{m\_\{k\}\}i\}\{\(\\tilde\{F\}\_\{m\_\{k\}\}T\_\{k\}\)\}\_\{p\_\{m\_\{k\}\}i\}\\right\)=\(∏m=1m≠n,mkN\(𝟏1×Pm​Hm\)i\)​\(𝟏1×Pmk​\(Hmk⊙\(F~mk​Tk\)\)\)i\.\\displaystyle=\\left\(\\prod\_\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n,m\_\{k\}\\end\{subarray\}\}^\{N\}\{\\left\(\{\\mathbf\{1\}\_\{1\\times P\_\{m\}\}\}H\_\{m\}\\right\)\}\_\{i\}\\right\)\{\\left\(\{\\mathbf\{1\}\_\{1\\times P\_\{m\_\{k\}\}\}\}\(H\_\{m\_\{k\}\}\\odot\(\\tilde\{F\}\_\{m\_\{k\}\}T\_\{k\}\)\)\\right\)\}\_\{i\}\.Therefore, by reordering products and broadcasting along the first axis of the matrix as \([3\.17](https://arxiv.org/html/2607.06976#S3.E17)\) does not depend on the first axis variablepnp\_\{n\}we have

F^k,n​\(Gn∗Tk\)=𝟏Pn×1​\(⊙m=1m≠n𝑁​𝐡~mT\)\.~\\hat\{F\}\_\{k,n\}\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)=\\mathbf\{1\}\_\{P\_\{n\}\\times 1\}\\left\(\\overset\{N\}\{\\underset\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}\{\\odot\}\}\\tilde\{\\mathbf\{h\}\}^\{T\}\_\{m\}\\right\)\.
On the other hand, ifn=mkn=m\_\{k\}, \([3\.16](https://arxiv.org/html/2607.06976#S3.E16)\) becomes

\(F^k,n​\(Gn∗Tk\)\)pn​i=\(∏m=1m≠nN∑pm\(Hm\)pm​i\)​\(∑q\(F~n\)pn​q​\(Tk\)q​i\)=\(∏m=1m≠nN\(𝟏1×Pm​Hm\)i\)​\(F~n​Tk\)pn​i\.~\\begin\{aligned\} \{\\left\(\\hat\{F\}\_\{k,n\}\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)\\right\)\}\_\{p\_\{n\}i\}&=\\left\(\\prod\_\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}^\{N\}\\sum\_\{p\_\{m\}\}\{\(H\_\{m\}\)\}\_\{p\_\{m\}i\}\\right\)\\left\(\\sum\_\{q\}\{\(\\tilde\{F\}\_\{n\}\)\}\_\{p\_\{n\}q\}\{\(T\_\{k\}\)\}\_\{qi\}\\right\)\\\\ &=\\left\(\\prod\_\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}^\{N\}\{\\left\(\{\\mathbf\{1\}\_\{1\\times P\_\{m\}\}\}H\_\{m\}\\right\)\}\_\{i\}\\right\)\{\\left\(\\tilde\{F\}\_\{n\}T\_\{k\}\\right\)\}\_\{p\_\{n\}i\}\.\\end\{aligned\}\(3\.18\)Therefore, by reordering products and broadcasting terms along the first axis of the matrix which does not depend on the variablepnp\_\{n\}for \([3\.18](https://arxiv.org/html/2607.06976#S3.E18)\), we have

F^k,n​\(Gn∗Tk\)=\[𝟏Pn×1​\(⊙m=1m≠n𝑁​𝐡~mT\)\]⊙\(F~n​Tk\),~\\hat\{F\}\_\{k,n\}\\left\(G\_\{n\}\\ast T\_\{k\}\\right\)=\\left\[\\mathbf\{1\}\_\{P\_\{n\}\\times 1\}\\left\(\\overset\{N\}\{\\underset\{\\begin\{subarray\}\{c\}m=1\\\\ m\\neq n\\end\{subarray\}\}\{\\odot\}\}\\tilde\{\\mathbf\{h\}\}^\{T\}\_\{m\}\\right\)\\right\]\\odot\\left\(\\tilde\{F\}\_\{n\}T\_\{k\}\\right\),which concludes the proof\. ∎

This theorem enables us to efficiently compute each component of the matrix equation \([3\.15](https://arxiv.org/html/2607.06976#S3.E15)\)\. Finally, we can use[Proposition2\.2](https://arxiv.org/html/2607.06976#S2.Thmtheorem2)to find the minimizerCnC\_\{n\}of the LS problem \([3\.13](https://arxiv.org/html/2607.06976#S3.E13)\) for thenn\-th last layer parameter\.

Now, we present the entire LS step for MIONet \([Algorithm1](https://arxiv.org/html/2607.06976#alg1)\)\. Since the loss is the sum of squares of multilinear functions with last layer parameters\(θ1L,…,θNL\)\(\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\)ofNNbranches, 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π∈SN\\pi\\in S\_\{N\}\. Form=1,…,Nm=1,\\dots,N, fix all parameters exceptθπ​\(m\)L\\theta\_\{\\pi\(m\)\}^\{L\}, then solve the LS problem in terms ofθπ​\(m\)L\\theta\_\{\\pi\(m\)\}^\{L\}to updateθπ​\(m\)L\\theta\_\{\\pi\(m\)\}^\{L\}\. Note that a unique minimizer is guaranteed ifλπ​\(m\)\>0\\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 \([Algorithm2](https://arxiv.org/html/2607.06976#alg2)\) is similar to the LSGD algorithm for DeepONet introduced in\[Choi2025, Algorithm 2\]\. To speed up the ALS step, one may introduce enhanced line search \(ELS\)\[Rajih2008\]or partitioned ALS \(PALS\)\[Tichavsky2016\]\.

Algorithm 1Alternating Least Squares Step for MIONetOutput:Optimized last layer parametersθ1L,…,θNL\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}

1:functionALS\(

θ1B,…,θNB,θT,θ1L,…,θNL\\theta\_\{1\}^\{B\},\\dots,\\theta\_\{N\}^\{B\},\\theta^\{T\},\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\)

2:whileStopping criterion is not metdo

3:Choose permutation

π∈SN\\pi\\in S\_\{N\}
4:for

m=1,…,Nm=1,\\dots,Ndo

5:

θπ​\(m\)L←L​Sπ​\(m\)​\(θ1B,…,θNB,θT,θ1L,…,θNL\)\\theta\_\{\\pi\(m\)\}^\{L\}\\leftarrow LS\_\{\\pi\(m\)\}\(\\theta\_\{1\}^\{B\},\\dots,\\theta\_\{N\}^\{B\},\\theta^\{T\},\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\)⊳\\trianglerightSolve the LS problem in terms of

θπ​\(m\)L\\theta\_\{\\pi\(m\)\}^\{L\}and update

6:endfor

7:endwhile

8:endfunction

Algorithm 2Hybrid Least Squares/Gradient Descent for MIONetΘB=\(θ1B,…,θNB\)\\Theta^\{B\}=\(\\theta\_\{1\}^\{B\},\\dots,\\theta\_\{N\}^\{B\}\),ΘL=\(θ1L,…,θNL\)\\Theta^\{L\}=\(\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\)

1:

\(ΘB,θT,ΘL\)←\(Θ0B,θ0T,Θ0L\)\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)\\leftarrow\(\\Theta\_\{0\}^\{B\},\\theta\_\{0\}^\{T\},\\Theta\_\{0\}^\{L\}\): Initial parameters for the branches and the trunk

2:

ΘL←A​L​S​\(ΘB,θT,ΘL\)\\Theta^\{L\}\\leftarrow ALS\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)⊳\\trianglerightSolve the ALS problem for eachθmL\\theta\_\{m\}^\{L\}in the full batch

3:for

i=1,…i=1,\\dotsdo

4:

\(ΘB,θT\)←G​D​\(ΘB,θT,ΘL\)\(\\Theta^\{B\},\\theta^\{T\}\)\\leftarrow GD\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)⊳\\trianglerightUse a gradient descent type optimizer to findΘB\\Theta^\{B\}andθT\\theta^\{T\}

5:

ΘL←A​L​S​\(ΘB,θT,ΘL\)\\Theta^\{L\}\\leftarrow ALS\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)
6:endfor

## 4Experimental 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\[Choi2025\]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[Algorithm3](https://arxiv.org/html/2607.06976#alg3)\.

Algorithm 3ALS\+Adam for MIONetΘB=\(θ1B,…,θNB\)\\Theta^\{B\}=\(\\theta\_\{1\}^\{B\},\\dots,\\theta\_\{N\}^\{B\}\),ΘL=\(θ1L,…,θNL\)\\Theta^\{L\}=\(\\theta\_\{1\}^\{L\},\\dots,\\theta\_\{N\}^\{L\}\)

1:

\(ΘB,θT,ΘL\)←\(Θ0B,θ0T,Θ0L\)\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)\\leftarrow\(\\Theta\_\{0\}^\{B\},\\theta\_\{0\}^\{T\},\\Theta\_\{0\}^\{L\}\): Initial parameters

2:for

i=1,…,I0i=1,\\dots,I\_\{0\}do

3:

\(ΘB,θT,ΘL\)←A​d​a​m​\(ΘB,θT,ΘL\)\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)\\leftarrow Adam\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)⊳\\trianglerightInitial Adam stage for all parameters

4:endfor

5:

ΘL←A​L​S​\(ΘB,θT,ΘL\)\\Theta^\{L\}\\leftarrow ALS\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)⊳\\trianglerightSolve the ALS problem for eachθmL\\theta\_\{m\}^\{L\}in the full data batch

6:for

i=1,…i=1,\\dotsdo⊳\\trianglerightWork unit block

7:for

j=1,…,J0j=1,\\dots,J\_\{0\}do

8:

\(ΘB,θT\)←A​d​a​m​\(ΘB,θT,ΘL\)\(\\Theta^\{B\},\\theta^\{T\}\)\\leftarrow Adam\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)⊳\\trianglerightUse Adam for hidden layer parameters

9:endfor

10:

ΘL←A​L​S​\(ΘB,θT,ΘL\)\\Theta^\{L\}\\leftarrow ALS\(\\Theta^\{B\},\\theta^\{T\},\\Theta^\{L\}\)
11:endfor

For 2\-branch MIONet training in supervised learning, we use the mean square error \(MSE\) loss

ϵ1​Ldata\+λ1​‖C1‖F2\+λ2​‖C2‖F2,~\\epsilon\_\{1\}L\_\{\\text\{data\}\}\+\\lambda\_\{1\}\{\\\|C\_\{1\}\\\|\}\_\{F\}^\{2\}\+\\lambda\_\{2\}\{\\\|C\_\{2\}\\\|\}\_\{F\}^\{2\},\(4\.1\)and for the unsupervised learning, we use PI\-loss

ϵ1​Ldata\+ϵ2​Lphysics\+λ1​‖C1‖F2\+λ2​‖C2‖F2,~\\epsilon\_\{1\}L\_\{\\text\{data\}\}\+\\epsilon\_\{2\}L\_\{\\text\{physics\}\}\+\\lambda\_\{1\}\{\\\|C\_\{1\}\\\|\}\_\{F\}^\{2\}\+\\lambda\_\{2\}\{\\\|C\_\{2\}\\\|\}\_\{F\}^\{2\},\(4\.2\)where∥⋅∥F\{\\\|\\cdot\\\|\}\_\{F\}denotes the Frobenius norm,ϵ1=1\\epsilon\_\{1\}=1,LdataL\_\{\\text\{data\}\}is theL2L^\{2\}MSE on the data pairs\(u^d1,y^d1\)\(\\hat\{u\}\_\{d\_\{1\}\},\\hat\{y\}\_\{d\_\{1\}\}\)wherey^d1\\hat\{y\}\_\{d\_\{1\}\}corresponds to the given data points of the governing PDE, andLphysicsL\_\{\\text\{physics\}\}is theL2L^\{2\}MSE of the PDE residuals\(u^d2,y^d2\)\(\\hat\{u\}\_\{d\_\{2\}\},\\hat\{y\}\_\{d\_\{2\}\}\)wherey^d2\\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\)\)\(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 tensorFkF\_\{k\}depends only on one input argument of the target operator\. According to[Theorem3\.4](https://arxiv.org/html/2607.06976#S3.Thmtheorem4), 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\[Choi2025\]\. We use the Adam optimizer withlr=10−3\\text\{lr\}=10^\{\-3\}and\(β1,β2\)=\(0\.99,0\.999\)\(\\beta\_\{1\},\\beta\_\{2\}\)=\(0\.99,0\.999\)\. The Adam momentums are maintained between the LS steps\. He normal initialization\[He2015\]is used for parameter initialization, and the Swish functionx/\(1\+e−x\)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\.,λ1=λ2=λ\\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[Algorithm3](https://arxiv.org/html/2607.06976#alg3)\. 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 \(branchff\) 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×100100\\times 100data block is used as one batch in each Adam iteration\. On the other hand, training with ALS\+Adam uses the batch size of50×5050\\times 50in 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[Table2](https://arxiv.org/html/2607.06976#S4.T2)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\.

Table 2:Inputs for networks, network structures and weight of MIONet models\.IC and BC stand for initial condition and boundary condition, respectively\. FCN and CNN stand for fully connected network and convolutional neural network, respectively\. CNN structures are described in the corresponding subsections\. Swish activation is used on all branches and trunks\.### 4\.1Reaction\-diffusion with variable source and diffusivity

In this section, we consider a reaction\-diffusion equation with variable source and diffusivity as inputs:

∂u∂t=∂∂x​\(D​\(x\)​∂u∂x\)\\displaystyle\\frac\{\\partial u\}\{\\partial t\}=\\frac\{\\partial\}\{\\partial x\}\\left\(D\(x\)\\frac\{\\partial u\}\{\\partial x\}\\right\)\+R​\(u\)\+f​\(x\),\\displaystyle\+R\(u\)\+f\(x\),\\qquad\(x,t\)∈\(0,1\)×\(0,1\],\\displaystyle\(x,t\)\\in\(0,1\)\\times\(0,1\],u​\(x,0\)\\displaystyle u\(x,0\)=0,\\displaystyle=0,\\qquadx∈\(0,1\),\\displaystyle x\\in\(0,1\),u​\(0,t\)=u​\(1,t\)\\displaystyle u\(0,t\)=u\(1,t\)=0,\\displaystyle=0,\\qquadt∈\(0,1\],\\displaystyle t\\in\(0,1\],whereR​\(u\)=0\.01​u2R\(u\)=0\.01u^\{2\}\. Here, we aim to learn a solution operator that maps the sourceffand diffusivityDDto the solutionuuvia MIONet\. The inputs are generated from a Gaussian process \(GP\) in the interval\[0,1\]\[0,1\]with zero mean and a squared exponential covariance kernel

k​\(x1,x2\)=σ2​exp​\(−\|x1−x2\|22​l2\),k\(x\_\{1\},x\_\{2\}\)=\\sigma^\{2\}\\text\{exp\}\\left\(\-\\frac\{\{\|x\_\{1\}\-x\_\{2\}\|\}^\{2\}\}\{2l^\{2\}\}\\right\),\(4\.3\)with a pair of scale and variance parameters\(l,σ2\)=\(0\.2,1\)\(l,\\sigma^\{2\}\)=\(0\.2,1\)forffand\(l,σ2\)=\(0\.2,0\.35\)\(l,\\sigma^\{2\}\)=\(0\.2,0\.35\)forDD\. The input functionsffandDDare discretized at3333equidistant grid points of\[0,1\]\[0,1\], and the output function is evaluated at33×3333\\times 33equidistant grid points of\[0,1\]2\{\[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 size129×257129\\times 257in the space domain and time domain\. For the training data, we note that 1,000,000 reference solutions from every pair of\(f,D\)\(f,D\)need to be prepared, which takes about 8 GB of memory\.

### 4\.2Constant coefficient advection with variable source and IBC

In this section, we solve a 1D constant coefficient advection equation with a source term via MIONet:

∂u∂t\+a​∂u∂x=f​\(x\),\(x,t\)∈\(0,1\]2,u​\(x,0\)=P​\(x\),x∈\[0,1\],u​\(0,t\)=Q​\(t\),t∈\[0,1\],~\\begin\{aligned\} \\frac\{\\partial u\}\{\\partial t\}\+a\\frac\{\\partial u\}\{\\partial x\}&=f\(x\),\\qquad&&\(x,t\)\\in\{\(0,1\]\}^\{2\},\\\\ u\(x,0\)&=P\(x\),\\qquad&&x\\in\[0,1\],\\\\ u\(0,t\)&=Q\(t\),\\qquad&&t\\in\[0,1\],\\end\{aligned\}whereaais a fixed constant andP​\(0\)=Q​\(0\)P\(0\)=Q\(0\)\. We aim to learn a solution operator which maps the sourceff, BCPP, and ICQQto the solutionuu\.

Note that the analytical solution is given as follows:

u∗​\(x,t\)=\{1a​\(F​\(x\)−F​\(x−a​t\)\)\+P​\(x−a​t\),x−a​t≥0,1a​\(F​\(x\)−F​\(0\)\)\+Q​\(t−xa\),x−a​t<0,~u^\{\*\}\(x,t\)=\\begin\{cases\}\\frac\{1\}\{a\}\\left\(F\(x\)\-F\(x\-at\)\\right\)\+P\(x\-at\),\\qquad&x\-at\\geq 0,\\\\ \\frac\{1\}\{a\}\\left\(F\(x\)\-F\(0\)\\right\)\+Q\(t\-\\frac\{x\}\{a\}\),\\qquad&x\-at<0,\\end\{cases\}whereF​\(x\)=∫0xf​\(s\)​𝑑sF\(x\)=\\int\_\{0\}^\{x\}f\(s\)\\,dsis the antiderivative offf\. Since this solution may have non\-differentiable cusps along the linex−a​t=0x\-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\)=0f\(0\)=0andP′​\(0\)=−1a​Q′​\(0\)P^\{\\prime\}\(0\)=\-\\frac\{1\}\{a\}Q^\{\\prime\}\(0\)— to ensure differentiability of the solution\. Also, instead of having separate IC and BC as input functions,PPandQQ, they can be naturally concatenated into one input function along the domains of IC and BC\. Let𝐩=\[P​\(0\)​⋯​P​\(1\)\]\\mathbf\{p\}=\[P\(0\)\\cdots P\(1\)\]and𝐪=\[Q​\(0\)​⋯​Q​\(1\)\]\\mathbf\{q\}=\[Q\(0\)\\cdots Q\(1\)\]be the discretizations ofPPandQQalong their domains,\[0,1\]×\{0\}\[0,1\]\\times\\\{0\\\}and\{0\}×\[0,1\]\\\{0\\\}\\times\[0,1\], respectively\. SinceP​\(0\)=Q​\(0\)P\(0\)=Q\(0\), we can concatenate𝐩\\mathbf\{p\}and𝐪\\mathbf\{q\}by𝐫=\[Q​\(1\)​⋯​Q​\(0\)=P​\(0\)​⋯​P​\(1\)\]\\mathbf\{r\}=\[Q\(1\)\\cdots Q\(0\)=P\(0\)\\cdots P\(1\)\], where𝐪\\mathbf\{q\}is flipped and one of the duplicate valuesP​\(0\)P\(0\)orQ​\(0\)Q\(0\)is removed\.𝐫\\mathbf\{r\}will be used for the second input argument of the operator\.

To generate input functionsff,PPandQQ, we use the GP with zero mean and a squared exponential covariance kernel in[4\.3](https://arxiv.org/html/2607.06976#S4.E3)\. First, we sampleggfrom the GP in the interval\[0,1\]\[0,1\]with scale factorl=0\.2l=0\.2and varianceσ2=1\\sigma^\{2\}=1, andhhfrom the GP in the interval\[−a,1\]\[\-a,1\]withl=0\.2l=0\.2andσ2=1\\sigma^\{2\}=1\. Then, we set

f​\(x\)=g​\(x\)−g​\(0\),x∈\[0,1\],P​\(x\)=h​\(x\),x∈\[0,1\],Q​\(t\)=h​\(−a​t\),t∈\[0,1\]\.~\\begin\{aligned\} f\(x\)&=g\(x\)\-g\(0\),\\qquad&&x\\in\[0,1\],\\\\ P\(x\)&=h\(x\),\\qquad&&x\\in\[0,1\],\\\\ Q\(t\)&=h\(\-at\),\\qquad&&t\\in\[0,1\]\.\\end\{aligned\}The generatedff,PPandQQsatisfyf​\(0\)=0f\(0\)=0,P​\(0\)=Q​\(0\)P\(0\)=Q\(0\)andP′​\(0\)=−1a​Q′​\(0\)P^\{\\prime\}\(0\)=\-\\frac\{1\}\{a\}Q^\{\\prime\}\(0\)\. In this problem, we choosea=0\.5a=0\.5\. The original input functionsff,PPandQQare discretized at3333equidistant grid points of\[0,1\]\[0,1\], and the concatenated input𝐫\\mathbf\{r\}is a vector of length6565\. The output function is evaluated on33×3333\\times 33equidistant grid points of\[0,1\]2\{\[0,1\]\}^\{2\}\. The antiderivative of the exact solution is computed by the composite trapezoidal rule in finer grids of size129129instead of3333\.

### 4\.32D 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:

−∇⋅\(κ​∇u\)=f,\(x,y\)∈Ω=\(0,1\)2,u=g,\(x,y\)∈∂Ω,~\\begin\{aligned\} \-\\nabla\\cdot\\left\(\\kappa\\nabla u\\right\)&=f,\\qquad&&\(x,y\)\\in\\Omega=\{\(0,1\)\}^\{2\},\\\\ u&=g,\\qquad&&\(x,y\)\\in\\partial\\Omega,\\end\{aligned\}\(4\.4\)whereκ≡1\\kappa\\equiv 1and the MIONet takes two inputs, sourceffand Dirichlet BCgg, and generates a solutionuu\. The reference solutions are generated by the finite difference method on finer grids of size129×129129\\times 129\.

The 2D input sourceffis generated from a GP with zero mean and a 2D squared exponential covariance kernel

k​\(x1,x2,y1,y2\)=σ2​exp​\(−\|x1−x2\|22​lx2−\|y1−y2\|22​ly2\),~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\),with scale factorslx=ly=0\.2l\_\{x\}=l\_\{y\}=0\.2and varianceσ2=0\.1\\sigma^\{2\}=0\.1\. We take33×3333\\times 33equidistant grid points in\[0,1\]2\{\[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 with3×33\\times 3,2×22\\times 2, and2×22\\times 2kernels each with2×22\\times 2strides\. Since the channel sizes are\[1,16,32,64\]\[1,16,32,64\], the output of the CNN is a4×44\\times 4image of6464channels\.

For the BC input functiongg, we consider a flattened 1D functiong~\\tilde\{g\}on the interval\[0,4\]\[0,4\], generated by a GP with zero mean and periodic covariance kernel

k​\(x1,x2\)=σ2​exp​\(−2l2​sin2⁡\(π​\|x1−x2\|p\)\),~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\),where the scale factor, the period, and the variance arel=0\.3l=0\.3,p=4p=4, andσ2=0\.1\\sigma^\{2\}=0\.1, respectively, so thatg​\(𝐡​\(t\)\)=g~​\(t\)g\(\\mathbf\{h\}\(t\)\)=\\tilde\{g\}\(t\)where𝐡:\[0,4\]→∂Ω\\mathbf\{h\}\\colon\[0,4\]\\to\\partial\\Omegais the arc length parametrization of∂Ω\\partial\\Omegasuch that

𝐡​\(t\)=\{\(t,0\),t∈\[0,1\),\(1,t−1\),t∈\[1,2\),\(3−t,1\),t∈\[2,3\),\(0,4−t\),t∈\[3,4\]\.~\\begin\{aligned\} \\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\{aligned\}The input functiong~\\tilde\{g\}is discretized on129129equidistant grid points in\[0,4\]\[0,4\], and the output function is evaluated on33×3333\\times 33equidistant grid points in\[0,1\]2\{\[0,1\]\}^\{2\}\.

### 4\.4Results

![Refer to caption](https://arxiv.org/html/2607.06976v1/pics/4_new/Loss_Error_1_HHH_SSS_1.png)Figure 2:Solving various PDE problems via MIONet: Adam\-only \(Red\) vs\. ALS\+Adam \(Blue\)\.The top row denotes the training loss \(without regularization terms in[4\.1](https://arxiv.org/html/2607.06976#S4.E1)and[4\.2](https://arxiv.org/html/2607.06976#S4.E2)\) for different seeds, and the bottom row shows the mean relativeL2L^\{2\}error for the validation data over time by the seeds\. The shaded areas represent the maximum and minimum values among the seeds, and the solid lines represent their averages\. The plots are drawn on log\-log scales\.![Refer to caption](https://arxiv.org/html/2607.06976v1/pics/4_new/Data_Test_Eval_MIONet_HHH_SSS_1_Vert.png)Figure 3:Test data evaluation for trained MIONet: Adam\-only vs\. ALS\+Adam\.From the top, the reaction\-diffusion, advection, and Poisson examples are illustrated\. The models are evaluated with parameters trained for 10,000 seconds\.As shown in[Figure2](https://arxiv.org/html/2607.06976#S4.F2), training with ALS\+Adam outperforms classical Adam training in terms of training loss decay and model performance \(relativeL2L^\{2\}error for unseen data\) in both supervised and unsupervised learning\.

[Figure3](https://arxiv.org/html/2607.06976#S4.F3)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 asttincreases\. In the advection example, the errors tend to form lines parallel to the linex−0\.5​t=0x\-0\.5t=0\. For the Poisson’s equation, the errors are concentrated near the boundary\.

## 5Conclusion

In this paper, we propose a novel method to improve the training of vanilla MIONets, which generalizes the LSGD method for vanilla DeepONets\[Choi2025\]\. By interpreting the generalL2L^\{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