On the Optimization Landscape of Neural Collapse under MSE Loss: Global Optimality with Unconstrained Features

Abstract

When training deep neural networks for classification tasks, an intriguing empirical phenomenon has been widely observed in the last-layer classifiers and features, where (i) the class means and the last-layer classifiers all collapse to the vertices of a Simplex Equiangular Tight Frame (ETF) up to scaling, and (ii) cross-example within-class variability of last-layer activations collapses to zero. This phenomenon is called Neural Collapse (NC), which seems to take place regardless of the choice of loss functions. In this work, we justify NC under the mean squared error (MSE) loss, where recent empirical evidence shows that it performs comparably or even better than the de-facto cross-entropy loss. Under a simplified unconstrained feature model, we provide the first global landscape analysis for vanilla nonconvex MSE loss and show that the (only!) global minimizers are neural collapse solutions, while all other critical points are strict saddles whose Hessian exhibit negative curvature directions. Furthermore, we justify the usage of rescaled MSE loss by probing the optimization landscape around the NC solutions, showing that the landscape can be improved by tuning the rescaling hyperparameters. Finally, our theoretical findings are experimentally verified on practical network architectures.

Neural Collapse with Normalized Features: A Geometric Analysis over the Riemannian Manifold

“…Prior arts and related works on N C. The empirical N C phenomenon has inspired a recent line of theoretical studies on understanding why it occurs [17,21,44,47,71,86,87]. Like ours, most of these works studied the problem under the UFM.…”

“…Such a phenomenon is referred to as Neural Collapse (N C) [50], which has been shown empirically to persist across a broad range of canonical classification problems, on different loss functions (e.g., cross-entropy (CE) [16,50,87], mean-squared error (MSE) [71,86], and supervised contrasive (SC) losses [21]), on different neural network architectures (e.g., VGG [63], ResNet [24], and DenseNet [28]), and on a variety of standard datasets (such as MNIST [39], CIFAR [35], and ImageNet [12], etc). Recently, in independent lines of research, many works are devoted to learning maximally compact and separated features; see, e.g., [13,42,43,51,52,60,72,73,76].…”

“…We consider the commonly used CE loss and formulate the problem as a Riemannian optimization problem over products of unit spheres (i.e., the oblique manifold). Our study is also based upon the assumption of the so-called unconstrained feature model (UFM) [47,86,87] or layer-peeled model [17], where the last-layer features of the deep network are treated as free optimization variables to simplify the nonlinear interactions across layers. The underlying reasoning is that modern deep networks are often highly overparameterized with the capacity of learning any representations [26,45,62], so that the last-layer features can approximate, or interpolate, any point in the feature space.…”

“…More specifically, we prove that every local minimizer is a global solution satisfying the N C properties, and all the other critical points exhibit directions with negative curvature. Our analysis for the manifold setting is based upon a nontrivial extension of recent studies for the N C with penalized formulations [23,71,86,87], which could be of independent interest. Our work brings new tools from Riemannian optimization for analyzing optimization landscapes of training deep networks with an increasingly common practice of feature normalization.…”

“…Depsite the tremendous success of deep learning in engineering and scientific applications over the past decades, the underlying mechanism of deep neural networks (DNNs) still largely remains mysterious. Towards the goal of understanding the learned deep representations, a recent seminal line of works [16,23,50,86,87] presents an intriguing phenomenon that persists across a range of canonical classification problems during the terminal phase of training. Specifically, it has been widely observed that last-layer features (i.e., the output of the penultimate layer) and last-layer linear classifiers of a trained DNN exhibit simple but elegant mathematical structures, in the sense that • (NC1) Variability Collapse: the individual features of each class concentrate to their classmeans.…”