Hybrid Least Squares/Gradient Descent Methods for MIONets
This work addresses the slow training of MIONets, a type of neural operator, by providing a more efficient optimization method that is compatible with common loss functions.
The paper proposes a hybrid least squares/gradient descent (LSGD) method for MIONets that accelerates training by using alternating least squares for last-layer parameters, achieving faster convergence than standard gradient descent.
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 of $L^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.