Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path

Abstract

Recent work [Papyan, Han, and Donoho, 2020] discovered a phenomenon called Neural Collapse (NC) that occurs pervasively in today's deep net training paradigm of driving cross-entropy loss towards zero. In this phenomenon, the last-layer features collapse to their class-means, both the classifiers and class-means collapse to the same Simplex Equiangular Tight Frame (ETF), and the behavior of the last-layer classifier converges to that of the nearest-class-mean decision rule. Since then, follow-ups-such as Mixon et al. [2020] and Poggio and Liao [2020a,b]formally analyzed this inductive bias by replacing the hard-to-study cross-entropy by the more analytically tractable mean squared error (MSE) loss. But, these works stopped short of demonstrating the empirical reality of MSE-NC on benchmark datasets and canonical networks-as had been done in Papyan, Han, and Donoho [2020] for the cross-entropy loss. In this work, we establish the empirical reality and reproducibility of MSE-NC by reporting experimental observations for three prototypical networks and five canonical datasets-with code for reproducing NC in PyTorch for both MSE and cross-entropy in the following Google Colaboratory notebook: here. Following this, we develop three main contributions inspired by MSE-NC. Firstly, we show a new theoretical decomposition of the MSE loss into (A) a term assuming the last-layer classifier is exactly the least-squares or Webb and Lowe [1990] classifier and (B) a term capturing the deviation from this least-squares classifier. Secondly, we exhibit experiments on canonical datasets and networks demonstrating that, during training, term-(B) is negligible. This motivates a new theoretical construct: the central path, where the linear classifier stays MSE-optimal-for the given feature activations-throughout the dynamics. Finally, through our study of continually renormalized gradient flow along the central path, we produce closed-form dynamics that predict full Neural Collapse in an unconstrained features model.

Extended Unconstrained Features Model for Exploring Deep Neural Collapse

“…While most existing papers consider cross-entropy loss, in this paper we focus on the mean squared error (MSE) loss, which has been recently shown to be powerful also for classification tasks (Hui & Belkin, 2020). (We note that the occurrence of neural collapse when training practical DNNs with MSE loss, and its positive effects on their performance, have been shown empirically in a very recent paper (Han et al, 2021)). We start with analyzing the (plain) UFM, showing that for the regularized MSE loss the collapsed features can be more structured than in the cross-entropy case (e.g., they may possess also orthogonality), which affects also the structure of the weights.…”

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

“…Like ours, most of these works studied the problem under the UFM. In particular, despite the nonconvexity, recent works showed that the only global solutions are N C solutions for a variety of nonconvex training losses (e.g., CE [17,44,87], MSE [71,86], SC losses [21]) and different problem formulations (e.g., penalized, constrained, and unconstrained) [23,44,71,86,87]. Recently, this study has been extended to deeper models with the MSE training loss [71].…”

“…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.…”

Robust Training under Label Noise by Over-parameterization

“…Robustness of learned representations against label noise. Recently, a line of work showed an intriguing and universal phenomenon of learned deep representations under natural setting [90][91][92][93], that the last-layer representations of each class collapse to a single dimension. However, the collapsed representation loses the variability of the data and is vulnerable to corruptions such as label noise.…”