Optimizer choice matters for the emergence of Neural Collapse

AI summary

Overview

• Research area: Deep learning theory and optimization — specifically the geometry of learned representations known as Neural Collapse (NC), and how the choice of optimizer shapes it.
• Technical level: Advanced. The paper combines a large empirical study with convergence proofs on sign-gradient descent and unconstrained feature models; it assumes familiarity with weight decay coupling (Adam vs. AdamW), class-mean geometry, and simplex equiangular tight frames.
• Scope: The paper argues and demonstrates that Neural Collapse is not optimizer-agnostic: coupled weight decay (as in Adam) permits NC, while decoupled weight decay (as in AdamW) prevents it, supported by a new diagnostic metric NC0, four theorems, and roughly 3,900+ training runs.

What This Paper Is About

Neural Collapse is the tendency of a network's last-layer features and classifier weights to arrange themselves into a highly symmetric configuration in the terminal phase of training. Prior theoretical explanations largely ignored the optimizer, implicitly treating NC as universal across optimization methods. This paper shows that the optimizer — and specifically how weight decay is applied — decides whether NC can emerge at all, and backs that claim with both theory and extensive experiments.

Key Contributions

1. A new diagnostic metric, NC0. The paper introduces NC0, the row sum of the last-layer weight (formally α_t = (1/K)‖W_tᵀ1‖²), whose convergence to zero is a necessary (though not sufficient) condition for NC. Proposition 2.1 proves that NC2 and NC3 together imply NC0. Unlike the standard NC metrics, NC0 is analytically tractable: if NC0 diverges, NC cannot have occurred, even if other metrics plateau at small positive values.
2. A demonstration that weight-decay coupling matters. Across a wide range of experiments, coupled weight decay — L2 regularization applied directly inside the gradient update — is identified as a necessary condition for NC to emerge in adaptive optimizers such as Adam and Signum. AdamW and SignumW, despite their algorithmic similarity to Adam and Signum, fail to produce NC.
3. The first theoretical result on momentum in NC. The paper shows momentum has an accelerating effect on NC with SGD, beyond its effect on the convergence of the training loss.
4. Four theorems on NC0 dynamics. SGD with decoupled weight decay (Theorem 3.1), SGD with coupled weight decay (Theorem 3.2), SignGD with decoupled weight decay (Theorem 3.3), and SignGD with coupled weight decay (Theorem 3.4) exhibit qualitatively different NC0 behaviour, supported by over 3,900 training runs.

Main Findings

• NC0 as a necessity test: Since NC2 and NC3 imply NC0, a diverging NC0 rules out NC, providing a clearer criterion than existing metrics, which typically plateau at small but nonzero values under finite training and learning-rate decay.
• SGD converges exponentially: Under SGD with momentum β ∈ [0,1), weight decay λ ∈ [0,1), and learning rate η > 0, NC0 decays exponentially to zero when 0 < ηλ < 2 (decoupled weight decay) or 0 < ηλ < 2(1+β) (coupled weight decay). The rate is proportional to the weight decay, and momentum accelerates it.
• Decoupled weight decay under SignGD does not vanish: With SignGD on the unconstrained feature model (UFM) with the feature matrix fixed to an NC solution, NC0 increases monotonically from zero to the limit (K−2)²/λ². It does not go to zero.
• Coupled weight decay under SignGD can vanish: With SignGD, coupled weight decay, W₀ = 0, and a suitable learning-rate decay scheme η(t) → 0, NC0 does converge to 0. The mechanism is that as ‖W‖ grows, the coupled decay term can flip the sign of the gradient, producing a non-monotonic NC0 trajectory that rises before falling.
• Empirical confirmation at scale: AdamW and SignumW show NC0 and NC3 metrics orders of magnitude larger than Adam and Signum, even at weight decay values several orders of magnitude higher. In one setting (Table 2), SignumW's final NC0 is 3185.69 (≫100% of its initialization value) versus 0.78 for Signum, and AdamW's is 5.33 versus 0.34 for Adam. SGD achieves the lowest values (NC0 = 2.14e-04, a change of < −99.5%).
• Weight decay is essential in practice: Larger weight decay leads to a stronger decrease of NC metrics across all optimizers. The authors note their experiments cannot exclude eventual NC without weight decay in an asymptotic limit, but argue weight decay is essential in practical finite-length training on realistic models.
• Momentum accelerates NC beyond loss convergence: Two runs with weight decay 0.005 and momentum 0.7 versus 0.9 converge to almost exactly the same training loss yet have considerably different final NC metrics, indicating they reach solutions with very different geometric structure.
• Interpolating Adam and AdamW: Holding total weight decay fixed at 0.0005 and momentum at 0.9, increasing the coupled component smoothly improves NC0, NC2 and NC3 while validation accuracy stays largely unaffected. Coupled weight decay = 0 is equivalent to AdamW; coupled weight decay = 0.0005 is equivalent to Adam.
• SGD and SGDW stay close: Consistent with Theorems 3.1 and 3.2, the gap between coupled and decoupled weight decay is much more pronounced in adaptive optimizers than in SGD.
• Singular-value evidence: With AdamW and SignumW, the smallest singular value of W increases during training, indicating failure to satisfy NC3; NC0 and the nonzero singular values of the centered class means M grow and show high variance, indicating NC2 is also less well fulfilled.
• Partial neural collapse: AdamW can reach minimal NC1 and NC2 values among all optimizers while NC0 diverges and NC3 is not satisfied, suggesting NC properties do not always emerge jointly.
• NC4 is redundant in this study: NC4 is consistently satisfied whenever training accuracy approaches 100%, regardless of NC1–NC3, and is largely uncorrelated with the other metrics, so it is excluded from the main analysis.

Methodology in Plain English

The authors attack the problem from two sides. On the empirical side, they train ResNet9 and VGG9 on MNIST, FashionMNIST and Cifar10 using six optimizers — Adam, AdamW, SGD, SGD with decoupled weight decay (SGDW), Signum, and Signum with decoupled weight decay (SignumW). Each optimizer is run with three learning rates, six momentum values, and six weight decay values, giving 2 × 3 × 6 × 108 = 3,888 training runs; with ablations the total exceeds 3,900. All networks train for 200 epochs with batch size 128, and the learning rate decays by a factor of 10 after one-third and two-thirds of training. Runs with too-large weight decay are discarded because the model no longer trains. On the theoretical side, they define NC0, prove it is implied by NC2 and NC3, and then study NC0's trajectory under SGD and SignGD, using the tractable unconstrained feature model with features fixed to an NC solution (SignGD is the special case of Adam and AdamW when β₁, β₂ and ε are set to zero). They also run an interpolation experiment that mixes coupled and decoupled weight decay at a fixed total, and a small-scale MLP experiment on a separable dataset to visualize NC0 dynamics across SGD, Adam, AdamW, Adam0 and AdamW0.

Why This Matters

• Impact on research: Existing NC theory largely studies gradient flow or gradient descent and assumes NC is optimizer-independent. This work identifies weight-decay coupling as part of an optimizer's implicit bias, invalidates the equivalence often assumed between L2 regularization on gradients and weight decay on parameters for adaptive optimizers and momentum, and provides a new analysis-friendly metric (NC0) alongside the first result linking momentum to NC.
• Practical connections reported in the paper:
  • Inducing NC has been used to improve generalization in class-imbalanced training.
  • The emergence of NC has been linked to improved transfer learning.
  • The presence of NC has been connected to better out-of-distribution detection.
  • Their interpolation experiment shows that NC strength can vary while validation accuracy stays largely unchanged, supporting skepticism about using NC to explain generalization.
• Industry relevance: Adam and AdamW are default choices in modern training pipelines, and the choice between them is usually treated as a minor implementation detail. This paper says the two produce geometrically different solutions even at equal total regularization, which is directly relevant to anyone tuning optimizers, weight decay, or momentum for classification models.

Future Directions

1. Broader architectures and datasets. The authors call for expanding experiments to larger models such as Vision Transformers and DenseNets and more diverse datasets; their preliminary results on ViTs largely confirm the findings extend to Transformers.
2. Beyond the last layer. The study only analyzes NC properties in the last layer. Prior work suggests these properties may also appear in intermediate layers, so examining NC across depths could illuminate hierarchical feature representations.
3. Newer optimizers. Lion and Mars (first-order) and Shampoo, SOAP and Muon (second-order) show promising convergence and generalization benefits, but their effects on NC remain largely unexplored.
4. Theory beyond NC0. The theoretical results cover SignGD on the unconstrained feature model and only track NC0; a full theoretical understanding of NC1–NC3 under realistic optimization dynamics, and theoretical conditions under which only a subset of NC properties holds, remain open.

Target Audience

This paper suits optimization and deep learning theory researchers, particularly those working on implicit bias, weight decay, and representational geometry; practitioners who need to justify Adam versus AdamW choices; and graduate students already comfortable with NC literature who want a rigorous account of why the optimizer cannot be abstracted away. Readers without a background in convergence analysis or class-mean geometry will find the theoretical sections challenging, though the empirical figures and the NC0 diagnostic are broadly accessible.

Authors’ abstract

Neural Collapse (NC) refers to the emergence of highly symmetric geometric structures in the representations of deep neural networks during the terminal phase of training. Despite its prevalence, the theoretical understanding of NC remains limited. Existing analyses largely ignore the role of the optimizer, thereby suggesting that NC is universal across optimization methods. In this work, we challenge this assumption and demonstrate that the choice of optimizer plays a critical role in the emergence of NC. The phenomenon is typically quantified through NC metrics, which, however, are difficult to track and analyze theoretically. To overcome this limitation, we introduce a novel diagnostic metric, NC0, whose convergence to zero is a necessary condition for NC. Using NC0, we provide theoretical evidence that NC cannot emerge under decoupled weight decay in adaptive optimizers, as implemented in AdamW. Concretely, we prove that SGD, SignGD with coupled weight decay (a special case of Adam), and SignGD with decoupled weight decay (a special case of AdamW) exhibit qualitatively different NC0 dynamics. Also, we show the accelerating effect of momentum on NC (beyond convergence of train loss) when trained with SGD, being the first result concerning momentum in the context of NC. Finally, we conduct extensive empirical experiments consisting of 3,900 training runs across various datasets, architectures, optimizers, and hyperparameters, confirming our theoretical results. This work provides the first theoretical explanation for optimizer-dependent emergence of NC and highlights the overlooked role of weight-decay coupling in shaping the implicit biases of optimizers.