A global convergence theory for deep ReLU implicit networks via over-parameterization

1 INTRODUCTION . 1 ) Background and Motivation : In the last decade , implicit deep learning ( El Ghaoui et al. , 2019 ) have attracted more and more attention . Its popularity is mainly because it generalizes the recursive rules of many widely used neural network architectures . A line of recent works ( Bai et al. , 2019 ; El Ghaoui et al. , 2019 ; Bai et al. , 2020 ) have shown that the implicit neural network architecture is a wider class that includes most current neural network architectures as special cases , such as feed-forward neural networks , convolution neural networks , residual networks , and recurrent neural networks . Moreover , implicit deep learning is also well known for its competitive performance compared to other regular deep neural networks but using significantly fewer computational resources ( Dabre & Fujita , 2019 ; Dehghani et al. , 2018 ; Bai et al. , 2018 ) . Although a line of literature has been shown the superior performance of implicit neural networks experimentally , the theoretical understanding is still limited . To date , it is still unknown if a simple first-order optimization method such as ( stochastic ) gradient descent can converge on an implicit neural network activated by a nonlinear function . Unlike a regular deep neural network , an implicit neural network could have infinitely many layers , resulting in the possibility of divergence of the forward propagation ( El Ghaoui et al. , 2019 ; Kawaguchi , 2021 ) . The main challenge in establishing the convergence of implicit neural network training lies in the fact that , in general , the equilibrium equation of implicit neural networks can not be solved in closed-form . What exacerbates the problem is the well-posedness of the forward propagation . In other words , the equilibrium equation may have zero or multiple solutions . A line of recent studies have suggested a number of strategies to handle this well-posedness challenge , but they all involved reformulation or solving subproblems in each iteration . For example , El Ghaoui et al . ( 2019 ) suggested to reformulate the training as the Fenchel divergence formulation and solve the reformulated optimization problem by projected gradient descent method . However , this requires solving a projection subproblem in each iteration and convergence was only demonstrated numerically . By using an extra softmax layer , Kawaguchi ( 2021 ) established global convergence result of gradient descent for a linear implicit neural network . Unfortunately , their result can not be extended to nonlinear activations , which are critical to the learnability of deep neural networks . This paper proposes a global convergence theory of gradient descent for implicit neural networks activated by nonlinear Rectified Linear Unit ( ReLU ) activation function by using overparameterization . Specifically , we show that random initialized gradient descent with fixed stepsize converges to a global minimum of a ReLU implicit neural network at a linear rate as long as the implicit neural network is overparameterized . Recently , over-parameterization has been shown to be effective in optimizing finite-depth neural networks ( Zou et al. , 2020 ; Nguyen & Mondelli , 2020 ; Arora et al. , 2019 ; Oymak & Soltanolkotabi , 2020 ) . Although the objective function in the training is nonsmooth and non-convex , it can be shown that GD or SGD converge to a global minimum linearly if the width m of each layer is a polynomial of the number of training sample n and the number of layers h , i.e. , m = poly ( n , h ) . However , these results can not be directly applied to implicit neural networks , since implicit neural networks have infinitely many hidden layers , i.e. , h → ∞ , and the well-posedness problem surfaces during the training process . In fact , Chen et al . ( 2018 ) ; Bai et al . ( 2019 ; 2020 ) have all observed that the time and number of iterations spent on forward propagation are gradually increased with the the training epochs . Thus , we have to ensure the unique equilibrium point always exists throughout the training given that the width m is only polynomial of n. 2 ) Preliminaries of Implicit Deep Learning : In this work , we consider an implicit neural network with the transition at the ` -th layer in the following form ( El Ghaoui et al. , 2019 ; Bai et al. , 2019 ) : z ` = σ ( γ√ m Az ` −1 + φ ( x ) ) , ( 1 ) where φ : Rd → Rm is a feature mapping function that transforms an input vector x ∈ Rd to a desired feature vector φ , φ ( x ) , z ` ∈ Rm is the output of the ` -th layer , A ∈ Rm×m is a trainable weight matrix , σ ( u ) = max { 0 , u } is the ReLU activation function , and γ ∈ ( 0 , 1 ) is a fixed scalar to scaleA . As will be shown later in Section 2.1 , γ plays the role of ensuring the existence of the limit z∗ = lim ` →∞ z ` . In general , the feature mapping function φ is a nonlinear function , which extracts features from the low-dimensional input vector x . In this paper , we consider a simple nonlinear feature mapping function φ given by φ ( x ) , 1√ m σ ( Wx ) , ( 2 ) where W ∈ Rm×d is a trainable parameter matrix . As ` → ∞ , an implicit neural network can be considered an infinitely deep neural network . Consequently , z∗ is not only the limit of the sequence { z ` } ∞ ` =0 with z0 = 0 , but it is also the equilibrium point ( or fixed point ) of the equilibrium equation : z∗ = σ ( γ̃Az∗ + φ ) , ( 3 ) where γ̃ , γ/ √ m. In implicit neural networks , the prediction ŷ for the input vector x is the combination of the fixed point z∗ and the feature vector φ , i.e. , ŷ = uTz∗ + vTφ , ( 4 ) where u , v ∈ Rm are trainable weight vectors . For simplicity , we use θ , vec ( A , W , u , v ) to group all training parameters . Given a training data set { ( xi , yi ) } ni=1 , we want to minimize L ( θ ) = n∑ i=1 1 2 ( ŷi − yi ) 2 = 1 2 ‖ŷ − y‖2 , ( 5 ) where ŷ and y are the vectors formed by stacking all the prediction and labels . 3 ) Main Results : Our results are based on the following observations . We first analyze the forward propagation and find that the unique equilibrium point always exists if the scaled matrix γ̃A in Eq . ( 3 ) has an operator norm less than one . Thus , the well-posedness problem is reduced to finding a sequence of scalars { γk } ∞k=1 such that γ̃kA ( k ) is appropriately scaled . To achieve this goal , we show that the operator norm A ( k ) is uniformly upper bounded by a constant over all iterations . Consequently , a fixed scalar γ is enough to ensure the well-posedness of Eq . ( 3 ) . Our second observation is from the analysis of the gradient descent method with infinitesimal step-size ( gradient flow ) . By applying the chain rule with the gradient flow , we derive the dynamics of prediction ŷ ( t ) which is governed by the spectral property of a Gram matrix . In particular , if the smallest eigenvalue of the Gram matrix is lower bounded throughout the training , the gradient descent method enjoys a linear convergence rate . Along with some basic functional analysis results , it can be shown that the smallest eigenvalue of the Gram matrix at initialization is lower bounded if no two data samples are parallel . Although the Gram matrix varies in each iteration , the spectral property is preserved if the Gram matrix is close to its initialization . Thus , the convergence problem is reduced to showing the Gram matrix in latter iterations is close to its initialization . Our third observation is that we find random initialization , over-parameterization , and linear convergence jointly enforce the ( operator ) norms of parameters upper bounded by some constants and close to their initialization . Accordingly , we can use this property to show that the operator norm ofA is upper bounded and the spectral property of the Gram matrix is preserved throughout the training . Combining all these insights together , we can conclude that the random initialized gradient descent method with a constant step-size converges to a global minimum of the implicit neural network with ReLU activation . The main contributions of this paper are summarized as follows : ( i ) By scaling the weight matrix A with a fixed scalar γ , we show that the unique equilibrium point z∗ for each x always exists during the training if the parameters are randomly initialized , even for the nonlinear ReLU activation function . ( ii ) We analyze the gradient flow of implicit neural networks . Despite the non-smooth and nonconvexity of the objective function , the convergence to a global minimum at a linear rate is guaranteed if the implicit neural network is over-parameterized and the data is non-degenerate . ( iii ) Since gradient descent is discretized version of gradient flow , we can show gradient descent with fixed stepsize converges to a global minimum of implicit neural networks at a linear rate under the same assumptions made by the gradient flow analysis , as long as the stepsize is chosen small enough . Notation : For a vector x , ‖x‖ is the Euclidean norm of x . For a matrix A , ‖A‖ is the operator norm of A . If A is a square matrix , then λmin ( A ) and λmax ( A ) denote the smallest and largest eigenvalue ofA , respectively , and λmax ( A ) ≤ ‖A‖ . We denote [ n ] , { 1 , 2 , · · · , n } . 2 WELL-POSEDNESS AND GRADIENT COMPUTATION . In this section , we provide a simple condition for the equilibrium equation ( 3 ) to be well-posed in the sense that the unique equilibrium point exists . Instead of backpropagating through all the intermediate iterations of a forward pass , we derive the gradients of trainable parameters by using the implicit function theorem . In this work , we make the following assumption on parameter initialization . Assumption 1 ( Random Initialization ) . The entries Aij and Wij are randomly initialized by the standard Gaussian distribution N ( 0 , 1 ) , and ui and vi are randomly initialized by the symmetric Bernoulli or Rademacher distribution . Remark 2.1 . This initialization is similar to the approaches widely used in practice ( Glorot & Bengio , 2010 ; He et al. , 2015 ) . The result obtained in this work can be easily extended to the case where the distributions forAij , Wij , ui , and vi are replaced by sub-Gaussian random variables . 2.1 FORWARD PROPAGATION AND WELL-POSEDNESS . In a general implicit neural network , Eq . ( 3 ) is not necessarily well-posed , since it may admit zero or multiple solutions . In this work , we show that scaling the matrix A with γ̃ = γ/ √ m guarantees the existence and uniqueness of the equilibrium point z∗ with random initialization . This follows from a foundational result in random matrix theory as restated in the following lemma . Lemma 2.1 ( Vershynin ( 2018 ) , Theorem 4.4.5 ) . Let A be an m × n random matrix whose entries Aij are independent , zero-mean , and sub-Gaussian random variables . Then , for any t > 0 , we have ‖A‖ ≤ CK ( √ m + √ n + t ) with probability at least 1 − 2e−t2 . Here C > 0 is a fixed constant , and K = maxi , j ‖Aij‖ψ2 . Under Assumption 1 , Lemma 2.1 implies that , with exponentially high probability , ‖A‖ ≤ c √ m for some constant c > 0 . By scalingA by a positive scalar γ̃ , we show that the transition Eq . ( 1 ) is a contraction mapping . Thus , the unique equilibrium point exists with detailed proof in Appendix A.1 . Lemma 2.2 . If ‖A‖ ≤ c √ m for some c > 0 , then for any γ0 ∈ ( 0 , 1 ) , the scalar γ , min { γ0 , γ0/c } uniquely determines the existence of the equilibrium z∗ for every x , and ‖z ` ‖ ≤ 11−γ0 ‖φ‖ for all ` . Lemma 2.2 indicates that equilibria always exist if we can maintain the operator norm of the scaled matrix ( γ/ √ m ) A less than 1 during the training . However , the operator norms of matrix A are changed by the update of the gradient descent . It is hard to use a fixed scalar γ to scale the matrix A over all iterations , unless the operator norm of A is bounded . In Section 3 , we will show ‖A‖ ≤ 2c √ m always holds throughout the training , provided that ‖A ( 0 ) ‖ ≤ c √ m at initialization and the width m is sufficiently large . Thus , by using the scalar γ = min { γ0 , γ0/ ( 2c ) } for any γ0 ∈ ( 0 , 1 ) , equilibria always exist and the equilibrium equation Eq . ( 3 ) is well-posed .

Summary:  In this paper, the authors theoretically analyze the convergence of gradient descent for an implicit neural network with infinite layers with ReLU activation.  The authors show the unique fixed point of the infinite-layered mapping when the weight matrix $\boldsymbol{A}$ has a properly bounded spectral norm.  Using implicit differentiation, the authors show the partial gradient at the fixed point.  Furthermore, the authors show the linear convergence rate by proving the strictly positive-definite of the Gram matrix $\boldsymbol{G}(t)$  (and $\boldsymbol{H}(t)$).