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

We thank the reviewer for the constructive feedback. We respond to your questions/concerns below:

Extension to deep nets We have acknowledged the limitations of our results in the paper through multiple remarks. However, as the reviewer pointed out, the discussion on extensions to deep nets was not sufficiently detailed.. We will revise our conclusion section to elaborate more. Specifically, it is our belief that the key technical ideas (alignment dynamics in transient phase, and max-margin bias in asymptotic phase), whose connection to NC was not well explored in prior work, will be useful to study the convergence to NC in deep nets as both ideas appear in deep nets with positively homogenous activations (Kumar and Haupt, 2025; Ji and Telgarsky, 2020; Lyu and Li, 2019). Therefore, following the same high-level proof and utilizing these existing results on alignment and max-margin for deep nets might be the way to extend the current work to a more practical setting.

Assumption in Proposition 1 We will add a remark to explain the limitation of this assumption on the data correlations. Specifically, this assumption is required to show that the subsets of the parameter space defined in Assumption 4 are invariant under GF. We believe such an assumption is not needed in practice, but our limited understanding of the ROAs and their invariant subsets, as we have discussed from line 279 to line 289, has led to this additional technical condition to ensure directional convergence. Future research on better characterizations of the ROAs and their invariant sets will naturally relax or even potentially remove this requirement.

Features in ResNet We clarify that the main purpose of our ResNet experiments is mostly to examine the role of the normalization layer in determining the NC level in deep networks. It is unlikely the features in ResNet would satisfy our separability assumption (in fact, our MNIST experiment already shows that it is only satisfied approximately for MNIST). Nonetheless, the fact that our results align with some of the empirical observations (RMSNorm improves NC) shows that results with the separability assumption are not uninformative for practice, and relaxing this assumption should lead to more practically relevant results.

Clarifying Footnote 1 Footnote 1 formally says that there exists no dataset satisfying equation 1 with some $0<\mu_s\leq 1$ and $\mu_d>\frac{1}{\sqrt{K-1}}$, we will add its proof to the appendix. Note that the simple dataset the reviewer mentioned does not invalidate this statement, as this dataset satisfies equation 1 with $\mu_d=1\leq \frac{1}{\sqrt{K-1}}$ (note that $K=2$ here)

Claims in Section 4.1 The Claims in Section 4.1 are directly implied by existing results, which we have referenced at the beginning of the claims. We understand that this makes the paper not self-contained (which is also pointed out by Reviewer ti2o). We will, in Section 4.1, add high-level explanations to the proofs with illustrative figures, and in the appendix, explain how the claims are derived from previous results.

Definitions of u vectors $u_+$ and $u_-$ vectors are max-margin vectors for positive-class and negative-class data, respectively, defined in Theorem 1, right after equation (4). In Theorem 1, we have stated two cases, and the "case one" is when $\mathcal{K}=\lbrace +,-\rbrace$, and $u_k, \mathcal{K}=\lbrace +,-\rbrace$ defined for this case gives rise to $u_+, u_-$

Lastly, we will fix the typos the reviewer mentioned (in particular, V is in $\mathbb{R}^{d_y\times h}$, as the reviewer pointed out), and add the definitions for identity matrix and acronym ETF (Equiangular Tight Frame).

We hope our rebuttal addresses your questions. If you have additional questions/concerns, please feel free to post them during the discussion phase.

References:

Kumar and Haupt, Early Directional Convergence in Deep Homogeneous Neural Networks for Small Initializations, TMLR, 2025

Ji and Telgarsky, Directional convergence and alignment in deep learning, NeurIPS, 2020

Lyu and Li, Gradient descent maximizes the margin of homogeneous neural networks, ICLR, 2019

We thank the reviewer for the constructive feedback. We respond to your questions/concerns below:

Section 4 is too technical We thank the reviewer for the suggestions. We will, in Section 4.1, add high-level explanations to the proofs with illustrative figures, and also improve the appendix to better explain the technical aspects of the proofs.

Requirement on weight balancedness Indeed, the analysis can be extended to the case when $|v_{j0}|>\lVert w_{j0}\rVert$. We will check how much editing is required in the proof to determine whether to modify the main statement or to add a remark.

Requirement on homogenous activations We believe it is possible to extend the current analysis to networks with approximately homogenous activations. Specifically, we studied the dynamical behaviors of GF in different phases: alignment in the transient phase and max-margin bias in the asymptotic phase. For the alignment, it is mainly derived from the linearization of the network around the origin; thus, similar analysis can be done with non-homogeneous activations, albeit losing weight balancedness, potentially making the alignment dynamics harder to work with. For the max-margin, it requires positively homogeneous activations. Nonetheless, we note that recently there has been work on extension to approximately homogeneous activations (Cai et al., 2025). We are happy to add the discussion here to the appendix if needed.

Adding bias We will add a discussion on the network bias term in the appendix. Specifically, since $\sigma(\langle w, x\rangle +b)=\sigma(\langle [w;b],[x;1]\rangle)$, adding a bias term effectively adds one homogenous coordinate to the entire dataset. Notably, homogenous coordinate increases data correlation: $\langle [x_i;1],[x_j;1]\rangle=\langle x_i,x_j\rangle +1$, thus it mainly affects the negative correlation between data points with different labels. In the case when $\min_i \|x_i\|\gg 1$, the orthogonal separability (among $[x_i;1]$) still holds with the bias term.

Time scale of $T^\*$ We first note that $T^\*$ is the time to achieve inter-class separation, but not necessarily NC. For binary problems, its scale is characterized in (Min et al., 2024), which mainly depends on the data separability $\mu_s,\mu_d$ and dataset size. For multi-class problems, our current Proposition 1 did not provide a scale for $T^*$, but we believe it should have a similar scale as the binary case. We will work on adding the scale to the revision.

We hope our rebuttal addresses your questions. If you have additional questions/concerns, please feel free to post them during the discussion phase.

References:

Y. Cai, et al., Implicit Bias of Gradient Descent for Non-Homogeneous Deep Networks, ICML 2025

H. Min, et al., Early Neuron Alignment in Two-layer ReLU Networks with Small Initialization, ICLR 2024

We thank the reviewer for the constructive feedback. We respond to your questions/concerns below:

Why neuron dynamics decouple During the alignment phase (any time t prior to $T=O(\log(1/\epsilon))$), the neuron dynamics decouple as (16)(17) even with the presence of softmax: The neuron coupling is always through $\hat{y}_i$, and Lemma 1 suggests that during alignment phase, $\hat{y}_i\simeq \mathbf{1}/K$, making the dependence on other neurons negligible up to some $\epsilon$ error. Moreover, this approximation error is well controlled throughout the alignment phase. Specifically, at initialization, the norm of the weights satisfies $\lVert w_j(0)\rVert^2=O(\epsilon^2)$, then one can show that the norm only grows linearly. Then a Gronwall argument shows that for any t before $T=O(\log(1/\epsilon))$, the norm can only grow up to $\lVert w_j(t)\rVert^2=O(\epsilon)$, which gives rise to Lemma 2. The Gronwall argument is used in line 874 in its proof. Finally, $\lVert w_j(t)\rVert^2$ is used to control the error as shown in Lemma 3, and all the subsequent results (Lemmas 4 and 5) on the dynamical behaviors of neurons have taken this error into account.

What happens in the transition from decouple dynamics to max-margin behavior? We note that our results adopt two different analyses for two phases. For the alignment phase prior to $T^\*$, the approximation of decoupled neuron dynamics remains valid; for the dynamic phase after $T^\*$, there is no neuron decoupling, nor does the convergence analysis rely on an approximate decoupled dynamics. The asymptotic stability of GF, i.e., the limit point $\theta/\lVert\theta\rVert$ exists, is shown through prior work (Lyu and Li, 2019) with the residual errors fully taken into account, and our results in Proposition show its connection to NC. We will clarify the different analyses for the two phases in the revision.

Quantitative analysis of the attraction basins As we have discussed in Section 4.2, ROAs of multiple stationary points of the decoupled neuron dynamics pose unique challenges (to our best knowledge, our paper discussed them first) in understanding the directional convergence of neurons in the alignment phase in multi-class problems, and our current analysis can only provide invariant subsets of these ROAs, as defined by Assumption 4. Admittedly, the invariant subsets are relatively small when compared with the actual ROAs, because random initialization does not satisfy Assumption 4, yet empirically, the MNIST experiments show that the alignment with the attractors from random initialization. The quantitative analysis of these ROAs, as well as the relaxation of Assumption 4, has broader importance in understanding the implicit bias of GF on shallow networks in multi-class problems and is left to future research.

Accuracy gap between SOTA We thank the reviewer for the observation. Similar gap between SOTA exists for ResNet18 in the empirical results in our main paper, for which we believe the main source of the gap is the lack of data augumentation: we rerun the ResNet-CIFAR10 experiments with data agumentation and report the NC and Acc. of the trained network after 50 epochs in the following table, from which we see the Acc. improvement from agumentation, but less prominent NC (it is harder to collapse data under augmentation).

For ResNet50-CIFAR10, there was an additional issue of training instability: we used a modified ResNet50, for which the SGD with an initial step size of 0.1 sometimes diverges. This led us to use an initial step size of 0.01. Given that empirically a larger step size leads to better generalization, the small step size in our experiments with ResNet50 further enlarges the gap between SOTA. We will refine the experiments to better align with practice.

When does Assumption 4 hold Assumption 4 does not hold under random initialization. We will add this to the remark on the limitations of Assumption 4. We believe the analysis can be extended to all cases when the weight initialization does not exactly land on the boundary between two regions of attraction, as mentioned in Section 4.2, but showing the inter-class separation time $T^*$ exists requires additional technical efforts, and thus is left to future research.

How to extend to deep nets We believe a similar two-phase analysis is possible for three-layer networks, given certain initialization and data assumptions. Existing work (Kumar and Haupt, 2025) has shown that weight alignment also happens at the early phase of GF when training deep networks with positive homogeneous activations, and likewise, max-margin bias also appears in the asymptotic phase (Ji and Telgarsky, 2020; Lyu and Li, 2019). We believe our analysis suggests a recipe for connecting these results to NC in deep networks in future works.

Does Lemma 3 account for error As we have discussed in previous point, the error accumulates but well controlled until $T=O(\log(1/\epsilon))$, after which we no longer need to control the error as the analysis in the second phase does not require decoupled neuron dynamics.

We hope our rebuttal addresses your questions. If you have additional questions/concerns, please feel free to post them during the discussion phase.

References:

Kumar and Haupt, Early Directional Convergence in Deep Homogeneous Neural Networks for Small Initializations, TMLR, 2025

Ji and Telgarsky, Directional convergence and alignment in deep learning, NeurIPS, 2020

Lyu and Li, Gradient descent maximizes the margin of homogeneous neural networks, ICLR, 2019

We thank the reviewer for the constructive feedback. We respond to your questions/concerns below:

Comparison with DUFM We agree with the reviewer that DUFM studies multi-layer networks, but we do not think the DUFM captures practical NC more faithfully due to its two main limitations: (1) it does not incorporate any input data, and (2) as a consequence, it exhibits exact collapse even at the shallowest layer, which is inconsistent with practical networks, where the shallowest-layer features typically show little to no collapse, and collapse becomes progressively more pronounced from shallow to deep layers. Furthermore, the analysis of DUFM lacks convergence guarantees beyond a local result. Our contribution is precisely to (1) incorporate input data into the analysis and (2) bridge the gap between landscape analysis (neural collapse under UFM/DUFM) and convergence analysis (implicit bias in shallow networks), establishing convergence to neural collapse.

As such, although DUFM studies a more practical model than ours, our analysis offers advantages over (D)UFM in several aspects (data dependence, convergence, etc.) Moreover, since the dynamic behaviors (alignment, and max-margin bias) also appear in GF on deep networks with positive homogenous activations(Kumar and Haupt, 2025; Ji and Telgarsky, 2020; Lyu and Li, 2019), we believe our work, although on shallow networks, opens up oppotunities to tackle the problem of convergence to NC in deep networks.

Can two-layer ReLU achieve singleton collapse Based on our remark "Intra-class directional collapse" from line 144 to line 154, we believe directional collapse is the global optimal, and normalization along the collapsed direction is one way to achieve singleton collapse as in prior works. First of all, we note that shallow ReLU does not have a universal approximation capability in terms of mapping from input to hidden layer. The universal approximation generally refers to the ability to approximate any input-to-output function mapping $f^*: \mathbb{R}^{d_x}\rightarrow\mathbb{R}^{d_y}$. Here, to achieve NC, the first layer mapping $\sigma(W^\top x)$ must approximate a function $F: \mathbb{R}^{d_x}\rightarrow\mathbb{R}^{h}$ that maps training data to ETF features at hidden layer, but no universal approximation theorem exists in this case. Therefore, we believe the directional collapse is due to the limited representation power of $\sigma(W^\top x)$ for mapping orthogonally separable dataset to collapsed features. Then, our discussion on RMSNorm only suggests that there exists a shallow network with RMSNorm achieving full NC on orthogonally separable dataset, whether RMSNorm is necessary does not immediately follow from this statement. In fact, the understanding of the role of normalization on NC is mainly on an empirical level. Besides our experiments on ResNet with RMSNorm, (Pan and Cao, 2023) demonstrated that batch normalization has a critical influence on the emergence of NC, and their experiments show that models exhibit a stronger tendency toward NC when batch normalization is used. Theoretically proving such a relation warrants a separate paper in future research.

We hope our rebuttal addresses your questions. If you have additional questions/concerns, please feel free to post them during the discussion phase.

References:

Kumar and Haupt, Early Directional Convergence in Deep Homogeneous Neural Networks for Small Initializations, TMLR, 2025

Ji and Telgarsky, Directional convergence and alignment in deep learning, NeurIPS, 2020

Lyu and Li, Gradient descent maximizes the margin of homogeneous neural networks, ICLR, 2019

Pan and Cao. Towards understanding neural collapse: The effects of batch normalization and weight decay." arXiv 2023.