Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias

1 Introduction . One major mystery in deep learning is why deep neural networks generalize despite overparameterization ( Zhang et al. , 2017 ) . To tackle this issue , many recent works turn to study the implicit bias of gradient descent ( GD ) — what kind of theoretical characterization can we give for the low-loss solution found by GD ? The seminal works by Soudry et al . ( 2018a , b ) revealed an interesting connection between GD and margin maximization : for linear logistic regression on linearly separable data , there can be multiple linear classifiers that perfectly fit the data , but GD with any initialization always converges to the maxmargin ( hard-margin SVM ) solution , even when there is no explicit regularization . Thus the solution found by GD has the same margin-based generalization bounds as hard-margin SVM . Subsequent works on linear models have extended this theoretical understanding of GD to SGD ( Nacson et al. , 2019b ) , other gradient-based methods ( Gunasekar et al. , 2018a ) , other loss functions with certain poly-exponential tails ( Nacson et al. , 2019a ) , linearly non-separable data ( Ji and Telgarsky , 2018 , 2019b ) , deep linear nets ( Ji and Telgarsky , 2019a ; Gunasekar et al. , 2018b ) . ∗Equal contribution †Most of the work is done when Kaifeng Lyu and Runzhe Wang were at Tsinghua University . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Given the above results , a natural question to ask is whether GD has the same implicit bias towards max-margin solutions for machine learning models in general . Lyu and Li ( 2020 ) studied the relationship between GD and margin maximization on deep homogeneous neural network , i.e. , neural network whose output function is ( positively ) homogeneous with respect to its parameters . For homogeneous neural networks , only the direction of parameter matters for classification tasks . For logistic and exponential loss , Lyu and Li ( 2020 ) assumed that GD decreases the loss to a small value and achieves full training accuracy at some time point , and then provided an analysis for the training dynamics after this time point ( Theorem 3.1 ) , which we refer to as late phase analysis . It is shown that GD decreases the loss to 0 in the end and converges to a direction satisfying the Karush-Kuhn-Tucker ( KKT ) conditions of a constrained optimization problem ( P ) on margin maximization . However , given the non-convex nature of neural networks , KKT conditions do not imply global optimality for margins . Several attempts are made to prove the global optimality specifically for two-layer nets . Chizat and Bach ( 2020 ) provided a mean-field analysis for infinitely wide two-layer Squared ReLU nets showing that gradient flow converges to the solution with global max margin , which also corresponds to the max-margin classifier in some non-Hilbertian space of functions . Ji and Telgarsky ( 2020a ) extended the proof to finite-width neural nets , but the width needs to be exponential in the input dimension ( due to the use of a covering condition ) . Both works build upon late phase analyses . Under a restrictive assumption that the data is orthogonally separable , i.e. , any data point xi can serve as a perfect linear separator , Phuong and Lampert ( 2021 ) analyzed the full trajectory of gradient flow on two-layer ReLU nets with small initialization , and established the convergence to a piecewise linear classifier that maximizes the margin , irrespective of network width . In this paper , we study the implicit bias of gradient flow on two-layer neural nets with Leaky ReLU activation ( Maas et al. , 2013 ) and logistic loss . To avoid the lazy or Neural Tangent Kernel ( NTK ) regime where the weights are initialized to large random values and do not change much during training ( Jacot et al. , 2018 ; Chizat et al. , 2019 ; Du et al. , 2019b , a ; Allen-Zhu et al. , 2018 , 2019 ; Zou et al. , 2018 ; Arora et al. , 2019b ) , we use small initialization to encourage the model to learn features actively , which is closer to real-life neural network training . When analyzing convergence behavior of training on neural networks , one can simplify the problem and gain insights by assuming that the data distribution has a simple structure . Many works particularly study the case where the labels are generated by an unknown teacher network that is much smaller/simpler than the ( student ) neural network to be trained . Following Brutzkus et al . ( 2018 ) ; Sarussi et al . ( 2021 ) and many other works , we consider the case where the dataset is linearly separable , namely the labels are generated by a linear teacher , and study the training dynamics of two-layer Leaky ReLU nets on such dataset . 1.1 Our Contribution . Among all the classifiers that can be represented by the two-layer Leaky ReLU nets , we show any global-max-margin classifier is exactly linear under one more data assumption : the dataset is symmetric , i.e. , if x is in the training set , then so is −x . Note that such symmetry can be ensured by simple data augmentation . Still , little is known about what kind of classifiers neural network trained by GD learns . Though Lyu and Li ( 2020 ) showed that gradient flow converges to a classifier along KKT-margin direction , we note that this result is not sufficient to guarantee the global optimality since such classifier can have nonlinear decision boundaries . See Figure 1 ( left ) for an example . In this paper , we provide a multi-phase analysis for the full trajectory of gradient flow , in contrast with previous late phase analyses which only analyzes the trajectory after achieving 100 % training accuracy . We show that gradient flow with small initialization converges to a global-max-margin linear classifier ( Theorem 4.2 ) . The proof leverages power iteration to show that neuron weights align in two directions in an early phase of training , inspired by Li et al . ( 2021 ) . We further show the alignment at any constant training time by associating the dynamics of wide neural net with that of two-neuron neural net , and finally , extend the alignment to the infinite time limit by applying Kurdyka-Łojasiewicz ( KL ) inquality in a similar way as Ji and Telgarsky ( 2020a ) . The alignment at convergence implies that the convergent classifier is linear . The above results also justify a recent line of works studying the so-called simplicity bias : GD first learns linear functions in the early phase of training , and the complexity of the solution increases as training goes on ( Kalimeris et al. , 2019 ; Hu et al. , 2020 ; Shah et al. , 2020 ) . Indeed , our result establishes a form of extreme simplicity bias of GD : if the dataset can be fitted by a linear classifier , then GD learns a linear classifier not only in the beginning but also at convergence . On the pessimistic side , this paper suggests that such global margin maximization result could be fragile . Even for linearly separable data , global-max-margin classifiers may be nonlinear without the symmetry assumption . In particular , we show that for any linearly separable dataset , gradient flow can be led to converge to a linear classifier with suboptimal margin by adding only 3 extra data points ( Theorem 6.2 ) . See Figure 1 ( right ) for an example . 2 Related Works . Generalization Aspect of Margin Maxmization . Margin often appears in the generalization bounds for neural networks ( Bartlett et al. , 2017 ; Neyshabur et al. , 2018 ) , and larger margin leads to smaller bounds . Jiang et al . ( 2020 ) conducted an empirical study for the causal relationships between complexity measures and generalization errors , and showed positive results for normalized margin , which is defined by the output margin divided by the product ( or powers of the sum ) of Frobenius norms of weight matrices from each layer . On the pessimistic side , negative results are also shown if Frobenius norm is replaced by spectral norm . In this paper , we do use the normalized margin with Frobenius norm ( see Section 3 ) . Learning on Linearly Separable Data . Some works studied the training dynamics of ( nonlinear ) neural networks on linearly separable data ( labels are generated by a linear teacher ) . Brutzkus et al . ( 2018 ) showed that SGD on two-layer Leaky ReLU nets with hinge loss fits the training set in finite steps and generalizes well . Frei et al . ( 2021 ) studied online SGD ( taking a fresh sample from the population in each step ) on the two-layer Leaky ReLU nets with logistic loss . For any data distribution , they proved that there exists a time step in the early phase such that the net has a test error competitive with that of the best linear classifier over the distribution , and hence generalizes well on linearly separable data . Both two papers reveal that the weight vectors in the first layer have positive correlations with the weight of the linear teacher , but their analyses do not imply that the learned classifier is linear . In the NTK regime , Ji and Telgarsky ( 2020b ) ; Chen et al . ( 2021 ) showed that GD on shallow/deep neural nets learns a kernel predictor with good generalization on linearly separable data , and it suffices to have width polylogarithmic in the number of training samples . Still , they do not imply that the learned classifier is linear . Pellegrini and Biroli ( 2020 ) provided a mean-field analysis for two-layer ReLU net showing that training with hinge loss and infinite data leads to a linear classifier , but their analysis requires the data distribution to be spherically symmetric ( i.e. , the probability density only depends on the distance to origin ) , which is a more restrictive assumption than ours . Sarussi et al . ( 2021 ) provided a late phase analysis for gradient flow on two-layer Leaky ReLU nets with logistic loss , which establishes the convergence to linear classifier based on an assumption called Neural Agreement Regime ( NAR ) : starting from some time point , for any training sample , the outputs of all the neurons have the same sign . However , it is unclear why this can happen a priori . Comparing with our work , we analyze the full trajectory of gradient flow and establish the convergence to linear classifier without assuming NAR . Phuong and Lampert ( 2021 ) analyzed the full trajectory for gradient flow on orthogonally separable data , but every KKT-margin direction attains the global max margin ( see Appendix H ) in their setting , which it is not necessarily true in general . In our setting , KKT-margin direction with suboptimal margin does exist . Simplicity Bias . Kalimeris et al . ( 2019 ) empirically observed that neural networks in the early phase of training are learning linear classifiers , and provided evidence that SGD learns functions of increasing complexity . Hu et al . ( 2020 ) justified this view by proving that the learning dynamics of two-layer neural nets and simple linear classifiers are close to each other in the early phase , for dataset drawn from a data distribution where input coordinates are independent after some linear transformation . The aforementioned work by Frei et al . ( 2021 ) can be seen as another theoretical justification for online SGD on aribitrary data distribution . Shah et al . ( 2020 ) pointed out that extreme simplicity bias can lead to suboptimal generalization and negative effects on adversarial robustness . Small Initialization . Several theoretical works studying neural network training with small initialization can be connected to simplicity bias . Maennel et al . ( 2018 ) uncovered a weight quantization effect in training two-layer nets with small initialization : gradient flow biases the weight vectors to a certain number of directions determined by the input data ( independent of neural network width ) . It is hence argued that gradient flow has a bias towards “ simple ” functions , but their proof is not entirely rigorous and no clear definition of simplicity is given . This weight quantization effect has also been studied under the names of weight clustering ( Brutzkus and Globerson , 2019 ) , condensation ( Luo et al. , 2021 ; Xu et al. , 2021 ) . Williams et al . ( 2019 ) studied univariate regression and showed that two-layer ReLU nets with small initialization tend to learn linear splines . For the matrix factorization problem , which can be related to training neural networks with linear or quadratic activations , we can measure the complexity of the learned solution by rank . A line of works showed that gradient descent learns solutions with gradually increasing rank ( Li et al. , 2018 ; Arora et al. , 2019a ; Gidel et al. , 2019 ; Gissin et al. , 2020 ; Li et al. , 2021 ) . Such results have been generalized to tensor factorization where the complexity measure is replaced by tensor rank ( Razin et al. , 2021 ) . Beyond small initialization of our interest and large initialization in the lazy or NTK regime , Woodworth et al . ( 2020 ) ; Moroshko et al . ( 2020 ) ; Mehta et al . ( 2021 ) studied feature learning when the initialization scale transitions from small to large scale .

Summary: The paper studies the implicit bias for finite two layers Leaky ReLU networks. For linearly separable symmetric data, the authors show that the global max-margin direction corresponds to a linear classifier. In addition, under additional regularity conditions and for a sufficiently small initialization scale, the authors show gradient flow achieves the global max-margin. For non-symmetric yet linearly separable data, the authors show that gradient flow can be led to converge to a linear classifier with suboptimal margin by adding 3 data points.