Convergence analysis of wide shallow neural operators within the framework of Neural Tangent Kernel

Partial Differential Equations (PDEs) are essential for modeling a wide range of phenomena in physics, biology, and engineering. Nonetheless, the numerical solution of PDEs has always been a significant challenge in the field of scientific computation. Traditional numerical approaches, such as finite difference, finite element, finite volume, and spectral methods, can encounter difficulties due to the curse of dimensionality when applied to PDEs with a high number of dimensions. In recent years, the impressive achievements of deep learning in various domains, including computer vision, natural language processing, and reinforcement learning, have led to an increased interest in utilizing machine learning techniques to tackle PDE-related problems.

For scientific problems, neural network-based methods are primarily divided into two categories: neural solvers and neural operators. Neural solvers, such as PINNs [1], DRM [2], utilize neural networks to represent the solutions of PDEs, minimizing some form of residual to enable the neural networks to approximate the true solutions closely. There are two potential advantages of neural solvers. First, this is an unsupervised learning approach, which means it does not require the costly process of obtaining a large number of labels as in supervised learning. Second, as a powerful represention tool, neural networks are known to be effictive for approximating continuous functions [3], smooth functions [4], Sobolev functions [5]. This presents a potentially viable avenue for addressing the challenges of high-dimensional PDEs. Nevertheless, existing neural solvers face several limitations when compared to classical numerical solvers like FEM and FVM, particularly in terms of accuracy, and convergence issues. In addition, neural solvers are typically limited to solving a fixed PDE. If certain parameters of the PDE change, it becomes necessary to retrain the neural network.

Neural operator (also called operator learning) aims to approximate unknown operator, which often takes the form of the solution operator associated with a differential equation. Unlike most supervised learning methods in the field of machine learning, where both inputs and outputs are of finite dimensions, operator learning can be regarded as a form of supervised learning in function spaces. Because the inputs and outputs of neural networks are of finite dimensions, some operator learning methods, such as PCA-net and DeepONet, use an encoder to convert infinite-dimensional inputs into finite-dimensional ones, and a decoder to convert finite-dimensional outputs back into infinite-dimensional outputs. The PCA-Net architecture was proposed as an operator learning framework in [6], where principal component analysis (PCA) is employed to obtain data-driven encoders and decoders, combining with a neural network mapping between the finite-dimensional latent spaces. Building on early work by [7], DeepONet [8] consists of a deep neural network for encoding the discrete input function space and another deep neural network for encoding the domain of the output functions. The encoding network is conventionally referred to as the “branch-net”, while the decoding network is referred to as the “trunk-net”. In contrast to the PCA-Net and DeepONet architectures mentioned above,, Fourier Neural Operator (FNO), introduced in [9], does not follow the encoder-decoder-net paradigm. Instead, FNO is a composition of linear integral operators and nonlinear activation functions, which can be seen as a generalization the structure of finite-dimensional neural networks to a function space setting.

The theoretical research on neural operators mostly focuses on the study of approximation errors and generalization errors. As we know, the theoretical basis for the application of neural networks lies in the fact that neural networks are universal approximators. This also holds for neural operators. Regarding the analysis of approximation errors in neural operators, the aim is to identify whether neural operator also possess a universal approximation property, i.e. the ability to approximate a wide class of operators to any given accuracy. As shown in [7], (shallow) operator networks can approximate continuous operators mapping between spaces of continuous functions with arbitrary accuracy. Building on this result, DeepONets have also been proven to be universal approximators. For neural operators following the encoder-decoder paradigm like DeepONet and PCA-Net, Lemma 22 in [10] provides a consistent approximation result, which states that if two Banach spaces have the approximation property, then continuous maps between them can be approximated in a finite-dimensional manner. The universal approximation capability of the FNO was initially established in [11], drawing on concepts from Fourier analysis, and specifically leveraging the density of Fourier series to demonstrate the FNO’s ability to approximate a broad spectrum of operators. For a more quantitative analysis of the approximation error of neural operators, see [12]. In addition to approximation errors, the error analysis of encoder-decoder style neural operators also includes encoding and reconstruction errors. [13] has provided both lower and upper bounds on the total error for DeepONets by using the spectral decay properties of the covariance operators associated with the underlying measures. By employing tools from non-parametric regression, [14] has provided an analysis of the generalization error for neural operators with basis encoders and decoders. The results in [14] holds on neural operators with some popular encoders and decoders, such as those using Legendre polynomials, trigonometric functions, and PCA. For more details on the recent advances and theoretical research in operator learning, refer to the review [12].

Up to this point, the theoretical exploration of the convergence and optimization aspects of neural operators has received relatively little attention. To our best knowledge, only [15] and [16] have touched upon the optimization of neural operators. Based on restricted strong convexity (RSC), [15] has presented a unified framework for gradient descent and apply the framework to DeepONets and FNOs, establishing convergence guarantees for both. [16] has briefly analyzed the training of physics-informed DeepONets and derived a weighting scheme guided by NTK theory to balance the data and the PDE residual terms in the loss function. In this paper, we focus on the training error of shallow neural operator in [7] with the framework of NTK, showing that gradient descent converges at a global linear rate to the global optimum.

The neural operator considered in this paper was originally introduced in [4], aimining to approximate a non-linear operator. Specifically, suppose that $\sigma$ is a continuous and non-polynomial function, $X$ is a Banach space, $K_{1}\subset X$, $K_{2}\subset\mathbb{R}^{d}$ are two compact sets in $X$ and $\mathbb{R}^{d}$, respectively, $V$ is a compact set in $C(K_{1})$. Assume that $G^{*}:V\rightarrow C(K_{2})$ is a nonlinear continuous operator. Then, an operator net can be formulated in terms of two shallow neural networks. The first is the so-called branch net $\beta(u)=(\beta_{1}(u),\cdots,\beta_{m}(u))$, defined for $1\leq r\leq m$ as

where $\{x_{s}\}_{1\leq s\leq q}\subset K_{1}$ are the so-called sensors and $a_{rk},\xi_{rk}^{s},\theta_{rk}$ are weights of the neural network.

The second neural network is the so-called trunk net $\tau(y)=(\tau_{1}(y),\cdots,\tau_{m}(y))$, defined as

where $y\in K_{2}$ and $w_{r},\zeta_{r}$ are weights of the neural network. Then the branch net and trunk net are combined to approximate the non-linear operator $G^{*}$, i.e.,

As shown in [4], (shallow) operator networks can approximate, to arbitrary accuracy, continuous operators mapping between spaces of continuous functions. Specifically, for any $\epsilon>0$, there are positive integers $p,m$ and $q$, constants $a_{rk},\xi_{rk}^{s},\theta_{rk},\zeta_{r}\in\mathbb{R}$, $w_{r}\in\mathbb{R}^{d}$, $x_{s}\in K_{1},s=1,\cdots q$, $r=1,\cdots,m$ and $k=1,\cdots,m$, such that

holds for all $u\in V$ and $y\in K_{2}$.

The training of neural networks is performed using a supervised learning process. It involves minimizing the mean-squared error between the predicted output $G(u)(y)$ and the actual output $G^{*}(u)(y)$. Specifically, assume we have samples $\{(u_{i},G^{*}(u_{i}))\}_{i=1}^{N},u_{i}\sim\mu$, where $\mu$ is a probability supported on $V$. The aim is to minimize the following loss function:

In this paper, we primarily focus on the shallow neural operators with ReLU activation functions. Formally, we consider a shallow operator of the following form.

where we equate the function $u$ with its value vector $(u(x_{1}),\cdots,u(x_{q}))^{T}\in\mathbb{R}^{q}$ at the points $\{x_{s}\}_{s=1}^{q}$.

We denote the loss function by $L(W,\tilde{W},a)$. The main focus of this paper is to analyze the gradient descent in training the shallw neural operator. We fix the weights $\{a_{rk}\}_{r=1,k=1}^{m,p}$ and apply gradient (GD) to optimize the weights $\{w_{r}\}_{r=1}^{m},\{\tilde{w}_{rk}\}_{r=1,k=1}^{m,p}$. Specifically,

where $\eta>0$ is the learning rate and $L(W,\tilde{W})$ is an abbreviation of $L(W,\tilde{W},a)$.

At this point, the loss function is

where $z_{j}^{i}=G^{*}(u_{i})(y_{j})$. Throughout this paper, we consider the initialization

and assume that $\|u_{i}\|_{2}=\mathcal{O}(1),\|y_{j}\|_{2}=\mathcal{O}(1)$ for all $j\in[n]$. Note that here we treat vector $u_{i}=(u_{i}(x_{1}),\cdots,u_{x_{q}})$ and function $u_{i}$ as equivalent.

In this section, we present our result for gradient flow, which can be viewed as a continuous form of gradient descent with an infinitesimal time step size. The analysis of gradient flow in continuous time serves as a foundational step for comprehending discrete gradient descent algorithms. We prove that the gradient flow converges to the global optima of the loss under over-parameterization and some mild conditions on training samples. The time continuous form can be characterized as the following dynamics

for $r\in[m],k\in[p]$. We denote $G^{t}(u_{i})(y_{j})$ the prediction on $y_{j}$ at time $t$ under $u_{i}$, i.e., with weights $w_{r}(t),\tilde{w}_{rk}(t)$.

Thus, we can deduce that

and

Then, the dynamics of each prediction can be calculated as follows.

We let $G^{t}(u_{i}):=(G^{t}(u_{i})(y_{1}),\cdots,G^{t}(u_{i})(y_{n_{2}}))\in\mathbb{R}^{n_{2}}$ and $G^{t}(u):=(G^{t}(u_{1}),\cdots,G^{t}(u_{n_{1}}))\in\mathbb{R}^{n_{1}n_{2}}$. Then, we have

where $H_{j}^{i}(t)\in\mathbb{R}^{n_{2}\times n_{2}}$, whose $(i_{1},j_{1})$-th entry is defined as

and $\tilde{H}_{j}^{i}(t)\in\mathbb{R}^{n_{2}\times n_{2}}$, whose $(i_{1},j_{1})$-th entry is defined as

Thus, we can write the dynamics of predictions as follows.

where $H(t),\tilde{H}(t)\in\mathbb{R}^{n_{1}n_{2}\times n_{2}n_{2}}$. We can divide $H(t),\tilde{H}(t)$ into $n_{1}\times n_{1}$ blocks, the $(i,j)$-th block of $H(t)$ is $H_{j}^{i}(t)$ and the $(i,j)$-th block of $\tilde{H}(t)$ is $\tilde{H}_{j}^{i}(t)$. From the form of $H(t),\tilde{H}(t)$, we can derive the Gram matrices induced by the random initialization, which we denote by $H^{\infty}$ and $\tilde{H}^{\infty}$. Note that although $H^{\infty}$ and $\tilde{H}^{\infty}$ are large matrices, we can divide $H^{\infty}$ and $\tilde{H}^{\infty}$ into $n_{1}\times n_{1}$ blocks, where each block is a $n_{2}\times n_{2}$ matrix. Following the notation above, the $(i_{1},j_{1})$-th entry of $(i,j)$-th block of $H^{\infty}$ is

Thus, the $(i,j)$-th block can be written as

where $H_{2}^{\infty}\in\mathbb{R}^{n_{2}\times n_{2}}$ and the $(i_{1},j_{1})$-th entry of $H_{2}^{\infty}$ is

Thus, $H^{\infty}$ can be seen as a Kronecker product of matrices $H_{1}^{\infty}$ and $H_{2}^{\infty}$, where $H_{1}^{\infty}\in\mathbb{R}^{n_{1}\times n_{1}}$ and the $(i,j)$-th entry of $H_{1}^{\infty}$ is $\mathbb{E}[\sigma(\tilde{w}^{T}u_{i})\sigma(\tilde{w}^{T}u_{j})]$. Similarly, we have that $\tilde{H}^{\infty}$ is a Kronecker product of matrices $\tilde{H}_{1}^{\infty}$ and $\tilde{H}_{2}^{\infty}$, where $\tilde{H}_{1}^{\infty}\in\mathbb{R}^{n_{1}\times n_{1}}$ and the $(i,j)$-th entry of $\tilde{H}_{1}^{\infty}$ is $\mathbb{E}[u_{i}^{T}u_{j}I\{\tilde{w}^{T}u_{i}\geq 0,\tilde{w}^{T}u_{j}\geq 0\}]$, the $(i_{1},j_{1})$-th entry of $\tilde{H}_{2}^{\infty}$ is $\mathbb{E}[\sigma(w^{T}y_{i_{1}})\sigma(w^{T}y_{j_{1}})]$ .

Similar to the situation in $L^{2}$ regression, we can show two essential facts: (1) $\|H(0)-H^{\infty}\|_{2}=\mathcal{O}(1/\sqrt{m})$, $\|\tilde{H}(0)-\tilde{H}^{\infty}\|_{2}=\mathcal{O}(1/\sqrt{m})$ and (2) for all $t\geq 0$, $\|H(t)-H(0)\|_{2}=\mathcal{O}(1/\sqrt{m})$, $\|\tilde{H}(t)-\tilde{H}(0)\|_{2}=\mathcal{O}(1/\sqrt{m})$.

Therefore, roughly speaking, as $m\to\infty$, the dynamics of the predictions can be written as

which results in the linear convergence.

We first show that the Gram matrices are strictly positive definite under mild assumptions.

Then we can verify the two facts that $H(0),\tilde{H}(0)$ are close to $H^{\infty},\tilde{H}^{\infty}$ and $H(t),\tilde{H}(t)$ are close to $H(0),\tilde{H}(0)$ by following two lemmas.

With these preparations, we come to the final conclusion.

Proof Sketch: Note that

thus if $\lambda_{min}(H(t))\geq\lambda_{0}/2$ and $\lambda_{min}(\tilde{H}(t))\geq\tilde{\lambda}_{0}/2$, we have

This yields that $\frac{d}{dt}\left(exp((\lambda_{0}+\tilde{\lambda}_{0})t)\|z-G^{t}(u)\|_{2}^{2}\right)\leq 0$, i.e., $exp((\lambda_{0}+\tilde{\lambda}_{0})t)\|z-G^{t}(u)\|_{2}^{2}$ is non-increasing, thus we have

On the other hand, roughly speaking, the continous dynamics of $w_{r}(t)$ and $\tilde{w}_{rk}(t)$, i.e., (3) and (4), show that

Thus if the prediction decays like (10), we can deduce that

Combining this with the stability of the descrete Gram matrices, i.e., Lemma 3, we have $\|H(t)-H^{\infty}\|_{2}\leq\lambda_{0}/4$, $\|\tilde{H}(t)-\tilde{H}^{\infty}\|_{2}\leq\tilde{\lambda}_{0}/4$ and $\lambda_{min}(H(t))\geq\lambda_{0}/2$, $\lambda_{min}(\tilde{H}(t))\geq\tilde{\lambda}_{0}/2$, when $m$ is sufficiently large.

From such equivalence, we can arrive at the desired conclusion.

In this section, we are going to demonstrate that randomly initialized gradient in training shallow neural operators converges to the golbal minimum at a linear rate. Unlike the continuous time case, the discrete time case requires a more refined analysis. In the following, we first present our main result and then outline the proof’s approach.

For the $L^{2}$ regression problems, [17] has demonstrated that if the learning rate $\eta=\mathcal{O}(\lambda_{0}/n^{2})$, then randomly initialized gradient descent converges to a globally optimal solution at a linear convergence rate when $m$ is large enough. The requirement of $\eta$ is derived from the decomposition for the residual in the $(k+1)$-th iteration, i.e.,

where $u(t)$ is the prediction vector at $t$-iteration under the shallow neural network and $y$ is the true prediction. Although using the method from [17] can also yield linear convergence of gradient descent in training shallow neural operators, the requirements for the learning rate $\eta$ would be very stringent due to the dependency on $\lambda_{0},\tilde{\lambda}_{0}$ and $n$. Thus, instead of decomposing the residual into the two terms as above, we write it as follows, which serves as a recursion formula.

Just as in the case of $L^{2}$ regression, we prove our conclusion by induction. From the recursive formula above, it can be seen that both the estimation of $H(t)$ and $\tilde{H}(t)$, as well as the estimation of the residual $I(t)$, depend on $\|w_{r}(t)-w_{r}(0)\|_{2}$ and $\|\tilde{w}_{rk}(t)-\tilde{w}_{rk}(0)\|_{2}$. Therefore, our inductive hypothesis is the following differences between weights and their initializations.

This condition can lead to the linear convergence of gradient descent, i.e., result in Theorem 2.