A multiscale neural network based on hierarchical matrices
Yuwei Fan, Lin Lin, Lexing Ying, Leonardo Zepeda-Nunez
Introduction
In the past decades, there has been a great combined effort in developing efficient algorithms to solve linear problems issued from discretization of integral equations (IEs), and partial differential equations (PDEs). In particular, multiscale methods such as multi-grid methods , the fast multipole method , wavelets , and hierarchical matrices , have been strikingly successful in reducing the complexity for solving such systems. In several cases, for operators of pseudo-differential type, these algorithms can achieve linear or quasi-linear complexity. In a nutshell, these methods aim to use the inherent multiscale structure of the underlying physical problem to build efficient representations at each scale, thus compressing the information contained in the system. The gains in complexity stem mainly from processing information at each scale, and merging it in a hierarchical fashion.
Even though these techniques have been extensively applied to linear problems with outstanding success, their application to nonlinear problems has been, to the best of our knowledge, very limited. This is due to the high complexity of the solution maps. In particular, building a global approximation of such maps would normally require an extremely large amount of parameters, which in return, is often translated to algorithms with a prohibitive computational cost. The development of algorithms and heuristics to reduce the cost is an area of active research . However, in general, each method is application-dependent, and requires a deep understanding of the underlying physics.
On the other hand, the surge of interest in machine learning methods, in particular, deep neural networks, has dramatically improved speech recognition , visual object recognition , object detection, etc. This has fueled breakthroughs in many domains such as drug discovery , genomics , and automatic translation , among many others . Deep neural networks have empirically shown that it is possible to obtain efficient representations of very high-dimensional functions. Even though the universality theorem holds for neural networks , i.e., they can approximate arbitrarily well any function with mild regularity conditions, how to efficiently build such approximations remains an open question. In particular, the degree of approximation depends dramatically on the architecture of the neural network, i.e. how the different layers are connected. In addition, there is a fine balance between the number of parameters, the architecture, and the degree of approximation .
This paper aims to combine the tools developed in deep neural networks with ideas from multiscale methods. In particular, we extend hierarchical matrices (-matrices) to nonlinear problems within the framework of neural networks. Let
be a nonlinear generalization of pseudo-differential operators, issued from an underlying physical problem, described by either an integral equation or a partial differential equation, where can be considered as a parameter in the equation, is either the solution of the equation or a function of it, and is the number of variables.
We build a neural network with a novel multiscale structure inspired by hierarchical matrices. We interpret the application of an -matrix to a vector using a neural network structure as follows. We first reduce the dimensionality of the vector, or restrict it, by multiplying it by a short and wide structured matrix. Then we process the encoded vector by multiplying it by a structured square matrix. Then we return the vector to its original size, or interpolate it, by multiplying it by a structured tall and skinny matrix. Such operations are performed separately at different spatial scales. The final approximation to the matrix-vector multiplication is obtained by adding the contributions from all spatial scales, including the near-field contribution, which is represented by a near-diagonal matrix. Since every step is linear, the overall operation is also a linear mapping. This interpretation allows us to directly generalize the -matrix to nonlinear problems by replacing the structured square matrix in the processing stage by a structured nonlinear network with several layers. The resulting artificial neural network, which we call multiscale neural network, only requires a relatively modest amount of parameters even for large problems.
We demonstrate the performance of the multiscale neural network by approximating the solution to the nonlinear Schrödinger equation , as well as the Kohn-Sham map . These mappings are highly nonlinear, and are still well approximated by the proposed neural network, with a relative accuracy on the order of . Furthermore, we do not observe overfitting even in the presence of a relatively small set of training samples.
Although machine learning and deep learning literature is vast, the application of deep learning to numerical analysis problems is relatively new, though that is rapidly changing. Research using deep neural networks with multiscale architectures has primarily focused on image and video processing.
Deep neural networks have been recently used to solve PDEs and classical inverse problems . For general applications of machine learning to nonlinear numerical analysis problems, the work of Raissi and Karnidiakis used machine learning, in particular, Gaussian processes, to find parameters in nonlinear equations ; Chan, and Elsheikh predicted the basis function on the coarse grid in multiscale finite volume method by neural network; Khoo, Lu and Ying used neural network in the context of uncertainty quantification ; Zhang et al used neural network in the context of generating high-quality interatomic potentials for molecular dynamics . Wang et al. applied non-local multi-continuum neural network on time-dependent nonlinear problems . Khrulkov et al. and Cohen et al. developed deep neural network architectures based on tensor-train decomposition . In addition, we note that deep neural networks with related multi-scale structures have been proposed mainly for applications such as image processing, however, we are not aware of any applications of such architectures to solving nonlinear differential or integral equations.
2 Organization
The reminder of the paper is organized as follows. Section 2 reviews the -matrices and interprets the -matrices using the framework of neural networks. Section 3 extends the neural network representation of -matrices to the nonlinear case. Section 4 discusses the implementation details and demonstrates the accuracy of the architecture in representing nonlinear maps, followed by the conclusion and future directions in Section 5.
Neural network architecture for ℋ\mathcal{H}-matrices
In this section, we aim to represent the matrix-vector multiplication of -matrices within the framework of neural networks. For the sake of clarity, we succinctly review the structure of -matrices for the one dimensional case in Section 2.1. We interpret -matrices using the framework of neural networks in Section 2.2, and then extend it to the multi-dimensional case in Section 2.3.
Hierarchical matrices (-matrices) were first introduced by Hackbusch et al. in a series of papers as an algebraic framework for representing matrices with a hierarchically off-diagonal low-rank structure. This framework provides efficient numerical methods for solving linear systems arising from integral equations (IE) and partial differential equations (PDE) . In the sequel, we follow the notation in to provide a brief introduction to the framework of -matrices in a simple uniform and Cartesian setting. The interested readers are referred to for further details.
where and are periodic in and is smooth and numerically low-rank away from the diagonal. A discretization with an uniform grid with discretization points yields the linear system given by
We introduce a hierarchical dyadic decomposition of the grid in levels as follows. We start by the -th level of the decomposition, which corresponds to the set of all grid points defined as
Fig. 2 illustrates the block partition of induced by the dyadic partition, and the decomposition induced by the different interaction lists at each level that follows (2.4).
Thus the matrix-vector multiplication (2.2) can be expressed as
as illustrated in Fig. 3, which constitutes the basis for writing -matrices as a neural network.
2 Matrix-vector multiplication as a neural network
An artificial neural network, in particular, a feed-forward network, can be thought of the composition of several simple functions, usually called propagation functions, in which the intermediate one-dimensional variables are called neurons, which in return, are organized in vector, or tensor, variables called layers. For example, an artificial feed-forward neural network
with layer can be written using the following recursive formula
Within this context, training a network refers to finding the weights and biases, whose entries are collectively called parameters, in order to approximate a given map. This is usually done by minimizing a loss function using a stochastic optimization algorithm.
We interpret the structure of -matrices (2.6) using the framework of neural networks. The different factors in (2.7) possess a distinctive structure, which we aim to exploit by using locally connected (LC) network. LC networks are propagating functions whose weights have a block-banded constraint. For the one-dimensional example, we also treat as a 2-tensor of dimensions , where is the channel dimension and is the spatial dimension, and be a 2-tensor of dimensions . We say that is connected to by a LC networks if
Each LC network requires parameters, , , , , and to be characterized. Next, we define three types of LC network by specifying some of their parameters,
Restriction network: we set and in LC. This network represents the multiplication of a block diagonal matrix with block sizes and a vector with size , as illustrated by Fig. 4 (a). We denote this network using . The application of is depicted in Fig. 4 (a).
Kernel network: we set and . This network represents the multiplication of a cyclic block band matrix of block size and band size times a vector of size , as illustrated by the upper portion of Fig. 4 (b). To account for the periodicity we pad the input layer on the spatial dimension to the size . We denote this network by . This network has two steps: the periodic padding of on the spatial dimension, and the application of (2.10). The application of is depicted in Fig. 4 (b).
Interpolation network: we set and in LC. This network represents the multiplication of a block diagonal matrix with block size , times a vector of size , as illustrated by the upper figure in Fig. 4 (c). We denote the network , which has two steps: the application of (2.10), and the reshaping of the output to a vector by column major indexing. The application of is depicted in Fig. 4 (c).
2.2 Neural network representation
Following (2.7), in order to construct the neural network (NN) architecture for (2.7), we need to represent the following four operations:
From 1.1 and the definition of and , the operations (2.11a) and (2.11c) are equivalent to
respectively. Analogously, 1.3 indicates that (2.11b) is equivalent to
Combining (2.12), (2.13) and (2.14), we obtain Algorithm 1, whose architecture is illustrated in Fig. 5. In particular, Fig. 5 is the translation to the neural network framework of (2.7) (see Fig. 3) using the building blocks depicted in Fig. 4.
Moreover, the memory footprints of the neural network architecture and -matrices are asymptotically the same with respect the spatial dimension of and . This can be readily shown by computing the total number of parameters. For the sake of simplicity, we only count the parameters in the weights, ignoring those in the biases. A direct calculation yields the number of parameters in , and :
respectively. Hence, the number of parameters in Algorithm 1 is
3 Multi-dimensional case
Following the previous section, the extension of Algorithm 1 to the -dimensional case can be easily deduced using the tensor product of one-dimensional cases. Consider below for instance, and the generalization to higher dimensional case will be straight-forward. Suppose that we have an IE in 2D given by
we discretize the domain with a uniform grid with () discretization points, and let be the resulting matrix obtained from discretizing (2.17). We denote the set of all grid points as
where is called the band size for tensor. Thus 1 can be generalized to tensors yielding the following properties.
Next, we characterize LC networks for the 2D case. An NN layer for 2D can be represented by a 3-tensor of size , in which is the channel dimension and , are the spatial dimensions. If a layer with size is connected to a locally connected layer with size , then
where . As in the 1D case, the channel dimension corresponds to the rank , and the spatial dimensions correspond to the grid points of the discretized domain. Analogously to the 1D case, we define the LC networks , and and use them to express the four operations in (2.11) which constitute the building blocks of the neural network. The extension is trivial, the parameters , and in the one-dimensional LC networks are replaced by their 2-dimensional counterpart , and , respectively. We point out that for the 1D case is replaced by , for the 2D case in the definition of LC.
Using the notations above we extend Algorithm 1 to the 2D case in Algorithm 2. We crucially remark that the function in Algorithm 2 is not the usual major column based reshaping. It reshapes a 2-tensor with size to a 3-tensor with size , by treating the former as a block tensor with block size , and reshaping each block as a vector following the formula with , for , and . Fig. 6 provides an example for the case . The is its inverse.
Multiscale neural network
In this section, we extend the aforementioned NN architecture to represent a nonlinear generalization of pseudo-differential operators of the form
Due to its multiscale structure, we refer to the resulting NN architecture as the multiscale neural network (MNN). We consider the one-dimensional case below for simplicity, and the generalization to higher dimensions follows directly as in Section 2.3.
NN can represent nonlinearities by choosing the activation function, , to be nonlinear, such as ReLU or sigmoid. The range of the activation function also imposes constraints on the output of the NN. For example, the range of “ReLU” in and the range of the sigmoid function is $\mathsf{LCK}\mathsf{LCR}\mathsf{LCI}$ networks in Algorithm 1 are still treated as restriction and interpolation operations between coarse grid and fine grid, respectively, so we use the linear activation functions in these layers. Particularly, we also use the linear activation function for the last layer of the adjacent part, which is marked in line 13 in Algorithm 3.
As in the linear case, we calculate the number of parameters of MNN and obtain (neglecting the number of parameters in in (2.10))
2 Translation-invariant case
Compared to the LC network, the only difference is that the parameters and are independent of . Hence, inheriting the definition of , and , we define the layers , and , respectively. By replacing the LC layers in Algorithm 1 by the corresponding CNN layers, we obtain the neural network architecture for the translation invariant kernel.
For the nonlinear case, the translation invariant kernel for the linear case can be extended to kernels that are equivariant to translation, i.e. for any translation ,
For this case, all the LC layers in Algorithm 3 can be replaced by its corresponding CNN layers. The number of parameters of , and are
Thus, the number of parameters in Algorithm 3 using CNN is
Numerical results
In this section we discuss the implementation details of MNN. We demonstrate the accuracy of the MNN architecture using two nonlinear problems: the nonlinear Schrödinger equation (NLSE), and the Kohn-Sham map (KS map) in the Kohn-Sham density functional theory (KSDFT).
where is the target solution generated by a numerical discretization of the PDEs and is the predicted solution by MNN. The optimization is performed using the NAdam optimizer . The weights and biases in MNN are initialized randomly from the normal distribution and the batch size is always set between th and th of the number of train samples.
2 NLSE with inhomogeneous background potential
The nonlinear Schrödinger equation (NLSE) is widely used in quantum physics to describe the single particle properties of the Bose-Einstein condensation phenomenon . Here we study the NLSE with inhomogeneous background potential :
with period boundary condition. We aim to find its ground state denoted by . We take a strongly nonlinear case in this work and thus consider a defocusing cubic Schrödinger equation. Due to the cubic term, an iterative method is required to solve (4.2) numerically. We employ the method in for the numerical solution, which solves a time-dependent NLSE by a normalized gradient flow. The MNN is used to learn the map from the background potential to the ground state
This map is equivariant to translation, and thus MNN is implemented using the CNN layers. The constraints in (4.2) can be guaranteed by adding a post-correction in the network as
where is the prediction of the neural network. In the following, we study the performance of MNN on 1D and 2D cases.
For the one-dimensional case, the number of discretization points is , and we set and . The potential is chosen as
We also compare the MNN with an instance of the classical convolutional neural networks (CNN). The CNN architecture used in is adopted as the reference architecture. Table 4 presents the setup of the networks and the training and validation errors for CNN. Clearly, by comparing the results of MNN in Tables 2 and 3 with the results of CNN in Table 4, one can observe that the MNN not only reduces the number of parameters, but also improve the accuracy.
throughout the results shown in Tables 2 and 3, the validation errors are very close to the corresponding training errors, thus no overfitting is observed. Fig. 8 presents a sample for the potential and its corresponding solution and prediction solution by MNN. We can observe that the prediction solution agrees with the target solution very well.
2.2 Two-dimensional case
For the two-dimensional case, the number of discretization points is , and we set and . The potential is chosen as
Tables 5 and 6 present the numerical results for different number of channels, , and different number of layers ,, respectively. Similarly to the 1D case, the choice of parameters and also yield accurate results in the 2D case. Fig. 9 presents a sample of the potential in the test set and its corresponding solution, prediction solution and its error with respect to the reference solution.
3 Kohn-Sham map
Kohn-Sham density functional theory is the most widely used electronic structure theory. It requires the solution of the following set of nonlinear eigenvalue equations (real arithmetic assumed for all quantities):
Here is the number of electrons (spin degeneracy omitted), is the spatial dimension, and stands for the Kronecker delta. In addition, all eigenvalues are real and ordered non-decreasingly, and is the electron density, which satisfies the constraint
The Kohn-Sham equations (4.7) need to be solved self-consistently, which can also viewed as solving the following fixed point map
Here the mapping from to is called the Kohn-Sham map, which for a fixed potential is reduced to a linear eigenvalue problem, and it constitues the most computationally intensive step for solving (4.7). We seek to approximate the Kohn-Sham map using a multiscale neural network, whose output was regularized so it satisfies (4.8).
In the following numerical experiments the potential, , is given by
where is the dimension and . We set or for 1D and for 2D. The coefficients are randomly chosen following the uniform distribution , and the centers of the Gaussian wells , are chosen randomly under the constraint that , unless explicitely specified. The Kohn-Sham map is discretized using a pseudo-spectral method , and solved by a standard eigensolver.
We set and we generated data sets using different number of wells, , which in this case is also equal to the number of electrons , ranging from to .
The number of discretization points is . We trained the architecture defined in Section 3 for each , setting the number of levels , using different values for and .
Table 7 shows that there is no overfitting, even at this level of accuracy and number of parameters. This behavior is found in all the numerical examples, thus we only report the test error in what follows.
From Table 8 we can observe that as we increase the error decreases sharply. Fig. 10 depict this behavior. In Fig. 10 we have that if , then the network output , fails to approximate accurately; however, by modestly increasing , the network is able to accurately approximate .
However, the accuracy of the network stagnates rapidly. In fact, increasing beyond does not provide any considerable gains. In addition, Table 8 shows that the accuracy of the network is agnostic to the number of Gaussian wells present in the system.
In addition, we studied the relation between the quality of the approximation and . We fixed , and we trained several networks using different values of , ranging from , i.e., a very shallow network, to . The results are summarized in Table 9. We can observe that the error decreases sharply as the depth of the network increases, but it rapidly stagnates as becomes large.
The Kohn-Sham map is a very non-linear mapping, and we demonstrate below that the non-linear activation functions in the network play a crucial role. Consider a linear network, in which we have completely eliminated the non-linear activation functions. This linear network can be easily implemented by setting , and . We then train the resulting linear network using the same data as before and following the same optimization procedure. Table 10 shows that the approximation error can be very large, even for relatively small .
In addition, we note that the favorable behavior of MNN with respect to shown in the prequel is rooted in the fact that the band gap (here the band gap is equal to ) remains approximately the same as we increase . According to the density functional perturbation theory (DFPT), the Kohn-Sham map is relatively insensitive to the change of the external potential. This setup mimics an insulating system. On the other hand, we may choose in (4.10) to be , and relax the constraint between Gaussian centers to . The rest of the coefficients are randomly chosen using the same distributions as before. In this case, we generate a new data set, in which the average gap for the generated data sets are , , and for equal to , , , and , respectively. The decrease of the band gap with respect to the increase of resembles the behavior of a metallic system, in which the the Kohn-Sham map becomes more sensitive to small perturbations of the potential. After generating the samples, we trained two different networks: the MNN with , and a regular CNN with layers, channels, window size . The results are shown in Tables 11 and 12. From Tables 11 and 12 we can observe that MNN outperforms CNN and that as the band gap decreases, the performance gap between MNN and CNN widens. We point out that it is possible to partially alleviate this adverse dependence on the band gap in the MNN by introducing a nested hierarchical structure to the interpolation and restriction operators as shown in .
3.2 Two-dimensional case
The discretization is the standard extension to 2D using tensor products, using a grid. In this case we only used and we followed the same number of training and test samples as that in the D case. We fixed , , and we trained the network for different number of channels, . The results are displayed in Table 13, which shows the same behavior as for the 1D case, in which the error decays sharply and then stagnates, and there is no over fitting. In particular, the network is able to effectively approximate the Kohn-Sham map as shown in Fig. 11. Fig. 11a shows the output of neural network for a test sample and Fig. 11b shows the approximation error with respect to the reference.
Conclusion
We have developed a multiscale neural network (MNN) architecture for approximating nonlinear mappings, such as those arising from the solution of integral equations (IEs) or partial differential equations (PDEs). In order to control the number of parameters, we first rewrite the widely used hierarchical matrix into the form of a neural network, which mainly consists of three sub-networks: restriction network, kernel network, and interpolation network. The three sub-networks are all linear, and correspond to the components of a singular value decomposition. We demonstrate that such structure can be directly generalized to nonlinear problems, simply by replacing the linear kernel network by a multilayer kernel network with nonlinear activation functions. Such “nonlinear singular value decomposition operation” is performed at different spatial scales, which can be efficiently implemented by a number of locally connected (LC) networks, or convolutional neural networks (CNN) when the mapping is equivariant to translation. Using the parameterized nonlinear Schrödinger equation and the Kohn-Sham map as examples, we find that MNN can yield accurate approximation to such nonlinear mappings. When the mapping has degrees of freedom, the complexity of MNN is only . Thus the resulting MNN can be further used to accelerate the evaluation of the mapping, especially when a large number of evaluations are needed within a certain range of parameters.
Acknowledgements
The authors thank Yuehaw Khoo for constructive discussions. The work of Y.F. and L.Y. is partially supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Computing (SciDAC) program and the National Science Foundation under award DMS-1818449, and the GCP Research Credits Program from Google. The work of Y.F. is also partially supported by AWS Cloud Credits for Research program from Amazon. The work of L.L and L. Z. is partially supported by the Department of Energy under Grant No. DE-SC0017867, by the Department of Energy under the CAMERA project, and by the Air Force Office of Scientific Research under award number FA9550-18-1-0095.