PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, Zhihui Zhu

arXiv:2203.01238v2cs.LGcs.AIcs.ITmath.ITmath.OCstat.ML

TL;DR

MSE has shown performance comparable to or better than cross-entropy, but its neural-collapse landscape was less fully characterized than cross-entropy's. Using an unconstrained-feature model, the paper analyzes MSE's global landscape and rescaled variant. It finds neural-collapse global solutions and strict saddles at other critical points.

  • Problem

    MSE's neural-collapse theory was limited, with prior analyses often addressing gradient flow rather than the global optimization landscape.

    Source: §1 · p. 2

  • Method

    The authors analyze regularized MSE under an unconstrained-feature model, characterizing global solutions and critical points.

    Source: §1 · p. 3, §2.1 · p. 4

  • Results

    Every critical point is either a global solution exhibiting neural collapse or a strict saddle with negative curvature.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The analysis extends theoretical support for neural collapse under MSE beyond results restricted to particular training algorithms.

    Source: §1 · p. 2, §5 · p. 13

    The theoretical analysis uses an unconstrained-feature model that treats last-layer features as free optimization variables.

    Source: §2.1 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Although MSE has shown performance comparable to or better than cross-entropy, its neural-collapse optimization landscape was less understood. This work analyzes MSE with unconstrained features, characterizing neural-collapse global solutions and a benign landscape, and examines rescaled MSE.

1 Introduction

  • Neural collapse describes within-class feature variability collapsing, class means forming a Simplex ETF, classifier self-duality, and a nearest-class-center decision rule.The passage set also states that the class means have equal length and are maximally distant.Source: §1 · p. 2
  • Earlier MSE work reported neural collapse and faster collapse, but theoretical analysis remained limited and often studied gradient flow rather than iterative training with nontrivial stepsizes.The cited prior analyses include small-initialization linearization and continual feature renormalization.Source: §1 · p. 2
  • The authors characterize all global solutions as satisfying neural-collapse properties across all feature dimensions, unlike prior analyses limited to particular gradient-flow solutions or d > K.Prior MSE analyses cited here studied gradient flow; one relied on small initializations, while another considered continually renormalized features.Source: §1 · pp. 2–3
  • Every critical point is either a neural-collapse global solution or a strict saddle with negative curvature, ruling out spurious local minima.Methods that escape strict saddles, such as stochastic gradient descent, converge to a global solution exhibiting neural collapse.Source: §1 · p. 3
  • Rescaled MSE also exhibits neural collapse, and its scaling aligns gradients perpendicular to class decision boundaries, potentially aiding learning discriminative features.The authors report that rescaling is empirically important for competitive performance, particularly with many classes.Source: §1 · p. 3
  • For d ≥ K, MSE and cross-entropy learn similar neural-collapse features and classifiers; for d < K, their performance differs, and MSE solutions depend importantly on bias.Without bias, MSE minimizers are orthonormal matrices rather than Simplex ETFs when d ≥ K; cross-entropy can achieve neural collapse without bias.Source: §1 · pp. 3–4

2 The Problem Setup

The setup models balanced multiclass classification with a regularized MSE objective and treats last-layer features as free variables for analysis. It also motivates rescaled MSE for tasks with many classes and investigates its geometric rationale.

2 The Problem Setup

  • The classification setup uses a feature mapping followed by a linear classifier, with balanced class samples trained to approximate one-hot labels using regularized MSE.The setup denotes the number of classes by K, uses n samples per class, and specifies a positive weight-decay parameter.Source: §2 · p. 4

2.1 Basic Problem Formulation Based on Unconstrained Feature Models

  • The unconstrained feature model treats last-layer features as free optimization variables, motivated by overparameterized networks' capacity to approximate arbitrary representations.In this formulation, weight decay is applied to both classifier weights and features; the authors note that feature energy can be implicitly bounded when inputs are bounded.Source: §2.1 · pp. 4–5

2.2 Rescaled MSE Loss under Unconstrained Features

  • For large-class tasks, the authors motivate rescaled MSE as a way to emphasize the correct class and examine its geometric rationale for loss design.They state that vanilla MSE has reported good performance mainly for smaller class counts, while rescaling is often used for larger ones.Source: §2.2 · p. 5

3 Main Theoretical Results

Under the unconstrained feature model, the paper characterizes global minimizers of vanilla MSE and analyzes its optimization landscape. It also uses landscape visualizations to explain how rescaling may improve optimization around neural-collapse solutions.

3.1 Global Optimality Conditions

  • With balanced class sizes, the only global minimizers of vanilla MSE satisfy neural-collapse properties, subject to the theorem’s regularization conditions.When λWλH ≥ 1/(NK), W and H are zero; below that threshold, the minimizers satisfy NC1 and NC3.Source: §3.1 · pp. 5–6
  • For d ≥ K−1, class means at a global solution form a simplex ETF for appropriate regularization choices; for d < K−1, they are its best rank-d approximation.In the lower-dimensional case, the class means need not have equal lengths or be maximally pairwise-distanced.Source: §3.1 · p. 6
  • The paper’s analysis of vanilla MSE establishes global optimality conditions across feature dimensions and weight-decay choices, whereas the corresponding rescaled-MSE analysis is left for future work.The authors say handling both rescaling parameters α and M requires additional technical work.Source: §3.1 · p. 7

3.2 Characterizations of The Benign Global Landscape

  • For d > K, the full MSE loss is strict saddle, so every critical point is a global minimizer or has a negative-curvature direction.The paper states that strict-saddle geometry with no spurious local minima supports convergence to a global solution for methods that can escape strict saddles.Source: §3.2 · pp. 7–8

3.3 Delving Deeper into Optimization Landscapes: Why Rescaling Helps?

  • The visualization fixes W as a simplex ETF and plots feature norm s against angle θ to the class classifier, separating directions parallel and perpendicular to the decision boundary.The plotted landscape shows vanilla MSE steep in s but flat in θ, while CE and rescaled MSE with large M are steeper in θ over the described region.Source: §3.3 · pp. 8–9, Figure 2 · p. 9

4 Experiments

Experiments evaluate neural collapse (NC) on practical networks, introduce measures of feature diversity and margins, and test algorithm choice, feature dimension, fixed classifiers, and rescaled MSE.

New metrics for evaluating NC

  • The experiments use NC1–NC3 and add numerical rank to measure feature dimensionality and cosine margin to assess individual-feature margins.Numerical rank is averaged across classes; smaller values indicate stronger feature collapse toward class means.Source: §4 · p. 10

The prevalence of NC across different optimization algorithms

  • NC1, NC2, and NC3 converge to zero with SGD, Adam, and L-BFGS on CIFAR10, while cosine-margin distributions align with test accuracy across algorithms.The experiments use ResNet18 with vanilla MSE; the reported generalization performances differ by algorithm.Source: Figure 3 · p. 10, §4 · p. 11

Improving network efficiency via fixing classifiers as simplex ETFs

  • Fixing the last-layer classifier as a Simplex ETF and reducing feature dimension to d = K can reduce trainable parameters without sacrificing generalization performance.The reported comparison is on CIFAR10-ResNet18 with data augmentation and SGD.Source: §4 · p. 11, Figure 5 · p. 12

Choice of the feature dimension d

  • For MSE, d ≥ K−1 is crucial for fitting training data and generalizing to test data; as d decreases below K−1, more samples lie on the decision boundary.By contrast, CE produces similar performance across the tested feature dimensions.Source: §4 · pp. 11–12, Figure 4 · p. 11

Experiments of the rescaled MSE loss

  • With one rescaling parameter fixed, increasing the other correlates positively with NC and training and test performance on miniImageNet.The experiment varies α with fixed M and M with fixed α using ResNet18.Source: §4 · p. 12, Figure 6 · p. 12

5 Conclusion

The paper analyzes the global landscape of MSE under the unconstrained feature model and concludes that global solutions exhibit neural collapse without spurious local minimizers. It also notes that rescaling is empirically useful but heuristic, and that its landscape explanation relies on a strong simplification.

5 Conclusion

  • Under the unconstrained feature model, all global solutions exhibit NC, and the landscape has no spurious local minimizers.The conclusion extends the studied scope of NC under MSE beyond particular training algorithms and, with sufficiently expressive architectures, broadens the prevalence claim across CE and MSE.Source: §5 · p. 13
  • Rescaled MSE improves NC and test performance over vanilla MSE, but the rescaling is heuristic and its landscape explanation simplifies correlated variables H, W, and b to θ and s.The paper leaves deriving an optimal loss function for inducing NC to future work.Source: §5 · p. 13

Appendices

The appendices define the Hessian representation and Simplex ETF notation, then organize supporting proofs and experimental details.

Appendices

  • The appendix represents the Hessian as a bilinear form, avoiding a tensor representation or vectorizing the variable.This bilinear form is used throughout the appendix.Source: Appendices · p. 18
  • The appendix formally defines a standard K-Simplex ETF and extends the notation to a general Simplex ETF in R^d.The general definition uses an orthonormal matrix, with a specified choice when d = K−1.Source: Appendices · p. 18
  • Appendix A gives dataset, architecture, and training settings; Appendix B proves Theorem 3.1; Appendix C details the loss visualizations.Source: Appendices · p. 18

A Technical Details of the Experimental Setup in Section 4

The experimental appendix specifies datasets and training settings, defines NC measures, and includes a visualization comparing learned features under MSE and CE.

A Technical Details of the Experimental Setup in Section 4

  • Experiments use ResNet18; CIFAR10 is trained for 200 epochs with SGD, Adam, or LBFGS, while miniImageNet uses rescaled MSE and SGD.CIFAR10 uses weight decay 5 × 10^-4 and data augmentation unless stated otherwise; miniImageNet is split into 50,000 training and 10,000 validation images across all 100 classes.Source: A Technical Details of the… · p. 18

Three NC measures NC1-NC3

Visual comparison of features learned by MSE and CE losses with feature dimension d = 3

  • With feature dimension d = 3, the CIFAR10 visualization shows CE-learned features as more diverse and discriminative than MSE-learned features.The comparison uses ResNet18 trained with SGD and displays randomly selected training samples from each class.Source: A Technical Details of the… · p. 19, Figure 7 · p. 19

B Proof of Theorem 3.1 in Section 3.1

The appendix restates the global optimality conditions for the MSE objective, including a threshold case in which the global minimizer has zero weights and features.

B Proof of Theorem 3.1 in Section 3.1

B.1 Main Proof

The main proof reduces global optimization to characterizing the weight and feature matrices for a given bias, then optimizing over the bias.

B.1 Main Proof

  • The proof first characterizes W and H in terms of the bias b, using the centered label matrix eY = Y − b1⊤.Source: §B.1 · p. 20
  • For fixed bias energy, minimizing the stated objective is equivalent to minimizing the singular values of the centered label matrix.Source: §B.1 · pp. 20–21
  • The proof restricts the bias norm to c ≤ 1/√K and considers separate cases for feature dimension d < K and d ≥ K.Source: §B.1 · p. 21
  • After optimizing over the bias in those cases, the proof invokes Theorem B.4 to characterize W⋆ and H⋆.Source: §B.1 · p. 21

B.2 Supporting Lemmas

The supporting lemmas characterize critical points and global solutions through singular vectors, then show that every non-minimizing critical point has negative curvature.

B.2 Supporting Lemmas

  • At a critical point, the gradients imply a balance relation between W and H, allowing an orthogonalized representation without changing the product WH or Hessian information.Source: §B.2 · pp. 21–22
  • Any critical point that is not a local minimizer is a strict saddle with at least one negative Hessian eigenvalue.Source: §B.2 · p. 22
  • In the orthogonalized representation, each paired nonzero column of W and row of H corresponds to scaled left and right singular vectors of eY.Source: §B.2 · p. 23
  • For d ≥ K and d < K, the minimizers are characterized by covering the eligible singular directions, with the rank constraint differing across cases.Source: §B.2 · pp. 23–24
  • The global minimum objective is attained by the stated forms of W⋆ and H⋆, including a best-rank-d approximation when d < K.Source: §B.2 · pp. 22, 24, 27
  • The negative-curvature argument considers an uncovered singular direction and uses the directional Hessian to establish strict saddles in both dimension cases.Source: §B.2 · pp. 24–25
  • With balanced class counts, the centered label matrix has at least K−2 singular values equal to √n; the remaining two depend on b.Source: §B.2 · p. 25
  • When λWλH ≥ 1/NK, W⋆ and H⋆ are zero; below that threshold, the theorem gives dimension-dependent global-solution forms.Source: §B.2 · p. 27

C.1 Details of the Visualization Technique

The appendix derives a visualization of classifier outputs and losses using a feature vector's norm and angle in a plane spanned by two ETF classifier weights.

C.1 Details of the Visualization Technique

  • The feature vector is parameterized by radial coordinate s and angular coordinate θ, and the classifier output is expressed from these coordinates.Source: §C.1 · pp. 28–29
  • The visualization fixes classifier weights to a K-Simplex ETF with zero bias and places h in the plane spanned by wk and wk′.Source: §C.1 · p. 28
  • Figure 2 plots the rescaled MSE and CE losses as functions of the feature coordinates.Source: §C.1 · p. 29

C.2 Visualization of the Gradient Vector Field

The gradient-field visualization shows how rescaling MSE changes descent directions around a feature when the classifier is fixed as a simplex ETF. Increasing α or M strengthens the gradient component along θ while leaving the s component fixed.

C.2 Visualization of the Gradient Vector Field

  • Figure 8 fixes W as a simplex ETF and plots the landscape for one feature, with its norm s and angle θ as axes.The figure visualizes gradient fields for MSE and CE.Source: §C.2 · p. 29, Figure 8 · p. 30
  • The arrows point in gradient-descent directions, with lengths proportional to gradient norm.Source: §C.2 · p. 30
  • Increasing α or M strengthens the MSE gradient component along θ while leaving its s component fixed.For α or M in (1, ∞), the θ-direction component increases while the s-direction component stays fixed.Source: §C.2 · p. 30
Loading 2203.01238v2…