On the Optimization Landscape of Neural Collapse under MSE Loss: Global Optimality with Unconstrained Features
This work addresses a foundational problem in deep learning theory by rigorously analyzing optimization landscapes, offering insights for improving training with MSE loss, though it is incremental as it builds on prior empirical observations of Neural Collapse.
The paper tackles the problem of justifying the Neural Collapse phenomenon under MSE loss by providing a global landscape analysis, showing that global minimizers are neural collapse solutions and other critical points are strict saddles, with experimental verification on practical networks.
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.