Understanding Square Loss in Training Overparametrized Neural Network Classifiers

1 INTRODUCTION . The pursuit of better classifiers has fueled the progress of machine learning and deep learning research . The abundance of benchmark image datasets , e.g. , MNIST , CIFAR , ImageNet , etc. , provides test fields for all kinds of new classification models , especially those based on deep neural networks ( DNN ) . With the introduction of CNN , ResNets , and transformers , DNN classifiers are constantly improving and catching up to the human-level performance . In contrast to the active innovations in model architecture , the training objective remains largely stagnant , with cross-entropy loss being the default choice . Despite its popularity , cross-entropy has been shown to be problematic in some applications . Among others , Yu et al . ( 2020 ) argued that features learned from cross-entropy lack interpretability and proposed a new loss aiming for maximum coding rate reduction . Pang et al . ( 2019 ) linked the use of cross-entropy to adversarial vulnerability and proposed a new classification loss based on latent space matching . Guo et al . ( 2017 ) discovered that the confidence of most DNN classifiers trained with cross-entropy is not well-calibrated . Recently , several alternative losses have seen revived interests for deep classifiers . In particular , many existing works have presented empirical evidence promoting the use of square loss over crossentropy . Hui & Belkin ( 2020 ) conducted large-scale experiments comparing the two and found that square loss tends to perform better in natural language processing related tasks while cross-entropy usually yields slightly better accuracy in image classification . Similar comparisons are also made in Demirkaya et al . ( 2020 ) . Kornblith et al . ( 2020 ) compared a variety of loss functions and output layer regularization strategies on the accuracy and out-of-distribution robustness , and found that square loss has greater class separation and better out-of-distribution robustness . In comparison to the empirical investigation , theoretical understanding of square loss in training deep learning classifiers is still lacking . Through our lens , square loss has its uniqueness among classic classification losses , and we argue that it has great potentials for modern classification tasks . Below we list our motivations and reasons why . Explicit feature modeling Deep learning ’ s success can be largely attributed to its superior ability as feature extractors . For classification , the ideal features should be separated between classes and concentrated within classes . However , when optimizing cross-entropy loss , it ’ s not clear what the learned features should look like ( Yu et al. , 2020 ) . In comparison , square loss uses the label codings ( one-hot , simplex etc . ) as features , which can be modeled explicitly to control class separations . Model Calibration An ideal classifier should not only give the correct class prediction , but also with the correct confidence . Calibration error measures the closeness of the predicted confidence to the underlying conditional probability ⌘ . Using square loss in classification can be essentially viewed as regression where it treats discrete labels as continuous code vectors . It can be shown that the optimal classifier under square loss is 2⌘ 1 , linear with the ground truth . This distinguishing property allows it to easily recover ⌘ . In comparison , the optimal classifiers under the hinge loss and cross-entropy are sign ( 2⌘ 1 ) and log ( ⌘1 ⌘ ) , respectively . Therefore , hinge loss doesn ’ t provide reliable information on the prediction confidence , and cross-entropy can be problematic when ⌘ is close to 0 or 1 ( Zhang , 2004 ) . Hence , in terms of model calibration , square loss is a natural choice . Connections to popular approaches Mixup ( Zhang et al. , 2017 ) is a popular data augmentation technique where augmented data are constructed via convex combinations of inputs and their labels . Like in square loss , mixup treats labels as continuous and is shown to improve the generalization of DNN classifiers . In knowledge distillation ( Hinton et al. , 2015 ) , where a student classifier is trying to learn from a trained teacher , Menon et al . ( 2021 ) proved that the “ optimal ” teacher with the ground truth conditional probabilities provides the lowest variance in student learning . Since classifiers trained using square loss is a natural consistent estimator of ⌘ , one can argue that it is a better teacher . In supervised contrastive learning ( Khosla et al. , 2020 ) , the optimal features are the same as those from square loss with simplex label coding ( Graf et al. , 2021 ) ( details in Section 4 ) . Despite its lack of popularity in practice , square loss has many advantages that can be easily overlooked . In this work , we systematically investigate from a statistical estimation perspective , the properties of deep learning classifiers trained using square loss . The neural networks in our analysis are required to be sufficiently overparametrized in the neural tangent kernel ( NTK ) regime . Even though this restricts the implication of our results , it is a necessary first step towards a deeper understanding . In summary , our main contributions are : • Generalization error bound : We consider two cases , according to whether classes are separable or not . In the general non-separable case , we adopt the classical binary classification setting with smooth conditional probability . Fast rate of convergence is established for overparametrized neural network classifiers with Tsybakov ’ s noise condition . If two classes are separable with positive margins , we show that overparametrized neural network classifiers can provably reach zero misclassification error with probability exponentially tending to one . To the best of our knowledge , this is the first such result for separable but not linear separable classes . Furthermore , we bridge these two cases and offer a unified view by considering auxiliary random noise injection . • Robustness ( margin property ) : When two classes are separable , the decision boundary is not unique and large-margin classifiers are preferred . In the separable case , we further show that the decision boundary of overparametrized neural network classifiers trained by square loss can not be too close to the data support and the resulting margin is lower bounded away from zero , providing theoretical guarantees for robustness . • Calibration error : We show that classifiers trained using square loss are inherently well-calibrated , i.e. , the trained classifier provides consistent estimation of the ground-truth conditional probability in L1 norm . Such property doesn ’ t hold for cross-entropy . • Empirical evaluation : We corroborate our theoretical findings with empirical experiments in both synthetic low-dimensional data and real image data . Comparing to cross-entropy , square loss has comparable generalization error but noticeable advantages in robustness and model calibration . This work contributes towards the theoretical understanding of deep classifiers , from an estimation point of view , which has been a classic topic in statistics literature . Among others , Mammen & Tsybakov ( 1999 ) established the optimal convergence rate for 0-1 loss excess risk when the decision boundary is smooth . Zhang ( 2004 ) ; Bartlett et al . ( 2006 ) extended the analysis to various surrogate losses . Audibert & Tsybakov ( 2007 ) ; Kohler & Krzyzak ( 2007 ) studied the convergence rates for plug-in classifiers from local averaging estimators . Steinwart et al . ( 2007 ) investigated the convergence rate for support vector machine using Gaussian kernels . We build on and extend classic results to neural networks in the NTK regime . Comparing to existing works on deep learning classification , e.g. , Kim et al . ( 2018 ) derived fast convergence rates of ReLU DNN classifiers that minimize the empirical hinge loss , our results incorporate the training algorithm and apply to trained classifiers . We require the neural network to be overparametrized , which has been extensively studied recently , under the umbrella term NTK . Most such results are in the regression setting with a handful of exceptions . Ji & Telgarsky ( 2019 ) showed that only polylogarithmic width is sufficient for gradient descent to overfit the training data using logistic loss . Hu et al . ( 2020 ) proved generalization error bound for regularized NTK in classification . Cao & Gu ( 2019 ; 2020 ) provided optimization and generalization guarantees for overparametrized network trained with cross-entropy . In comparison , our results are sharper in the sense that we take the ground truth data assumptions into consideration . This allows a faster convergence rate , especially when the classes are separable , where the exponential convergence rate is attainable . The NTK framework greatly reduces the technical difficulty for our theoretical analysis . However , our results are mainly due to properties of the square loss itself and we expect them to hold for a wide range of classifiers . There are other works investigating the use of square loss for training ( deep ) classifiers . Han et al . ( 2021 ) uncovered that the “ neural collapse ” phenomenon also occurs under square loss where the last-layer features eventually collapse to their simplex-style class-means . Muthukumar et al . ( 2020 ) compared classification and regression tasks in the overparameterized linear model with Gaussian features , illustrating different roles and properties of loss functions used at the training and testing phases . Poggio & Liao ( 2019 ) made interesting observations on effects of popular regularization techniques such as batch normalization and weight decay on the gradient flow dynamics under square loss . These findings support our theoretical results ’ implication , which further strengthens our beliefs that the essence comes from the square loss and our analysis can go beyond NTK regime . The rest of this paper is arranged as follows . Section 2 presents some preliminaries . Main theoretical results are in Section 3 . The simplex label coding is discussed in Section 4 followed by numerical studies in Section 5 and conclusions in Section 6 . Technical proofs and details of the numerical studies can be found in the Appendix . 2 PRELIMINARIES . Notation For a function f : ⌦ ! R , let kfk1 = supx2⌦ |f ( x ) | and kfkp = ( R ⌦ |f ( x ) | pdx ) 1/p . For a vector x , kxkp denotes its p-norm , for 1 p 1 . Lp and lp are used to distinguish function norms and vector norms . For two positive sequences { an } n2N and { bn } n2N , we write an . bn if there exists a constant C > 0 such that an Cbn for all sufficiently large n. We write an ⇣ bn if an . bn and bn . an . Let [ N ] = { 1 , . . . , N } for N 2 N , I be the indicator function , and Id be the d⇥ d identity matrix . N ( µ , ⌃ ) represents Gaussian distribution with mean µ and covariance ⌃ . Classification problem settings Let P be an underlying probability measure on ⌦ ⇥ Y , where ⌦ ⇢ Rd is compact and Y = { 1 , 1 } . Let ( X , Y ) be a random variable with respect to P . Suppose we have observations { ( xi , yi ) } ni=1 ⇢ ( ⌦ ⇥ Y ) n i.i.d . sampled according to P . The classification task is to predict the unobserved label y given a new input x 2 ⌦ . Let ⌘ defined on ⌦ denote the conditional probability , i.e. , ⌘ ( x ) = P ( y = 1|x ) . Let PX be the marginal distribution of P on X . The key quantity of interest is the misclassification error , i.e. , 0-1 loss . In the population level , the 0-1 loss can be written as L ( f ) = E ( X , Y ) ⇠P I { sign ( f ( X ) ) 6= Y } =EX⇠PX [ ( 1 ⌘ ( X ) ) I { f ( X ) 0 } + ⌘ ( X ) I { f ( X ) < 0 } ] , ( 2.1 ) where the expectation is taken with respect to the probability measure P . Clearly , an optimal classifier with the minimal 0-1 loss is 2⌘ 1 . According to whether labels are deterministic , there are two scenarios of interest . If ⌘ only takes values from { 0 , 1 } , i.e. , labels are deterministic , we call this case the separable case1 . Let ⌦1 = 1In the separable case we consider , the classes are not limited to linearly separable but can be arbitrarily complicated . { x|⌘ ( x ) = 1 } , ⌦2 = { x|⌘ ( x ) = 0 } and ⌦ = ⌦1 [ ⌦2 . If the probability measure of { x|⌘ ( x ) 2 ( 0 , 1 ) } is non-zero , i.e. , the labels contain randomness , we call this case the non-separable case . In the separable case , we further assume that there exists a positive margin , i.e. , dist ( ⌦1 , ⌦2 ) 2 > 0 , where is a constant , and dist ( ⌦1 , ⌦2 ) = infx2⌦1 , x02⌦2 kx x0k2 . In the non-separable case , to quantify the difficulty of classification , we adopt the well-established Tsybakov ’ s noise condition ( Audibert & Tsybakov , 2007 ) , which measures how large the “ difficult region ” is where ⌘ ( x ) ⇡ 1/2 . Definition 2.1 ( Tsybakov ’ s noise condition ) . Let 2 [ 0,1 ] . We say P has Tsybakov noise exponent if there exists a constant C , T > 0 such that for all 0 < t < T , PX ( |2⌘ ( X ) 1| < t ) C ·t . A large value of implies the difficult region to be small . It is expected that a larger leads to a faster convergence rate of a neural network classifier . This intuition is verified for the overparametrized neural network classifier trained by square loss and ` 2 regularization . See Section 3 for more details . Neural network setup We mainly focus on the one-hidden-layer ReLU neural network family F with m nodes in the hidden layer , denoted by fW , a ( x ) = 1p m mX r=1 ar ( W > r x ) , where x 2 ⌦ , W = ( W1 , · · · , Wm ) 2 Rd⇥m is the weight matrix in the hidden layer , a = ( a1 , · · · , am ) > 2 Rm is the weight vector in the output layer , ( z ) = max { 0 , z } is the rectified linear unit ( ReLU ) . The initial values of the weights are independently generated from Wr ( 0 ) ⇠ N ( 0 , ⇠2Im ) , ar ⇠ unif { 1 , 1 } , 8r 2 [ m ] . Based on the observations { ( xi , yi ) } ni=1 , the goal of training a neural network is to find a solution to min W nX i=1 l ( fW , a ( xi ) , yi ) + µR ( W , a ) , ( 2.2 ) where l is the loss function , R is the regularization , and µ 0 is the regularization parameter . Note in Equation 2.2 that we only consider training the weights W . This is because a · ( z ) = sign ( a ) · ( |a|z ) , which allows us to reparametrize the network to have all ai ’ s to be either 1 or 1 . In this work , we consider square loss associated with ` 2 regularization , i.e. , l ( fW , a ( xi ) , yi ) = ( fW , a ( xi ) yi ) 2 and R ( W , a ) = kW k22 . A popular way to train the neural network is via gradient based methods . It has been shown that the training process of DNNs can be characterized by the neural tangent kernel ( NTK ) ( Jacot et al. , 2018 ) . As is usually assumed in the NTK literature ( Arora et al. , 2019 ; Hu et al. , 2020 ; Bietti & Mairal , 2019 ; Hu et al. , 2021 ) , we consider data on the unit sphere Sd 1 , i.e. , kxik2 = 1 , 8i 2 [ n ] , and the neural network is highly overparametrized ( m n ) and trained by gradient descent ( GD ) . For details about NTK and GD in one-hidden-layer ReLU neural networks , we refer to Appendix A . In the rest of this work , we use fW ( k ) , a to denote the GD-trained neural network classifier under square loss associated with ` 2 regularization , where k is the iteration number satisfying Assumption D.1 and W ( k ) is the weight matrix after k-th iteration .

Summary: The paper studies wide neural networks trained with the square loss in the neural tangent kernel regime from a theoretical perspective. In particular, it provides generalization error bounds, robustness and calibration results and makes a connection to supervised contrastive learning. Furthermore, an empirical comparison between cross entropy and square loss is provided.