Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Neural Collapse under Gradient Flow on Shallow ReLU Networks for Orthogonally Separable Data

Hancheng Min, Zhihui Zhu, René Vidal

TL;DR

This work analyzes gradient flow on a shallow two-layer ReLU network trained for classification under orthogonally separable data. It proves that, with small initialization, GF converges to solutions exhibiting Neural Collapse, where last-layer features collapse directionally within classes and class means become orthogonal, with a projected self-duality between features and classifier weights. The results connect NC to the implicit bias of GF, showing how data structure and nonlinear activation shape NC differently from unconstrained feature models, and extend the analysis from binary to multi-class settings. Numerical experiments on MNIST and real networks corroborate the theoretical NC characterizations and demonstrate practical effects of normalization layers on NC. Overall, the paper bridges NC theory with implicit bias in gradient-based training and highlights how training dynamics and data structure jointly drive structured representations in shallow networks.

Abstract

Among many mysteries behind the success of deep networks lies the exceptional discriminative power of their learned representations as manifested by the intriguing Neural Collapse (NC) phenomenon, where simple feature structures emerge at the last layer of a trained neural network. Prior works on the theoretical understandings of NC have focused on analyzing the optimization landscape of matrix-factorization-like problems by considering the last-layer features as unconstrained free optimization variables and showing that their global minima exhibit NC. In this paper, we show that gradient flow on a two-layer ReLU network for classifying orthogonally separable data provably exhibits NC, thereby advancing prior results in two ways: First, we relax the assumption of unconstrained features, showing the effect of data structure and nonlinear activations on NC characterizations. Second, we reveal the role of the implicit bias of the training dynamics in facilitating the emergence of NC.

Neural Collapse under Gradient Flow on Shallow ReLU Networks for Orthogonally Separable Data

TL;DR

This work analyzes gradient flow on a shallow two-layer ReLU network trained for classification under orthogonally separable data. It proves that, with small initialization, GF converges to solutions exhibiting Neural Collapse, where last-layer features collapse directionally within classes and class means become orthogonal, with a projected self-duality between features and classifier weights. The results connect NC to the implicit bias of GF, showing how data structure and nonlinear activation shape NC differently from unconstrained feature models, and extend the analysis from binary to multi-class settings. Numerical experiments on MNIST and real networks corroborate the theoretical NC characterizations and demonstrate practical effects of normalization layers on NC. Overall, the paper bridges NC theory with implicit bias in gradient-based training and highlights how training dynamics and data structure jointly drive structured representations in shallow networks.

Abstract

Among many mysteries behind the success of deep networks lies the exceptional discriminative power of their learned representations as manifested by the intriguing Neural Collapse (NC) phenomenon, where simple feature structures emerge at the last layer of a trained neural network. Prior works on the theoretical understandings of NC have focused on analyzing the optimization landscape of matrix-factorization-like problems by considering the last-layer features as unconstrained free optimization variables and showing that their global minima exhibit NC. In this paper, we show that gradient flow on a two-layer ReLU network for classifying orthogonally separable data provably exhibits NC, thereby advancing prior results in two ways: First, we relax the assumption of unconstrained features, showing the effect of data structure and nonlinear activations on NC characterizations. Second, we reveal the role of the implicit bias of the training dynamics in facilitating the emergence of NC.
Paper Structure (26 sections, 17 theorems, 85 equations, 5 figures)

This paper contains 26 sections, 17 theorems, 85 equations, 5 figures.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Main result: Neural Collapse under GF on Two-layer ReLU Networks
  4. Detailed Discussions: Connecting Neural Collapse with Implicit Bias of GF
  5. Proof of Neural Collapse in Binary Classification
  6. Proof Sketch of Neural Collapse in Multi-class Classification
  7. Numerical Experiments
  8. Conclusion
  9. Additional Remarks and Related Works
  10. Extension to ReLU nets with biases
  11. Additional related works
  12. Matrix factorization in deep learning theory
  13. Gradient descent on shallow neural networks
  14. Experimental Details
  15. Experimental details on classifying MNIST digits
  16. ...and 11 more sections

Key Result

Theorem 1

Given orthogonally separable data (Assumption assump_data), $\epsilon$-small and balanced initialization (Assumption assump_init) for a sufficiently small $\epsilon$, and some additional case-dependent assumptions on the data and initialization shape, the limit $\bar{\bm{\theta}}:=\lim_{t\rightarrow

Figures (5)

  • Figure 1: Visualization of the NC phenomenon.
  • Figure 2: Convergence and implicit bias of GF in two-layer ReLU networks: (a) An example of orthogonally separable data (Assumption \ref{['assump_data']}), gray region indicates the "dead region" for neuron weights: $\{\bm{w}: \left\langle \bm{w}, x_i\right\rangle\leq 0, \forall i\in[n]\}$; (b) Weight initialization following Assumption \ref{['assump_init']}, input neuron weights have small norm and random directions; (c) During alignment phase, inter-class separation is achieved through directional alignment between input neuron weights and data points, as described in \ref{['eq_class_sep']}; (d) Asymptotically, neuron weights diverge to infinity while their directions align with the class-wise max-margin directions, as described in \ref{['eq_asym_dir_conv']}.
  • Figure 3: Validating Theorem \ref{['thm_nc_conv']} in classifying MNIST digits $\{0,1,2\}$. (a) Normalized correlation matrix of subsampled 500 MNIST digits; (b) (For the trained network) visualization of output neuron weights (as crosses, gray dashed line represents $\tilde{\bm{e}}_k$ directions for references), the average input neuron weights (as grayscale image, surrounded by colored box), and the average of the digits (for comparison, next to the neuron weights); (c) PCA of raw digits data $\bm{X}$, keeping the top 3 principal components (1000 points visualized) results in a $\sim\!61\%$ relative approximation error for $\bm{X}$; (d) (For the trained network) PCA of last-layer feature $\bm{\phi}_{\bm{\theta}}(\bm{X})$ and classifiers (rows of $\bm{V}$), keeping the top 3 components (1000 points visualized) results in a $\sim\!0.2\%$ relative approximation error for $\bm{\phi}_{\bm{\theta}}(\bm{X})$.
  • Figure 4: Measuring NC in trained modified ResNet18 on MNIST (top) and CIFAR10 (bottom)
  • Figure 5: Measuring NC in trained modified ResNet50 on MNIST (top) and CIFAR10 (bottom)

Theorems & Definitions (36)

  • Remark 1
  • Remark 2
  • Theorem 1: NC of GF on Two-layer ReLU Networks
  • Claim : Inter-class separation via alignment, based the analyses from phuong2021inductivemin2023early
  • Claim : Directional collapse and self-duality via max-margin bias, based on the analyses from phuong2021inductiveji2019gradient
  • Proposition 1: Inter-class separation in multi-class problems
  • Proposition 2: Variability collapse and self-duality in multi-class problems
  • Lemma 1
  • proof
  • Lemma 2
  • ...and 26 more