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

AI summary

Authors: Hancheng Min (Shanghai Jiao Tong University), Zhihui Zhu (Ohio State University), René Vidal (University of Pennsylvania) arXiv: 2510.21078v1 [cs.LG], 24 Oct 2025

Overview

Research area: Deep learning theory — the theoretical analysis of Neural Collapse (NC) and the implicit bias of gradient-based training, with proofs for two-layer ReLU networks.

Technical level: Advanced. The paper is a theorem-and-proof paper requiring comfort with gradient flow, Clarke sub-differentials, max-margin arguments, and matrix/optimization notation.

One-sentence scope: The paper proves that gradient flow on a two-layer ReLU network trained on orthogonally separable data provably converges to a Neural Collapse configuration, and explains how the implicit bias of gradient flow drives that collapse.

What This Paper Is About

Neural Collapse is the empirical phenomenon in which, at the end of training, last-layer features from the same class collapse together, class means become maximally separated, and classifier weights align with class means. Most theory on why this happens studies the Unconstrained Feature Model (UFM), which treats the last-layer features as free, unconstrained optimization variables and analyzes the resulting matrix-factorization-like landscape. This paper instead asks what happens when the features are actually produced by a two-layer ReLU network fed with real input data, and how the training dynamics — not just the loss landscape — produce NC. The goal is to prove NC for gradient flow with small initialization on orthogonally separable data and to expose the implicit bias mechanism that creates it.

Key Contributions

1. A provable NC result on a real network rather than a free-feature model. Theorem 1 shows that gradient flow with small, balanced initialization on a width-h two-layer ReLU network converges (its normalized weight direction has a limit) and that the limit exhibits NC, for both binary and multi-class orthogonally separable data. This relaxes the unconstrained-feature assumption, so the effect of the input data structure and the ReLU nonlinearity on the NC characterization becomes visible.

2. A concrete account of how implicit bias produces NC in binary classification. In the proof of Theorem 1 for the binary case, the paper traces NC to two phases of gradient flow: an early alignment phase that makes inter-class features mutually orthogonal, and a late asymptotic max-margin phase that drives intra-class directional collapse and self-duality.

3. An extension of implicit-bias analysis from binary to multi-class settings. The multi-class proof sketch extends prior implicit-bias results to the multi-class case, addressing new difficulties caused by the multi-dimensional network output and the cross-entropy loss.

4. A sharper NC vocabulary for shallow ReLU networks. The paper replaces the classical NC properties with three modified ones: intra-class directional collapse, orthogonal class means, and projected self-duality, each tied to the ReLU activation or the input data.

Main Findings

• Intra-class directional collapse replaces collapse to a singleton. For the limiting weight direction, the last-layer features of class-k data satisfy φ(x_i) = ⟨s_k u_k, x_i⟩ · φ̄_k for all i in class k. So instead of collapsing to a single point, features of the same class collapse into a one-dimensional subspace spanned by φ̄_k, with lengths that can vary. The residual within-class variability is governed by the projections {⟨s_k u_k, x_i⟩}, a large reduction relative to the variability of the raw inputs {x_i}. If features are normalized to unit norm (for example by RMSnorm), they collapse exactly to their class means.

• Class means are orthogonal, not maximally separated. The class-mean directions satisfy φ̄_k ≥ 0 and ⟨φ̄_k, φ̄_{k'}⟩ = 0 for all k ≠ k'. Normalized, they form a non-negative orthogonal frame. The paper explains the orthogonal structure arises because ReLU features are non-negative, and notes any orthogonal frame can be turned into a simplex ETF by removing its global mean — consistent with UFM results that use ReLU as the activation.

• Projected self-duality replaces plain self-duality. In the binary case the limiting classifier satisfies V̄ = s_+ φ̄_+ᵀ − s_- φ̄_-ᵀ. In the multi-class case V̄ = sqrt(K/(K−1)) (I − (1/K)𝟙𝟙ᵀ)[s_1 φ̄_1, …, s_K φ̄_K]ᵀ, which gives V̄V̄ᵀ = (K/(K−1))(I − (1/K)𝟙𝟙ᵀ) Φ̄Φ̄ᵀ (I − (1/K)𝟙𝟙ᵀ). Since Φ̄Φ̄ᵀ is diagonal with positive entries, V̄ forms a scaled simplex ETF and attains maximum margin; when the diagonal scales s_k are all equal, V̄ is an exact simplex ETF and each classifier converges to the corresponding projected class mean, up to scaling.

• Scale factors come from per-class max margins. Each s_k is written as s_k = sqrt( γ_k^{-1} / (2 Σ_{k∈𝒦} γ_k^{-1}) ), where γ_k = max_{u ∈ S^{D−1}} min_{i ∈ ℐ_k} ⟨x_i, u⟩ is the maximum margin achievable using only class-k data and u_k is the corresponding max-margin direction. The scaling factors s_+, s_- are tied, via Lyu and Li's result, to a max-margin problem defined on the entire dataset.

• The binary proof decomposes NC into two phases. In the early alignment phase, neuron norms start at scale O(ε) and stay at scale O(ε^{1/2}) for a time of length Θ(log(1/ε)); during this window the directional dynamics are approximately decoupled and driven by directional alignment between input neuron weights and data. Stationary points with all activations zero form a "dead region" (zero gradient). Extremal vectors (those with some ξ_ij ≠ 0) are, for binary orthogonally separable data, the class mean directions x̄_+ and x̄_-. For neurons whose output weight is positive, x̄_+ is an attractor and x̄_- is a repeller, and the opposite holds for neurons whose output weight is negative.

• Inter-class separation is a provable intermediate state. Under a non-degenerate initialization (Assumption 3), there exists a time T* and non-empty subsets of neurons such that after T*, W_+(t)ᵀx_i > 0 for positive-class data and ≤ 0 for negative-class data (and symmetrically for W_-(t)); all remaining neurons are in the dead region. This immediately gives inter-class separation, ⟨φ(x_i), φ(x_{i'})⟩ = 0 for i in class + and i' in class −, for all t ≥ T*.

• After separation, gradient flow splits into independent linear-network problems. The loss decomposes into a positive-class sum and a negative-class sum, so the flow on {W_+, V_+} decouples entirely from that on {W_-, V_-}. Each piece amounts to training a two-layer linear network on same-label, positively correlated data, where the asymptotic max-margin bias of gradient flow applies.

• The limiting directions have a rank-one, aligned form. The limit satisfies W̄_+ = s_+ u_+ g_+ᵀ, V̄_+ = s_+ g_+ᵀ, W̄_- = s_- u_- g_-ᵀ, V̄_- = −s_- g_-ᵀ for some unit vectors g_+, g_-. The paper notes that the alignment condition (V̄_+W̄_+ᵀ ∝ u_+) and the balancedness condition (V̄_+ᵀV̄_+ = W̄_+ᵀW̄_+) can only hold simultaneously when W̄_+ᵀ has rank 1, which forces NC.

• The result applies under weaker conditions than prior convergence work. Prior convergence-to-NC results use the mean squared loss and require extra conditions such as initialization near a global optimum, weight decay regularization, or large width. This work uses the exponential loss or logistic loss (binary) and cross-entropy loss (multi-class), without explicit regularization or width over-parametrization.

• Setup and assumptions. The data must be orthogonally separable: same-label pairs of normalized inputs have inner product ≥ μ_s with 0 < μ_s ≤ 1, and different-label pairs have inner product ≤ −μ_d with 0 < μ_d ≤ 1/sqrt(K−1) (the paper notes no dataset can satisfy the assumption with μ_d > 1/sqrt(K−1)). Initialization must be ε-small and balanced: w_j(0) = ε w_j0, v_j(0) = ε v_j0, ‖w_j0‖ = ‖v_j0‖. Balancedness is preserved for all t along gradient flow. Networks are bias-free, with width h ≥ K.

• The data assumption is acknowledged as strong, with three justifications given. Shallow networks have limited expressive power for collapsing features, so more structured data is needed; precise characterization of asymptotic weight convergence is limited to structurally simple data; and the authors state that simple real data satisfies orthogonal separability approximately, giving NC characters that match the theorem.

• Experimental details are not reported in the available content. Section 5 is referenced as showing that simple real data approximately satisfies orthogonal separability and that the authors' experiments do not require balanced initialization, but no dataset names, dataset sizes, benchmark numbers, or model configurations appear in the provided text.

Methodology in Plain English

The authors analyze gradient flow (the continuous-time limit of gradient descent) rather than running experiments as the primary tool. They consider a two-layer ReLU network that is bias-free, with the last-layer features being the ReLU activations of the hidden layer and the last-layer classifier being the output weight matrix.

They put structure on the data — orthogonal separability, meaning same-class inputs point roughly the same way and different-class inputs point roughly opposite ways — and on initialization: start every neuron with a very small weight of scale ε, and keep the input and output weight norms of each neuron equal. Because the network is positively homogeneous in its parameters, the loss drives the weights to grow without bound, so the object of study is the direction of the weight vector, and the proof shows that this direction has a limit.

The proof then splits the training trajectory into two regimes. In the early regime, the neuron norms are still tiny (order ε^{1/2}) for a long stretch of time (order log(1/ε)), so only the directions of the input neurons move, and those directions rotate toward their class-average data direction while avoiding a "dead region" where all activations are zero. This rotation alone makes the features of different classes orthogonal to each other. Once that happens, the loss separates into one term per class, and the network effectively becomes two independent linear networks trained on positively correlated same-class data. For those, known max-margin results say the limiting weight direction must align with the max-margin direction of a linear classifier on the whole dataset; combining this with the balance condition forces the weight matrices to be rank one and identifies the limit explicitly, which is exactly the NC configuration. The multi-class proof must additionally handle the vector-valued output and the cross-entropy loss.

Why This Matters

Impact on research. Prior NC theory mostly established that global minima of a free-feature surrogate loss have NC structure, which leaves open whether actual training dynamics land there. This paper closes part of that gap by showing that gradient flow itself — via its implicit bias and small-initialization behavior — reaches an NC configuration on a real network with real inputs. It also refines the NC vocabulary for shallow ReLU networks: directional collapse rather than point collapse, orthogonal rather than maximally separated means, and projected rather than exact self-duality. These refinements clarify which NC properties are artifacts of the unconstrained-feature idealization and which survive contact with data and nonlinearity.

Real-world applications. The paper does not enumerate application domains, but the findings bear on:

• Representation learning and feature reuse — understanding when last-layer features carry class-discriminative, low-dimensional structure informs how features from a trained classifier are reused downstream.
• Normalization layer design — the paper notes that unit-norm normalization (e.g., RMSnorm) turns directional collapse into exact collapse to class means, which speaks to choices of normalization in classifier heads.
• Separable or near-separable classification tasks — the authors state that simple real data approximately satisfies orthogonal separability, so the predicted NC characters are claimed to transfer approximately to such data.
• Interpreting trained classifiers and their margins — the result connects the trained classifier to class-wise max-margin directions, giving a geometric reading of what a trained head encodes.

Industry relevance. Practitioners who train deep classifiers on well-separated data often observe that last-layer features become very low-dimensional and that classifier weights mirror class means. This paper gives a theoretical reason rooted in the training algorithm rather than explicit regularization, which matters for settings where weight decay or margin-based losses are not used and for anyone relying on those features for transfer, retrieval, or model inspection.

Future Directions

• Relax orthogonal separability to realistic data. The authors themselves frame the data assumption as strong and motivate it by the limited expressive power of a single hidden layer and by the tractability of asymptotic weight analysis. Extending the characterization beyond approximately separable data is the natural next step.
• Handle deeper networks and non-ReLU activations. The orthogonal class means and non-negative frame arise specifically from ReLU; whether analogous characterizations hold for other activations, or whether the limit is still rank one and margin-aligned with more layers, is open.
• Remove the balanced-initialization requirement from the theory. Balanced initialization is maintained exactly along gradient flow and is needed for the analysis, while the authors note their experiments do not require it. A proof that does not need it would broaden the scope.
• Explain the gap between the loss function and the implicit bias in multi-class settings. The paper states that the multi-class analysis requires new technical machinery because of the vector-valued output and cross-entropy loss; sharpening that analysis — and testing how much the characterizations depend on the specific loss (exponential, logistic, cross-entropy) — remains an open direction.

Target Audience

This paper is written for machine learning theorists and mathematically inclined graduate students working on optimization dynamics, implicit bias, and the theory of deep learning. It will be most useful to readers already familiar with Neural Collapse, the Unconstrained Feature Model, gradient flow, and max-margin analyses of homogeneous networks. Practitioners without a theory background will find the informal takeaways about directional collapse, orthogonal class means, and projected self-duality useful, but the theorem and proof structure are aimed at a research audience.

Authors’ 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.