PaperlayerView exact revision on arXiv

Source-linked AI summary

Extended Unconstrained Features Model for Exploring Deep Neural Collapse

Tom Tirer, Joan Bruna

arXiv:2202.08087v3cs.LG

TL;DR

Neural collapse has been observed in trained networks, and prior UFM analyses largely considered regularized cross-entropy. This paper analyzes regularized MSE and extends the UFM with another weight layer and ReLU; its nonlinear extension is empirically useful for modeling practical-network neural collapse.

  • Problem

    Prior UFM analyses largely considered regularized cross-entropy, leaving the regularized-MSE minimizers and models of neural-collapse behavior across depth to analyze.

    Source: §1 · p. 1, §4.1 · p. 5

  • Method

    The paper characterizes regularized-MSE UFM minimizers and extends the model with another weight layer and ReLU nonlinearity.

    Source: §1 · p. 2, §7 · p. 8

  • Results

    The nonlinear extended UFM empirically models neural collapse observed in practical networks, and ReLU better distinguishes feature behavior across two levels in collapse rate and structure.

    Source: §6 · p. 8, §7 · p. 9

  • Takeaways & Limitations

    The extended model avoids the stated failure of adding ReLU to the plain UFM, where negative first-layer entries cannot be modified by gradient-based optimization after random initialization.

    Source: §7 · p. 9

    The analysis treats feature matrices as free variables rather than tying them to a data distribution, and the perturbation result is only asymptotic.

    Source: §5 · p. 6, §7 · p. 9

Abstract

from arXiv · show

The modern strategy for training deep neural networks for classification tasks includes optimizing the network's weights even after the training error vanishes to further push the training loss toward zero. Recently, a phenomenon termed "neural collapse" (NC) has been empirically observed in this training procedure. Specifically, it has been shown that the learned features (the output of the penultimate layer) of within-class samples converge to their mean, and the means of different classes exhibit a certain tight frame structure, which is also aligned with the last layer's weights. Recent papers have shown that minimizers with this structure emerge when optimizing a simplified "unconstrained features model" (UFM) with a regularized cross-entropy loss. In this paper, we further analyze and extend the UFM. First, we study the UFM for the regularized MSE loss, and show that the minimizers' features can have a more delicate structure than in the cross-entropy case. This affects also the structure of the weights. Then, we extend the UFM by adding another layer of weights as well as ReLU nonlinearity to the model and generalize our previous results. Finally, we empirically demonstrate the usefulness of our nonlinear extended UFM in modeling the NC phenomenon that occurs with practical networks.

1. Introduction

Neural collapse describes features converging within classes while class means form a structured frame aligned with classifier weights. This paper analyzes regularized MSE in the unconstrained features model and extends that model with another weight layer and ReLU.

1. Introduction

  • Neural collapse includes within-class feature convergence, centered class means forming a simplex ETF, and alignment of classifier weights with that frame.After collapse, classification is based on the nearest class center in feature space.Source: §1 · p. 1
  • The unconstrained features model treats training features as free optimization variables, motivated by the adaptability of overparameterized networks.The paper situates its work in a training regime that continues weight optimization after training error vanishes to push loss toward zero.Source: §1 · p. 1
  • Regularized MSE minimizers can have more delicate feature structure, including orthogonality, which also affects the classifier weights.The paper contrasts this with the simplex ETF structure found for regularized cross-entropy.Source: §1 · p. 1
  • The nonlinear extended UFM empirically models neural collapse observed in practical networks.The contributions also include an asymptotic analysis when features are a fixed perturbation around structured collapse.Source: §1 · p. 2
  • The paper extends the UFM with another weight layer and ReLU, generalizing its results while addressing behavior across depth.A linear extended UFM is also analyzed, with a stated limitation in modeling practical depthwise neural collapse.Source: §1 · p. 2

2. Background and Related Work

The section introduces neural collapse and the unconstrained features model, then situates the paper's regularized-MSE analysis and model extensions against prior work.

2. Background and Related Work

  • The network's feature vector is the output of its feature mapping, and the network parameters minimize empirical risk with a loss and regularizer.The model includes classifier weights and bias, as well as parameters for a nonlinear compositional feature mapping.Source: §2 · p. 2
  • Prior empirical work describes neural collapse as within-class feature convergence, structured class means, classifier-weight alignment, and nearest-center classification.This phenomenon has been studied during training after the training error reaches zero.Source: §2 · p. 2
  • The UFM treats training features as free variables, motivated by the capacity of overparameterized networks to adapt feature mappings to training data.The model pairs a classifier and features with a loss and regularization term.Source: §2 · p. 2
  • Prior regularized MSE work considered the plain UFM, while this paper also extends it with another layer and nonlinearity to address its one-level scope.Earlier regularized or constrained cross-entropy UFM results characterize global minimizers as simplex ETFs; unregularized MSE can have noncollapse minimizers.Source: §2 · p. 3

3. NC for Unconstrained Features Model with Regularized MSE Loss

For regularized MSE, the minimizing feature and weight structures depend on bias treatment. The bias-free case yields an orthogonal frame, while the unregularized-bias case yields a simplex ETF.

3. NC for Unconstrained Features Model with Regularized MSE Loss

  • The paper characterizes global minimizers for bias-free and unregularized-bias settings, and notes that matrix-factorization optimization has no local minima, only global minima and strict saddle points.The characterizations also inform conclusions for optimizable bias with positive regularization.Source: §3 · p. 3

3.1. The Bias-Free Case

  • For bias-free regularized MSE, when c ≤ 1 the collapsed features form an orthogonal frame (H^T H = ρ I_K) and weights align with feature columns; when c > 1, the minimizer is (W*, H*) = (0, 0).Source: §3.1 · pp. 3–4

3.2. The Unregularized-Bias Case

  • With unregularized optimizable bias, minimizers have simplex-ETF features and aligned weights when c ≤ 1; when c > 1, the minimizer has zero weights and features.The paper notes that this result resembles prior cross-entropy results and reports it as a theorem for MSE.Source: §3.2 · p. 4
  • For positive finite bias regularization, minimizers may be neither orthogonal frames nor simplex ETFs, although experiments found centered features forming a simplex ETF and weights aligned with features.The bias-free and unregularized-bias settings provide the two limiting cases for this observation.Source: §3.2 · pp. 4–5

4. Extended Unconstrained Features Model

The one-level UFM cannot represent behavior across depth, so the paper adds an intermediate weight layer and then ReLU. The nonlinear extension preserves an orthogonal-frame structure under a stated condition, while the linear extension still misses depthwise behavior seen in practical networks.

4.1. Unconstrained Features Model With an Additional Layer

  • With d>K, global minimizers make both H1 and W1H1 collapse to orthogonal frames, with their Gram matrices proportional to corresponding weight products.To the authors’ knowledge, this is the first characterization of minimizers for a multilayer UFM.Source: §4.1 · p. 5
  • With cross-entropy loss, the authors empirically observe that H1 and W1H1 at global minimizers collapse to a simplex ETF.The theorem’s orthogonal-frame result for MSE differs from this observed cross-entropy structure.Source: §4.1 · p. 5
  • The linear extension does not match practical depthwise collapse: both levels collapse simultaneously, and their NC1 variability is similar when H1's columns avoid W1's null space.The authors introduce ReLU to break this similarity between feature levels.Source: §4.1 · pp. 5–6

4.2. Non-Linear Unconstrained Features Model

  • Under nuclear-norm equality, the ReLU-activated deepest features collapse to an orthogonal frame aligned with the output weights.The theorem requires this equality for its proof; the authors report observing it in their numerical experiments.Source: §4.2 · p. 6

5. Toward Generalizing the UFMs Results to Other Models

Because existing theoretical analyses treat features as free optimization variables, the paper takes an initial step toward connecting them to data. Its asymptotic analysis shows that certain optimal-weight properties return even when features are fixed and non-collapsed.

5. Toward Generalizing the UFMs Results to Other Models

  • Unlike standard UFM analyses with freely optimized features, this analysis fixes a non-collapsed feature matrix to demonstrate a step toward data-linked models.The paper identifies connecting features to a data distribution as an important direction for future work.Source: §5 · p. 6
  • As sample count tends to infinity, optimal weights recover properties including orthogonal structure and alignment with the fixed, non-collapsed feature matrix.The authors frame this as a feasibility demonstration, while deferring a comprehensive study linking the models to data.Source: §5 · p. 6

6. Numerical Results

The experiments use gradient descent to corroborate theorems by tracking objective values and neural-collapse metrics. They report that ReLU better distinguishes feature behavior across levels, while the nonlinear extended model’s metrics correlate with those of a practical ResNet18.

6. Numerical Results

  • The theorem-setting experiments use gradient descent and demonstrate theorem verification through convergence of the NC metrics to zero.Metrics are computed every 5e3 iterations.Source: §6 · p. 6
  • The no-bias MSE case measures NC2 against orthogonal frames, while the unregularized-bias MSE case measures it against simplex ETFs.Figures 1 and 2 are described as corroborating Theorems 3.1 and 3.2, respectively.Source: Figure 1 · p. 7, Figure 2 · p. 7, §6 · p. 7
  • In the linear two-level model, both W1H1 and H1 collapse to orthogonal frames, as measured by NC1 and NC2.Figure 3 reports metrics for the two-level model.Source: §6 · p. 7, Figure 3 · p. 8
  • Adding ReLU better distinguishes the two feature levels’ collapse rates and structures than the linear extended UFM does.The comparison is between Figures 3 and 4.Source: §6 · p. 8, Figure 3 · p. 8, Figure 4 · p. 8
  • ResNet18’s NC metric behavior correlates with the nonlinear extended UFM under both MSE without bias and cross-entropy with bias.The comparison uses a ResNet18 trained on MNIST.Source: §6 · p. 8, Figure 5 · p. 8

7. Conclusion

The paper characterizes regularized-MSE UFM minimizers and extends the model with another weight layer and ReLU to capture collapse across depth. Experiments demonstrate the nonlinear extension’s usefulness for modeling neural collapse in practical networks.

7. Conclusion

  • The paper characterizes global minimizers of the regularized-MSE UFM, including distinctions from cross-entropy results such as the effect of bias.It also extends the UFM with another weight layer and ReLU to generalize the results.Source: §7 · p. 8
  • Experiments verify the theorems and demonstrate the nonlinear extended UFM’s usefulness for modeling neural collapse in practical networks.Source: §7 · p. 9
  • A single-level UFM with ReLU can fail after random initialization because gradient-based optimization cannot modify negative entries in its first layer; the extended model avoids this limitation.Source: §7 · p. 9
  • The plain UFM has limited expressiveness when features are fixed, and Theorem 5.1’s result is only asymptotic.The authors note that showing positive neural-collapse effects on generalization with limited training data may require departing from the plain UFM.Source: §7 · p. 9

A. Proof of Theorem 3.1

  • When c=K√nλHλW≤1, minimizers have zero within-class variability and aligned class weights and features; when c>1, the minimizer is (W,H)=(0,0).Source: A. Proof of Theorem 3.1 · pp. 11–12

A.1. Alternative proof for Theorem 3.1

  • The alternative proof reduces the objective to singular values, whose minimizing values are equal, giving H^T H proportional to I_K.Source: §A.1 · p. 13

B. Proof of Theorem 3.2

The proof shows that the unregularized bias centers the features and, below a regularization threshold, the class features have simplex-equiangular geometry.

B. Proof of Theorem 3.2

C. Proof of Theorem 4.1

The proof establishes zero within-class variability and characterizes the linear extended UFM’s minimizers through orthogonal structure and layer alignment. It also identifies when the first layer preserves the within-class variability metric.

C. Proof of Theorem 4.1

  • Jensen’s inequality forces every minimizer to have zero within-class variability, so each class’s feature columns are equal.Source: C. Proof of Theorem 4.1 · p. 18
  • The minimizers have H = σ_H R and W2 = σ_W R⊤, with nonnegative scalars and arbitrary orthonormal R; therefore W2 is proportional to H.The scalar values depend on K, n, and the regularization parameters and are obtained numerically from a fourth-degree polynomial.Source: C. Proof of Theorem 4.1 · p. 20
  • Factoring the collapsed features across layers gives W1H1 as an orthogonal frame, with W1 and W2 aligned to the corresponding feature representations.Source: C. Proof of Theorem 4.1 · p. 20
  • The same minimizer analysis for the second two-factor subproblem yields H1 = σ_H1 R̃ and W = σ_W R̃⊤, so W is proportional to H1.Source: C. Proof of Theorem 4.1 · p. 20

D. On the Within-Class Variability Metric NC1

E. Proof of Theorem 4.2

The ReLU proof relates the nonlinear objective to a constrained problem and establishes orthogonal collapse under a stated assumption. The accompanying proof for Theorem 5.1 solves a fixed-feature weight subproblem in closed form.

E. Proof of Theorem 4.2

  • The proof carries the orthogonal-collapse and alignment properties from the intermediate minimization to global minimizers of the nonlinear model.Source: E. Proof of Theorem 4.2 · p. 22
  • Under the assumption that the nuclear norms of H* and σ(H*) agree, the ReLU minimizer can be mapped to one whose activated features have orthogonal collapse.The authors conjecture the assumption may hold necessarily, but state it as an assumption of the theorem.Source: E. Proof of Theorem 4.2 · pp. 22–23

F. Proof of Theorem 5.1

  • For fixed H, the proof solves the strongly convex weight-only subproblem in closed form and recovers W* when H is fixed to H*n.Source: F. Proof of Theorem 5.1 · p. 24

F.1. Intuitive explanation of the result

For features perturbed around a collapsed solution, the asymptotic effect is interpreted as scaling the corresponding weight minimizer. The discussion extends this intuition to the extended UFM but leaves a rigorous analysis for future work.

F.1. Intuitive explanation of the result

  • With zero-mean feature perturbations, the asymptotic weight minimizer is expected to resemble the perfectly collapsed-feature minimizer up to scaling.Source: §F.1 · p. 25
  • The same heuristic extends to fixed perturbed features in the extended UFM, whose objective retains orthogonal invariances and a singular-value formulation.Source: §F.1 · p. 25
  • A rigorous study of fixed-feature effects in the extended UFM is deferred to future work.Source: §F.1 · p. 25

G. More Numerical Results for the Unconstrained Features Model

The numerical results show similar neural-collapse metrics for the nonlinear extended UFM and a trained ResNet18 on CIFAR10. Additional experiments corroborate the stated UFM minimizer structures and show collapse at both feature levels in the linear extended model.

G. More Numerical Results for the Unconstrained Features Model

Related papers

Loading 2202.08087v3…