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 studies the implicit bias of gradient flow in the lens of margin maximization. A previous work by Lyu & Li (2020) shows that gradient flow directionally converges to a KKT point of the margin maximization problem. The authors then focus on understanding what kind of local/global optimality of margin can be guaranteed by applying the KKT convergence result. The authors consider various settings on 2-layer/deep linear/ReLU neural nets, and for each setting they prove either local/global optimality of margin or an impossibility result.  As the positive results could be more preferable by the ML community, I summarize the positive results below:  * Global optimality for (deep) fully-connected linear nets; * Global optimality for general two-layer linear nets with no weight sharing, assuming all the weights are non-zero; * Local optimality for general two-layer ReLU nets with no weight sharing, assuming all neurons have non-zero inputs; * Per-layer global optimality for general (deep) linear nets (i.e., every layer weights achieve the global max margin if all other layers are fixed.) * Per-layer local optimality for general (deep) ReLU nets (i.e., every layer weights achieve a local max margin if all other layers are fixed.), assuming all neurons have non-zero inputs.  Here "general" nets can have arbitrary connections between neurons in adjacent layers.