A Geometric Analysis of Neural Collapse with Unconstrained Features

Zhihui Zhu; Tianyu Ding; Jinxin Zhou; Xiao Li; Chong You; Jeremias Sulam; Qing Qu

Source-linked AI summary

A Geometric Analysis of Neural Collapse with Unconstrained Features

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

arXiv:2105.02375v1May 6, 2021cs.LGcs.AIcs.ITmath.OCstat.ML

TL;DR

The paper asks how Neural Collapse can be theoretically understood and efficiently optimized despite the nonconvexity of neural-network training. It studies a top-down unconstrained feature model and shows that regularized cross-entropy has a benign landscape whose global solutions are Simplex ETFs, with practical implications for classifier design and training cost.

Problem
The paper addresses the limited theoretical understanding of Neural Collapse and the optimization landscape governing last-layer features and classifiers in deep networks.
Method
The paper models last-layer features as free optimization variables and analyzes a cross-entropy objective with bias and weight decay.
Results
Simplex ETFs are the only global minimizers, while all other critical points are strict saddles; experiments observe Neural Collapse across SGD, Adam, and LBFGS.
Takeaways & Limitations
Fixing the classifier as a Simplex ETF and setting feature dimension near the class count can maintain performance while reducing computational costs.
Takeaways & Limitations
The analysis treats last-layer features as free variables and applies weight decay only to the classifier and features rather than all network parameters.

Abstract

from arXiv · show

We provide the first global optimization landscape analysis of $Neural\;Collapse$ -- an intriguing empirical phenomenon that arises in the last-layer classifiers and features of neural networks during the terminal phase of training. As recently reported by Papyan et al., this phenomenon implies that ($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. We study the problem based on a simplified $unconstrained\;feature\;model$, which isolates the topmost layers from the classifier of the neural network. In this context, we show that the classical cross-entropy loss with weight decay has a benign global landscape, in the sense that the only global minimizers are the Simplex ETFs while all other critical points are strict saddles whose Hessian exhibit negative curvature directions. In contrast to existing landscape analysis for deep neural networks which is often disconnected from practice, our analysis of the simplified model not only does it explain what kind of features are learned in the last layer, but it also shows why they can be efficiently optimized in the simplified settings, matching the empirical observations in practical deep network architectures. These findings could have profound implications for optimization, generalization, and robustness of broad interests. For example, our experiments demonstrate that one may set the feature dimension equal to the number of classes and fix the last-layer classifier to be a Simplex ETF for network training, which reduces memory cost by over $20\%$ on ResNet18 without sacrificing the generalization performance.

1 Introduction

Neural Collapse describes last-layer features and classifiers converging to highly structured, maximally separable configurations. This paper analyzes that phenomenon through an unconstrained feature model and proves a benign optimization landscape with practical training and cost implications.

Neural Collapse: Neural Collapse includes vanishing within-class feature variability and class means arranged as a Simplex ETF.The phenomenon also includes self-duality between classifiers and class means and a nearest-class decision rule.
Analytical approach: The unconstrained feature model treats last-layer features as free optimization variables, isolating the topmost layers for analysis.This top-down simplification addresses the difficulty created by nonlinear interactions across network layers.
Main contributions: The regularized cross-entropy loss has a benign global landscape: Simplex ETFs are global solutions, while every other critical point is a strict saddle.Strict saddles have negative-curvature Hessian directions, so the landscape has no spurious local minimizers.
Main contributions: Optimization methods that escape strict saddles converge to global solutions exhibiting Neural Collapse, consistent with experiments across diverse algorithms.The paper reports this behavior for practical overparameterized networks trained with different optimization algorithms.
Significance: The feature-learning perspective connects the analysis to network design, generalization, robustness, and empirical Neural Collapse observations.The paper contrasts this top-down interpretation with prior landscape analyses that offer limited practical guidance.

2 The Problem Setup

The problem setup replaces learned penultimate representations with free feature variables in a simplified last-layer classification model. It uses balanced classes, cross-entropy, bias, and weight decay while acknowledging the simplification relative to full networks.

Deep-network background: A deep network maps inputs through composed affine and nonlinear layers, with the penultimate output serving as the learned feature representation.The final network output can be viewed as a linear classifier acting on these features.
Deep-network background: For K-class classification, training optimizes network parameters using a loss that compares model outputs with one-hot class labels.The setup introduces class sample counts and cross-entropy as the relevant training objective.
Unconstrained feature model: The unconstrained feature model treats each last-layer feature h = φθ(x) as a free decision variable, peeling away the preceding network layers.This isolates the last-layer classifier and features while preserving the objects characterized by Neural Collapse.
Objective and assumptions: The formulation assumes balanced classes and includes a classifier W, feature matrix H, bias b, and positive weight-decay penalties.The resulting problem is designed to more closely resemble practical cross-entropy training than earlier constrained, bias-free variants.
Model simplifications: Weight decay is applied to the last-layer classifier and features rather than all network parameters, an idealization motivated by bounded feature energy.Experiments cited by the paper report similar Neural Collapse and comparable performance under both approaches.
Model simplifications: Treating features as optimization variables remains a simplification, but the paper argues that overparameterization and experiments preserve the core Neural Collapse properties.The theory identifies Simplex ETFs as the only global minimizers and the loss as a strict saddle function.

3 Main Theoretical Results

The unconstrained feature model has Simplex ETF global solutions, and under d > K its regularized cross-entropy loss has a benign strict-saddle landscape. Thus, optimization methods that escape strict saddles can reach neural-collapse solutions, while bias and nonnegativity constraints shape how collapse is represented.

Global Optimality Conditions: Simplex ETFs are the only global minimizers under balanced classes and feature dimension d ≥ K.The classifier forms a scaled and rotated K-Simplex ETF; global solutions exhibit neural collapse.
Global Optimality Conditions: Neural collapse maximally separates different classes while equalizing their pairwise angles and eliminating within-class feature variability.The class structure is equiangular, and features within each class collapse together.
Bias and Feature Constraints: With unconstrained features, the optimal bias can collapse to identical entries or be removed without affecting classifier performance; nonnegative features change this behavior.Under nonnegativity constraints such as ReLU, the feature mean and bias need not vanish, while centered features can still form a Simplex ETF.
Benign Global Landscape: When d > K, every local minimizer is global, while every other critical point is a strict saddle with a negative Hessian eigenvalue.The theorem rules out spurious local minima and identifies negative-curvature escape directions.
Proof Strategy: The analysis connects the nonconvex factorized problem to a low-rank convex counterpart to derive global optimality and critical-point properties.The formulation also relates unconstrained feature training to low-rank matrix factorization.
Benign Global Landscape: Methods capable of escaping strict saddles, including stochastic gradient descent with random initialization, converge to global neural-collapse solutions almost surely.The result supports convergence across a diverse set of optimization algorithms in overparameterized networks.
Scope and Limitations: The strict-saddle proof uses the null space of W and therefore requires d > K, although the authors conjecture the result may extend to d = K.This boundary is identified as a direction for future work.

4 Experiments

Experiments test whether the unconstrained feature formulation captures Neural Collapse in practical networks and whether NC supports simpler training designs. Across architectures, datasets, optimizers, and weight-decay setups, the experiments examine NC metrics, accuracy, model capacity, and fixed classifiers.

4.1 The Prevalence of NC Across Different Optimization Algorithms: NC1–NC4 consistently collapse to zero under SGD, Adam, and LBFGS on MNIST and CIFAR10 with ResNet18.This indicates NC occurs regardless of the training algorithm, with maximally linearly separable features and a corresponding perfect linear classifier.
4.1 The Prevalence of NC Across Different Optimization Algorithms: Different training algorithms achieve Neural Collapse but exhibit notably different generalization performances on unseen data.Because NC characterizes the training data, generalization also depends on algorithmic bias and the learned weights for the feature representation.
4.2 The Validity of (4) Based on Unconstrained Feature Models for NC: Larger overparameterized networks more readily memorize random labels and tend to exhibit more severe Neural Collapse.Experiments use random-label CIFAR10 with ResNet18 and four-layer MLPs of varying widths; sufficiently large models achieve zero training error.
4.2 The Validity of (4) Based on Unconstrained Feature Models for NC: NC behavior is comparable when weight decay is applied to network parameters Θ or to the peeled features and classifier.The peeled setup reaches 99.57% test accuracy on MNIST and 77.92% on CIFAR10, compared with 99.60% and 78.42% for the classical formulation.
4.3 Insights from NC for Improving Network Designs: Fixing the last-layer classifier as a Simplex ETF matches the learned classifier in NC1, training accuracy, and test accuracy while reducing training parameters and computational complexity.The comparison covers MNIST-ResNet18 and CIFAR10-ResNet50 settings, including data augmentation for the latter experiment.
4.3 Insights from NC for Improving Network Designs: 8.01%, 11.76%, and 52.56% of total learning parameters can be reduced for ImageNet with fixed classifiers on ResNet50, DenseNet169, and ShuffleNet, respectively.For large class counts, Simplex ETFs are also close to orthonormal matrices because pairwise column inner products equal −1/(K−1).

5 Conclusion

The analysis establishes a benign optimization landscape for last-layer features and classifiers, while identifying boundaries concerning deeper layers, robustness, generalization, and large-class settings.

Simplex ETFs are the only global minimizers of the regularized cross-entropy loss, while all other critical points are strict saddles with negative curvature.The result applies to last-layer features and classifiers in the unconstrained feature model.
Methods that escape strict saddles converge to global solutions exhibiting Neural Collapse, consistent with empirical convergence across optimization algorithms.The empirical observation concerns practical overparameterized networks and diverse optimization algorithms.
Investigating Deeper Models: The analysis treats last-layer features as free optimization variables, leaving feature-learning landscapes in shallower layers for future investigation.Empirical evidence suggests shallower layers exhibit less severe variability collapse than deeper layers.
Study of the Relationship Between NC and Robustness: The robustness implications of Neural Collapse remain unresolved because maximum-dimensional linearly separable subspaces may improve robustness to random corruptions compared with ETF collapse.The passage calls for further empirical and theoretical investigation.
Study of the Relationship Between NC and Generalization: Neural Collapse concerns training data and cannot fully explain generalization, since different optimization algorithms can yield different test performance despite NC.The overparameterized setting permits infinitely many networks with NC for one training dataset.
Study of the Relationship Between NC and Network Training: The theory assumes arbitrary feature production before the final layer, whereas deep networks can have difficult training landscapes involving vanishing and exploding gradients.Residual and isometric learning are described as important for effective training beyond a few layers.
Dealing with a Large Number of Classes K ≫ d: The classification analysis focuses on K < d, a condition used to establish negative curvature, while many applications involve very large numbers of classes.Recommendation and document-retrieval systems are cited as examples where large K may make feature dimension costly.

Appendices

The appendices introduce gradient and Hessian notation, then provide proofs of global optimality and the strict-saddle property central to the paper.

The gradient of f(Z) is represented as a K × N matrix of entrywise partial derivatives.The appendix indexes entries by the coordinates z_ij.
The Hessian is alternatively expressed as a bilinear form, avoiding explicit vectorization of the matrix variable.This representation is used in the proof of the strict-saddle theorem.
Appendix B proves that Simplex ETFs are the only global minimizers, while Appendix C proves strict saddles and excludes spurious local minimizers.These correspond to the paper’s two major theoretical results.

A Basics

The appendix defines Simplex ETFs and develops nuclear-norm and matrix-analysis tools used in the proofs.

A general Simplex ETF is treated as a collection of points in R^d, while the standard construction is specified through matrix notation.The appendix also records identities involving I_K, 1_K, and P^⊤P = I_K.
The factorization argument uses Young’s inequality, trace identities, orthogonality, and Frobenius-norm relations.The equality case is characterized through the corresponding norm-balancing conditions.
Lemma A.3 extends the variational form of the nuclear norm for fixed Z and α > 0.The proof constructs factors from the SVD and identifies an equality-achieving factorization.
The nuclear norm of Z is connected to its singular values and singular value decomposition Z = UΣV^⊤.These definitions support the variational argument developed in Lemma A.3.
A symmetric positive semidefinite matrix has a unique real symmetric positive semidefinite kth root for each fixed integer k ≥ 2.This matrix-analysis result is listed as Lemma A.4.

B Proof of Theorem 3.1

Appendix B restates and proves the global optimality condition for the regularized cross-entropy model under a feature-dimension assumption.

When d > K, any global minimizer of f(W,H,b) = g(WH + b1^⊤) + λ… satisfies the theorem’s stated structural conditions.The displayed objective includes the regularized cross-entropy formulation.
At a global minimizer, either b⋆ = 0 or λ_b = 0, and the normalized classifier matrix forms a K-simplex ETF.The ETF characterization is given in the sense specified by the paper’s definition.

B.1 Main Proof

The proof reduces the objective to a scalar function of the classifier energy ρ, establishes its behavior at the boundaries, and concludes that a finite minimizer exists.

Main Proof: The objective is coercive because weight decay grows at infinity while cross-entropy remains non-negative, so every global minimizer is finite.Coercivity is defined as f(x) tending to +∞ when ||x||_2 tends to infinity.
Main Proof: For fixed classifier energy ρ, the proof lower-bounds the cross-entropy term and characterizes when the bound is attained.The equality conditions identify the parameter forms associated with the global optimum.
Main Proof: The scalar objective ξ(ρ; λW, λH) is increasing in ρ through the relevant auxiliary quantity and decreasing in c1 through the corresponding coefficient.These monotonicity properties determine the limiting behavior used to locate the minimum.
Main Proof: As ρ approaches 0, c1 approaches 1/(K−1) and c2 approaches log K.This gives the small-energy boundary behavior of the reduced objective.
Main Proof: As ρ approaches +∞, c1 approaches +∞, c2 approaches 0, and ξ(ρ; λW, λH) approaches +∞.The reduced objective therefore diverges at large classifier energy.
Main Proof: Continuity on [0,+∞), together with divergence at infinity, implies that ξ(ρ; λW, λH) achieves its minimum at a finite ρ.Figure 8 provides an example of this reduced one-dimensional behavior.

B.2 Supporting Lemmas

The supporting lemmas derive balance and symmetry conditions at critical points, lower-bound the cross-entropy loss, and specify the equality conditions leading to the target structure.

Supporting Lemmas: Critical-point equations imply a balance relation between W and H, expressed through their weighted Gram matrices.The gradients with respect to W and H are combined to obtain this relation.
Supporting Lemmas: Applying the lower bound to WH+b1^⊤ yields a reduced lower bound involving the classifier energy ρ and the regularization parameters.The argument uses a common c1 across samples and classes, followed by AM-GM inequalities.
Supporting Lemmas: The lower bound is attained only under the stated equality conditions, which constrain the logits and the relationships among W, H, and b.The proof combines conditions from the cross-entropy and balance arguments.
Supporting Lemmas: The equality conditions further enforce classwise symmetry in the logits and pairwise relations among classifier coordinates.These relations are derived for all distinct class indices k and ℓ.
Supporting Lemmas: For one-hot labels, Jensen’s inequality lower-bounds cross-entropy, with equality requiring equal non-target logits.The bound is parameterized by c1 and c2, whose definitions are supplied in the lemma.

C Proof of Theorem 3.2

For feature dimension d>K, the theorem establishes that the loss is a strict saddle function with no spurious local minima, and that its local minima have the prescribed global form.

C Proof of Theorem 3.2: When d>K, the function f(W,H,b) is a strict saddle function with no spurious local minimum.The theorem applies under the stated feature-dimension condition.
C Proof of Theorem 3.2: Every local minimizer is a global minimum of the form characterized in Theorem B.1.Thus local minima correspond to the paper’s target global solutions.
C Proof of Theorem 3.2: Every critical point that is not a local minimum has a negative-curvature direction, meaning its Hessian has at least one negative eigenvalue.This is the strict-saddle component of the theorem.

C.1 Main Proof

The main proof connects the factorized nonconvex problem to a convex formulation, transfers optimality conditions, and constructs negative-curvature directions for the remaining critical points.

C.1 Main Proof: Critical points are partitioned into two sets: one already consists of global optima, while the other must be shown to have negative curvature.This reduces the theorem to analyzing the second subset.
C.1 Main Proof: The factorized problem is connected to a convex program through Z=HW and a rescaled regularization parameter.A global minimizer of the convex program supplies a lower bound for the original factorized objective.
C.1 Main Proof: Optimality conditions from the convex problem are transferred to critical points of the nonconvex problem.When the transferred conditions hold, the corresponding factorized point is a global solution.
C.1 Main Proof: For points in the second set, the proof evaluates the Hessian bilinear form along a selected perturbation direction.The direction is designed using singular vectors of the loss Hessian.
C.1 Main Proof: Because d>K, W has a nonzero null-space vector a, and the balance relation implies H^⊤a=0.These orthogonality properties simplify the Hessian expression along the constructed direction.
C.1 Main Proof: The constructed direction yields strictly negative curvature for every critical point in the second set, completing the strict-saddle proof.The conclusion follows from the singular-value condition and the resulting Hessian inequality.

C.2 Supporting Lemmas

The supporting lemmas characterize global optimality for the convex reformulation and connect qualifying critical points of the nonconvex program to global minima.

C.2 Supporting Lemmas: Any critical point of the convex program is global, with optimality characterized by the stated subdifferential condition.
C.2 Supporting Lemmas: The convex program's global solutions are characterized by a necessary and sufficient optimality condition involving the gradient and nuclear-norm subdifferential.The nuclear-norm subdifferential is described using the compact SVD Z = UΣV^⊤ and an orthogonality-constrained residual.
C.2 Supporting Lemmas: A critical point of the nonconvex program satisfying the specified condition yields a product Z⋆ = W⋆H⋆ that satisfies the convex program's global optimality condition.The proof derives this using the factorized critical-point equations and the compact SVD of Z⋆.
C.2 Supporting Lemmas: Such a qualifying nonconvex critical point is therefore a global minimum of the convex reformulation and of the original factorized objective.The factorized objective attains the convex objective at (W⋆, H⋆, b⋆) and is no larger than every competing factorization.