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 paper studies the gradient flow training dynamics for deep and wide ResNet with the mean-field scaling. For a similar setting as (Lu et al., 2020), they report a limiting model with finite-dimensional approximation, and its global convergence result. The authors also justify the well-posedness of the resulting distributional dynamics, which seems only heuristically defined in prior work.