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: This paper considers margin maximization and simplicity bias properties of gradient flow on a two-layer leaky-ReLU network with linearly separable data. If the data is also symmetric (i.e., whenever (x,y) is in the dataset, (-x,-y) is also in the dataset), this paper proves that the max-margin network basically represents the max-margin linear classifier. Moreover, while prior results ensure convergence to KKT points which may not give the maximum margin, this paper shows that under a few technical conditions gradient flow can indeed maximize the margin. On the other hand, if the data is not symmetric, this paper provides an example where gradient flow cannot globally maximize the margin.