On the Theory of Implicit Deep Learning: Global Convergence with Implicit Layers

1 INTRODUCTION . A feedforward deep neural network consists of a stack of H layers , where H is the depth of the network . The value for the depthH is typically a hyperparameter and is chosen by network designers ( e.g. , ResNet-101 in He et al . 2016 ) . Each layer computes some transformation of the output of the previous layer . Surprisingly , several recent studies achieved results competitive with the state-ofthe-art performances by using the same transformation for each layer with weight tying ( Dabre & Fujita , 2019 ; Bai et al. , 2019b ; Dehghani et al. , 2019 ) . In general terms , the output of the l-th layer with weight tying can be written by z ( l ) = h ( z ( l−1 ) ; x , θ ) for l = 1 , 2 , . . . , H − 1 , ( 1 ) where x is the input to the neural network , z ( l ) is the output of the l-th layer ( with z ( 0 ) = x ) , θ represents the trainable parameters that are shared among different layers ( i.e. , weight tying ) , and z ( l−1 ) 7→ h ( z ( l−1 ) ; x , θ ) is some continuous function that transforms z ( l−1 ) given x and θ . With weight tying , the memory requirement does not increase as the depth H increases in the forward pass . However , the efficient backward pass to compute gradients for training the network usually requires to store the values of the intermediate layers . Accordingly , the overall computational requirement typically increases as the finite depth H increases even with weight tying . Instead of using a finite depth H , Bai et al . ( 2019a ) recently introduced the deep equilibrium model that is equivalent to running an infinitely deep feedforward network with weight tying . Instead of running the layer-by-layer computation in equation ( 1 ) , the deep equilibrium model uses rootfinding to directly compute a fixed point z∗ = liml→∞ z ( l ) , where the limit can be ensured to exist by a choice of h. We can train the deep equilibrium model with gradient-based optimization by analytically backpropagating through the fixed point using implicit differentiation ( e.g. , Griewank & Walther , 2008 ; Bell & Burke , 2008 ; Christianson , 1994 ) . With numerical experiments , Bai et al . ( 2019a ) showed that the deep equilibrium model can improve performance over previous state-ofthe-art models while significantly reducing memory consumption . Despite the remarkable performances of deep equilibrium models , our theoretical understanding of its properties is yet limited . Indeed , immense efforts are still underway to mathematically understand deep linear networks , which have finite values for the depth H without weight tying ( Saxe et al. , 2014 ; Kawaguchi , 2016 ; Hardt & Ma , 2017 ; Laurent & Brecht , 2018 ; Arora et al. , 2018 ; Bartlett et al. , 2019 ; Du & Hu , 2019 ; Arora et al. , 2019a ; Zou et al. , 2020b ) . In deep linear networks , the function h at each layer is linear in θ and linear in x ; i.e. , the map ( x , θ ) 7→ h ( z ( l−1 ) ; x , θ ) is bilinear . Despite this linearity , several key properties of deep learning are still present in deep linear networks . For example , the gradient dynamics is nonlinear and the objective function is nonconvex . Accordingly , understanding gradient dynamics of deep linear networks is considered to be a valuable step towards the mathematical understanding of deep neural networks ( Saxe et al. , 2014 ; Arora et al. , 2018 ; 2019a ) . In this paper , inspired by the previous studies of deep linear networks , we initiate a theoretical study of gradient dynamics of deep equilibrium linear models as a step towards theoretically understanding general deep equilibrium models . As we shall see in Section 2 , the function h at each layer is nonlinear in θ for deep equilibrium linear models , whereas it is linear for deep linear networks . This additional nonlinearity is essential to enforce the existence of the fixed point z∗ . The additional nonlinearity , the infinite depth , and weight tying are the three key proprieties of deep equilibrium linear models that are absent in deep linear networks . Because of these three differences , we can not rely on the previous proofs and results in the literature of deep linear networks . Furthermore , we analyze gradient dynamics , whereas Kawaguchi ( 2016 ) ; Hardt & Ma ( 2017 ) ; Laurent & Brecht ( 2018 ) studied the loss landscape of deep linear networks . We also consider a general class of loss functions for both regression and classification , whereas Saxe et al . ( 2014 ) ; Arora et al . ( 2018 ) ; Bartlett et al . ( 2019 ) ; Arora et al . ( 2019a ) ; Zou et al . ( 2020b ) analyzed gradient dynamics of deep linear networks in the setting of the square loss . Accordingly , we employ different approaches in our analysis and derive qualitatively and quantitatively different results when compared with previous studies . In Section 2 , we provide theoretical and numerical observations that further motivate us to study deep equilibrium linear models . In Section 3 , we mathematically prove convergence of gradient dynamics to global minima and the exact relationship between the gradient dynamics of deep equilibrium linear models and that of the adaptive trust region method . Section 5 gives a review of related literature , which strengthens the main motivation of this paper along with the above discussion ( in Section 1 ) . Finally , Section 6 presents concluding remarks on our results , the limitation of this study , and future research directions . 2 PRELIMINARIES . We begin by defining the notation . We are given a training dataset ( ( xi , yi ) ) ni=1 of n samples where xi ∈ X ⊆ Rmx and yi ∈ Y ⊆ Rmy are the i-th input and the i-th target output , respectively . We would like to learn a hypothesis ( or predictor ) from a parametric family H = { fθ : Rmx → Rmy | θ ∈ Θ } by minimizing the objective function L ( called the empirical loss ) over θ ∈ Θ : L ( θ ) = ∑n i=1 ` ( fθ ( xi ) , yi ) , where θ is the parameter vector and ` : Rmy × Y → R≥0 is the loss function that measures the difference between the prediction fθ ( xi ) and the target yi for each sample . For example , when the parametric family of interest is the class of linear models as H = { x 7→Wφ ( x ) |W ∈ Rmy×m } , the objective function L can be rewritten as : L0 ( W ) = n∑ i=1 ` ( Wφ ( xi ) , yi ) , ( 2 ) where the feature map φ is an arbitrary fixed function that is allowed to be nonlinear and is chosen by model designers to transforms an input x ∈ Rmx into the desired features φ ( x ) ∈ Rm . We use vec ( W ) ∈ Rmym to represent the standard vectorization of a matrix W ∈ Rmy×m . Instead of linear models , our interest in this paper lies on deep equilibrium models . The output z∗ of the last hidden layer of a deep equilibrium model is defined by z∗ = lim l→∞ z ( l ) = lim l→∞ h ( z ( l−1 ) ; x , θ ) = h ( z∗ ; x , θ ) , ( 3 ) where the last equality follows from the continuity of z 7→ h ( z ; x , θ ) ( i.e. , the limit commutes with the continuous function ) . Thus , z∗ can be computed by solving the equation z∗ = h ( z∗ ; x , θ ) without running the infinitely deep layer-by-layer computation . The gradients with respect to parameters are computed analytically via backpropagation through z∗ using implicit differentiation . 2.1 DEEP EQUILIBRIUM LINEAR MODELS . A deep equilibrium linear model is an instance of the family of deep equilibrium models and is defined by setting the function h at each layer as follows : h ( z ( l−1 ) ; x , θ ) = γσ ( A ) z ( l−1 ) + φ ( x ) , ( 4 ) where θ = ( A , B ) with two trainable parameter matrices A ∈ Rm×m and B ∈ Rmy×m . Along with a positive real number γ ∈ ( 0 , 1 ) , the nonlinear function σ is used to ensure the existence of the fixed point and is defined by σ ( A ) ij = exp ( Aij ) ∑m k=1 exp ( Akj ) . The class of deep equilibrium linear models is given by H = { x 7→ B ( liml→∞ z ( l ) ( x , A ) ) | A ∈ Rm×m , B ∈ Rmy×m } , where z ( l ) ( x , A ) = γσ ( A ) z ( l−1 ) + φ ( x ) . Therefore , the objective function for deep equilibrium linear models can be written as L ( A , B ) = n∑ i=1 ` ( B ( lim l→∞ z ( l ) ( xi , A ) ) , yi ) . ( 5 ) The outputs of deep equilibrium linear models fθ ( x ) = B ( liml→∞ z ( l ) ( x , A ) ) are nonlinear and non-multilinear in the optimization variable A . This is in contrast to linear models and deep linear networks . From the optimization viewpoint , linear modelsWφ ( x ) are called linear because they are linear in the optimization variables W . Deep linear networks W ( H ) W ( H−1 ) · · ·W ( 1 ) x are multilinear in the optimization variables ( W ( 1 ) , W ( 2 ) , . . . , W ( H ) ) ( this holds also when we replace x by φ ( x ) ) . This difference creates a challenge in the analysis of deep equilibrium linear models . Following previous works on gradient dynamics of different machine learning models ( Saxe et al. , 2014 ; Ji & Telgarsky , 2020 ) , we consider the process of learning deep equilibrium linear models via gradient flow : d dt At = − ∂L ∂A ( At , Bt ) , d dt Bt = − ∂L ∂B ( At , Bt ) , ∀t ≥ 0 , ( 6 ) where ( At , Bt ) represents the model parameters at time t with an arbitrary initialization ( A0 , B0 ) . Throughout this paper , a feature map φ and a real number γ ∈ ( 0 , 1 ) are given and arbitrary ( except in experimental observations ) and we omit their universal quantifiers for the purpose of brevity . 2.2 PRELIMINARY OBSERVATION FOR ADDITIONAL MOTIVATION . Our analysis is chiefly motivated as a step towards mathematically understanding general deep equilibrium models ( as discussed in Sections 1 and 5 ) . In addition to the main motivation , this section provides supplementary motivations through theoretical and numerical preliminary observations . In general deep equilibrium models , the limit , liml→∞ z ( l ) , is not ensured to exist ( see Appendix C ) . In this view , the class of deep equilibrium linear models is one instance where the limit is guaranteed to exist for any values of model parameters as stated in Proposition 1 : Proposition 1 . Given any ( x , A ) , the sequence ( z ( l ) ( x , A ) ) l in Euclidean space Rm converges . Proof . We use the nonlinearity σ to ensure the convergence in our proof in Appendix A.5 . Proposition 1 shows that we can indeed define the deep equilibrium linear model with liml→∞ z ( l ) = z∗ ( x , A ) . Therefore , understanding this model is a sensible starting point for theory of general deep equilibrium models . As our analysis has been mainly motivated for theory , it would be of additional value to discuss whether the model would also make sense in practice , at least potentially in the future . Consider an ( unknown ) underling data distribution P ( x , y ) = P ( y|x ) P ( x ) . Intuitively , if the mean of the P ( y|x ) is approximately given by a ( true unknown ) deep equilibrium linear model , then it would make sense to use the parametric family of deep equilibrium linear models to have the inductive bias in practice . To confirm this intuition , we conducted numerical simulations . To generate datasets , we first drew uniformly at random 200 input images for input data points xi from a standard image dataset — CIFAR-10 , CIFAR-100 or Kuzushiji-MNIST ( Krizhevsky & Hinton , 2009 ; Clanuwat et al. , 2019 ) . We then generated targets as yi = B∗ ( liml→∞ z ( l ) ( xi , A∗ ) ) + δi where δi i.i.d.∼ N ( 0 , 1 ) . Each entry of the true ( unknown ) matrices A∗ and B∗ was independently drawn from the standard normal distribution . For each dataset generated in this way , we used stochastic gradient descent ( SGD ) to train linear models , fully-connected feedforward deep neural networks with ReLU nonlinearity ( DNNs ) , and deep equilibrium linear models . For all models , we fixed φ ( x ) = x . See Appendix D for more details of the experimental settings . The results of this numerical test are presented in Figure 1 . In the figure , the plotted lines indicate the mean values over five random trials whereas the shaded regions represent error bars with one standard deviation . The plots for linear models and deep equilibrium linear models are shown with the best and worst learning rates ( separately for each model in terms of the final test errors at epoch = 5000 ) from the set of learning rates SLR = { 0.01 , 0.005 , 0.001 , 0.0005 , 0.0001 , 0.00005 } . The plots for DNNs are shown with the best learning rates ( separately for each depthH ) from the set SLR . As can be seen , all models preformed approximately the same at initial points , but deep equilibrium linear models outperformed both linear models and DNNs in test errors after training , confirming our intuition above . Moreover , we confirmed qualitatively same behaviors with four more datasets as well as for DNNs with and without bias terms in Appendix D. These observations additionally motivated us to study deep equilibrium linear models to obtain our main results in the next section . The purpose of these experiments is to provide a secondary motivation for our theoretical analyses .

Summary: This submission studies the dynamics and convergence properties of "deep equilibrium models", which are parametric fixed-point iterations corresponding to the infinite depth limit of "weight-tied" neural networks. As the authors point out, these networks differ from deep linear networks and networks in the NTK scaling in that the optimization remains nonlinear w/r/t the parameters. The authors prove two results: first, they establish linear convergence to the global minimum under the relatively strict assumption of a "local" PL-inequality; secondly, they show that the dynamics of the deep equilibrium models differs from gradient descent dynamics and, in fact, is related to a trust region Newton method.