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: The paper presents a proof of exponential convergence to global optimality in the over-parametrization settings for an implicit model with scaled weights parameters. Although existing work has established similar proofs for feedforward explicit neural networks, such methods don't work with non-linearly activated implicit models where the well-posedness issue poses challenges to the training process. The authors shows that by scaling the weights, well-posedness can be ensured. The convergence result is obtained first on continuous settings and is then extended to discrete settings. Numerical experiments on real datasets confirms the finding.