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

X. Y. Han; Vardan Papyan; David L. Donoho

Source-linked AI summary

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

X. Y. Han, Vardan Papyan, David L. Donoho

arXiv:2106.02073v4Jun 3, 2021cs.LGcs.AImath.DGmath.OCstat.ML

TL;DR

Neural Collapse is well documented under CE training, but its behavior and dynamics under MSE require investigation. This paper verifies MSE-NC empirically, decomposes the MSE loss around the least-squares classifier, and studies renormalized gradient flow on a central path. The resulting closed-form dynamics imply Neural Collapse, while experiments show the classifier-deviation term is negligible during training.

Problem
The paper addresses the limited theoretical understanding of Neural Collapse under MSE loss despite comparable test performance between MSE- and CE-trained networks.
Method
The paper decomposes MSE loss into NC-interpretable least-squares terms and classifier deviation, then analyzes renormalized gradient flow along the central path.
Results
Experiments show MSE-NC across five benchmark datasets and three canonical networks, while the classifier-deviation term is negligible during training; closed-form central-path dynamics imply (NC1)–(NC4).
Takeaways & Limitations
MSE loss provides an analytically tractable setting in which empirical Neural Collapse can be connected to explicit training dynamics.
Takeaways & Limitations
CIFAR100 and ImageNet are omitted because prior work required an additional MSE scaling heuristic whose rigorous study was outside this paper’s scope.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper investigates Neural Collapse in MSE-trained deep networks, establishing its empirical occurrence and developing a central-path theory whose closed-form dynamics imply NC.

MSE classification: MSE-trained networks achieve test accuracies comparable to CE-trained networks, motivating MSE as an analytically tractable basis for studying NC.The paper emphasizes that MSE offers more mathematical opportunities than the harder-to-analyze CE loss.
MSE decomposition: The MSE loss decomposes into NC-interpretable terms plus a deviation term measuring the classifier’s departure from the least-squares optimum.The decomposition is L = LNC1 + LNC2/3 + L⊥LS; the first terms assume the least-squares classifier, while L⊥LS captures deviation from it.
Central path: Empirically, the deviation term becomes negligible during training, so classifier–feature pairs remain approximately on the central path.The approximation begins at an early epoch and persists into the Terminal Phase of Training.
Dynamics: Renormalized features and continually renormalized gradient flow yield exact closed-form dynamics that imply NC, including (NC1)–(NC4).The dynamics are derived for the central-path model and expressed through the SNR matrix and its initialization singular vectors.
Empirical motivation: MSE-trained networks exhibit Neural Collapse across five benchmark datasets and three canonical architectures, with (NC1)–(NC4) observed during training.The experiments use MNIST, FashionMNIST, CIFAR10, SVHN, and STL10 with VGG, ResNet, and DenseNet networks.

2 DECOMPOSITION OF MSE LOSS

The paper decomposes MSE loss into a least-squares component interpretable through Neural Collapse and a nonnegative deviation measuring classifier suboptimality. Experiments show the deviation is negligible early in training and through the Terminal Phase, motivating the central path.

MSE loss decomposition: Theorem 1 decomposes MSE loss into an optimal least-squares term depending only on features and a deviation term measuring classifier mismatch.The deviation is nonnegative and vanishes exactly when the classifier equals the least-squares classifier.
NC interpretation: The least-squares component further separates activation collapse from convergence of class-means and classifiers toward a shared Simplex ETF.These correspond respectively to NC1 and NC2–NC3.
Empirical decomposition: The classifier-deviation term becomes negligible compared with the least-squares component from an early epoch and remains so through the Terminal Phase of Training.Thus, the full MSE loss approximately equals the least-squares component during this regime.
Central path: The central path consists of feature-classifier pairs whose real classifier exactly equals the MSE-optimal least-squares classifier for the features.On this path, full MSE loss equals its least-squares component, allowing analysis to focus on that component.

3 EXACT CLOSED-FORM ANALYSIS ON CENTRAL PATH

The paper analyzes continually renormalized gradient flow restricted to the central path under explicit assumptions. It derives closed-form SNR dynamics whose limiting behavior implies all four Neural Collapse properties.

3 EXACT CLOSED-FORM ANALYSIS ON CENTRAL PATH: The continually renormalized flow projects the least-squares gradient onto the tangent space of the manifold while requiring the within-class covariance to remain full-rank.This flow maintains the renormalized feature constraint used in the theory.
3 EXACT CLOSED-FORM ANALYSIS ON CENTRAL PATH: Figure 2 indicates that classifier deviation is negligible from an early epoch, making feature-classifier pairs effectively central-path during the Terminal Phase.The decomposition further suggests Simplex ETF organization precedes continued within-class activation compression.
3 EXACT CLOSED-FORM ANALYSIS ON CENTRAL PATH: Theorem 3 assumes central-path restriction, zero global feature mean, and renormalization to identity within-class covariance.These assumptions make the closed-form Neural Collapse analysis tractable.
3.1 INVARIANCE PROPERTY: The central-path invariance permits invertible coordinate changes that produce renormalized features with spherical covariance and preserve the relevant classifier structure.The signal-to-noise ratio matrix is identified with the class-means matrix of the renormalized features.
3.2 SIGNAL-TO-NOISE RATIO MATRIX ON THE CENTRAL PATH: SNR singular values quantify class separation along singular-vector directions, with the smallest singular value representing the smallest class separation.Larger smallest nonzero singular values correspond to greater immunity of linear separation to noise.
3.3 DYNAMICS OF SNR: Proposition 2 gives exact linear dynamics for every nonzero SNR singular value: a_j(t)=a_j+t.The result provides explicit time evolution for the SNR spectrum.
3.3 DYNAMICS OF SNR: The SNR dynamics drive signal relative to noise toward Neural Collapse, with the limiting matrix and renormalized features implying NC1–NC4.Corollary 2 states convergence of the SNR matrix, while the accompanying results establish the four NC properties.

4 RELATED WORKS

Related analyses have established mathematical frameworks for Neural Collapse, including work on MSE loss and CE-loss models. The paper distinguishes its approach by studying renormalized gradient-flow dynamics on the central path.

Prior Neural Collapse analyses: Several prior works proposed mathematical frameworks for explaining Neural Collapse after its initial publication.The paper surveys these works in Appendix E.
MSE-NC analyses: Existing MSE-NC analyses considered either a linearized unconstrained-features ODE or homogeneous weight-normalized networks.These approaches differ from the dynamics studied in this paper.
CE-loss analyses: CE-loss studies examined NC in unconstrained-features or layer-peeled models, focusing on loss landscapes or global minima rather than training dynamics.The paper positions its analysis as dynamic rather than solely equilibrium-based.

5 CONCLUSION

The paper verifies Neural Collapse in MSE-trained canonical networks and datasets, then develops a decomposition, central-path construct, and renormalized flow whose closed-form dynamics imply Neural Collapse.

Empirical findings: Experiments verify Neural Collapse in prototypical classification deep nets trained with MSE loss on canonical datasets.The paper then derives and measures a new decomposition of MSE loss.
Theoretical framework: The observed tendency of classifiers toward least-squares optimality motivates the central path, where classifiers remain exactly MSE-optimal for the features.Renormalized features and continually renormalized gradient flow are then analyzed on this path.
Theoretical framework: The resulting closed-form dynamics imply the occurrence of Neural Collapse.This connects the empirical MSE-NC observation to an explicit theoretical dynamical analysis.

REPRODUCIBILITY STATEMENT

The paper documents reproducible MSE-NC experiments across five datasets and canonical network architectures, with shared training procedures and public code and measurements.

The paper provides PyTorch code, experimental details, and deposited measurements for reproducing MSE-NC and CE-NC results.
Experiments use MNIST, FashionMNIST, CIFAR10, SVHN, and STL10 with preprocessing matching Papyan, Han, and Donoho (2020).
The study trains VGG, ResNet, and DenseNet models using the same depth-selection and architecture-specification procedures as Papyan, Han, and Donoho (2020).
The selected architectures vary by dataset, including VGG11/13, ResNet18/34/50, and DenseNet40/100/201/250.
Optimization, parameters, and hyperparameter tuning match Papyan, Han, and Donoho (2020).
Scaling-heuristic experimentation was considered prohibitively expensive computationally.

A.6 EXPERIMENTAL RESULTS

MSE-trained networks exhibit multiple Neural Collapse properties and improve beyond zero training error, while some test-data and STL10 configurations converge more slowly or behave atypically.

The median test-accuracy improvement from achieving zero error to the end of training is 0.962 percentage points, versus 1.833 percentage points for the mean.
MSE-trained features and classifiers approach equinormness, equiangularity, maximal-equiangularity, and self-duality.
Last-layer features undergo variability collapse, and classifier decisions converge toward the nearest class center decision rule.
Networks trained with MSE tend to achieve faster activation collapse and better adversarial robustness than networks trained with CE.
STL10-ResNet and STL10-DenseNet sometimes converge more slowly to Neural Collapse or show trends inconsistent with other dataset-network pairs.
The authors hypothesize that STL10’s larger images create a higher-dimensional, harder non-convex optimization problem that may slow convergence.
The authors note that MSE modifications proposed for more challenging classification problems may have been appropriate for the STL10 outliers.
On test data, variability collapse remains visible but proceeds much more slowly than on training data, with larger final plotted measures.

B THEOREM 1

Theorem 1 decomposes MSE loss into a least-squares component and a classifier-deviation term; the least-squares component further separates NC-related contributions.

The MSE loss decomposes as L(f_W, f_H) = L_LS(f_H) + L_⊥LS(f_W, f_H).
The decomposition isolates deviation of the trained classifier from the least-squares solution in the second term.
For fixed feature activations, the least-squares classifier satisfies the relevant first-order optimality condition.
The least-squares component further decomposes into L_NC1(f_H) and L_NC2/3(f_H).
The further decomposition follows from orthogonality between perturbations affecting class or global means and those affecting within-class covariance.

C.2 ADDITIONAL INTUITIONS FOR THEOREM 2 TERMS

The loss terms isolate activation collapse and convergence toward the Simplex ETF, while invariance motivates analyzing normalized features and fiber-to-fiber optimization on the central path.

LNC1 is a variance term that vanishes only under activation collapse (NC1).
With class means fixed, LNC1 approaches zero only when within-class covariance ΣW approaches zero.
LNC2/3 measures the distance of least-squares class-mean predictions from the Simplex ETF associated with NC2 and NC3.
Features related by an invertible transformation generate identical least-squares predictions and MSE loss, so optimization moves between equivalent fibers.
The normalized features manifold consists of activations with identity within-class covariance, allowing each full-rank feature matrix to be represented through a positive-definite transformation.
Continual renormalization preserves class predictions and MSE performance while repeatedly projecting gradient steps back to the normalized manifold.

D.4 PROOF OF PROPOSITION 2 (GRADIENT FLOW IN ALIGNED SNR COORDINATES)

In aligned SNR coordinates, the proof differentiates singular values through feature rows and obtains their gradient-flow dynamics, yielding an explicit closed-form solution.

The derivative of each SNR singular value ωj depends only on the corresponding aligned feature row xj.
The derivation applies the chain rule to the feature-row dynamics and uses the central-path loss representation and label identities.
Under continually renormalized gradient flow, the dynamics of each non-zero SNR singular value are given explicitly by Lemma 5.
Proposition 2 solves these dynamics in closed form as ωj(t)=aj+t for j=1,...,C−1.

D.5 PROOF OF COROLLARY 1

The closed-form dynamics imply that all non-zero SNR singular values diverge while becoming asymptotically equal relative to one another.

All non-zero SNR singular values diverge to infinity under the Equation 9 dynamics.
Their normalized fourth powers converge to one for every j=1,...,C−1.

D.6 PROOF OF COROLLARY 2

Because SNR singular vectors remain fixed and singular values equalize, the renormalized SNR matrix and features converge to the Neural Collapse geometry under the stated assumptions.

The left and right singular vectors of the SNR matrix remain constant during continually renormalized gradient flow.
Corollary 2 states that the SNR matrix converges to a form determined by its initial singular vectors and largest singular value.
Equal non-zero singular values, zero-mean columns, and rank C−1 imply that renormalized class means converge to a Simplex ETF.
NC3 and NC4 follow from NC1 and NC2 on the central path by the cited theorem of Papyan, Han, and Donoho.
The related infinite-width counterexample shows that Neural Collapse need not occur in two-layer infinite-width networks, while other compared analyses omit closed-form dynamics or convergence rates.

E.7 ZHU, DING, ZHOU, LI, YOU, SULAM, AND QU (2021)

Zhu et al. characterize Neural Collapse in an unconstrained features model and report additional empirical phenomena and architecture substitutions. Their analysis establishes strong loss-landscape results, while the experiments remain limited to selected network–dataset combinations.

Model and theory: Zhu et al. study an unconstrained features model with classifier weights, biases, features, and weight-decay parameters.The model minimizes cross-entropy loss with weight decay over W, H, and b.
Model and theory: All minima in their model exhibit Neural Collapse, and all local minima are global minima.
Empirical phenomena: NC emerges independently of whether training uses SGD, ADAM, or L-BFGS.
Empirical phenomena: NC also emerges with completely shuffled one-hot target vectors, while increasing width expedites NC under random labels.
Architecture substitutions: Their ablations suggest replacing global weight decay with last-layer feature-and-classifier decay, or trained classifiers with fixed Simplex ETF classifiers, without affecting performance.
Scope: The experiments cover limited network–dataset combinations, motivating conjectures about generalization and architecture design rather than establishing those conjectures broadly.