[Paper Review] Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path

X. Y. Han, Vardan Papyan+1 more|arXiv (Cornell University)|Jun 3, 2021

Domain Adaptation and Few-Shot LearningComputer Science33 references19 citations

TL;DR

This paper investigates Neural Collapse (NC) in deep networks trained with mean squared error (MSE) loss, demonstrating that NC emerges similarly to cross-entropy training. By decomposing the MSE loss into components tied to NC and introducing the 'central path'—a theoretical trajectory where the classifier remains MSE-optimal—it derives exact gradient flow dynamics that predict NC, offering a mathematically tractable framework for understanding NC's emergence.

ABSTRACT

The recently discovered Neural Collapse (NC) phenomenon occurs pervasively in today's deep net training paradigm of driving cross-entropy (CE) loss towards zero. During NC, last-layer features collapse to their class-means, both classifiers and class-means collapse to the same Simplex Equiangular Tight Frame, and classifier behavior collapses to the nearest-class-mean decision rule. Recent works demonstrated that deep nets trained with mean squared error (MSE) loss perform comparably to those trained with CE. As a preliminary, we empirically establish that NC emerges in such MSE-trained deep nets as well through experiments on three canonical networks and five benchmark datasets. We provide, in a Google Colab notebook, PyTorch code for reproducing MSE-NC and CE-NC: at https://colab.research.google.com/github/neuralcollapse/neuralcollapse/blob/main/neuralcollapse.ipynb. The analytically-tractable MSE loss offers more mathematical opportunities than the hard-to-analyze CE loss, inspiring us to leverage MSE loss towards the theoretical investigation of NC. We develop three main contributions: (I) We show a new decomposition of the MSE loss into (A) terms directly interpretable through the lens of NC and which assume the last-layer classifier is exactly the least-squares classifier; and (B) a term capturing the deviation from this least-squares classifier. (II) We exhibit experiments on canonical datasets and networks demonstrating that term-(B) is negligible during training. This motivates us to introduce a new theoretical construct: the central path, where the linear classifier stays MSE-optimal for feature activations throughout the dynamics. (III) By studying renormalized gradient flow along the central path, we derive exact dynamics that predict NC.

Motivation & Objective

• To establish that Neural Collapse (NC) occurs in deep networks trained with mean squared error (MSE) loss, mirroring its behavior under cross-entropy (CE) loss.
• To decompose the MSE loss into interpretable components related to NC, distinguishing between NC-compatible terms and deviation terms.
• To introduce the 'central path'—a theoretical construct where the linear classifier remains MSE-optimal throughout training—enabling exact dynamic analysis.
• To derive exact gradient flow dynamics along the central path that predict the emergence of NC, providing a mathematically rigorous foundation for NC.

Proposed method

• Decompose the MSE loss into (A) terms interpretable via NC theory when the classifier is the least-squares classifier, and (B) a deviation term from this idealized classifier.
• Empirically show that term (B) is negligible during training on canonical networks and datasets, justifying the central path approximation.
• Define the central path as the trajectory where the classifier remains MSE-optimal for the feature activations at each training step.
• Analyze renormalized gradient flow along the central path to derive exact dynamics governing the convergence to NC.
• Use the central path and renormalized dynamics to analytically predict the limiting behaviors of NC: feature collapse, convergence to Simplex ETF, and classifier alignment.
• Leverage the analytical tractability of MSE loss to provide a theoretical framework for NC that is more accessible than the complex CE loss.

Experimental results

Research questions

• RQ1 Does Neural Collapse emerge in deep networks trained with mean squared error (MSE) loss, similar to its behavior under cross-entropy (CE) loss?
• RQ2 Can the MSE loss be decomposed into components that directly reflect the mechanisms of Neural Collapse?
• RQ3 Is there a theoretical trajectory—such as the 'central path'—where the classifier remains MSE-optimal, enabling exact dynamic analysis of NC?
• RQ4 Can exact gradient flow dynamics along the central path predict the convergence to Neural Collapse?
• RQ5 What are the mathematical conditions under which NC emerges under MSE training, and how do they compare to CE-based training?

Key findings

• Neural Collapse emerges in deep networks trained with MSE loss across three canonical architectures and five benchmark datasets, confirming its robustness beyond CE training.
• The MSE loss decomposes into a component aligned with NC (when the classifier is least-squares optimal) and a small deviation term, which is empirically negligible during training.
• The central path is introduced as a theoretical construct where the classifier remains MSE-optimal throughout training, enabling exact dynamic analysis.
• Exact dynamics derived from renormalized gradient flow along the central path predict the convergence to Neural Collapse, including feature collapse and Simplex ETF alignment.
• The theoretical framework based on MSE loss provides a more analytically tractable path to understanding NC than the complex CE loss, enabling precise derivation of NC emergence.
• The results suggest that NC is not an artifact of CE loss but a fundamental property of overparameterized deep learning under standard optimization, even with MSE.