Research map / Optimization / Implicit bias and generalization
Neural Collapse under Gradient Flow on Shallow ReLU Networks for Orthogonally Separable Data
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that gradient flow from small balanced initialization on a two-layer ReLU network trained on orthogonally separable data provably converges to a neural-collapse structure (directional collapse, orthogonal class means, projected self-duality), linking NC to the implicit bias of GF.
Problem
Neural collapse is usually analyzed for unconstrained last-layer features, which ignores input data and nonlinear activations and does not explain how training dynamics produce NC.
Approach
Analyzes GF on a bias-free two-layer ReLU network under small balanced initialization for orthogonally separable data. Early alignment separates classes, then the asymptotic max-margin bias of GF yields intra-class directional collapse and projected self-duality.
Key results
- On MNIST digits {0,1,2}, PCA of last-layer features keeps about 0.2% relative approximation error with 3 components, versus about 61% for raw data, consistent with directional collapse.
- In modified ResNet18 experiments over 5 runs, RMSNorm before the ReLU classifier markedly improves intra-class directional collapse (lower NC1) compared with no normalization or LayerNorm.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add RMSNorm before the ReLU classifier | Modified ResNet18 · MNIST and CIFAR10 | NC1 | Helped |
| Use top 3 PCA components of trained last-layer features instead of raw digits | Two-layer ReLU network · MNIST digits 0,1,2 | relative approximation error: ~61% -> ~0.2% | Helped |
Surprising findings
- Class-mean features are orthogonal and non-negative rather than forming a simplex ETF, because ReLU features are non-negative.
Limitations
- Requires strong orthogonal separability and small balanced initialization assumptions.
- Multi-class proof needs a stricter separability condition (Xmax^2/(Xmin^2 mu_d mu_s^2) < 2K-3) and a semi-local initialization assumption.
- Experiments are small-scale (MNIST, CIFAR10 with ResNets); theory is limited to shallow networks.
Experimental setup
Datasets: MNIST CIFAR10
Benchmarks: MNIST CIFAR10
Metrics: NC1 NC2 NC3 test accuracy
Baselines: unconstrained feature model analyses of neural collapse
Builds on: unconstrained feature model max-margin bias of gradient flow in linear networks alignment phase of small-initialization training
Closest papers by meaning
- Wide Neural Networks Trained with Weight Decay Provably Exhibit Neural Collapse (ICLR 2025 · similarity 0.92)
- Neural Collapse Beyond the Unconstrained Features Model: Landscape, Dynamics, and Generalization in the Mean-Field Regime (ICML 2025 · similarity 0.90)
- Neural Collapse is Globally Optimal in Deep Regularized ResNets and Transformers (NeurIPS 2025 · similarity 0.90)
- Generalized Neural Collapse for a Large Number of Classes (ICML 2024 · similarity 0.89)
- Neural Collapse in Cumulative Link Models for Ordinal Regression: An Analysis with Unconstrained Feature Model (NeurIPS 2025 · similarity 0.89)
- Neural collapse vs. low-rank bias: Is deep neural collapse really optimal? (NeurIPS 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.