Deep Classifiers trained with the Square Loss

Here we consider a model of the dynamics of gradient flow under the square loss in overparametrized ReLU networks. We show that convergence to a solution with the absolute minimum ρ , which is the product of the Frobenius norms of each layer weight matrix, is expected when normalization by a Lagrange multiplier (LN) is used together with Weight Decay (WD). We prove that SGD converges to solutions that have a bias towards 1) large margin (i.e. small ρ ) and 2) low rank of the weight matrices. In addition, we predict the occurence of Neural Collapse without ad hoc assumptions such as the unconstrained features hypothesis. Abstract Recent results of [1] suggest that square loss performs on par with cross-entropy loss in classification tasks with deep networks. While the theoretical understanding of training deep networks with the cross-entropy loss has been growing ([2] and [3]) in terms of margin maximization, the study of square loss for classification has been lagging behind. Here we consider a specific model of the dynamics of, first, gradient flow (GF), and then SGD, in overparametrized ReLU networks under the square loss. Under the assumption of convergence to zero loss minima, we show that solutions have a bias toward small ρ , defined as the product of the Frobenius norms of each layer unnormalized weight matrix. We assume that during training there is normalization using a Lagrange multiplier (LM) of each layer weight matrix but the last one, together with Weight Decay (WD). For λ → 0 the solution would be the interpolating solution with minimum ρ . In the absence of LM+WD, good solutions for classification may still be achieved because of the implicit bias towards small norm solutions in the GD dynamics introduced by carefully chosen close-to-zero initial conditions on the norms of the weights, similar to the case of overparametrized linear networks (see Appendix E). However, for λ = 0 we often observe solutions with large ρ that are suboptimal and probably in the NTK regime. We show that convergence to an ideal equilibrium of SGD ( ˙ V k = 0 for all minibatches) with λ > 0 would imply rank one weight matrices. This is impossible generically, implying that SGD never converges to the same set of V k across all minibatches. We claim that this is the origin of