On Margin Maximization in Linear and ReLU Networks

1 INTRODUCTION . A central question in the theory of deep learning is how neural networks generalize even when trained without any explicit regularization , and when there are far more learnable parameters than training examples . In such optimization problems there are many solutions that label the training data correctly , and gradient descent seems to prefer solutions that generalize well ( Zhang et al. , 2016 ) . Hence , it is believed that gradient descent induces an implicit bias ( Neyshabur et al. , 2014 ; 2017 ) , and characterizing this bias has been a subject of extensive research in recent years . A main focus in the theoretical study of implicit bias is on homogeneous neural networks . These are networks where scaling the parameters by any factor α > 0 scales the predictions by αL for some constant L. For example , fully-connected and convolutional ReLU networks without bias terms are homogeneous . Lyu & Li ( 2019 ) proved that in linear and ReLU homogeneous networks trained with the exponential or the logistic loss , if gradient flow converges to zero loss1 , then the direction to which the parameters of the network converge can be characterized as a first order stationary point ( KKT point ) of the maximum margin problem in the parameter space . Namely , the problem of minimizing the ` 2 norm of the parameters under the constraints that each training example is classified correctly with margin at least 1 . They also showed that this KKT point satisfies necessary conditions for optimality . However , the conditions are not known to be sufficient even for local optimality . It is analogous to showing that some unconstrained optimization problem converges to a point with gradient zero , without proving that it is either a global or a local minimum . In this work we consider several architectures of homogeneous neural networks with linear and ReLU activations , and study whether the aforementioned KKT point is guaranteed to be a global optimum of the maximum margin problem , a local optimum , or neither . Perhaps surprisingly , our results imply that in many cases , such as depth-2 fully-connected ReLU networks and depth-2 diagonal linear networks , the KKT point may not even be a local optimum of the maximum-margin problem . On the flip side , we identify multiple settings where a local or global optimum can be guaranteed . We now describe our results in a bit more detail . We denote by N the class of neural networks without bias terms , where the weights in each layer might have an arbitrary sparsity pattern , and 1They also assumed directional convergence , but ( Ji & Telgarsky , 2020 ) later showed that this assumption is not required . Table 1 : Results on depth-2 networks . Linear ReLU Fully-connected Global ( Thm . 3.1 ) Not local ( Thm . 3.2 ) Nno-share Not local ( Thm . 4.1 ) Not local ( Thm . 3.2 ) Nno-share assuming non-zero weights vectors Global ( Thm . 4.2 ) Not local ( Thm . 4.2 ) Nno-share assuming non-zero inputs to all neurons Global ( Thm . 4.2 ) Local , Not global ( Thm . 4.3 ) N assuming nonzero inputs to all neurons Not local ( Thm . 4.4 ) Not local ( Thm . 4.4 ) Table 2 : Results on deep networks . Linear ReLU Fully-connected Global ( Thm . 3.1 ) Not local ( Thm . 3.2 ) Nno-share assuming non-zero inputs to all neurons Not local ( Thm . 5.1 ) Not local ( Thm . 5.1 ) N - max margin for each layer separately Global ( Thm . 5.2 ) Not local ( Thm . 5.3 ) N - max margin for each layer separately , assuming non-zero inputs to all neurons Global ( Thm . 5.2 ) Local , Not global ( Thm . 5.4 ) weights might be shared2 . The class N contains , for example , convolutional networks . Moreover , we denote by Nno-share the subclass of N that contains only networks without shared weights , such as fully-connected networks and diagonal networks ( cf . Gunasekar et al . ( 2018b ) ; Yun et al . ( 2020 ) ) . We describe our main results below , and also summarize them in Tables 1 and 2 . Fully-connected networks : . • In linear fully-connected networks of any depth the KKT point is a global optimum3 . • In fully-connected depth-2 ReLU networks the KKT point may not even be a local optimum . Moreover , this negative result holds with constant probability over the initialization , i.e. , there is a training dataset such that gradient flow with random initialization converges with constant probability to the direction of a KKT point which is not a local optimum . Depth-2 networks in N : • The positive result on fully-connected linear networks does not extend to networks with sparse weights : In linear diagonal networks the KKT point may not be a local optimum . • In our proof for the above negative result , the KKT point contains a neuron whose weights vector is zero . However , in practice gradient descent often converges to networks that do not contain such zero neurons . We show that for linear networks in Nno-share , if the KKT point has only non-zero weights vectors , then it is a global optimum . We also show that even for the simple case of depth-2 diagonal linear networks , the optimality of the KKT points can be unexpectedly subtle , in the context of margin maximization in the predictor space ( see Remark 4.1 ) . • For ReLU networks in Nno-share , in order to obtain a positive result we need a stronger assumption . We show that if the KKT point is such that for every input in the dataset the input to every hidden neuron in the network is non-zero , then it is guaranteed to be a local optimum ( but not necessarily a global optimum ) . • For linear or ReLU convolutional networks , even if the above assumptions hold , the KKT point may not be a local optimum . Deep networks in N : . • We show that the positive results on depth-2 linear and ReLU networks in Nno-share ( under the assumptions described above ) do not apply to deeper networks . • We study a weaker notion of margin maximization : maximizing the margin for each layer separately . For linear networks of depthm ≥ 2 inN ( including networks with shared weights ) , we show that the KKT point is a global optimum of the per-layer maximum margin problem . For ReLU networks the KKT point may not even be a local optimum of this problem , but under the assumption on non-zero inputs to all neurons it is a local optimum . 2See Section 2 for the formal definition . 3We note that margin maximization for such networks in the predictor space is already known ( Ji & Telgarsky , 2020 ) . However , margin maximization in the predictor space does not necessarily imply margin maximization in the parameter space . In the paper , our focus is on understanding what can be guaranteed for the KKT convergence points specified in Lyu & Li ( 2019 ) . Accordingly , in most of our negative results , the construction assumes some specific initialization of gradient flow , and does not quantify how “ likely ” they are to be reached under some random initialization . An exception is our negative result for depth-2 fully-connected ReLU networks ( Thm . 3.2 ) , which holds with constant probability under reasonable random initializations . Understanding whether this can be extended to the other settings we consider is an interesting problem for future research . Finally , we consider non-homogeneous networks , for example , networks with skip connections or bias terms . Lyu & Li ( 2019 ) showed that a smoothed version of the normalized margin is monotonically increasing when training homogeneous networks . They observed empirically that the normalized margin is monotonically increasing also when training non-homogeneous networks , but did not provide a proof for this phenomenon and left it as an open problem . We give an example for a simple non-homogeneous network where the normalized margin ( as well as the smoothed margin ) is strictly decreasing ( see Thm . 6.1 ) . The paper is structured as follows : In Section 2 we provide necessary notations and definitions , and discuss the most relevant prior results . Additional related works are discussed in Appendix A . In Sections 3 , 4 and 5 we state our results on fully-connected networks , depth-2 networks in N and deep networks in N respectively , and provide some proof ideas . In Section 6 we state our result on non-homogeneous networks . All formal proofs are deferred to Appendix C . 2 PRELIMINARIES . Notations . We use bold-faced letters to denote vectors , e.g. , x = ( x1 , . . . , xd ) . For x ∈ Rd we denote by ‖x‖ the Euclidean norm . We denote by 1 ( · ) the indicator function , for example 1 ( t ≥ 5 ) equals 1 if t ≥ 5 and 0 otherwise . For an integer d ≥ 1 we denote [ d ] = { 1 , . . . , d } . Neural networks . A fully-connected neural network Φ of depth m ≥ 2 is parameterized by a collection θ = [ W ( l ) ] ml=1 of weight matrices , such that for every layer l ∈ [ m ] we have W ( l ) ∈ Rdl×dl−1 . Thus , dl denotes the number of neurons in the l-th layer ( i.e. , the width of the layer ) . We assume that dm = 1 and denote by d : = d0 the input dimension . The neurons in layers [ m− 1 ] are called hidden neurons . A fully-connected network computes a function Φ ( θ ; · ) : Rd → R defined recursively as follows . For an input x ∈ Rd we set h′0 = x , and define for every j ∈ [ m − 1 ] the input to the j-th layer as hj = W ( j ) h′j−1 , and the output of the j-th layer as h ′ j = σ ( hj ) , where σ : R→ R is an activation function that acts coordinate-wise on vectors . Then , we define Φ ( θ ; x ) = W ( m ) h′m−1 . Thus , there is no activation function in the output neuron . When considering depth-2 fully-connected networks we often use a parameterization θ = [ w1 , . . . , wk , v ] where w1 , . . . , wk are the weights vectors of the k hidden neurons ( i.e. , correspond to the rows of the first layer ’ s weight matrix ) and v are the weights of the second layer . We also consider neural networks where some weights can be missing or shared . We define a class N of networks that may contain sparse and shared weights as follows . A network Φ in N is parameterized by θ = [ u ( l ) ] ml=1 where m is the depth of Φ , and u ( l ) ∈ Rpl are the parameters of the l-th layer . We denote by W ( l ) ∈ Rdl×dl−1 the weight matrix of the l-th layer . The matrix W ( l ) is described by the vector u ( l ) , and a function gl : [ dl ] × [ dl−1 ] → [ pl ] ∪ { 0 } as follows : W ( l ) ij = 0 if gl ( i , j ) = 0 , and W ( l ) ij = uk if gl ( i , j ) = k > 0 . Thus , the function gl represents the sparsity and weight-sharing pattern of the l-th layer , and the dimension pl of u ( l ) is the number of free parameters in the layer . We denote by d : = d0 the input dimension of the network and assume that the output dimension dm is 1 . The function Φ ( θ ; · ) : Rd → R computed by the neural network is defined recursively by the weight matrices as in the case of fully-connected networks . For example , convolutional neural networks are in N . Note that the networks in N do not have bias terms and do not allow weight sharing between different layers . Moreover , we define a subclass Nno-share of N , that contains networks without shared weights . Formally , a network Φ is in Nno-share if for every layer l and every k ∈ [ pl ] there is at most one ( i , j ) ∈ [ dl ] × [ dl−1 ] such that gl ( i , j ) = k. Thus , networks in Nno-share might have sparse weights , but do not allow shared weights . For example , diagonal networks ( defined below ) and fully-connected networks are in Nno-share . A diagonal neural network is a network in Nno-share such that the weight matrix of each layer is diagonal , except for the last layer . Thus , the network is parameterized by θ = [ w1 , . . . , wm ] where wj ∈ Rd for all j ∈ [ m ] , and it computes a function Φ ( θ ; · ) : Rd → R defined recursively as follows . For an input x ∈ Rd set h0 = x . For j ∈ [ m − 1 ] , the output of the j-th layer is hj = σ ( diag ( wj ) hj−1 ) . Then , we have Φ ( θ ; x ) = w > mhm−1 . In all the above definitions the parameters θ of the neural networks are given by a collection of matrices or vectors . We often view θ as the vector obtained by concatenating the matrices or vectors in the collection . Thus , ‖θ‖ denotes the ` 2 norm of the vector θ . The ReLU activation function is defined by σ ( z ) = max { 0 , z } , and the linear activation is σ ( z ) = z . In this work we focus on ReLU networks ( i.e. , networks where all neurons have the ReLU activation ) and on linear networks ( where all neurons have the linear activation ) . We say that a network Φ is homogeneous if there exists L > 0 such that for every α > 0 and θ , x we have Φ ( αθ ; x ) = αLΦ ( θ ; x ) . Note that in our definition of the class N we do not allow bias terms , and hence all linear and ReLU networks inN are homogeneous . With the exception of Section 6 which studies non-homogeneous networks , all networks considered in this work are homogeneous . Optimization problem and gradient flow ( GF ) . Let S = { ( xi , yi ) } ni=1 ⊆ Rd× { −1 , 1 } be a binary classification training dataset . Let Φ be a neural network parameterized by θ ∈ Rm . For a loss function ` : R → R the empirical loss of Φ ( θ ; · ) on the dataset S is L ( θ ) : = ∑n i=1 ` ( yiΦ ( θ ; xi ) ) . We focus on the exponential loss ` ( q ) = e−q and the logistic loss ` ( q ) = log ( 1+e−q ) , and consider gradient flow ( GF ) on the objective L ( θ ) . This setting captures gradient descent with an infinitesimal step size . Let θ ( t ) be the trajectory of GF . Starting from an initial point θ ( 0 ) , the dynamics of θ ( t ) is given by the differential equation dθ ( t ) dt = −∇L ( θ ( t ) ) . Note that the ReLU function is not differentiable at 0 . Practical implementations of gradient methods define the derivative σ′ ( 0 ) to be some constant in [ 0 , 1 ] . We note that the exact value of σ′ ( 0 ) has no effect on our results . Convergence to a KKT point of the maximum-margin problem . We say that a trajectory θ ( t ) converges in direction to θ̃ if limt→∞ θ ( t ) ‖θ ( t ) ‖ = θ̃ ‖θ̃‖ . In this work we rely on the following theorem : Theorem 2.1 ( Paraphrased from Lyu & Li ( 2019 ) ; Ji & Telgarsky ( 2020 ) ) . Let Φ be a homogeneous linear or ReLU neural network . Consider minimizing either the exponential or the logistic loss over a binary classification dataset { ( xi , yi ) } ni=1 using GF . Assume that there exists time t0 such that L ( θ ( t0 ) ) < 1 , namely , Φ classifies every xi correctly . Then , GF converges in direction to a first order stationary point ( KKT point ) of the following maximum margin problem in parameter space : min θ 1 2 ‖θ‖2 s.t . ∀i ∈ [ n ] yiΦ ( θ ; xi ) ≥ 1 . ( 1 ) Moreover , L ( θ ( t ) ) → 0 and ‖θ ( t ) ‖ → ∞ as t→∞ . In the case of ReLU networks , Problem 1 is non-smooth . Hence , the KKT conditions are defined using Clarke ’ s subdifferential , which is a generalization of the differential for non-differentiable functions . See Appendix B for a formal definition . We note that Lyu & Li ( 2019 ) proved the above theorem under the assumption that θ converges in direction , and Ji & Telgarsky ( 2020 ) showed that such a directional convergence occurs and hence this assumption is not required . Lyu & Li ( 2019 ) also showed that the KKT conditions of Problem 1 are necessary for optimality . In convex optimization problems , necessary KKT conditions are also sufficient for global optimality . However , the constraints in Problem 1 are highly non-convex . Moreover , the standard method for proving that necessary KKT conditions are sufficient for local optimality , is by showing that the KKT point satisfies certain second order sufficient conditions ( SOSC ) ( cf . Ruszczynski ( 2011 ) ) . However , even when Φ is a linear neural network it is not known when such conditions hold . Thus , the KKT conditions of Problem 1 are not known to be sufficient even for local optimality . A linear network with weight matrices W ( 1 ) , . . . , W ( m ) computes a linear predictor x 7→ 〈β , x〉 where β = W ( m ) · . . . ·W ( 1 ) . Some prior works studied the implicit bias of linear networks in the predictor space . Namely , characterizing the vector β from the aforementioned linear predictor . Gunasekar et al . ( 2018b ) studied the implications of margin maximization in the parameter space on the implicit bias in the predictor space . They showed that minimizing ‖θ‖ ( under the constraints in Problem 1 ) implies : ( 1 ) Minimizing ‖β‖2 for fully-connected networks ; ( 2 ) Minimizing ‖β‖2/L for depth-L diagonal networks ; ( 3 ) Minimizing ‖β̂‖2/L for depth-L convolutional networks with full-dimensional filters , where β̂ are the Fourier coefficients of β . However , these implications may not hold if GF converges to a KKT point which is not a global optimum of Problem 1 . For some classes of linear networks , positive results were obtained directly in the predictor space , without assuming convergence to a global optimum of Problem 1 in the parameter space . Most notably , for fully-connected linear networks ( of any depth ) , Ji & Telgarsky ( 2020 ) showed that under the assumptions of Thm . 2.1 , GF maximizes the ` 2 margin in the predictor space . Note that margin maximization in the predictor space does not necessarily imply margin maximization in the parameter space . Moreover , some results on the implicit bias in the predictor space of linear convolutional networks with full-dimensional convolutional filters are given in Gunasekar et al . ( 2018b ) . However , the architecture and set of assumptions are different than what we focus on .

Summary: This paper is inspired by recent results which have shown that homogeneous networks, when trained on losses with exponential decay converge in direction to KKT points, i.e. points which are first order stationary w.r.t. the opimization of the margin of classification. Since this margin maximization is non-convex, KKT points may not be global maximizers of the margin or not even local maximizers. The authors therefore study a wide range of architecture of linear and ReLU networks, some shallow, some deep, with or without sparse and shared weights, and study in which case all KKT points are global margin maximizers, local margin maximizers or neither (for negative results, counterexamples are given).  Finally they study what happens for non-homogeneous networks, in particular in the presence of skip connection. While for homogeneous networks the margin increases during training, they give an example of a simple network with skip connection where the margin is strictly decreasing.