On the Global Convergence of Gradient Descent for multi-layer ResNets in the mean-field regime

1 INTRODUCTION . Training of multi-layer neural networks ( NN ) requires us to find weights in the network such that its outputs perfectly match the prescribed outputs for a given set of training data . The usual approach is to formulate this problem as a nonconvex minimization problem and solve it with a first-order optimization method based on gradient descent ( GD ) . Extensive computational experience shows that in the overparametrized regime ( where the total number of parameters in the NN far exceeds the minimum number required to fit the training data ) , GD methods run for sufficiently many iterations consistently find a global minimum achieving the zero-loss property , that is , a perfect fit to the training data . What is the mechanism that allows GD to perform so well on this large-scale nonconvex problem ? Part of the explanation is that in the overparametrized case , the parameter space contains many global minima , and some evidence suggests that they are distributed throughout the space , making it easier for the optimization process to find one such solution . Many approaches have been taken to characterize this phenomenon more rigorously , including landscape analysis , the neural tangent kernel approach , and mean-field analysis . All such viewpoints aim to give an idea of the structure and size of the NN required to ensure global convergence . Our approach in this paper is based on mean-field analysis and gradient-flow analysis , the latter being the continuous and mean-field limit of GD . We will examine residual neural networks ( ResNets ) , and study how deep and wide a ResNet needs to be to match the data with high accuracy and high confidence . To relax the assumptions on the activation function as far as possible , we follow the setup in ( Chizat & Bach , 2018 ) , which requires this function to be either 2-homogeneous or partially 1-homogeneous . We show that both depth and width of the NN depend algebraically on and η , which are the accuracy and confidence levels , respectively . Mean-field analysis translates the training process of the ResNet to a gradient-flow partial differential equation ( PDE ) . The training process evolves weights on connections between neurons . When dealing with wide neural networks , instead of tracing the evolution of each weight individually , one can record the evolution of the full distribution of the weight configuration . This perspective translates the coupled ordinary differential equation system ( ODE ) that characterizes evolution of individual weights into a PDE ( the gradient-flow equation ) characterizing the evolution of the distribution . The parameters in the PDE naturally depend on the properties of the activation functions . Gradient-flow analysis is used to show that the PDE drives the solution to a point where the loss function becomes zero . We obtain our results on zero-loss training of ResNet with GD by translating the zero-loss property of the gradient-flow PDE back to the discrete-step setting . This strategy of the proof was taken in an earlier paper ( Ding et al. , 2021 ) where multi-layer ResNets were also analyzed . The main difference in this current paper is that the assumptions on the activation function and the initial training state for obtaining the global convergence are both much relaxed . This paper adopts the setup from ( Chizat & Bach , 2018 ) of minimal Lipschitz continuity requirements on the activation function . Furthermore , the paper ( Ding et al. , 2021 ) required a dense support condition to be satisfied on the final parameter configuration has a support condition . This condition is hard to justify in any realistic setting , and is discared from current paper . Further details on these issues appear in Section 3 . We discuss the setup of the problem and formally derive the continuous and mean-field limits in Section 2 . In Section 3 , we discuss related work , identify our contribution and present the main theorem in its general terms . After precise definitions and assumptions are specified in Section 4 , we present the two main ingredients in the proof strategy . The mean-field limit is obtained by connecting the training process of the ResNet to a gradient-flow PDE in Section 5 , and the zero-loss property of the limiting PDE is verified in Section 6 . The main theorem is a direct corollary of Theorem 5.1 and Theorem 6.1 ( or Theorem 6.2 ) . 2 RESNET AND GRADIENT DESCENT . The ResNet can be specified as follows : zl+1 ( x ) = zl ( x ) + 1 ML M∑ m=1 f ( zl ( x ) , θl , m ) , l = 0 , 1 , . . . , L− 1 , ( 1 ) where M and L are the width and depth , respectively ; z0 ( x ) = x ∈ Rd is the input data ; and zL ( x ) is the output from the last layer . The configuration of the NN is encoded in parameters ΘL , M = { θl , m } L−1 , Ml=0 , m=1 , where each parameter θl , m is a vector in R k and f : Rd × Rk → Rd is the activation function . The formulation ( 1 ) covers “ conventional ” ResNets , which have the specific form zl+1 ( x ) = zl ( x ) + 1 ML M∑ m=1 Ul , mσ ( W > l , mzl ( x ) + bl , m ) , l = 0 , 1 , . . . , L− 1 , where Wl , m , Ul , m ∈ Rd , bl.m ∈ R , and σ is the ReLU activation function . In this example , we have θl , m = ( Wl , m , Ul , m , bl , m ) ∈ Rk , with k = 2d+ 1 . Denote by ZΘL , M ( l ; x ) the output of the ResNet defined by ( 1 ) . ( This quantity is the same as zL ( x ) defined above , but we use this alternative notation to emphasize the dependece on parameters ΘL , M . ) The goal of training ResNet is to seek parameters ΘL , M that minimize the following mismatch or loss function : E ( ΘL , M ) = Ex∼µ [ 1 2 ( g ( ZΘL , M ( L ; x ) ) − y ( x ) ) 2 ] , ( 2 ) where g ( x ) : Rd → R is a given measuring function , y ( x ) ∈ R is the label corresponding to x , and µ is the probability from which the data x is drawn . Classical gradient descent updates the parameters according to the formula Θn+1L , M = Θ n L , M − h∇ΘE ( ΘnL , M ) , where h is the step length . In the limit as h→ 0 , the updating process can be characterized by the following ODE ( Chizat & Bach , 2018 , Def 2.2 ) ( rescaled by L , M ) : dΘL , M ( s ) ds = −ML∇ΘE ( ΘL , M ) , for s ≥ 0 , ( 3 ) where s represents pseudo-time , the continuous analog of the discrete stepping process . 2.1 THE CONTINUOUS LIMIT AND THE MEAN-FIELD LIMIT . The continuous limit of ( 1 ) is obtained when the ResNet is infinitely deep , with L → ∞ . By reparametrizing the indices l = [ 0 , · · · , L− 1 ] with the continuous variable t ∈ [ 0 , 1 ] , we can view z in ( 1 ) as a function in t that satisfies a coupled ODE , with 1/L being the stepsize in t. Accordingly , θl , m can be recast as θm ( t = l/L ) , and denoting Θ ( t ) = { θm ( t ) } Mm=1 , we can write the continuous limit of ( 1 ) as dz ( t ; x ) dt = 1 M M∑ m=1 f ( z ( t ; x ) , θm ( t ) ) , t ∈ [ 0 , 1 ] , with z ( 0 ; x ) = x . ( 4 ) Extending ( 2 ) , we define the cost functional E as E ( Θ ) = Ex∼µ [ 1 2 ( g ( ZΘ ( 1 ; x ) ) − y ( x ) ) 2 ] , ( 5 ) where ZΘ ( t ; x ) solves ( 4 ) for a given collection Θ ( t ) of the M functions { θm ( t ) } . Similar to ( 3 ) , we can use GD to find the configuration of Θ ( t ) that minimizes ( 5 ) by making Θ ( t ) flow in the descending direction of E ( Θ ) . Denote s the pseudo-time of the training process , and Θ ( s ; t ) the collection of functions at the training time s : ∂Θ ∂s = −M δE δΘ ∣∣∣∣ Θ ( s ; · ) , s > 0 , t ∈ [ 0 , 1 ] ( 6 ) where δEδΘ is the functional derivative of E with respect to Θ , and thus a list of M functions of t for every fixed s. The mean-field limit is obtained by making the ResNet infinitely wide , that is , M →∞ . Considering that the right hand side of ( 4 ) has the form of an expectation , it approaches an integral in the limit , with respect to a certain probability density . Denoting this PDF by ρ ( θ , t ) ∈ C ( [ 0 , 1 ] ; P2 ) 1 , and assuming that the θm are drawn from it , the ODE for z translates to the following : dz ( t ; x ) dt = ∫ Rk f ( z ( t ; x ) , θ ) dρ ( θ , t ) , t ∈ [ 0 , 1 ] with z ( 0 ; x ) = x . ( 7 ) Mimicking ( 5 ) , we define the following cost function in the mean-field setting : E ( ρ ) = Ex∼µ [ 1 2 ( g ( Zρ ( 1 ; x ) ) − y ( x ) ) 2 ] , ( 8 ) where Zρ ( t ; x ) is the solution to ( 7 ) for a given ρ . Then , similar to the gradient flow for ΘL , M and Θ ( t ) , the probability distribution ρ that encodes the configuration of θ flows in the descending direction of E ( ρ ) in pseudo-time s. Since ρ ( θ , t , s ) needs to be a probability density for all s and t , its evolution in s is characterized by a gradient flow in the Wasserstein metric ( Chizat & Bach , 2018 ; Lu et al. , 2020 ; Ding et al. , 2021 ) : ∂ρ ∂s = ∇θ · ( ρ∇θ δE δρ ∣∣∣∣ ρ ( · , · , s ) ) , s > 0 , t ∈ [ 0 , 1 ] with ρ ( θ , t , 0 ) = ρini ( θ , t ) , ( 9 ) where δEδρ is the functional derivative with respect to ρ , and thus a function of ( θ , t ) for every fixed s. Using the classical calculus-of-variations method , this functional derivative can be computed as : δE δρ ∣∣∣∣ ρ ( θ , t ) = Ex∼µ ( p > ρ ( t ; x ) f ( Zρ ( t ; x ) , θ ) ) , ( 10 ) where pρ ( · ; x ) , parameterized by x , maps [ 0 , 1 ] → Rd , and is a vector solution to the following ODE : dp > ρ dt = −p > ρ ∫ Rk ∂zf ( Zρ , θ ) ρ ( θ , t ) dθ . ( 11 ) 1A collection of probability distribution that is continuous in t and has bounded second moment in θ for all t. The definition is to be made rigorous in Def 4.1. with pρ ( t = 1 ; x ) = ( g ( Zρ ( 1 ; x ) ) − y ( x ) ) ∇g ( Zρ ( 1 ; x ) ) . In the later sections , to emphasize the s dependence , we use δE ( Θ ( s ) ) δΘ and δE ( ρ ( s ) ) δρ to denote δE δΘ ∣∣ Θ ( s ; · ) and δE δρ ∣∣∣ ρ ( · , · , s ) respectively . As a summary , to update ρ ( θ , t , s ) to ρ ( θ , t , s + δs ) with an infinitesimal δs , we solve ( 7 ) for Zρ ( t ; x ) , using the given ρ ( θ , t , s ) , and compute pρ using ( 11 ) . This then allows us to compute δE ( ρ ( s ) ) δρ ( θ , t ) which , in turn , yields ρ ( θ , t , s + δs ) from ( 9 ) . In ( 11 ) , ∂zf is a d × d matrix that stands for the Jacobian of f with respect to its z argument .

Summary: This work analyzes infinitely deep and infinitely wide residual networks. The infinite depth limit is the neural ODE. The infinite width limit is the mean-field limit. The main contribution establishing global convergence of the infinitely wide and deep ResNet. Finally, the global convergence of the infinite limit is translated to the finite but large setup.