TransNet: Transferable Neural Networks for Partial Differential
  Equations
Transfer learning for partial differential equations (PDEs) is to develop a
pre-trained neural network that can be used to solve a wide class of PDEs.
Existing transfer learning approaches require much information of the target
PDEs such as its formulation and/or data of its solution for pre-training. In
this work, we propose to construct transferable neural feature spaces from
purely function approximation perspectives without using PDE information. The
construction of the feature space involves re-parameterization of the hidden
neurons and uses auxiliary functions to tune the resulting feature space.
Theoretical analysis shows the high quality of the produced feature space,
i.e., uniformly distributed neurons. Extensive numerical experiments verify the
outstanding performance of our method, including significantly improved
transferability, e.g., using the same feature space for various PDEs with
different domains and boundary conditions, and the superior accuracy, e.g.,
several orders of magnitude smaller mean squared error than the state of the
art methods.
Zezhong Zhang, Feng Bao, Lili Ju, Guannan Zhang
# TransNet: Transferable Neural Networks for Partial Differential Equations

###### Abstract

Transfer learning for partial differential equations (PDEs) is to develop a pre-trained neural network that can be used to solve a wide class of PDEs. Existing transfer learning approaches require much information of the target PDEs such as its formulation and/or data of its solution for pre-training. In this work, we propose to construct transferable neural feature spaces from purely function approximation perspectives without using PDE information. The construction of the feature space involves re-parameterization of the hidden neurons and uses auxiliary functions to tune the resulting feature space. Theoretical analysis shows the high quality of the produced feature space, i.e., uniformly distributed neurons. Extensive numerical experiments verify the outstanding performance of our method, including significantly improved transferability, e.g., using the same feature space for various PDEs with different domains and boundary conditions, and the superior accuracy, e.g., several orders of magnitude smaller mean squared error than the state of the art methods.

Machine Learning, Transferable Neural Networks, Partial Differential Equations

## 1 Introduction

Rapid advancement of deep learning has attracted significant attention of researchers to explore how to use deep learning to solve scientific and engineering problems. Since numerical solutions of partial differential equations (PDEs) sits at the heart of many scientific areas, there is a surge of studies on how to use neural networks to leverage data and physical knowledge to solve PDEs (Raissi et al., 2019; E and Yu, 2018; Long et al., 2018; Zang et al., 2020; Li et al., 2021; Lu et al., 2021; Gin et al., 2021; Zhang et al., 2021; Teng et al., 2022; Clark Di Leoni et al., 2023). The neural network-based methods have several advantages over traditional numerical methods (e.g., finite element, finite difference and finite volume), such as avoiding the need for numerical integration, generating differentiable solutions, exploiting advanced computing capabilities, e.g., GPUs. Nevertheless, a major drawback of these deep learning methods for solving PDEs is high computational cost associated with the neural network training/retraining using stochastic gradient descent (SGD). One of the popular strategies to alleviate this issue is transfer learning.

Transfer learning for PDEs is to develop a pre-trained neural network that can be effectively re-used to solve a PDE with multiple coefficients or in various domains, or to solve multiple types of PDEs. When transfer a pre-trained neural network from one scenario to another, the feature space, e.g., the hidden layers, are often frozen or slightly perturbed, which can dramatically reduce the training overhead by orders of magnitude. However, existing transfer learning approaches for PDEs, e.g., (Lu et al., 2021; Li et al., 2021; Chakraborty, 2020; Desai et al., 2021), require information/knowledge of the target family of PDEs to pre-train a neural network model. The needed information could be the analytical definitions of the PDEs including initial and boundary conditions, and/or measurement data of the PDE's solution. These requirement not only leads to time-consuming simulation data generation using other PDE solvers, but also limits the transferability of the pre-trained neural network (i.e., the pre-trained network is only transferable to the same or similar type of PDEs that are used for pre-training).

To overcome the above challenges, in this paper we propose a transferable neural network (TransNet) to improve the transferability of neural networks for solving PDEs. The key idea is construct a pre-trained neural feature space without using any PDE information, so that the pre-trained feature space could be transferred to a variety of PDEs with different
domains and boundary conditions. We limit our attention to single-hidden-layer fully-connected neural networks, which have sufficient expressive power for low-dimensional PDEs that are commonly used in science and engineering fields. Specifically, we treat each hidden neuron as a basis function and re-parameterize all the neurons to separate the parameters that determine the neuron's location and the ones that control the shape (i.e., the slope) of the activation function. Then, we develop a simple, yet very effective, approach to generate uniformly distributed neurons in the unit ball, and rigorously prove the uniform neuron distribution. Then, the shape parameters of the neurons are tuned using auxiliary functions, i.e., realizations of a Gaussian process. The entire feature space construction (determining the hidden neurons' parameters) does not require the PDE's formulation or data of the PDE's solution. When applying the constructed feature space to a PDE problem, we only need to solve for the parameters of the output layer by minimizing the standard PDE residual loss. This can be done by either solving a simple least squares problem for linear PDE or combining a least squares solver with a nonlinear iterative solver, e.g., Pichard iteration, for nonlinear PDEs.

The major contributions of this work are summarized as

* We develop transferable neural feature spaces that are independent of any PDE, and can be applied to effectively solve various linear and nonlinear PDE problems.
* We theoretically and computationally prove the uniform distribution of the hidden neurons, viewed as global non-orthogonal basis, for the proposed TransNet in the unit ball of any dimension.
* We demonstrate the superior accuracy and efficiency of the proposed TransNet for solving PDEs, e.g., the mean square errors of TransNet are several orders of magnitudes smaller than those by the state-of-the-art methods.

## 2 Related work

Studies on using neural networks for solving PDEs can be traced back to some early works, e.g., (Dissanayake and Phan-Thien, 1994; Lagaris et al., 1998). Recent advances mostly have been focused on physics-informed neural network (PINN). The general idea of PINN is to represent the PDE's solution by a neural network, and then train the network by minimizing certain measurement of the PDE's residual at a set of samples in the domain of computation. Several improvements on the training and sampling were proposed in (Lu et al., 2021; Anitescu et al., 2019; Zhao and Wright, 2021; Krishapriyan et al., 2021). Besides direct minimizing the PDE's residual, there are studies on how to combine traditional PDE solvers with neural networks. For example, the deep Ritz method (E and Yu, 2018) uses the variational form of PDEs and combines the stochastic gradient descent with numerical integration to train the network; the deep Galerkin method (Sirignano and Spiliopoulos, 2018) combines the Galerkin method with machine learning; the PDE-Net (Long et al., 2018, 2019) uses a stack of neural networks to approximate the PDE solutions over a multiple of time steps.

Another type of deep learning method for PDEs is to use neural networks to learn a family of PDE operators, instead of a single equation. The Fourier neural operator (FNO) (Li et al., 2021) parameterizes the integral kernel in Fourier space and is generalizable to different spatial/time resolutions. The DeepONet (Lu et al., 2021) extends the universal approximation theorem (Chen and Chen, 1995) to deep neural networks, and its variant (Wang et al., 2021) further reduces the amount of data needed for training. The physics-informed neural operator (PINO) (Li et al., 2021) combines operator learning with function approximation to achieve higher accuracy. MIONet (Jin et al., 2022) was proposed to learn multiple-input operators via tensor product based on low-rank approximation.

Random feature models have also been used to solve PDEs (Sun et al., 2018; Liu et al., 2022) or learn PDE operators (Nelsen and Stuart, 2021). The theory of random feature models for function approximation was developed due to its natural connection with kernel methods (Liu et al., 2022; Bach, 2017). The proposed TransNet can be viewed as an improved random feature model for PDEs from two perspectives: (1) the re-parameterization of the hidden neurons to separate the parameters that determine locations of the neurons and the ones that control the activation function slope, (2) the usage of auxiliary functions to tune the neural feature space, which makes a critical contribution to the improvement of the accuracy of TransNet in solving PDEs.

## 3 Transferable neural networks for PDEs

### Problem setting and background

We introduce the problem setup for using neural networks to solve partial differential equations. The PDE of interest can be presented in a general formulation, i.e.,

\[\begin{cases}\mathcal{L}(u(\mathbf{y}))=f(\mathbf{y})&\text{for}\ \mathbf{y}\in\Omega,\\ \mathcal{B}(u(\mathbf{y}))=g(\mathbf{y})&\text{for}\ \mathbf{y}\in\partial\Omega,\end{cases} \tag{1}\]

where \(\Omega\subset\mathbb{R}^{d}\) with the boundary \(\partial\Omega\) is the spatial-temporal bounded domain under consideration, \(\mathbf{y}:=(\mathbf{x},t)=(x_{1},\dots,x_{d-1},t)^{\top}\) is a column vector includes both spatial and temporal variables, \(u\) denotes the unknown solution of the PDE, \(\mathcal{L}(\cdot)\) is a differential operator, \(\mathcal{B}(\cdot)\) is the operator defining the initial and/or boundary conditions, \(f(\mathbf{y})\) and \(g(\mathbf{y})\) are the right hand sides associated with the operators \(\mathcal{L}(\cdot)\) and \(\mathcal{B}(\cdot)\), respectively. For notational simplicity, we assume that the solution is a scalar function; the proposed method can be extended to vector-valued functions without
any essential difficulty. We limit our attention to the single-hidden-layer fully-connected neural networks, denoted by

\[u_{\rm NN}(\mathbf{y}):=\sum_{m=1}^{M}\alpha_{m}\,\sigma(\mathbf{w}_{m}\mathbf{y}+b_{m})+ \alpha_{0}, \tag{2}\]

where \(M\) is the number of hidden neurons, the row vector \(\mathbf{w}_{m}=(w_{m,1},\ldots,w_{m,d})\) and the scalar \(b_{m}\) are the weights and bias of the \(m\)-th hidden neuron, the row vector \(\mathbf{\alpha}=(\alpha_{0},\alpha_{1},\ldots,\alpha_{M})\) includes the weights and bias of the output layer, and \(\sigma(\cdot)\) is the activation function. As demonstrated in Section 4, this type of neural networks have sufficient expressive power for solving a variety of PDEs with satisfactory accuracy.

A typical method (Karniadakis et al., 2021) for solving the PDE in Eq. (1) is to directly parameterize the solution \(u(\mathbf{y})\) as a neural network \(u_{\rm NN}(\mathbf{y})\) in Eq. (2) and optimize the neural network's parameters by minimizing the PDE residual loss, e.g., \(L(\mathbf{y})=\|\mathcal{L}(u(\mathbf{y}))-\mathcal{L}(u_{\rm NN}(\mathbf{y}))\|_{2}+\| \mathcal{B}(u(\mathbf{y}))-\mathcal{B}(u_{\rm NN}(\mathbf{y}))\|_{2}\), at a set of spatial-temporal locations. Despite the good performance of these approaches in solving PDE problems, its main drawback is the _limited transferability_ because of the high computational cost of gradient-based re-training and hyperparameter re-tuning. When there is any change to the operators \(\mathcal{L}(\cdot),\mathcal{B}(\cdot)\), the right-hand-side functions \(f(\mathbf{y}),g(\mathbf{y})\), or the shape of the domain \(\Omega\), the neural network \(u_{\rm NN}(\mathbf{y})\) often needs to be re-trained using gradient-based optimization (even though the current parameter values could provide a good initial guess for the re-training), or the hyperparameters associated with the network and the optimizer need to be re-tuned. In comparison, the random feature models require much lower re-training cost, which has been exploited in learning operators (Nelsen and Stuart, 2021) and dynamical systems (McDonald and Alvarez, 2021; Liu et al., 2022b).

### The neural feature space

We can treat each hidden neuron \(\sigma(\mathbf{w}_{m}\mathbf{y}+b_{m})\) as a non-linear feature map from the space of \(\mathbf{y}\in\mathbb{R}^{d}\) to the output space \(\mathbb{R}\). From the perspective of approximation theory, the set of hidden neurons \(\{\sigma(\mathbf{w}_{m}\mathbf{y}+b_{m})\}_{m=1}^{M}\) can be viewed as a globally supported basis in \(\mathbb{R}^{d}\). The _neural feature space_, denoted by \(\mathcal{P}_{\rm NN}\), can be defined by the linear space expanded by the basis \(\{\sigma(\mathbf{w}_{m}\mathbf{y}+b_{m})\}_{m=1}^{M}\), i.e.,

\[\mathcal{P}_{\rm NN}=span\Big{\{}1,\sigma(\mathbf{w}_{1}\mathbf{y}+b_{1}),\ldots, \sigma(\mathbf{w}_{M}\mathbf{y}+b_{M})\Big{\}}, \tag{3}\]

where the constant basis corresponds to the bias of the output layer. Then, the neural network in Eq. (2) lives in the linear space, i.e., \(u_{\rm NN}(\mathbf{y})\in\mathcal{P}_{\rm NN}\). In other words, the neural network approximation can be viewed as a spectral method with _non-orthogonal_ basis, and the parameters \(\mathbf{\alpha}\) in Eq. (2) of the output layer of \(u_{\rm NN}(\mathbf{y})\) contains the coefficients of the expansion in the neural feature space \(\mathcal{P}_{\rm NN}\).

In the PINN methods, the neural feature space \(\mathcal{P}_{\rm NN}\) and the coefficient \(\mathbf{\alpha}\) are trained simultaneously using stochastic gradient descent methods, which often leads to a non-convex and ill-conditioned optimization problem. It has been shown that the non-convexity and ill-conditioning in the neural network training are major reasons of unsatisfactory accuracy of the trained neural network. A natural idea to reduce the complexity of the training is to decouple the training of \(\mathcal{P}_{\rm NN}\) from that of \(\mathbf{\alpha}\). For example, in random feature models, \(\mathcal{P}_{\rm NN}\) is defined by randomly generating the parameters\(\{(\mathbf{w}_{m},b_{m})\}_{m=1}^{M}\) from a user-defined probability distribution; the coefficients \(\mathbf{\alpha}\) can then be obtained by solving a linear system when the operators \(\mathcal{L}\), \(\mathcal{B}\) in Eq. (1) are linear. However, the numerical experiments in Section 4 show that the random feature model based on Eq. (2) converges very slowly with the increase of the number of features. This drawback motivates us to develop a methodology to customize the neural feature space \(\mathcal{P}_{\rm NN}\) to improve the accuracy, efficiency and transferability of \(u_{\rm NN}\) in solving PDEs.

### Constructing the transferable neural feature space

This section contains the key ingredients of the proposed TransNet. The goal is to construct a single neural feature space \(\mathcal{P}_{\rm NN}\) that can be used to solve various PDEs in different domains.

#### 3.3.1 Re-parameterization of \(\mathcal{P}_{\rm NN}\)

The first step is to re-parameterize the hidden neuron \(\sigma(\mathbf{w}_{m}\mathbf{y}+b_{m})\), viewed as a basis function in \(\Omega\), to separate the components that determine the _location_ of the neuron and the components that control the _shape_ of the neuron.

The idea of handling the locations of the basis functions is inspired by the studies on activation patterns of ReLU networks. When \(\sigma\) is the ReLU function, there is a _partition hyperplane_ defined by

\[w_{m,1}y_{1}+w_{m,2}y_{2}+\cdots+w_{m,d}y_{d}+b_{m}=0 \tag{4}\]

that separates the activated and inactivated regions for this neuron. The intersections of multiple partition hyperplanes associated with different neurons define a linear region of ReLU network. Studies have shown that the expressive power of a ReLU network is determined by the number of linear regions and the distribution of those linear regions. In principle, the more _uniformly distributed_ linear regions in the domain \(\Omega\), the more expressive power the ReLU network has. For other activation functions, e.g., \(tanh(\cdot)\) that is widely used in solving PDEs due to its smoothness, the partition hyperplane in Eq. (4) can be used to describe the geometric property of the neuron.

Specifically, let us re-write Eq. (4) into the following point
slope form:

\[\gamma_{m}\big{(}a_{m,1}(y_{1}-r_{m}a_{m,1})+\cdots+a_{m,d}(y_{d}-r_{m}a_{m,d}) \big{)}=0, \tag{5}\]

where \(\mathbf{a}_{m}=(a_{m,1},\ldots,a_{m,d})\) is a unit vector, i.e., \(\|\mathbf{a}_{m}\|_{2}=1\), \(r_{m}>0\) and \(\gamma_{m}\in\mathbb{R}\) are two scalar parameters for the \(m\)-th neuron. We can relate Eq. (5) to Eq. (4) by

\[\left\{\begin{aligned} w_{m,i}&=\gamma_{m}a_{m,i}, \quad i=1,\cdots,d,\\ b_{m}&=-\gamma_{m}\sum_{i=1}^{d}a_{m,i}^{2}r_{m}, \end{aligned}\right. \tag{6}\]

which shows the desired geometric properties of the partition hyperplane in Eq. (4). In terms of the location, the unit vector \(\mathbf{a}_{m}\) is the normal direction of the partition hyperplane in \(\mathbb{R}^{d}\), the vector \((r_{m}a_{m,1},\ldots,r_{m}a_{m,d})\) indicates a point that the hyperplane passes, \(r_{m}\) is the distance between the origin and the partition hyperplane. An illustration is shown in Figure 1(a). In terms of the shape, the constant \(\gamma_{m}\) determines the steepness of the slope of the activation function along the normal direction \(\mathbf{a}_{m}\). Thus, the re-parameterization in Eq. (5) successfully separates the parameters determining location from the ones determining the shape.

#### 3.3.2 Generating uniformly distributed neurons for \(\mathcal{P}_{\mathrm{NN}}\)

The second step of constructing \(\mathcal{P}_{\mathrm{NN}}\) is to determine the parameters \(\{(\mathbf{a}_{m},r_{m})\}_{m=1}^{M}\) in Eq. (5), such that all the neurons are uniformly distributed in \(\Omega\). We assume \(\Omega\) is a _unit ball_, i.e., \(B_{1}(\mathbf{0})=\{\mathbf{y}:\|\mathbf{y}\|_{2}\leq 1\}\subset\mathbb{R}^{d}\) in this subsection. To proceed, we need to define a density function that measures the neuron distribution. For a given \(\mathbf{y}\in\Omega\), the distance between \(\mathbf{y}\) and the partition hyperplane in Eq. (5) is given by

\[dist(\mathbf{y},m)=|\mathbf{a}_{m}(\mathbf{y}-r_{m}\mathbf{a}_{m})|, \tag{7}\]

for \(m=1,\ldots,M\). We use this distance to define how close the point \(\mathbf{y}\) to the \(m\)-th neuron. The density function, denoted by \(D_{M}(\mathbf{y})\), is defined using the above distance, i.e.,

\[D_{M}(\mathbf{y})=\frac{1}{M}\sum_{m=1}^{M}\mathbf{1}_{dist(\mathbf{y},m)<\tau}(\mathbf{y}), \tag{8}\]

where \(\mathbf{1}_{dist(\mathbf{y},m)<\tau}(\mathbf{y})\) is the indicator function of the event that the distance between \(\mathbf{y}\) and the \(m\)-th neuron is smaller than a prescribed tolerance \(\tau>0\). Intuitively, \(D_{M}(\mathbf{y})\) measures the percentage of neurons whose partition hyperplane in Eq. (4) intersect the ball (with radius \(\tau\)) around \(\mathbf{y}\).

Next we propose the following approach, illustrated in Figure 1(b), to generate the parameters \(\{(\mathbf{a}_{m},r_{m})\}_{m=1}^{M}\). Specifically, we first generate the normal directions \(\{\mathbf{a}_{m}\}_{m=1}^{M}\) uniformly distributed on the \(d-1\)-dimensional unit sphere. Note that when \(d>2\), sampling uniformly in the angular space in the hyperspherical coordinate system does not lead to uniformly distributed samples on the unit sphere. This is known as the sphere point picking problem. To overcome this issue, we draw samples from the \(d\)-dimensional Gaussian distribution in the Cartesian coordinate system, and normalize the samples to unit vectors to obtain \(\{\mathbf{a}_{m}\}_{m=1}^{M}\). Then, we generate \(\{r_{m}\}_{m=1}^{M}\) uniformly from \([0,1]\) using the Monte Carlo method. The following theorem shows that our approach provides a set of uniformly distributed neurons in \(\Omega\), where the density is measured by \(D_{M}(\mathbf{y})\) in Eq. (8).

**Theorem 1** (Uniform neuron distribution): _Given the re-parameterization in Eq. (5), if \(\{\mathbf{a}_{m}\}_{m=1}^{M}\) are uniformly distributed random vectors on the \(d\)-dimensional unit sphere, i.e., \(\|\mathbf{a}_{m}\|_{2}=1\), and \(\{r_{m}\}_{m=1}^{M}\) are uniformly distributed random variables in \([0,1]\), then, for a fixed \(\tau\in(0,1)\),_

\[\mathbb{E}[D_{M}(\mathbf{y})]=\tau\ \text{ for any }\|\mathbf{y}\|_{2}\leq 1-\tau,\]

_where \(D_{M}(\mathbf{y})\) is the density function defined in Eq. (8)._

The proof is given in Appendix A; an illustration of the density function is given in Figure 1(c). This result is a little surprising that the distribution of \(\{r_{m}\mathbf{a}_{m}\}_{m=1}^{M}\), i.e., the red dots in Figure 1(b)-middle, are not uniformly distributed in the ball \(B_{1-\tau}(\mathbf{0})\), but the density function \(D_{M}(\mathbf{y})\) is a constant in the ball \(B_{1-\tau}(\mathbf{0})\).

**Remark 1** (The dimentionality): _Even though Theorem 1 holds for any dimension \(d\), the number of neurons required to cover a high-dimensional unit ball still could be intractable. On the other hand, the majority of PDEs commonly used in science and engineering are defined in low-dimensional domains, e.g., 3D spatial domain + 1D time domain. In this scenario, the proposed method is effective and easy to implement, as demonstrated in Section 4._

#### 3.3.3 Tuning the shape of the neurons in \(\mathcal{P}_{\mathrm{NN}}\) using auxiliary functions

The third step is to tune the shape parameters \(\{\gamma_{m}\}_{m=1}^{M}\) in Eq. (5) that controls the slope of the activation function. The experimental tests in Section 4.1 show that the slope parameters play a critical role in determining the accuracy of the neural network approximator \(u_{\mathrm{NN}}\). For simplicity, we assume the same shape parameter value for all neurons, i.e., \(\gamma=\gamma_{m}\) for \(m=1,\ldots,M\). Because we intend to construct a feature space \(\mathcal{P}_{\mathrm{NN}}\) that can be used in multiple scenarios, e.g., various PDEs with different domains and boundary conditions, we do not want to tune the shape parameter \(\gamma\) using any information about a specific PDE.

Our idea is to use auxiliary functions that have similar or more complicated spatial-temporal variation frequency as
the PDE solution to tune \(\gamma\). Specifically, we propose to use realizations of Gaussian processes to generate the auxiliary functions. The advantage of Gaussian process is that one can control the variation frequency of its realizations by adjusting the correlation length. Additionally, the Guassian process is independent of the coordinate system. Let us denote by \(G(\mathbf{y}|\omega,\eta)\) the Gaussian process, where \(\omega\) represents the abstract random variable and \(\eta\) is the correlation length. Given a correlation length, we first generate a set of realizations of the Gaussian process, denoted by \(\{G(\mathbf{y}|\omega_{k},\eta)\}_{k=1}^{K}\). For each realization, define the MSE loss as

\[\begin{split}\text{MSE}(u_{\mathrm{NN}}(\mathbf{y}),G(\mathbf{y}|\omega _{k},\eta))\\ =&\frac{1}{J}\sum_{j=1}^{J}\left[\sum_{m=1}^{M} \alpha_{m}\sigma(\mathbf{w}_{m}\mathbf{y}_{j}+b_{m})+\alpha_{0}-G(\mathbf{y}_{j}|w_{k}, \eta)\right]^{2},\end{split} \tag{9}\]

where the parameters \(\{\mathbf{w}_{m}\}_{m=1}^{M}\) and \(\{b_{m}\}_{m=1}^{M}\) are already determined using the strategy in Section 3.3.2 and Eq. (6), and \(J\) denotes the number of sample points. Unlike standard neural network training, the optimal coefficient \(\mathbf{\alpha}\) that minimizing the MSE loss can be efficiently achieved by solving the least squares problem. Hence, the shape parameter \(\gamma\) can be tuned by solving the following one-dimensional optimization problem

\[\min_{\gamma}\left\{\sum_{k=1}^{K}\min_{\mathbf{\alpha}}\left[\text{MSE}(u_{ \mathrm{NN}}(\mathbf{y}),G(\mathbf{y}|\omega_{k},\eta))\right]\right\}, \tag{10}\]

where for each candidate \(\gamma\), we solve \(K\) least squares problems to compute the total loss.

**Remark 2** (The choice of the correlation length): _There are two strategies to choose the correlation length \(\eta\). One is to use the prior knowledge about the PDE. For example, for the Naveier-Stokes equations with low Reynolds' number, we know the solution will not have very high-frequency oscillation. The other is to use an over-killing correlation length to ensure that the feature space has sufficient expressive power to solve the target PDE._

### Applying TransNet to linear and nonlinear PDEs

Once the neural feature space \(\mathcal{P}_{\mathrm{NN}}\) is constructed and tuned, we can readily use it to solve PDE problems. Even though \(\mathcal{P}_{\mathrm{NN}}\) is defined on the unit ball, i.e., \(B_{1}(\mathbf{0})\), we can always place the (bounded) domain \(\Omega\) for the target PDE in \(B_{1}(\mathbf{0})\) by simple translation and dilation. Thus, the feature space can be used to handle PDEs defined in various domains, as demonstrated in Section 4.

**Linear PDEs.** When \(\mathcal{L}\) and \(\mathcal{B}\) in Eq. (1) are linear operators, the unknown parameters \(\mathbf{\alpha}=(\alpha_{0},\ldots,\alpha_{M})\) in Eq. (2) can be easily determined by solving the following least squares problem, i.e.,

\[\begin{split}\min_{\mathbf{\alpha}}&\left\{\frac{1}{J_ {1}}\sum_{j=1}^{J_{1}}\left[\sum_{m=1}^{M}\alpha_{m}\,\mathcal{L}(\sigma(\mathbf{ w}_{m}\mathbf{y}_{j}+b_{m}))+\alpha_{0}-f(\mathbf{y}_{j})\right]^{2}\right.\\ &+\left.\frac{1}{J_{2}}\sum_{j=1}^{J_{2}}\left[\sum_{m=1}^{M} \alpha_{m}\,\mathcal{B}(\sigma(\mathbf{w}_{m}\mathbf{y}_{j}+b_{m}))+\alpha_{0}-g(\mathbf{y }_{j})\right]^{2}\right\}\end{split} \tag{11}\]

where the parameters \(\{\mathbf{w}_{m}\}_{m=1}^{M}\) and \(\{b_{m}\}_{m=1}^{M}\) are first computed using the strategy in Section 3.3.2 and Eq. (6).

**Nonlinear PDEs.** When one or both operators, \(\mathcal{L}\) and \(\mathcal{B}\), are nonlinear, there are two approaches to handle the situation. The first way is to wrap the least squares problem with a well established nonlinear iterative solver, e.g., Picard's methods, to solve the PDE. Within each iteration, the PDE is linearized such that we can update the coefficient \(\mathbf{\alpha}\) by solving the least squares problem as mentioned above. When there is sufficient knowledge to choose a

Figure 1: **(a)** Illustrates how the re-parameterization in Eq. (5) characterizes the location of a neuron. The blue line is the plane where \(\tanh(\cdot)=0\), \(\mathbf{a}_{m}\) (the arrow) is the normal direction of the plane, the red dot is the location \(r_{m}\mathbf{a}_{m}\) that the plane passes, \(r_{m}\) is the distance between the origin and the plane. **(b)** illustrates how to generate uniformly distributed neurons in the unit ball. The first step in **(b)-left** is to generate the normal directions \(\{\mathbf{a}_{m}\}_{m=1}^{M}\) uniformly distributed on unit sphere; the second step in **(b)-middle** is to generated \(\{r_{m}\}_{m=1}^{M}\) uniformly from \([0,1]\) defining the locations the neurons’ partition hyperplanes will pass; the blue lines in **(b)-right** show the distribution of the partition hyperplanes. **(c)** shows the density function \(D_{M}(\mathbf{y})\) with \(\tau=0.05\) in Eq. (8) for a set of neurons generated using our approach. We can see that our approach provides a uniformly distributed neurons in the ball \(B_{1-\tau}(\mathbf{0})\), which is consistent with Theorem 1.
proper nonlinear solver, we prefer this approach because the well-established theory on nonlinear solvers can ensure a good convergence rate. Thus, we in fcat adopt this approach for numerical experiments in this paper. The second feasible approach is to wrap a gradient descent optimizer around the total loss \(L(\mathbf{y})=\|\mathcal{L}(u(\mathbf{y}))-\mathcal{L}(u_{\mathrm{NN}}(\mathbf{y}))\|_{2}^{2 }+\|\mathcal{B}(u(\mathbf{y}))-\mathcal{B}(u_{\mathrm{NN}}(\mathbf{y}))\|_{2}^{2}\). Because the neural feature space \(\mathcal{P}_{\mathrm{NN}}\) is fixed, the optimization will be simpler than training the entire neural network from scratch. This approach is easier to implement and suitable for scenarios that standard nonlinear solvers do not provide a satisfactory solution.

**Remark 3** (Not using PDE's solution data): _In this work, we do not rely on any measurement data of the solution \(u(\mathbf{y})\) when using TransNet to solve PDEs, because the operators \(\mathcal{L}\) and \(\mathcal{B}\) in Eq. (1) are sufficient to ensure the existence and uniqueness of the PDE's solution. On the other hand, if any extra data of \(u(\mathbf{y})\) are available, TransNet can easily incorporate it into the least squares problem in Eq. (11) as a supervised learning loss._

### Complexity and accuracy of TransNet

The complexity of TransNet is greatly reduced compared to the scenario of using SGD to train the entire network. The construction of the neural feature space \(\mathcal{P}_{\mathrm{NN}}\) only involves random number generations and a simple one-dimensional optimization in Eq. (10). Moreover, these cost are completely offline, and the constructed \(\mathcal{P}_{\mathrm{NN}}\) is transferable to various PDE problems. The online operation for solving linear PDEs only requires solving one least squares problem, where the assembling of the least squares matrix can be efficiently done using the autograd function in Tensorflow or Pytorch. The numerical experiments in Section 4 show that that the accuracy and efficiency of TransNet is significantly improved compared with several baseline methods, because our method does not suffer from the slow convergence of SGD in neural network training.

## 4 Numerical experiments

We now demonstrate the performance of TransNet by testing several classic steady-state or time-dependent PDEs in two and three dimensional spaces. In Section 4.1, we illustrate how to construct the transferable feature space \(\mathcal{P}_{\mathrm{NN}}\). To test and demonstrate the transferability of our model, we build and test two neural features spaces, one for the 2D case and the other for the 3D case1. The constructed feature spaces are then used in Section 4.2 to solve the model PDE problems.

Footnote 1: Note that the dimension of the feature space is the sum of both space and time dimensions since it doesn’t differ them.

### Uniform neuron distribution

This experiment is to use and test the algorithm proposed in Section 3.3 to construct transferable neural feature spaces \(\mathcal{P}_{\mathrm{NN}}\) in the 2D and 3D unit balls. We tune the shape parameter \(\gamma=\gamma_{m}\) for \(m=1,\dots,M\) in Eq. (5) with \(K=50\) realizations of the Gaussian process. In addition, we also test the effect of the correlation length and the number of hidden neurons by setting different values for \(\eta\) and \(M\). For each setting of \(\eta\) and \(M\), the shape parameter \(\gamma\) is tuned separately. Additional information about the experiment setup is given in Appendix B.

Figure 2 illustrates the landscapes of the loss function \(\sum_{k=1}^{K}\min_{\mathbf{\alpha}}[\text{MSE}(u_{\mathrm{NN}}(\mathbf{y}),G(\mathbf{y}| \omega_{k},\eta))]\) of the optimization problem in Eq. (10) for 2D and 3D neural feature spaces. We report the results for two correlation lengths (\(\eta=0.5\) and \(\eta=1.0\)) combined with three numbers of hidden neurons (\(M=100,500,1000\) for 2D and \(M=500,1000,5000\) for 3D). We observe that the loss function behaves roughly like a parabolic curve for a fixed number of hidden neurons, so that the problem in Eq. (10) can be solved by a simple solver for one-dimensional optimization. More importantly, we observe that the optimal value for \(\gamma\) varies with the number of hidden neurons. This provides an important insight that tuning \(\gamma\) is a necessary operation to achieve optimal accuracy of \(u_{\mathrm{NN}}\) when changing the number of hidden neurons.

Figure 2: The loss landscapes of the optimizing problem in Eq. (10) for tuning the shape parameter \(\gamma\) of the feature space \(\mathcal{P}_{\mathrm{NN}}\) in two and three dimensional cases. The blue star is the optimal value for \(\gamma\) founded by our method. It shows that the optimal value for \(\gamma\) varies with the number of hidden neurons, meaning that tuning \(\gamma\) is a necessary operation to achieve optimal accuracy of \(u_{\mathrm{NN}}\) when changing the number of hidden neurons.
Figure 3 illustrates the error distribution when using TransNet to approximate three realizations of the Gaussian process with correlation length \(\eta=0.5\) in the 2D unit ball. Even though the purpose of TransNet is not to approximate the Gaussian process, it is interesting to check whether the uniform density \(D_{M}(\mathbf{y})\) (proved in Theorem 1) leads to uniform error distribution. We use 1000 hidden neurons and the shape parameter \(\gamma\) is set to 2. The bottom row of Figure 3 shows that the MSE error distributes uniformly in the unit ball, which demonstrates the effectiveness of the feature space generation method proposed in Section 3.3.

### PDE examples

We then use the constructed 2D and 3D neural feature spaces from Section 4.1 to solve two steady-state PDEs (i.e., the Poisson equation and the time-independent Navier-Stokes equation) and two time-dependent PDEs (i.e., the Fokker-Planck equation and the wave equation). The definitions of the PDEs under consideration are given in Appendix C. We perform the following testing cases:

* Poisson equation (2D space) in a box domain;
* Poisson equation (2D space) in a circular domain;
* Poisson equation (2D space) in an L-shaped domain;
* Poisson equation (2D space) in an annulus domain;
* Poisson equation (3D space) in a box domain;
* Steady-state Navier-Stokes equation (2D space);
* Fokker-Planck equation (1D space + 1D time);
* 2D Fokker-Planck equation (2D space + 1D time);
* 1D wave equation (1D space + 1D time)

to demonstrate the transferability of TransNet in solving various PDEs in different domains. Recall that for time-dependent PDEs, the temporal variable is simply treated as an extra dimension, so that we will use the 2D feature space to solve problems \((C_{7})\) and \((C_{9})\) and the 3D feature space to solve problem \((C_{8})\). We compare our method with two baseline methods, i.e., the random feature mode and the PINN. All the methods use the same network architecture, i.e., Eq. (2) with the \(tanh\) activation. Additional information about the setup of the experiments are given in Appendix D.

Figure 4 shows the MSE decay with the increasing of the number of the hidden neurons, where the number of hidden neurons are chosen as \(M\) = 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, respectively, for the 2D feature space, and \(M=\) 1000, 2000, 3000, 4000, 5000, respectively, for the 3D feature space. We observe that our TransNet achieves a superior performance for all the nine test cases, which demonstrates the outstanding transferability of TransNet. PINN with BFGS acceleration provides a good accuracy gain compared with PINN with Adam, which means the landscape of the PDE loss exhibits severe ill-conditioning as the SGD method approaches the minimizer2. In comparison, TransNet does not require SGD in solving the PDEs, so that TransNet does not suffer from the slow convergence of SGD used in PINN.

Footnote 2: BFGS can alleviate ill-conditioning by exploiting the second-order information, e.g., the approximate Hessian.

Figure 5 shows the density function \(D_{M}(\mathbf{y})\) in Eq. (8) of the feature spaces obtained by training PINN and the random feature models in solving the Poisson equation in the 2D space, i.e., case \((C_{1})\) - \((C_{4})\), where the constant \(\tau\) in Eq. (8) is set to 0.2. Compared with TransNet's uniform density shown in Figure 1(c), the feature spaces obtained by the baseline methods have highly non-uniform densities in the domain of computation. The random feature models tend to have higher density, i.e., more hidden neurons, near the center of the domain. The first row in Figure 5 can be viewed as the initial densities of the feature space for PINN; the second and the third rows are the final densities. We can see that the training of PINN does not necessarily lead to a more uniform density function \(D_{M}(\mathbf{y})\), which is one of the reasons why PINN cannot exploit the full expressive power of the neural network \(u_{\mathrm{NN}}\).

## 5 Conclusion

We propose a transferable neural network model to advance the state of the art of using neural networks to solve PDEs. The key ingredient is to construct a neural feature space independent of any PDE, which makes it easy to transfer the neural feature space to various PDEs in different domains. Moreover, because the feature space is in fact fixed when using TransNet to solve a PDE, we only need to solve linear least squares problems, which avoids the drawbacks of SGD

Figure 3: Top row: three realizations of the auxiliary Gaussian process with the correlation length \(\eta=0.5\). Bottom row: the distribution of the MSE of TransNet’s approximation with 1000 hidden neurons. Thanks to the feature space with the uniform density in the 2D unit ball (illustrated in Figure 1(c)), we obtain a TransNet approximation with very small MSE fluctuation.
based training algorithms, e.g., ill-conditioning. Numerical experiments show that the proposed TransNet can exploit more expressive power of a given neural network than the compared baselines. This work is the first scratch in this research direction, and there are multiple potential related topics that will be studied in our future work, including (1) _theoretical analysis of the convergence rate of TransNet in solving PDEs._ We observe in Figure 4 that the MSE of TransNet decays along with the increasing of the number of hidden neurons. A natural question to study is that whether TransNet can achieve the optimal convergence rate of the single-hidden-layer fully-connected neural network. (2) _Extension to multi-layer neural networks._ Even though the single-hidden-layer model has sufficient expressive power for the PDEs tested in this work, there are more complicated PDEs, e.g., turbulence models, that could require multi-layer models with much higher expressive power. (3) _The properties of the least squares problem._ In this work, we use the standard least squares solver of Pytorch in the numerical experiments. However, it is worth further investigation of the properties of this specific least squares problem. For example, since the set of neurons \(\{\sigma(\mathbf{w}_{m}\mathbf{y}+b_{m})\}_{m=1}^{M}\) forms a non-orthogonal basis, it is possible to have linearly correlated neurons which will reduce the column rank of the least squares problem.

\begin{table}
\begin{tabular}{c|c|c|c|c|c|c|c|c} \hline \hline  & \((C_{1})\) & \((C_{2})\) & \((C_{3})\) & \((C_{4})\) & \((C_{5})\) & \((C_{6})\) & \((C_{7})\) & \((C_{8})\) & \((C_{9})\) \\ \hline Random feature model & 0.25s & 0.22s & 0.22s & 0.19s & 0.96s & 12.85s & 0.92s & 1.21s & 0.47s \\ PINN:Adam & 29.69s & 25.34s & 24.57s & 22.24s & 110.59s & 69.73s & 61.45s & 97.12s & 49.25s \\ PINN:Adam+BFGS & 125.78s & 121.46s & 120.93s & 119.24s & 264.62s & 191.53s & 172.86s & 178.99s & 152.71s \\ TransNet & 0.27s & 0.20s & 0.20s & 0.17s & 1.03s & 11.14s & 0.97s & 1.27s & 0.51s \\ \hline \hline \end{tabular}
\end{table}
Table 1: The computing times of TransNet and the baselines in solving the nine PDE test cases with 1000 hidden neurons. TransNet and the random feature model are significantly faster than PINN because SGD is not required in them.

Figure 4: The MSE decay along with the increasing of the number of hidden neurons for \((C_{1})\) to \((C_{9})\), where all the methods use the same network architecture. Our TransNet significantly outperforms the baseline methods from two aspects: (i) _Transferability_: for a fixed number of hidden neurons, TransNet only need use one 2D feature space and one 3D feature space; (ii) _Accuracy_: TransNet achieves several orders of magnitude smaller MSE than PINN and the random feature models. TransNet does not suffer from the slow convergence in SGD-based neural network training, and can exploit more expressive power of a given neural network \(u_{\mathrm{NN}}\) to obtain more accurate PDE solutions.

Figure 5: The density function \(D_{M}(\mathbf{y})\) with \(\tau=0.2\) in Eq. (8) of the neural feature spaces obtained by training PINN and the random feature models in solving the Poisson equation in the 2D space, i.e., problems \((C_{1})\) - \((C_{4})\). Compared to the uniform density of TransNet in Figure 1(c), both PINN and the random feature model cannot provide feature spaces with uniform density, which is one explanation of their under-performance shown in Figure 4.
squares matrix, or even lead to an under-determined system. This will require the use of some regularization techniques, e.g., ridge regression, to stabilize the least squares system. Additionally, compressed sensing, i.e., \(\ell_{1}\) regularization, could be added to remove redundant neurons from the feature space as needed and obtain a sparse neural network.

## Acknowledgement

This work was supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Applied Mathematics Program, under the contract number ERKJ387. This work was accomplished at Oak Ridge National Laboratory (ORNL). ORNL is operated by UT-Battelle, LLC., for the U.S. Department of Energy under Contract DE-AC05-00OR22725.
