A S UBSPACE C ORRECTION M ETHOD FOR R E LU N EURAL N ETWORKS FOR S OLVING PDE S
Anonymous authors
Paper under double-blind review
A BSTRACT
In this paper, we propose a novel algorithm called Neuron-wise Parallel Subspace
Correction Method (NPSC) for training ReLU neural networks for numerical solution of partial differential equations (PDEs).
Despite of extremely extensive research activities in applying neural networks for numerical PDEs, there is still a serious lack of training algorithms that can be used to obtain approximation with adequate accuracy.
Based on recent results on the spectral properties of linear layers and landscape analysis for single neuron problems, we develop a special type of subspace correction method that deals with the linear layer and each neuron in the nonlinear layer separately.
An optimal preconditioner that resolves the ill-conditioning of the linear layer is presented, so that the linear layer is trained in a uniform number of iterations with respect to the number of neurons.
In each single neuron problem, a good local minimum is found by a superlinearly convergent algorithm, avoiding regions where the loss function is flat.
Performance of the proposed method is demonstrated through numerical experiments for function approximation problems and PDEs.
I NTRODUCTION
Neural networks, thanks to the universal approximation property (Cybenko, 1989; Pinkus, 1999), are promising tools for numerical solutions of partial differential equations (PDEs).
Moreover, it was shown in Siegel & Xu (2022b) that the approximation properties of neural networks have higher asymptotic approximation rates than that of traditional numerical algorithms such as finite element methods.
Such powerful approximation properties, however, can hardly be observed in numerical experiments even in simple tasks of function approximation.
It was shown in Hong et al.
that conventional training algorithms such as gradient descent converge too slowly to an accurate solution even for a very simple function.
Therefore, applying neural networks to solutions of PDEs must require novel training algorithms, different from the conventional ones for regression, image classification, and pattern recognition tasks.
In this viewpoint, there are many works on designing and analyzing training algorithms for neural networks to try to narrow the gap between the theoretical optimum and training results.
A hybrid least squares/gradient descent method was proposed in Cyr et al.
Plateau phenomena (Park et al., 2000) occurring in the gradient dynamics of
ReLU shallow neural networks were analyzed in a rigorous manner in Ainsworth & Shin (2021), and a training algorithm called active neuron least squares was designed to avoid plateau phenomena in Ainsworth & Shin (2022).
As a completely different approach, the orthogonal greedy algorithm was shown to achieve an optimal convergence rate (Siegel & Xu, 2022c) with respect to the number of neurons when it is applied to certain shallow neural networks in Siegel & Xu (2022a).
A surprising recent result on training of neural networks is that optimizing the linear layer parameters in a neural network is one bottleneck that leads to a large number of iterations of a gradient-based method.
More precisely, it was proven in Hong et al.
that optimizing the linear layer parameters in a ReLU shallow neural network requires solving a very ill-conditioned linear problem in general.
This work motivates us to separately design efficient solvers for the outer linear layer and the inner nonlinear layer, respectively, and train them alternately.
Meanwhile, there have been some encouraging works on learning a single neuron, a more simplified model than a nonlinear layer.
Convergence analyses of gradient methods for the single neuron problem with ReLU activation were presented in Soltanolkotabi (2017);
Tian (2017); Yehudai & Ohad
(2020) under various assumptions on input distributions.
The case of a single ReLU neuron with bias was analyzed in Vardi et al.
All of these results show that the global convergence of gradient methods for the single ReLU neuron problem can be attained under certain conditions.
Hence, learning a single neuron can be regarded as a much more hopeful task than learning a nonlinear layer with multiple neurons.
In this paper, we use the idea of subspace correction (Xu, 1992), which is well-known in the fields of numerical analysis and scientific computing.
Subspace correction methods provide a unified framework to design and analyze many modern iterative numerical methods such as block coordinate descent, multigrid, and domain decomposition methods.
In subspace correction methods, the solution space of a target problem is decomposed into a sum of subspaces, and local problems defined on subspaces are solved separately in either sequential or parallel manner.
Then a solution of the whole problem is constructed by assembling the solutions of the local problems appropriately.
Mathematical theory of subspace correction methods for convex optimization problems can be found in Park (2020); Tai & Xu (2002).
Subspace correction methods have been successfully applied to various nonlinear optimization problems appearing in engineering fields (see, e.g., Badea & Krause (2012);
Lee & Park (2019)).
In particular, there have been successful applications of block coordinate descent methods (Wright, 2015) for training of neural networks (Zeng et al., 2019; Zhang & Brand, 2017).
Therefore, we expect that the idea of the subspace correction method is also suitable for training of ReLU shallow neural networks.
Inspired by the above works, we propose a new training algorithm called Neuron-wise Parallel
Subspace Correction Method (NPSC), which is a special type of subspace correction method for a
ReLU shallow neural network.
The proposed method utilizes a space decomposition for the linear layer and each individual neuron.
In the first step of each epoch of the NPSC, the linear layer is fully trained by solving a relevant linear system.
This resolves ill-conditioning of the linear layer.
We design an optimal preconditioner for the linear layer problem in one-dimension, based on the relation between ReLU neural networks and linear finite elements investigated in He et al.
Hong et al.
In the second step, we train each single neuron in parallel, taking advantages of better convergence properties of learning a single neuron explained above.
Each neuron is trained by a superlinearly convergent algorithm (Marquardt, 1963).
Finally, an update for the parameters in the nonlinear layer is computed by assembling the corrections obtained in the local problems for each neuron.
Due to the intrinsic parallel structure, NPSC is suitable for parallel computation on distributed memory computers.
We present applications of NPSC to various function approximation problem and PDEs, and numerically verify that it outperforms conventional training algorithms.
The rest of this paper is organized as follows.
In Section 2, we summarize key features of ReLU shallow networks and describe our model problem.
The proposed NPSC is presented in Section 3.
Applications of NPSC to various function approximation problems and PDEs are presented in Section 4 to demonstrate the effectiveness of NPSC.
We conclude the paper with remarks in Section 5.
M ODEL P ROBLEM
We consider a shallow neural network with n neurons and d -dimensional inputs x ∈ R d given by
[Equation]
where θ is the collection of parameters consisting of a i ∈ R , ω i ∈ R d , and b i ∈ R for 1 ≤ i ≤ n , and σ : R → R is a ReLU activation function defined by σ ( x ) = max { 0 , x } .
The neural network (1) possesses total ( d + 2) n parameters.
Let Ω ⊂ R d be a bounded domain and let f ∈ L 2 (Ω) .
As a model problem, we consider the following minimization problem:
[Equation]
where a ( · , · ) is a continuous, coercive, and symmetric bilinear form defined on a Hilbert space
V ⊂ L 2 (Ω) .
The problem (2) can be seen as the Galerkin approximation of the problem
[Equation]
which can represent various elliptic PDEs (E & Yu, 2018; Xu, 2020), on the space
[Equation]
The space Σ 1 n enjoys the universal approximation property (Cybenko, 1989; Pinkus, 1999), namely, any function with sufficient regularity can be uniformly approximated by functions in Σ 1 n .
Recent results on the approximation property of Σ 1 n can be found in Siegel & Xu (2022c).
N EURON - WISE P ARALLEL S UBSPACE C ORRECTION M ETHOD (NPSC)
In this section, we introduce NPSC, a subspace correction method (Xu, 1992) that deals with each neuron in parallel for the problem (2).
We first present a space decomposition for NPSC.
Let Θ = R ( d +2) n denote the solution space of (2), i.e., the space for the parameter θ .
It admits a natural decomposition Θ = A ⊕ W , where A = R n and W = R ( d +1) n are the spaces for a and { ω, b } , respectively, and ⊕ denotes direct sum.
Since any { ω, b } ∈ W consists of the parameters ( { ω i , b i } ) ni
=1 from n neurons, W can be further decomposed as W = (cid:76) ni
=1 W i , where W i =
R d + is the space for { ω i , b i } .
Finally, we have the following space decomposition of Θ :
[Equation]
A graphical description for the space decomposition (3) is presented in Figure 1(a).
NPSC, our proposed method, is presented in Algorithm 1.
It is a subspace correction method (Xu, 1992) for (2) based on the space decomposition (3).
At the k th epoch, NPSC updates the parameter a first by minimizing the loss function with respect to a , then it updates the parameters in each neuron by minimizing E with respect to { ω i , b i } in parallel.
The update of { ω i , b i } is relaxed by an appropriate learning rate τ k > 0 as in the existing parallel subspace correction methods for optimization problems (Lee & Park, 2019;
Tai & Xu, 2002).
The overall structure of NPSC is depicted in Figure 1(b).
In the remainder of this section, we discuss parts of Algorithm 1 in detail.
3. A DJUSTMENT OF P ARAMETERS
We call that a neuron with parameter { ω i , b i } is dead if σ ( ω i · x + b i ) vanishes almost everywhere (a.e.) on the domain Ω .
In Vardi et al.
(2021, Theorem 3.1), it was shown that a neuron is initialized as dead with probability close to half if we employ a usual random initialization scheme.
A similar observation was made in Lu et al.
Such a dead neuron makes the functions { σ ( ω i · x + b i ) } ni =1 linearly dependent, i.e., there exist a i ∈ R , 1 ≤ i ≤ n , not all zero, such that (cid:80) ni =
a i σ ( ω i · x + b i ) = 0 .
Meanwhile, the linear dependence may occur by a combination of several non-dead neurons.
Theorem 1 characterizes some situations when { σ ( ω i · x + b i ) } ni = becomes linearly dependent; see Appendix A for a proof.
Theorem 1.
Consider the shallow neural network (1) .
The followings hold.
Choose an initial guess θ (0) = { a (0) , ω (0) , b (0) } and an initial learning rate τ 0 = 1 . for k = 1 , . . . , T do
Adjust { ω ( k − 1) , b ( k − 1) } to avoid linear dependence of the neurons (see Algorithm 2).
a ( k ) ∈ arg min
a
[Equation]
( k
[Equation]
[Equation]
, b ( k
[Equation]
[Equation]
) (see (6)) for i = 1 , . . . , n in parallel do
[Equation]
    a ( k )
[Equation]
⊕ n (cid:77) j =1 ,j ̸ = i ω ( k − 1) j , b i ⊕ n (cid:77) j =1 ,j ̸ = i b ( k − 1) j     (4) end for
Determine the learning rate τ k by backtracking (see Algorithm 3).
for i = 1 , . . . , n in parallel do
[Equation]
end for end for If a neuron satisfies b i ≤ − ess sup x ∈ Ω ω i · x , then it is dead. If there are more than d + 1 neurons such that b i ≥ − ess inf x ∈ Ω ω i · x , then the functions σ ( ω i · x + b i ) corresponding to these neurons are linearly dependent in Ω . Linear dependence of the functions { σ ( ω i · x + b i ) } ni =1 in the neural network (1) should be avoided because it means that several neurons are redundant and they are not utilized to improve the approximability of (1).
Theorem 1 motivates us to consider an adjustment procedure for { ω ( k − 1) , b ( k − 1) } when we enter the k th iteration of NPSC so that linear dependence does not occur.
Algorithmsummarizes the adjustment procedure for a given { ω, b } ∈ W .
Algorithm 2 Adjustment for { ω, b } in Algorithmfor i = 1 , . . . , n in parallel do Set ω i ← ω i | ω i | and b i ← b i | ω i | .
if b i ̸∈ (cid:18) − max x ∈ Ω ω i · x, − min x ∈ Ω ω i · x (cid:19) then Reset b i ∈ R such that b i ∼ Uniform (cid:18) − max x ∈ Ω ω i · x, − min x ∈ Ω ω i · x (cid:19) .
end if end for Since ω i determines the direction of a function σ ( ω i · x + b i ) , we may normalize { ω i , b i } so that | ω i | = 1 in Algorithm 2.
Then, for each b i that is not on the interval between the minimum and maximum of ω i · x , we relocate it on the interval.
Since the extrema in Ω agree with the essential extrema under mild conditions on Ω , this relocation step lets the neurons avoid the scenario of linear dependence described in Theorem 1.
In Algorithm 2, evaluations of the extrema of ω i · x are simple linear programs, so that they can be done efficiently by conventional algorithms for linear programming (Nocedal & Wright, 2006).
In particular, if Ω is a polyhedral domain, then we can evaluate the extrema of ω i · x by hands; we simply compute ω i · x at all the vertices of Ω and take the extrema among them.
3. a - MINIMIZATION P ROBLEMS
In Algorithm 1, a -minimization problems have the general form
[Equation]
where { ω, b } ∈ W .
The first-order optimality condition for (5) reads as
[Equation]
where K ∈ R n × n and β ∈ R n are given by
[Equation]
Assuming the basis functions { σ ( ω i · x + b i ) } ni =1 are linearly independent in Ω , the matrix K is symmetric and positive definite.
Hence, the system (6) can be solved efficiently by the conjugate gradient method, and the efficiency can be improved with an appropriate preconditioner.
Theorem 2 presents an optimal preconditioner in one-dimension.
While it was proven in Hong et al.
(2022) that the matrix K is usually ill-conditioned, the preconditioner in Theorem 2 can resolve this issue.
We refer readers to Appendix A for a proof of Theorem 2.
Computational aspects and numerical results regarding the preconditioner are presented in Appendix B.
Theorem 2.
Let Ω =
[0 , 1] ⊂ R .
Assume that 0 ≤ x 1 < x 2 < · · · < x n ≤ 1 under an appropriate reordering, where x i = − b
i /ω i for 1 ≤ i ≤ n .
Then we can find a matrix P ∈ R n × n explicitly such that the condition number κ ( PK ) of a preconditioned operator PK has an upper bound independent of n , ω , and b , i.e., κ ( PK ) =
O (1) .
3. { ω i , b i } - MINIMIZATION P ROBLEMS
Training the nonlinear layer is challenging because of its nonconvexity. Moreover, it was shown in Lu et al. (2020) that training of ReLU networks usually suffers from dead neurons.
In the following, we claim that training each neuron in the nonlinear layer separately has an advantage that some “good” local minima that avoid dead neurons can be found (cf.
Vardi et al.
We consider the { ω i , b i } -minimization problem (4) in Algorithm 1 for a fixed i .
We may assume that a i ̸ = 0 [-, then r i = 
 solves (7) since-]; otherwise, the minimization problem becomes trivial.
Under some elementary manipulations, (4) is rewritten as
[Equation]
where F ∈ L 2 (Ω) is a function determined in terms of f , a ( k ) , ω ( k − 1) j , and b ( k − 1) j for j ̸ = i .
A notable advantage of single neuron training (8) is that we can easily avoid regions where the loss function is flat.
The following theorem says that all the critical points of (8) on the flat regions make the neuron dead, and that all the critical points outside the flat regions have less loss values than those on the flat regions.
A proof of Theorem 3 can be found in Appendix A.
Theorem 3.
Consider the single neuron problem (8) .
The followings hold:
If { ω ∗ i , b ∗ i } ∈ R d +1 is a critical point of (8) such that σ ( ω ∗ i · x + b ∗ i ) vanishes on Ω , then we have E i ( ω ∗ i , b ∗ i ) =
E i ( 0 ) .
If { ω ∗ i , b ∗ i } ∈ R d +1 is a critical point of (8) such that σ ( ω ∗ i · x + b ∗ i ) does not vanish on a subset of Ω with nonzero measure, then we have E i ( ω ∗ i , b ∗ i ) <
E i ( 0 ) .
Theorem 3 implies that, if we choose an initial guess { ω (0) i , b (0) i } such that E i ( ω (0) i , b (0) i ) <
E i ( 0 ) for a monotone training algorithm, then the algorithm converges to a good local minimum that prevents the neuron to be dead.
That is, each { ω i , b i } -minimization problem of NPSC finds a good local solution.
Note that the condition E i ( ω (0) i , b (0) i ) <
E i ( 0 ) is satisfied by a random initialization of { ω (0) i , b (0) i } around {+
+} under some mild conditions; see Vardi et al.
(2021, Theorem 5.4).
Various optimization algorithms including first- and second-order methods can be used to solve (8). Among them, a good option is the Levenberg–Marquardt algorithm (Marquardt, 1963); PDEs usually arise from physics, so that the dimension d + 1 of { ω i , b i } is not very big in general.
The major computational cost of each iteration of the Levenberg–Marquardt algorithm is to solve a linear system with d + 1 unknowns; it is not time-consuming when d + 1 is small.
The Levenberg–Marquardt algorithm does not require explicit assembly of the Hessian, and it converges to a local minimum with the superlinear convergence rate (Yamashita & Fukushima, 2001), which is much faster than that of the first-order methods.
3. B ACKTRACKING FOR L EARNING R ATES
Recently, it was shown in Park (2022) that parallel subspace correction methods for convex optimization problems can be accelerated by adopting a backtracking scheme.
Hence, it is natural to utilize a backtracking scheme to find a suitable value for the learning rate τ k in Algorithm 1.
Due to the nonconvexity of the loss function E of (2), we are not able to adopt the backtracking scheme proposed in Park (2022) directly to Algorithm 1.
Moreover, since ReLU activation in (2) makes
E nonsmooth, conventional backtracking schemes such as Calatroni & Chambolle (2019); Scheinberg et al.
(2014) are not applicable.
A simple but effective backtracking scheme for finding τ k is presented in Algorithm 3.
By allowing adaptive increase and decrease of τ k along the epochs, the convergence rate of NPSC is improved.
Algorithm 3 Backtracking scheme to find τ k in Algorithm
Choose a minimum learning rate τ min = 10 − 12 .
τ k ← 2 τ k −repeatfor i = 1 , . . .
, n in parallel do
[Equation]
N UMERICAL R ESULTS
In this section, we present numerical results of NPSC applied to various function approximation problems and PDEs of the form (2).
The following algorithms are compared with NPSC in our numerical experiments: gradient descent (GD), Adam (Kingma & Ba, 2015), hybrid leastsquares/gradient descent (LSGD)
(Cyr et al., 2020), and the proposed NPSC.
All the algorithms were implemented in ANSI C with OpenMPI compiled by Intel C++ Compiler.
They were executed on a computer cluster equipped with multiple Intel Xeon SP-6148 CPUs (2.4GHz, 20C) and the operating system CentOS 7.4 64bit.
In all experiments, we use the He initialization (He et al., 2015) to initialize the parameters of (2).
That is, we set a i ∼ N (0 , 2 /n ) , ω i ∼ ( N (0 , 2 /d ))
d , and b i ∼ ( N (0 , 2 /d )) in (2).
All the numerical results presented in this section are averaged over 10 random initializations.
As shown in Appendix C, the performances of conventional training algorithms can be improved by utilizing the backtracking scheme presented in Algorithm 3.
Hence, for GD, Adam, and LSGD, we employ Algorithm 3 to find learning rates.
In LSGD and NPSC, a -minimization problems (6) are solved by the preconditioned conjugate gradient method with the preconditioner P in Theorem 2 and the stop criterion
[Equation]
Finally, { ω i , b i } -minimization problems (8) in NPSC are solved by the Levenberg–Marquardt algorithm (Marquardt, 1963) with the stop criterion
[Equation]
MPI parallelization is applied to NPSC in a way that each { ω i , b i } -minimization problem is assigned to a single processor and solved in parallel.
4. L 2 - FUNCTION APPROXIMATION PROBLEMS
If we set
V =
[Equation] [Equation]
in (2), then we obtain the L 2 -function approximation problem, which is the most elementary instance of (2).
That is, the solution of (2) in this case is the best L 2 -approximation of f found in Σ 1 n .
As our first example, we consider L 2 -approximation (10) for a sine function; we set
[Equation]
(11) in (2).
For numerical integration, we use the trapezoidal rule on 10,000 uniformly sampled points; see Appendix D. Figure 2(a) plots the relative energy errors
E ( θ ( k ) )
− E ∗ | E ∗ | obtained by NPSC, GD, Adam, and LSGD per epoch, where E ∗ is the loss corresponding to the exact solution.
It is clear from the loss decay that NPSC outperforms all the other methods.
The exact solution of (11) and its approximations obtained by 10 3 epochs of the training algorithms with 2 5 neurons are depicted in Figure 2(b).
One can readily observe that the NPSC result is most accurate among the approximations.
The L 2 -errors between the exact solution and its approximations for various numbers of neurons are presented in Figure 2(c). Only the NPSC result seems to be comparable to O ( n − 2 ) , the theoretical optimal rate derived in Siegel & Xu (2022c).
The second example is the following highly oscillatory instance of (10):
[Equation]
in which a similar problem appeared in Cai et al.
We note that it was demonstrated in Krishnapriyan et al.
(2021) that training for complex functions like (12) is a much harder task than training for simple functions.
Same as in (11), we use the trapezoidal rule on 10,000 uniformly sampled points for numerical integration.
Figures 2(d–f) present numerical results for the problem (12).
In Figure 2(d), NPSC shows stable decay of the loss while the losses of GD and Adam are stagnant along the epochs and that of LSGD blows up at in the beginning epochs and varies very rapidly.
In Figure 2(e), one can observe that the NPSC result captures both low- and high-frequency parts of the target function well, while the other ones captures the low-frequency part only (Xu et al., 2020).
Moreover, as shown in Figure 2(f), the L 2 -error of the NPSC decreases when the number of neurons increases, while those of the other algorithms seem to stagnate.
This highlights the robustness of NPSC for complex problems.
4. E LLIPTIC P ARTIAL D IFFERENTIAL E QUATIONS
Next, we consider the case
[Equation]
13) in the problem (2), where H 1 (Ω) is the usual Sobolev space consisting of L 2 (Ω) -functions with square-integrable gradients.
It is well-known that (2) becomes the Galerkin approximation of the weak formulation of the following elliptic PDE on Σ 1 n (Xu, 2020):
[Equation]
where ∂u/∂n denotes the normal derivative of u to ∂ Ω .
Hence, one can find a numerical approximation for the solution of the above PDE by solving (13).
We set
[Equation]
in (13); the exact solution is u ( x ) = cos 3 πx + cos 11 πx .
We can observe in Figure 3(a) that NPSC achieves a superior level of accuracy such that the other algorithms do not reach even with a large number of epochs.
As shown in Figures 3(b), L 2 –errors of the NPSC results decrease as the number of neuron increases, and their decreasing rates are comparable to the theoretical optimal rates in (Siegel & Xu, 2022c).
Finally, we consider the following two-dimensional example for (13):
[Equation]
(15) whose exact solution is u ( x 1 , x 2 ) = cos πx cos πy .
Since the preconditioner in Theorem 2 is applicable for one-dimension only, in this example, a -minimization problems are solved by n iterations of the unpreconditioned conjugate gradient method.
We use the quasi-Monte Carlo method (Caflisch, 1998) with 10,000 sampling points for numerical integration; see Appendix D for details.
It is verified by the numerical results for (15) presented in Figure 3(c, d) that NPSC outperforms the other methods and provides reasonable L 2 -errors in high-dimensional problems as well.
denote NPSC without the adjustment step and the Levenberg–Marquardt algorithm, respectively
Remark 1 .
When we train a neural network (1) for solving the PDE (13) with a gradient-based method, we have to evaluate the second-order derivative of σ ( x ) , which is the Dirac delta function.
In our experiments, we simply ignore Dirac delta terms in numerical integration.
This motivates us to consider high-order activation functions like ReLU k (Siegel & Xu, 2022b) as a future work.
4. A BLATION S TUDIES
Key components of the proposed NPSC are the adjustment step for parameters (Algorithm 2), the optimal preconditioner for a -minimization problems (Theorem 2), the Levenberg–Marquardt algorithm for { ω i , b i } -minimization problems, and the backtracking scheme for learning rates (Algorithm 3).
In order to validate the effects of these components, we conduct ablation studies for the adjustment step and the Levenberg–Marquardt algorithm; relevant results for the preconditioner and the backtracking scheme can be found in Appendices B and C, respectively.
Figure 4 depicts numerical comparisons among three algorithms: NPSC, NPSC without Algorithm 2, and NPSC without the Levenberg–Marquardt algorithm.
In the last algorithm, each { ω i , b i } - minimization problem is solved approximately by 20 iterations of GD.
Since the variants of NPSC achieve slower convergence rates than NPSC in all the examples, one can conclude that both Algorithm 2 and the Levenberg–Marquardt algorithm contribute to the fast convergence of NPSC.
We also note that Algorithm 2 helps NPSC to avoid unstable convergence behaviors like Figure 4(c, d).
C ONCLUSION
In this paper, we proposed NPSC for training ReLU neural networks for numerical solution of
PDEs.
Separately designing efficient solvers for the linear layer and each neuron and training them alternately, NPSC yields accurate results for function approximation problems and PDEs.
This paper leaves us several interesting and important topics for future research.
Although NPSC adopts the well-established framework of subspace correction methods, its rigorous convergence analysis is still open due to the nonconvexity of the model.
To solve high-order PDEs, we should generalize NPSC so that it can be applied to different kind of activation functions.
In addition, we should consider optimal preconditioners for optimizing linear layer in high dimensions.
Generalization of NPSC to many layers is not clear yet, and we consider this topic as a future work.
Reducing the gap between the theoretical approximation properties and training results is a challenging issue even for simple neural networks.
To the best of our knowledge, NPSC is the first result that deals with the ill-conditioning of layers and successfully reduces the gap.
We believe that this paper can play a role of a good preliminary work for efficient and accurate training of general network architectures.
A D EFERRED P ROOFS
In this appendix, we provide proofs of Theorems 1–3.
First, we provide the proof of Theorem 1.
Proof of Theorem 1.
If b i ≤ − ess sup x ∈ Ω ω i · x for some i , then we have
[Equation]
Hence, σ ( ω i · x + b i ) =
a.e. in Ω , which means that the corresponding neuron is dead.
Now, we suppose that b i ≥ − ess inf x ∈ Ω ω i · x for 1 ≤ i ≤ d + 2 .
Then we have
[Equation]
This implies that σ ( ω i · x + b i ) =
ω i · x + b i a.e. in Ω , i.e., each σ ( ω i · x + b i ) is affine in Ω .
Since the dimension of the space of all affine functions in Ω is d + 1 , { σ ( ω i · x + b i ) }
d +2 i =1 are linearly dependent in Ω .
Next, we introduce some preliminaries that will be used in the proof of Theorem 2.
For brevity, we only consider the case
< x 1 < x 2 < · · · < x n < 1 ; the cases x 1 = 0 and x n = 1 can be proven similarly.
By the Lax–Milgram theorem (Brezis, 2011, Corollary 5.8), there exists a bijective linear operator A :
V → V such that
[Equation]
We denote
[Equation]
and write them as a column vector Ψ( x ) =
[ ψ 1 ( x ) , . . .
, ψ n ( x )] T .
Then it follows that
[Equation]
where K was defined in (7) and A is applied entrywise.
In addition, we define ψ
[Equation]
=
[Equation]
[Equation]
and write Ψ( x )
=
[ ψ 0 ( x ) , ψ 1 ( x ) , . . .
, ψ n +1 ( x )] T .
We readily have Ψ( x ) =
R Ψ( x ) , where
[Equation]
In (17), 0 and I n denote the zero column vector of length n and the R n × n identity matrix, respectively.
If we write
[Equation]
then we get M = RKR T .
Invoking He et al.
(2020, Theorem 2.1) implies that { ψ i ( x ) } n +1 i =0 are linearly independent, so that they form a basis for the space V n of continuous and piecewise linear functions on the grid
= x 0 < x 1 < · · · < x n +1 = 1 , where x
i = − b i /ω i .
Meanwhile, V n admits another set of basis functions { ϕ i ( x ) } n +1 i =0 , which is given by
[Equation]
if
[Equation] [Equation]
if
[Equation]
∈ [ x
[Equation]
+1 )
[Equation]
0 , otherwise, 0 ≤ i ≤ n + 1 , with conventions x − 1 = 0 and x n +2 = 1 .
That is, { ϕ i ( x ) } n +1 i =0 is the collection of the standard hat basis functions defined on the grid { x i } n +1 i =
Similar to (18), we write Φ( x ) =
[ ϕ 0 ( x ) , ϕ 1 ( x ) , . . .
, ϕ n ( x )] T and set
[Equation]
The following lemma, which can be deduced by using elementary linear algebra, is an ingredient of the proof of Theorem 2.
Lemma 1.
For two positive integers m ≥ n , let A , B ∈ R m × m be nonsingular matrices and let
R ∈ R n × m be a surjective matrix.
Then RB − 1 R T is nonsingular and
[Equation]
Proof.
Since B − 1 is nonsingular and R T is injective, we get
[Equation]
for α ∈ R n , which implies that RB − 1 R T is nonsingular.
Let λ min and λ max stand for the minimum and maximum eigenvalues, respectively.
Then it follows that
[Equation]
In the same manner, we have
[Equation]
Combining (19) and (20) yields the desired result.
Now, we are ready to present our proof of Theorem 2.
Proof of Theorem 2.
First, we define Ψ + ( x ) =
[ ψ +0 ( x ) , ψ +1 ( x ) , . . .
, ψ + n +1 ( x )] T , where each ψ + i ( x ) is given by
[Equation]
It follows by direct calculation that
[Equation]
otherwise.
[Equation] [Equation]
Using (21)
[Equation]
it is straightforward to construct a matrix C 1 ∈ R ( n +2) ×
( n +2) satisfying
[Equation] [Equation]
The matrix C 1 is nonsingular since both { ψ i ( x ) } n +1 i
=
and { ψ + i ( x ) } n +1 i =0 forms bases for V n (He et al., 2020).
Meanwhile, one can verify the following relation between {{ ϕ i ( x ) } n +1 i =0 and { ψ + i ( x ) } n +1 i =0 by direct calculation:
[Equation]
Hence
[Equation]
we can construct a matrix C 2 ∈ R ( n +2) × ( n +2)
such that
[Equation]
24) explicitly using (23).
Combining (22) and (24) yields
[Equation]
where C = C 2 C − 1 1 .
See Appendix B for how to compute multiplication by C − 1efficiently.
Equation (25) implies that two matrices K and K ϕ are related as follows:
[Equation]
Since M = RKR T , invoking Lemma 1, setting P = C T K − 1 ϕ C and
[Equation]
27) completes the proof, where R was defined in (17).
A computationally cheap way of multiplication by P is presented in Appendix B.
Remark 2 .
In the case of the L 2 -function approximation problem (10), we can construct an alternative preconditioner P alt whose computational cost is a bit cheaper than P .
Let P alt =
C T diag( K ϕ ) − 1 C .
It follows from (26) that
[Equation]
In (10), K ϕ is a mass matrix for the linear finite element method defined on the grid { x i } n +1 i =0 , so that it satisfies κ (diag( K ϕ )
− 1 K ϕ )
=
O (1) .
Then we have
[Equation]
Therefore, P alt = ( RP − 1 alt R T ) − 1 satisfies κ ( P alt K )
=
O (1) by Lemma 1.
It is evident that the computational cost of diag( K ϕ ) − 1 is cheaper than that of K − 1 ϕ .
Finally, we present our proof of Theorem 3.
Proof of Theorem 3.
Let { ω ∗ i , b ∗ i } be a critical point of (8).
If σ ( ω ∗ i · x + b ∗ i ) vanishes on Ω , then we readily get {+E i ( ω ∗ i , b ∗ i ) = 
+}
= E i ( 0 ) .
Now, we suppose that σ ( ω ∗ i · x + b ∗ i ) is nontrivial, i.e., the set
[Equation]
has nonzero measure.
Note that σ ( ω ∗ i · x + b ∗ i ) =
ω ∗ i · x + b ∗ i and σ ′ ( ω ∗ i · x + b ∗ i )
= 1 on D .
Since ∇ E i ( ω ∗ i , b ∗ i )
= 0 , we obtain [Equation] i
[Equation] [Equation]
where the operator A was given in (16) and the integrals are done entrywise.
It follows that
[Equation]
which completes the proof.
Remark 3 .
In Vardi et al.
(2021, Theorem 4.2), a stronger result than Theorem 3 for the following special case was presented:
[Equation]
for some ˆ w i and ˆ b i .
It was proven that any critical point ω ∗ i of (8) such that σ ( ω ∗ i · x + b ∗ i ) is nontrivial is a global minimum.
B C OMPUTATIONAL A SPECTS OF THE P RECONDITIONER IN T HEOREM
In this appendix, we present computational aspects of the preconditioner P presented in Theorem 2.
Although the preconditioner P defined in (27) seems a bit complicated at the first glance, its computation requires only a cheap cost.
We provide a detailed explanation on how to deal with the preconditioner P in implementation.
In addition, we present numerical comparisons between the preconditioned and unpreconditioned conjugate gradient methods to support our theoretical result and highlight the computational efficiency of the preconditioner P .
B. I MPLEMENTATION I SSUES
Three nontrivial parts in the preconditioner P are the inverses of the matrices C 1 , RP − 1 R T , and
K ϕ .
First, we consider how to compute C − 1 1 α when a vector α ∈ R n +2 is given.
We solve a linear system C 1 β
= α in order to obtain C − 1 1 α .
Thanks to the sparsity pattern of C 1 , this linear system can be solved directly by the following O ( n ) elementary arithmetic operations:
[Equation]
Similarly, for α ∈ R n , one can compute ( RP − 1 R T ) − 1 α by solving a linear system
[Equation]
In (29), by the definition (17) of R , we get
[Equation]
for some p, q ∈ R .
Using the fact that the first and last entries of R T β are zero, two constants p and q can be determined by the following system of linear equations:
[Equation]
where 2 : n + 1 means “from the second to ( n + 1) th columns.”
Now, β is obtained by
[Equation]
Finally, we consider how to compute K − 1 ϕ α when a vector α ∈ R n +2 is given.
In most applications, the bilinear form a ( · , · ) is defined in terms of differential operators as in (13).
This implies that the matrix K ϕ is tridiagonal, so that K − 1 ϕ α can be obtained by the Thomas algorithm (see, e.g., Higham (2002, Section 9.6)), which requires only O ( n ) elementary arithmetic operations.
In conclusion, the computation of the preconditioner P can be done efficiently.
B. N UMERICAL R ESULTS
We turn our attention to numerical experiments for the preconditioner P .
We consider the aminimization problem (6) appearing in training of neural networks that solve the L 2 -function approximation problem (12) and the elliptic PDE (14).
In (6), we set the parameters ω and b as follows:
[Equation]
[Table]
We compare the numerical performances of the preconditioned and unpreconditioned conjugate gradient methods solving (6).
The initial guess a (0) is given by the zero vector.
Table 1 presents the number of iterations of the preconditioned and unpreconditioned conjugate gradient methods to meet the stop condition (9) with respect to various numbers of neurons n .
While the number of unpreconditioned iterations increases rapidly as n increases, the number of preconditioned iterations is uniformly bounded with respect to n .
This supports our theoretical result presented in Theorem 2, which claims that the condition number of the preconditioned operator PK is O (1) .
Since the computational cost of P is cheap, we can conclude that the preconditioner P is numerically efficient.
C E FFECT OF B ACKTRACKING FOR C ONVENTIONAL T RAINING A LGORITHMS
In this appendix, we present numerical results that show that the backtracking scheme presented in
Algorithm 3 is useful not only for the proposed NPSC but also for conventional training algorithms such as GD, Adam (Kingma & Ba, 2015), and LSGD (Cyr et al., 2020).
Figure 5 plots the relative energy error
E ( θ ( [Equation] E
[Equation] | E [Equation] | of GD, Adam, and LSGD for solving the problem (11), averaged over 10 random initializations, where k denotes the number of epochs and E ∗ is the energy corresponding to the exact solution of the problem.
The number of neurons used is 2 5 ; while we can observe similar results for the other numbers of neurons, we only provide the result of 2 5 neurons for brevity.
We observe that the algorithms equipped with the backtracking scheme outperform those with constant learning rates τ = 10 − 2 , 10 − 3 , and 10 − 4 in the sense of the convergence rate.
That is, Algorithm 3 seems to successfully find a good learning rate at each iteration of conventional training algorithms as well.
Hence, in Section 4, we employ Algorithm 3 to find learning rates of GD, Adam, and LSGD.
D N UMERICAL I NTEGRATION
This appendix is devoted to numerical integration schemes for computing the integral in our model problem (2).
If d = 1 , i.e., if the domain Ω =
[ x L , x R ] ⊂ R for some x L , x R ∈ R , then we approximate the integral of a function g ( x ) defined on Ω by the following simple trapezoidal rule:
[Equation]
where x L = x 1 < x 2 < · · · < x N = x
R are N uniform sampling points between x L and x R , i.e., x j = x L + j − 1 N − 1 ( x R − x L ) , 1 ≤ j ≤ N .
Approximation properties of the trapezoidal rule (30) can be found in standard textbooks on numerical analysis; see, e.g., Burden et al.
When d ≥ 2 , we adopt the quasi-Monte Carlo method (Caflisch, 1998) based on Halton sequences (Kocis & Whiten, 1997), which is known to overcome the curse of dimensionality in the sense that approximation error bounds independent of the dimension d are available.
In the quasiMonte Carlo method, the integral of a function g ( x ) defined on Ω is approximated by the average of the function evaluated at N sampling points:
[Equation]
where { x j } Nj =1 is a low-discrepancy sequence in Ω defined in terms of Halton sequences.
For the sake of description, we assume that Ω = [0 , 1] d ∈ R d .
Then the k th coordinate of x j ( 1 ≤ j ≤ N , 1 ≤ k ≤ d ) is the number j written in p k -ary representation, inverted, and written after the decimal point, where p k is the k th smallest prime number.
For example, if d = 3 , then the first four points of { x j } Nj =1 is given by
[Equation]
[Equation]
Approximation properties of (31) can be found in Caflisch (1998).