The Convex Geometry of Backpropagation: Neural Network Gradient Flows Converge to Extreme Points of the Dual Convex Program

1 INTRODUCTION . Neural networks ( NNs ) exhibit remarkable empirical performance in various machine learning tasks . However , a full characterization of the optimization and generalization properties of NNs is far from complete . Non-linear operations inherent to the structure of NNs , over-parameterization and the associated highly nonconvex training problem makes their theoretical analysis quite challenging . In over-parameterized models such as NNs , one natural question arises : Which particular solution does gradient descent/gradient flow find in unregularized NN training problems ? Suppose that X ∈ RN×d is the training data matrix and y ∈ { 1 , −1 } N is the label vector . For linear classification problems such as logistic regression , it is known that gradient descent ( GD ) exhibits implicit regularization properties , see , e.g. , ( Soudry et al. , 2018 ; Gunasekar et al. , 2018 ) . To be precise , under certain assumptions , GD converges to the following solution which maximizes the margin : arg min w∈Rd 1 2 ‖w‖22 , s.t . ynwTxn ≥ 1 , n ∈ [ N ] . ( 1 ) Here we denote [ N ] = { 1 , . . . , N } . Recently , there are several results on the implicit regularization of the ( stochastic ) gradient descent method for NNs . In ( Lyu & Li , 2019 ) , for the multi-layer homogeneous network with exponential or cross-entropy loss , with separable training data , it is shown that the gradient flow ( GF ) and GD finds a stationary point of the following non-convex max-margin problem : arg min θ 1 2 ‖θ‖22 , s.t . ynf ( θ ; xn ) ≥ 1 , n ∈ [ N ] , ( 2 ) where f ( θ ; x ) represents the output of the neural network with parameter θ given input x . In ( Phuong & Lampert , 2021 ) , by further assuming the orthogonal separability of the training data , it is shown that all neurons converge to one of the two max-margin classifiers . One corresponds to the data with positive labels , while the other corresponds to the data with negative labels . However , as the max-margin problem of the neural network ( 2 ) is a non-convex optimization problem , the existing results only guarantee that it is a stationary point of ( 2 ) , which can be a local minimizer or even a saddle point . In other words , the global optimality is not guaranteed . In a different line of work ( Pilanci & Ergen , 2020 ; Ergen & Pilanci , 2020b ; a ) , exact convex optimization formulations of two and three-layer ReLU NNs are developed , which have global optimality guarantees in polynomial-time when the data has a polynomial number of hyperplane arrangements , e.g. , in any fixed dimension or with convolutional networks of fixed filter size . In this work , we first derive an equivalent convex program corresponding to the maximal margin problem ( 2 ) . We then consider non-convex subgradient flow for unregularized logistic loss . We show that the limit points of non-convex subgradient flow can be identified via primal-dual correspondence in the convex optimization problem . We then present a sufficient condition on the dual variable to ensure that all stationary points of the non-convex max-margin problem are KKT points of the convex max-margin problem . For certain regular datasets including orthogonal separable data , we show that this sufficient condition on the dual variable holds , thus implies the convergence of gradient flow on the unregularized problem to the global optimum of the non-convex maximalo margin problem ( 2 ) . Consequently , this enables us to fully characterize the implicit regularization of unregularized gradient flow or gradient descent as convex regularization applied to a convex model . 1.1 RELATED WORK . There are several works studying the property of two-layer ReLU networks trained by gradient descent/gradient flow dynamics . The following papers study the gradient descent like dynamics in training two-layer ReLU networks for regression problems . Ma et al . ( 2020 ) show that for two-layer ReLU networks , only a group of a few activated neurons dominate the dynamics of gradient descent . In ( Mei et al. , 2018 ) , the limiting dynamics of stochastic gradient descent ( SGD ) is captured by the distributional dynamics from a mean-field perspective and they utlize this to prove a general convergence result for noisy SGD . Li et al . ( 2020 ) focus on the case where the weights of the second layer are non-negative and they show that the over-parameterized neural network can learn the ground-truth network in polynomial time with polynomial samples . In ( Zhou et al. , 2021 ) , it is shown that mildly over-parameterized student network can learn the teacher network and all student neurons converge to one of the teacher neurons . Beyond ( Lyu & Li , 2019 ) and ( Phuong & Lampert , 2021 ) , the following papers study the classification problems . In ( Chizat & Bach , 2018 ) , under certain assumptions on the training problem , with overparameterized model , the gradient flow can converge to the global optimum of the training problem . For linear separable data , utilizing the hinge loss for classification , Wang et al . ( 2019 ) introduce a perturbed stochastic gradient method and show that it can attain the global optimum of the training problem . Similarly , for linear separable data , Yang et al . ( 2021 ) introduce a modified loss based on the hinge loss to enable ( stochastic ) gradient descent find the global minimum of the training problem , which is also globally optimal for the training problem with the hinge loss . 1.2 PROBLEM SETTING . We focus on two-layer neural networks with ReLU activation , i.e. , f ( θ , X ) = ( XW1 ) +w2 , ( 3 ) where W1 ∈ Rd×m , w2 ∈ Rm and θ = ( W1 , w2 ) represents the parameter . Due to the ReLU activation , this neural network is homogeneous , i.e. , for any scalar c > 0 , we have f ( cθ ; X ) = c2f ( θ ; X ) . The training problem is given by min θ N∑ n=1 ` ( ynf ( θ ; xn ) ) , ( 4 ) where ` ( q ) : R → R+ is the loss function . We focus on the logistic , i.e , cross-entropy loss , i.e. , ` ( q ) = log ( 1 + exp ( −q ) ) . Then , we briefly review gradient descent and gradient flow . The gradient descent takes the update rule θ ( t+ 1 ) = θ ( t ) − η ( t ) g ( t ) , where g ( t ) ∈ ∂◦L ( θ ( t ) ) and ∂◦ represents the Clarke ’ s subdifferential . The gradient flow can be viewed as the gradient descent with infinitesimal step size . The trajectory of the parameter θ during training is an arc θ : [ 0 , +∞ ) → Θ , where Θ = { θ = ( W1 , w2 ) |W1 ∈ Rd×m , W2 ∈ Rm } . More precisely , the gradient flow is given by the differential inclusion d dt θ ( t ) ∈ −∂◦L ( θ ( t ) ) , for t ≥ 0 , a.e .. 2 MAIN RESULTS . In this section , we present our main results and defer the detailed analysis to the following sections . Consider the more general multi-class version of the problem withK classes . Suppose that ȳ ∈ [ K ] N is the label vector . Let Y = ( yn , k ) n∈ [ N ] , k∈ [ K ] ∈ RN×K be the encoded label matrix such that yn , k = { 1 , if ȳn = k , − 1 , otherwise . Similarly , we consider the following two-layer vector-output neural networks with ReLU activation : F ( Θ , X ) = f1 ( θ1 , X ) ... fK ( θK , X ) = ( XW ( 1 ) 1 ) +w ( 1 ) 2 ... ( XW ( K ) 1 ) +w ( K ) 2 , where we write Θ = ( θ1 , . . . , θK ) . For k = 1 , . . . , K , we have θk = ( W ( k ) 1 , w ( k ) 2 ) where W ( k ) 1 ∈ RN×m and w ( k ) 2 ∈ Rm . One can view each of the K outputs of F ( Θ , X ) as the output of a two-layer scalar-output neural network . Consider the following training problem : min Θ K∑ k=1 N∑ n=1 ` ( yn , kfk ( θk , xn ) ) . ( 5 ) According to ( Lyu & Li , 2019 ) , the gradient flow and the gradient descent finds a stationary point of the following non-convex max-margin problem : arg min Θ K∑ k=1 1 2 ‖θk‖22 , s.t . yn , kf ( θk ; xn ) ≥ 1 , n ∈ [ N ] , k ∈ [ K ] . ( 6 ) Denote the set of all possible hyperplane arrangement as P = { diag ( I ( Xw ≥ 0 ) ) |w ∈ Rd } , ( 7 ) and let p = |P| . We can also write P = { D1 , . . . , Dp } . From ( Cover , 1965 ) , we have an upper bound p ≤ 2r ( e ( N−1 ) r ) r where r = rank ( X ) . We first reformulate ( 6 ) as convex optimization . Proposition 1 The non-convex problem ( 6 ) is equivalent to the following convex program min K∑ k=1 p∑ j=1 ( ‖uj , k‖2 + ‖u′j , k‖2 ) , s.t . diag ( yk ) p∑ j=1 DjX ( uj , k − u′j , k ) ≥ 1 , ( 2Dj − I ) Xuj , k ≥ 0 , ( 2Dj − I ) Xu′j , k ≥ 0 , j ∈ [ p ] , k ∈ [ K ] . ( 8 ) where yk is the k-th column of Y . The dual problem of ( 8 ) is given by max tr ( ΛTY ) , s.t . diag ( yk ) λk 0 , max ‖w‖2≤1 |λTk ( XTw ) +| ≤ 1 , k ∈ [ K ] . ( 9 ) where λk is the k-th column of Λ . We present the detailed derivation of the convex formulation ( 8 ) and its dual problem ( 9 ) in the appendix . Given u ∈ Rd , we define D ( u ) = diag ( I ( Xu > 0 ) ) . ( 10 ) For two vectors u , v ∈ Rd , we define the cosine angle between u and v by cos∠ ( u , v ) = uTv ‖u‖2‖v‖2 . 2.1 OUR CONTRIBUTIONS . The following theorem illustrate that for neurons satisfying sign ( yTk ( Xw ( k ) 1 , i ) + ) = sign ( w ( k ) 2 , i ) at initialization , w ( k ) 1 , i align to the direction of ±XTD ( w ( k ) 1 , i ) yk at a certain time T , depending on sign ( w ( k ) 2 , ik , + ) at initialization . In Section 2.3 , we show that these are dual extreme points of ( 8 ) . Theorem 1 Consider the K-class classification training problem ( 5 ) for any dataset . Suppose that the neural network is scaled at initialization such that ‖w ( k ) 1 , i ‖2 = |w ( k ) 2 , i | for i ∈ [ m ] and k ∈ [ K ] . Assume that at initialization , for k ∈ [ K ] , there exists neurons ( w ( k ) 1 , ik , w ( k ) 2 , ik ) such that sign ( yTk ( Xw ( k ) 1 , ik ) + ) = sign ( w ( k ) 2 , ik ) = s , ( 11 ) where s ∈ { 1 , −1 } . Consider the subgradient flow applied to the non-convex problem ( 5 ) . Let δ ∈ ( 0 , 1 ) . Suppose that the initialization is sufficiently close to the origin . Then , for k ∈ [ K ] , there exist T = T ( δ , k ) such that cos∠ ( w ( k ) 1 , ik ( T ) , sXTD ( w ( k ) 1 , ik ( T ) ) yk ) ≥ 1− δ . Next , we impose conditions on the dataset to prove a stronger global convergence results on the flow . We say that the dataset ( X , ȳ ) is orthogonal separable among multiple classes if for all n , n′ ∈ [ N ] , xTnxn′ > 0 , if ȳn = ȳn′ , xTnxn′ ≤ 0 , if ȳn 6= ȳn′ . For orthogonal separable dataset among multiple classes , the subgradient flow for the non-convex problem ( 5 ) can find the global optimum of ( 6 ) up to a scaling constant . Theorem 2 Suppose that ( X , ȳ ) ∈ RN×d × [ K ] N is orthogonal separable among multiple classes . Consider the non-convex subgradient flow applied to the non-convex problem ( 5 ) . Suppose that the initialization is sufficiently close to the origin and scaled as in Theorem 1 . Then , the non-convex subgradient flow converges to the global optimum of the convex program ( 8 ) and hence the non-convex objective ( 6 ) up to scaling . Therefore , the above result characterizes the implicit regularization of unregularized gradient flow as convex regularization , i.e. , group ` 1 norm , in the convex formulation ( 8 ) . It is remarkable that group sparsity is enforced by small initialization magnitude with no explicit form of regularization .

Summary: The paper studies the subgradient flows when training a two-layer ReLU neural network. To this end, the non-convex max-margin problem is reformulated as a convex optimization problem. The authors then analyze the dual extreme points of the convex formulation and show the implicit regularization of unregularized gradient flow as convex regularization. Then, for the binary classification problem, it is proven that the KKT points of the non-convex max-margin problem correspond to the KKT points of the convex max-margin problem if the direction is dual feasible. The paper then demonstrates that this is the case under some conditions on spike-free matrices and orthogonal separable data. Finally it is shown that if the dataset is orthogonal separable and initialized sufficiently close to zero, the limiting point of the gradient flow is the global minimizer of the max-margin problem.