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

TL;DR

This paper analyzes the optimization landscape of neural collapse under MSE loss, proving that global minima correspond to neural collapse solutions and that all other critical points are saddles, with experimental validation on neural networks.

Contribution

It provides the first global landscape analysis for MSE loss in neural collapse, showing that only neural collapse solutions are global minima and others are saddle points.

Findings

Global minimizers are neural collapse solutions.

All other critical points are strict saddles with negative curvature.

Rescaling hyperparameters can improve the landscape around NC solutions.

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.