Neural Collapse Under Different Optimizers

An empirical study of Neural Collapse (Papyan et al, 2020) comparing SGD, Adam, AdamW, and SGDW on CIFAR-10 / ResNet-18.

---

Figure 1: PCA projections of penultimate-layer features at four training stages. SGD achieves near-perfect cluster separation by epoch 350 (within-class variance $= 0.013$); Adam, AdamW and SGDW retain visibly more ($0.047$, $0.133$ and $0.128$ respectively).

1. Introduction

Neural Collapse (NC) describes four geometric properties that emerge in the final-layer representations of deep classifiers when trained past zero training error: within-class features collapse to their class mean, the class means converge to a Simplex Equiangular Tight Frame, classifier weights align with the class means, and the network's predictions converge to a nearest-class-center rule. Papyan et al. (2020) established this phenomenon empirically using SGD. Recent theoretical work by Zhao et al. (ICLR 2026) proved that the choice of optimizer, specifically the coupling of weight decay, fundamentally determines whether NC emerges. This project provides a complementary geometric perspective: we reproduce the SGD baseline and extend the analysis to Adam, AdamW, and SGDW, using SVD spectrum analysis, norm tracking, and PCA projections to visualize how the representation geometry diverges across optimizers, and how much of that divergence traces back to decoupled weight decay versus adaptive gradients themselves.

2. Key Results

• SGD (Coupled) achieves near-complete Neural Collapse: within-class feature variance (NC1) drops to $0.013$, nearest-class-center agreement reaches $0.000$, and test accuracy peaks at $\mathbf{95.39%}$.
• Decoupled Weight Decay (SGDW & AdamW) degrades within-class collapse regardless of whether gradients are adaptive, and inflates classifier weight norms similarly (SGDW $3.47$, AdamW $4.26$, vs SGD's $1.80$). Moving from SGD to SGDW causes a $\sim 10\times$ bump in within-class variance ($0.013 \to 0.128$), placing it right alongside AdamW ($0.133$).
• Adaptive Gradient Scaling (Adam & AdamW) drives ETF Symmetry (NC2) by equalizing learning across feature dimensions. AdamW achieves the best NC2 ($0.351$, vs SGD's $0.615$, SGDW's $0.624$).
• Self-duality (NC3) shows a two-step degradation: decoupling alone takes NC3 from $0.497$ (SGD) to $1.112$ (SGDW); adding adaptive gradient scaling on top pushes it further to $1.496$ (AdamW). Unlike NC1, where decoupling alone accounts for nearly the entire effect, NC3 damage comes from decoupling and adaptivity independently.

3. Setup

3.1 Architecture

ResNet-18 adapted for CIFAR-10, following Papyan et al. (2020):

• First convolutional layer changed from $7 \times 7$ to $3 \times 3$ (stride $1$, padding $1$).
• Max-pooling layer replaced with identity.
• Output: $512$-dimensional feature space that goes through linear head with $10$ classes.

3.2 Config

All four optimizers are trained for $350$ epochs with identical data augmentation (RandomCrop with $4$-pixel padding, random horizontal flip) and batch size $128$. The learning rate follows a MultiStepLR schedule with $\gamma = 0.1$ at epochs $116$ and $233$.

Setting	SGD	Adam	AdamW	SGDW
Learning rate	$0.1$	$1 \times 10^{-3}$	$1 \times 10^{-3}$	$0.1$
Weight decay	$5 \times 10^{-4}$	$5 \times 10^{-4}$	$5 \times 10^{-2}$	$5 \times 10^{-4}$
Momentum	$0.9$	-	-	$0.9$
Betas	-	$(0.9, 0.999)$	$(0.9, 0.999)$	-

The step-decay schedule is chosen deliberately: the discrete LR drops create sharp phase transitions that make the onset of Neural Collapse easier to isolate and study, compared to cosine annealing where the transition is gradual.

4. Neural Collapse Metrics

All metrics are computed on training-set features extracted from checkpoints saved every $10$ epochs. Definitions follow Papyan et al. (2020):

• NC1 (Within-Class Variability Collapse): $\frac{1}{C} \text{tr}(\Sigma_W \Sigma_B^{\dagger})$, where $\Sigma_W$ and $\Sigma_B$ are the within-class and between-class covariance matrices. Approaches $0$ as features collapse to their class means.

• NC2 (Simplex ETF Convergence): $| \hat{M} \hat{M}^T - \frac{C}{C-1}(I_C - \frac{1}{C} \mathbf{1}\mathbf{1}^T) |_F$, where $\hat{M}$ contains the $\ell_2$-normalized centered class means. Approaches $0$ as the class means form a regular simplex.

• NC3 (Self-Duality): $| \hat{W} - \hat{M} |_F$, where $\hat{W}$ contains the $\ell_2$-normalized centered classifier weight vectors. Approaches $0$ as the classifier mirrors the feature geometry.

• NC4 (NCC Agreement): The fraction of training samples where the model's argmax prediction disagrees with a nearest-class-center classifier. Approaches $0$ as the two classifiers become equivalent.

Figure 2: Neural Collapse Metrics across 350 epochs. Vertical dashed lines mark LR drops at epochs 116 and 233. SGD converges to near-zero on all four metrics. AdamW achieves the best NC2 but diverges on NC3.

5. Experiments

5.1 Classification Performance

Optimizer	Weight Decay	Train Acc	Test Acc	Train Loss	Test Loss
SGD	Coupled	$100.00%$	$\mathbf{95.39%}$	$0.0018$	$0.1808$
Adam	Coupled	$99.93%$	$93.27%$	$0.0041$	$0.3310$
SGDW	Decoupled	$100.00%$	$94.53%$	$0.0004$	$0.2601$
AdamW	Decoupled	$100.00%$	$94.51%$	$0.0000$	$0.3830$

Figure 3: Training and test accuracy across 350 epochs.

SGD generalizes best. Adam exhibits the largest generalization gap ($6.66%$). AdamW and SGDW both partially recover through decoupled weight decay but still lag SGD by roughly $1%$. It is interesting to note that SGD generalizes the best ($95.39%$) with the tightest clusters, whereas AdamW has the best global symmetry but worse generalization ($94.51%$), and SGDW lands at nearly the same accuracy ($94.53%$) with neither. Here, NC1 tracks test performance far more closely than NC2 does.

5.2 Terminal NC Geometry

Metric	SGD (Coupled)	Adam (Coupled)	SGDW (Decoupled)	AdamW (Decoupled)	Ideal
NC1 (Variability Collapse)	$\mathbf{0.0133}$	$0.0471$	$0.1280$	$0.1335$	$\to 0$
NC2 (ETF Symmetry)	$0.6151$	$0.5875$	$0.6244$	$\mathbf{0.3513}$	$\to 0$
NC3 (Self-Duality)	$\mathbf{0.4965}$	$0.5058$	$1.1119$	$1.4957$	$\to 0$
NC4 (NCC Agreement)	$\mathbf{0.0000}$	$\mathbf{0.0000}$	$\mathbf{0.0000}$	$\mathbf{0.0000}$	$\to 0$

5.3 Cosine Similarity Structure

Figure 4: Cosine similarity $\hat{M}\hat{M}^T$ at epoch 350. All four optimizers develop the characteristic off-diagonal pattern of the Simplex ETF, with AdamW showing the most uniform off-diagonal values.

5.4 Singular Value Spectrum

Figure 5: Normalized singular value spectrum ($\sigma_i / \sigma_1$) of the centered class mean matrix at epoch 350. An ideal Simplex ETF has 9 equal singular values (dashed line). AdamW is the flattest (min ratio $0.83$), confirming its superior ETF structure; SGD is the steepest (min ratio $0.74$).

A perfect $C$-class Simplex ETF has rank $C - 1 = 9$ with all singular values equal. The SVD spectrum directly visualizes how close each optimizer's class-mean geometry is to this ideal structure. AdamW's variance-normalized gradient updates mechanically equalize the per-dimension learning dynamics, producing the most symmetric arrangement of class centroids, even though the individual clusters around those centroids are far less tight.

5.5 Norm Analysis

Figure 6: Average classifier weight norms $|w_c|$ and feature mean norms $|\mu_c|$ across training. AdamW inflates both to substantially larger magnitudes than SGD ($4.26$ vs $1.80$ for weights; $8.81$ vs $5.57$ for features), driven by its adaptive gradient scaling overpowering the weight decay. SGDW sits partway between ($3.47$ / $6.16$): closer to AdamW on weight norm, closer to SGD on feature norm.

Self-duality (NC3) requires $w_c \propto \mu_c$, the weight vectors and feature means must be proportionally aligned. AdamW's per-parameter adaptive learning rates push weights into directions that are gradient-efficient but geometrically misaligned with the feature means. The result is a $3\times$ worse ($1.50$ vs $0.50$) despite having the best global symmetry (NC2).

6. Discussion

• Papyan et al. showed NC metrics collapsing together under SGD. However, our results demonstrate that weight decoupling degrades NC1 almost entirely on its own (SGDW $\approx$ AdamW, with or without adaptive gradients), whereas NC2's improvement is specifically attributable to adaptive gradient scaling, not decoupling. NC3 is the exception, where both mechanisms contribute independently rather than one dominating.
• Connecting to Zhao et al. (2026): Their SignGD-based NC0 diagnostic explains why decoupled decay blocks collapse, matching our finding that NC1 damage is decoupling-only. Our SGD-to-SGDW gap on NC1 ($\sim 10\times$) is larger than the near-parity they report at the same nominal weight decay, plausibly due to our longer $350$-epoch schedule.
• Scope: These results are single-seed, single-architecture (ResNet-18), single-dataset (CIFAR-10) due to compute constraints. Extending to more architectures, datasets, and seeds would strengthen the conclusions.

7. Usage

# clone repo & install dependencies
git clone https://github.com/ahmadrazacdx/nc-optimizers.git
cd nc-optimizers
pip install -r requirements.txt
# prepare data (expects CIFAR-10 already downloaded to ./data; download=False by default)
python data.py
# run experiments (stored in `checkpoints/`)
python train.py --optimizer sgd
python train.py --optimizer adam
python train.py --optimizer adamw
python train.py --optimizer sgdw
# measure all nc metrics (stored in `artifacts/`)
python measure.py
# get all figures (stored in `figures/`)
python get_figures.py

train.py trains ResNet-18 with the specified optimizer and saves checkpoints every 10 epochs. The flag config is:

Flag	Requirement	Description	Default
--optimizer	required	Optimizer to use. Choices: sgd, adam, adamw, sgdw.	sgd
--epochs	optional	Total training epochs.	350
--batch_size	optional	Batch size.	128
--lr	optional	Learning rate (used default otherwise).	-
--weight_decay	optional	Weight decay (auto-set otherwise).	-

8. References

• Papyan, V., Han, X., & Donoho, D. L. (2020). Prevalence of Neural Collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40), 24652–24663. PNAS
• Zhao, J., Cheng, T. S., Masarczyk, W., & Lucchi, A. (2026). Optimizer choice matters for the emergence of Neural Collapse. ICLR 2026. arXiv:2602.16642