Wide Neural Networks Trained with Weight Decay Provably Exhibit Neural Collapse

## Abstract

Abstract Deep neural networks (DNNs) at convergence consistently represent the training data in the last layer via a highly symmetric geometric structure referred to as neural collapse. This empirical evidence has spurred a line of theoretical research aimed at proving the emergence of neural collapse, mostly focusing on the unconstrained features model. Here, the features of the penultimate layer are
free variables, which makes the model data-agnostic and, hence, puts into question its ability to capture DNN training. Our work addresses the issue, moving away from unconstrained features and studying DNNs that end with at least two linear layers. We first prove generic guarantees on neural collapse that assume (i) low training error and balancedness of the linear layers (for within-class variability collapse), and (ii) bounded conditioning of the features before the linear part (for orthogonality of class-means, as well as their alignment with weight matrices). We then show that such assumptions hold
for gradient descent training with weight decay: (i) for networks with a wide first layer, we prove low training error and balancedness, and (ii) for solutions that are either nearly optimal or stable under large learning rates, we additionally prove the bounded conditioning. Taken together, our results are the first to show neural collapse in the end-to-end training of DNNs.

## 1 Introduction

Among the many possible interpolators that a deep neural network (DNN) can find,
Papyan et al. () showed a strong bias of gradient-based training towards representations with a highly symmetric structure in the penultimate layer, which was dubbed neural collapse (NC). In particular, the feature vectors of the training data in the penultimate layer collapse to a single vector per class (NC1); these vectors form orthogonal or simplex equiangular tight frames (NC2), and they are aligned with the last layer’s row weight vectors (NC3). The question of why and how neural collapse emerges has been considered by
a popular line of research, see e.g. Lu & Steinerberger (); E & Wojtowytsch () and the discussion in Section 2. Many of these works focus on a simplified mathematical framework: the unconstrained features model (UFM) (Mixon et al., ; Han et al., ; Zhou et al., ), corresponding to the joint optimization over the last layer’s weights and the penultimate layer’s feature representations, which are treated as free variables. To account for the existence of the training data and of all the layers before the penultimate (i.e., the backbone of the network), some form of regularization on the free features is usually added. A number of papers has proved the optimality of NC in this model (Lu & Steinerberger, ; E & Wojtowytsch, ), its emergence with gradient-based methods (Mixon et al., ; Han et al., ) and a benign loss landscape (Zhou et al., ; Zhu et al., ). However, the major drawback of the UFM lies in its data-agnostic nature: it only acknowledges the presence of training data and backbone through a simple form of regularization (e.g., Frobenius norm or sphere constraint), which is far from being equivalent to end-to-end training. Moving beyond UFM, existing results are either only applicable to rather shallow networks (of at most three layers) (Kothapalli & Tirer, ; Hong & Ling, ) or hold under strong assumptions, such as symmetric quasi-interpolation (Xu et al., ; Rangamani & Banburski-Fahey, ), block-structured empirical NTK throughout training (Seleznova et al., ), or geodesic structure of the features across all layers (Wang et al., ).

In this paper, we provide the first end-to-end proof of within-class variability collapse (NC1) for a class of networks that end with at least two linear layers. Furthermore, we give rather weak sufficient conditions – either near-optimality or stability under large learning rates – for solutions to exhibit the orthogonality of class-means (NC2) and the alignment of class means with the last weight matrix (NC3). More precisely, our contributions can be summarized as follows:

- •
First, we show that within-class variability collapse (i.e., NC1) occurs, as long as the training error is low and the linear layers are approximately balanced, i.e., $\left\|W_{\ell+1}^{T}W_{\ell+1}-W_{\ell}W_{\ell}^{T}\right\|_{F}$ is small, where $W_{\ell+1},W_{\ell}$ are two consecutive weight matrices of the linear head. If, additionally, the conditioning of the linear head mapping is bounded, we bound the conditioning of the matrix of class means in the last layer, as well as that of the last weight matrix. This implies that, as the number of linear layers grows,
class means become orthogonal and, furthermore, they align with the last layer’s row vectors, which proves both NC2 and NC3.
- •
Next, we show that the sufficient conditions above for NC1 are satisfied by a class of deep networks with a wide first layer, smooth activations and pyramidal topology, after gradient training with weight decay. This provides the first guarantee of the emergence of NC1 for a deep network trained end-to-end via gradient descent.
- •
We further present two sufficient conditions under which the linear head is guaranteed to be
well-conditioned, hence NC1, NC2 and NC3 hold: either the network approaches a global optimum of the $\ell_{2}$-regularized square loss, or it nearly interpolates the data while being stable under large learning rates.
- •
Our numerical experiments on various architectures (fully connected, ResNet) and datasets (MNIST, CIFAR) confirm the insights coming from the theory: *(i)* NC2 is more prominent as the depth of the linear head increases, and *(ii)* the final linear layers are balanced at convergence. Furthermore, we show that, as the non-linear part of the network gets deeper, the non-negative layers become less non-linear and more balanced.

## 2 Related Work

#### Neural collapse.

Since its introduction by Papyan et al. (), neural collapse has been intensively studied, both from a theoretical and practical viewpoint. Practitioners use the NC for a number of applications, including transfer learning, OOD detection and generalization bounds (Galanti et al., ; Haas et al., ; Ben-Shaul & Dekel, ; Li et al., , ; Zhang et al., ). On the theoretical front, the most widely adopted framework to study the emergence of NC is the unconstrained features model (UFM) (Mixon et al., ; Fang et al., ). Under the UFM, the NC has been proved to be optimal with cross-entropy loss (E & Wojtowytsch, ; Lu & Steinerberger, ; Kunin et al., ), MSE loss (Zhou et al., ) and other losses (Zhou et al., ). Optimality guarantees for the generalization of NC to the class-imbalanced setting have been provided by Fang et al. (); Thrampoulidis et al. (); Hong & Ling (); Dang et al. (). Besides global optimality, a benign loss landscape around NC solutions has been proved in Zhu et al. (); Ji et al. (); Zhou et al. (), and the emergence of NC with gradient-based optimization under UFM has been studied by Mixon et al. (); Han et al. (); Ji et al. (); Wang et al. (). Jiang et al. () extend the analysis to large number of classes, Kothapalli et al. () to graph neural networks, Tirer et al. () generalize UFM with a perturbation to account for its imperfections, while Andriopoulos et al. () generalize UFM and NC to regression problems. A deep neural collapse is theoretically analyzed with deep UFM for the linear case in Dang et al. (); Garrod & Keating (), for the non-linear case with two layers in Tirer & Bruna (), and for the deep non-linear case in Súkeník et al. (, ), where, notably, the latter work provides the first negative result on deep NC with deep UFM.

A line of recent work aims at circumventing the data-agnostic nature of UFM, showing the emergence of neural collapse in settings closer to practice. Specifically, Seleznova et al. () assume a block structure in the empirical NTK matrix. Kernels are used by Kothapalli & Tirer () to analyze NC in wide two-layer networks, showing mostly negative results in the NTK regime. Beaglehole et al. () prove the emergence of deep neural collapse using kernel-based layer-wise training, and they show that NC is an optimal solution of adaptive kernel ridge regression in an over-parametrized regime. Sufficient (but rather strong) conditions for the emergence of NC beyond UFM are provided by Pan & Cao (). Hong & Ling () analyze two and three layer networks end-to-end, but only provide conditions under which the UFM optimal solutions are feasible. Wang et al. () consider the NC formation in residual networks, and the model is similar to the perturbed UFM in Tirer et al. (). However, the results crucially assume that the features lie on a geodesic in euclidean space, and proving that this is the case after training the ResNet is an open problem.
Rangamani & Banburski-Fahey (); Xu et al. () focus on homogeneous networks trained via gradient-based methods, making the strong assumption of symmetric quasi-interpolation. Then, Rangamani & Banburski-Fahey () do not prove this assumption, and the argument of Xu et al. () requires a regularization different from the one used in practice as well as interpolators with a given norm (whose existence is an open question). We also note that the quasi-interpolation property does not hold in practice exactly, which points to the need of a perturbation analysis.

#### Implicit bias.

Our approach leverages the (approximate) balancedness of the weights, which plays a central role in the analysis of the training dynamics of linear networks (Arora et al., ) and leads to a bias towards low-rank matrices (Arora et al., ; Tu et al., ). In the presence of weight decay, the low-rank bias can be made even more explicit (Dai et al., ), and it is reinforced by stochastic gradient descent (Wang & Jacot, ). Moving towards nonlinear models,
shallow networks with weight decay exhibit low-rank bias, as described by the variational norm (Bach, ), or Barron norm (E et al., ), and this bias provably emerges under (modified) GD dynamics (Abbe et al., ; Bietti et al., ; Lee et al., ). However, a single hidden layer appears to be insufficient to exhibit NC for most datasets. Moving towards deep nonlinear models, networks with weight decay are also known to exhibit low-rank bias (Galanti et al., ; Jacot et al., ), which can be related to neural collapse (Zangrando et al., ). The dimensionality and rank of the weights/representations varies between layers, exhibiting a bottleneck structure (Jacot, , ; Wen & Jacot, ). We remark that existing results apply to the global minima of the $\ell_{2}$-regularized loss. In contrast, our paper provides rather general sufficient conditions that are then provably satisfied by GD, thus showing that the training algorithm is responsible for neural collapse.

## 3 Balancedness and Interpolation Imply Neural Collapse

#### Notation and problem setup.

Given a matrix $A$ of rank $k$, we denote by $A_{i:}$ its $i$-th row,
by $A_{:i}$ its $i$-th column, by $s_{1}(A)\geq\cdots\geq s_{k}(A)$ its singular values in non-increasing order, and by $\kappa(A)$ the ratio $\frac{s_{1}(A)}{s_{k}(A)}$ between its largest and smallest non-zero singular values. We denote by $\|A\|_{F},\|A\|_{op}$ and $\sigma_{\textrm{min}}\left(A\right)$ its Frobenius norm, its operator norm and its smallest singular value, respectively.

We consider a neural network with $L_{1}$ non-linear layers with activation function $\sigma:\mathbb{R}\to\mathbb{R}$ followed by $L_{2}$ linear layers. Let $L:=L_{1}+L_{2}$ be the total number of layers, $W_{\ell}\in\mathbb{R}^{n_{\ell}\times n_{\ell-1}}$ the weight matrix at layer $\ell$, $X\in\mathbb{R}^{d\times N}$ the training data and $Y\in\mathbb{R}^{K\times N}$ the labels (for consistency, we set $n_{0}=d$ and $n_{L}=K$), where the output dimension $K$ corresponds to the number of classes and $N$ is the number of samples. We consider a one-hot encoding, i.e., the rows of $Y$ are elements of the canonical basis. Let $Z_{\ell}\in\mathbb{R}^{n_{\ell}\times N}$ be the output of layer $\ell$, given by

$$ $\displaystyle Z_{\ell}=\begin{cases}X&\ell=0,\\ \sigma\big{(}W_{\ell}Z_{\ell-1}\big{)}&\ell\in[L_{1}],\\ W_{\ell}Z_{\ell-1}&\ell\in\{L_{1}+1,\ldots,L_{1}+L_{2}\},\end{cases}$ (1) $$

where the activation function $\sigma$ is applied componentwise and, given an integer $n$, we use the shorthand $[n]:=\{1,\ldots,n\}$.
We write $\theta$ for the vector obtained by concatenating all parameters $\{W_{i}\}_{i\in[L]}$ and $W_{m:\ell}$ for the partial products of the weight matrices $W_{m}\cdots W_{\ell}$, so that $Z_{m}=W_{m:\ell+1}Z_{\ell}$ for all $m\geq\ell\in\{L_{1}+1,\dots,L_{1}+L_{2}\}$.

We index individual samples in a feature matrix $Z$ as $z_{ci},$ meaning the $i$-th sample of the $c$-th class, sometimes adding an upper-index to denote the layer or matrix to which the sample belongs. Let $\mu_{c}$ denote the mean of all samples from class $c$ and $\mu_{G}$ the global mean. Let $\bar{Z}$ be the matrix of class-means stacked into columns. The NC1 metric on a feature matrix $Z$ is given by $\frac{\text{tr}(\Sigma_{W})}{\text{tr}(\Sigma_{B})},$ where $\Sigma_{W}=\frac{1}{N}\sum_{c,i}(z_{ci}-\mu_{c})(z_{ci}-\mu_{c})^{\top}$ and $\Sigma_{B}=\frac{1}{K}\sum_{c=1}^{K}(\mu_{c}-\mu_{G})(\mu_{c}-\mu_{G})^{\top}.$ The NC2 metric on a feature matrix $Z$ is defined as $\kappa(\bar{Z})$, i.e., the conditioning number of the class-mean matrix of $Z.$ The NC3 metric on a feature matrix $Z$ and a weight matrix $W$ is defined as $\frac{1}{N}\sum_{c,i}\cos(z_{ci},W_{c:}),$ i.e., the average cosine similarity of features and weight vectors corresponding to the features’ class.

At this point, we state our result giving a set of sufficient conditions for NC1, NC2 and NC3.

###### Theorem 3.1 .

If the network satisfies

- •
approximate interpolation, i.e., $\left\|Z_{L}-Y\right\|_{F}\leq\epsilon_{1}$,
- •
approximate balancedness, i.e., $\left\|W_{\ell+1}^{\top}W_{\ell+1}-W_{\ell}W_{\ell}^{\top}\right\|_{op}\leq
\epsilon_{2}$, for $\ell\in\{L_{1}+1,\ldots,L-1\}$,
- •
bounded representations and weights, i.e., $\left\|Z_{L-2}\right\|_{op},\left\|Z_{L-1}\right\|_{op},\left\|W_{\ell}\right
\|_{op}\leq r$, for $\ell\in\{L_{1}+1,\ldots,L\}$,

then if $\epsilon_{1}\leq\mathop{\rm min}\nolimits\left(s_{K}(Y),\sqrt{\frac{(K-1)N}{4K
}}\right)$,

$$ $\text{NC1}(Z_{L-1})\leq\frac{r^{2}}{N}\frac{(\Psi(\epsilon_{1},\epsilon_{2},r) )^{2}}{\left(\sqrt{\frac{K-1}{K}}-\frac{2}{\sqrt{N}}\epsilon_{1}\right)^{2}},$ (2) $$

where $\Psi(\epsilon_{1},\epsilon_{2},r)=r\left(\frac{\epsilon_{1}}{s_{K}(Y)-\epsilon
_{1}}+\sqrt{n_{L-1}\epsilon_{2}}\right)$.
If we additionally assume that the linear part of the network is not too ill-conditioned, i.e., $\kappa(W_{L:L_{1}+1})\leq c_{3}$, then

$$ $\kappa(W_{L})\leq c_{3}^{\frac{1}{L_{2}}}(1+\epsilon)^{\frac{1}{L_{2}}}+{c_{3} ^{\frac{1}{L_{2}}-1}}\epsilon,$ (3) $$

with $\epsilon=\frac{\frac{L_{2}^{2}}{2}r^{2(L_{2}-1)}\epsilon_{2}}{\frac{(s_{K}(Y)-
\epsilon_{1})^{2}}{\left\|X\right\|_{op}^{2}r^{2L_{1}}}-\frac{L_{2}^{2}}{2}r^{
2(L_{2}-1)}\epsilon_{2}}$.
Finally,

$$ $\displaystyle\text{NC2}(Z_{L-1})$ $\displaystyle\leq\frac{\kappa(W_{L})+s_{K}(Y)^{-1}r\Psi(\epsilon_{1},\epsilon_ {2},r)}{1-s_{K}(Y)^{-1}r\Psi(\epsilon_{1},\epsilon_{2},r)}$ (4) $\displaystyle\text{NC3}(Z_{L-1},W_{L})$ $\displaystyle\geq\frac{(\sqrt{N}-\epsilon_{1})^{2}+N(\kappa(W_{L}))^{-2}-\left (r\Psi(\epsilon_{1},\epsilon_{2},r)+\sqrt{K}(\kappa(W_{L})^{2}-1)\right)^{2}}{ 2N\kappa(W_{L})(1+\epsilon_{1})}.$ (5) $$

The restrictions $\epsilon_{1}\leq s_{K}(Y)$ and $\epsilon_{1}\leq\sqrt{\frac{(K-1)N}{4K}}$ are mild, and we are interested in the regime in which $\epsilon_{1}$ is small. In words, (2) shows that, when $\epsilon_{1},\epsilon_{2}\approx 0$ (i.e., the network approximately interpolates the data in a balanced way), $\Psi(\epsilon_{1},\epsilon_{2},r)\approx 0$ and the within-class variability (which captures NC1) vanishes.
If in addition the depth of the linear part of the network grows, the RHS of (3) approaches 1, i.e., the last weight matrix $W_{L}$ is close to orthogonal. This implies that *(i)* $Z_{L-1}$ is also close to orthogonal (which captures NC2), and *(ii)*
the weights in the last layer align with $Z_{L-1}$ (which captures NC3). In fact, when
$\Psi(\epsilon_{1},\epsilon_{2},r)\approx 0$ and $\kappa(W_{L})\approx 1$, the RHS of both (4) and (5) is close to $1$. Below we give a proof sketch deferring the complete argument to Appendix B.

###### Theorem 3.1 .

#### Proof sketch.

We start with NC1. If $Z_{L}$ is already well-collapsed (which is guaranteed by the approximate interpolation) and $W_{L}$ is well-conditioned, then the only source of within-class variability in $Z_{L-1}$ is within the null space of $W_{L}.$ However, if $W_{L}$ and $W_{L-1}$ are balanced, the image of $Z_{L-1}$ must be approximately in a subspace of the row space of $W_{L}$ and, hence, $Z_{L-1}$ has little freedom within the kernel of $W_{L}$. More formally, consider first the case of perfect balancedness (i.e., $\epsilon_{2}=0$), and denote by $W_{L}^{+}$ the pseudo-inverse of $W_{L}$. Then, $\text{Im}(Z_{L-1})\subset\text{Im}(W_{L}^{\top})$ and

$$ $\left\lVert W_{L}^{+}W_{L}Z_{L-1}-W_{L}^{+}Y\right\rVert_{F}=\left\lVert Z_{L- 1}-W_{L}^{+}Y\right\rVert_{F}\leq\frac{\epsilon_{1}}{s_{K}(W_{L})}.$ (6) $$

As $s_{K}(W_{L})$ can be lower bounded by using the assumptions on approximate interpolation and boundedness of representations, the RHS of (6) is small and, therefore, $Z_{L-1}$ is close to a matrix with zero within-class variability. Moving to the case $\epsilon_{2}\neq 0$, we need to show that $W_{L}^{+}W_{L}Z_{L-1}$ is close to $Z_{L-1}.$ As $W_{L}^{+}W_{L}$ projects onto the row-space of $W_{L},$ only the part of $Z_{L-1}$ in the kernel of $W_{L}$ has to be considered. This part is controlled after writing $Z_{L-1}=W_{L-1}Z_{L-2}$, using the boundedness of $Z_{L-2}$ and the approximate balancedness between $W_{L}$ and $W_{L-1}.$ Finally, as $Z_{L-1}$ is close to $W_{L}^{+}Y,$ a direct computation yields the bound on NC1.

Next, to bound $\kappa(W_{L}),$ we notice that $(W_{L}W_{L}^{\top})^{L_{2}}-(W_{L:L_{1}+1}W_{L:L_{1}+1}^{\top})$ has small operator norm, since the weights of linear layers are approximately balanced. This allows to upper bound $\kappa(W_{L})^{2L_{2}}$ in terms of $\kappa(W_{L:L_{1}+1})$ (plus a small perturbation), which gives (3).

To lower bound the NC3 metric, we rescale $Z_{L-1}$ and $W_{L}$ to $Z^{\prime}_{L-1}$ and $W^{\prime}_{L}$, so that their columns and rows, respectively, have roughly equal size. Then, we reformulate the problem to proving that $\left\langle Z^{\prime}_{L-1},W^{\prime}_{L}Y\right\rangle$ is close to its theoretical maximum. We proceed to show this by arguing that, in this scaling and given that $W_{L}$ is sufficiently well-conditioned, $W^{\prime}_{L}$ can be replaced by $(W^{\prime}_{L})^{+}$. As $Z_{L-1}$ is close to $W_{L}^{+}Y$, we obtain (5). Finally, the bound on NC2 in (4) follows by combining (3) with the closeness between $Z_{L-1}$ and $W_{L}^{+}Y$ already obtained in the proof of NC1.

## 4 Gradient Descent Leads to No Within-Class Variability (NC1)

In this section, we show that NC1 holds for a class of neural networks with one wide layer followed by a pyramidal topology, as considered in Nguyen & Mondelli (). To do so, we show that the balancedness and interpolation conditions of Theorem 3.1 holds. We expect that similar conditions – and, therefore, NC1 – also hold under minimal over-parameterization and for ReLU networks, adapting e.g. the approach of Bombari et al. () and Zou & Gu (), respectively.

We consider a neural network as in (1), and we minimize the $\lambda$-regularized square loss
$C_{\lambda}(\theta)=\frac{1}{2}\left\lVert z_{L}(\theta)-y\right\rVert_{2}^{2}
+\frac{\lambda}{2}\left\lVert\theta\right\rVert_{2}^{2}$, where $z_{L}$ and $y$ are obtained by vectorizing $Z_{L}$ and $Y$, respectively, and $\theta$ collects all the parameters of the network. To do so, we consider the gradient descent (GD) update $\theta_{k+1}=\theta_{k}-\eta\nabla C_{\lambda}(\theta_{k})$, where $\eta$ is the step size and $\theta_{k}=(W_{\ell}^{k})_{\ell=1}^{L}$ contains all parameters at step $k.$ We also denote by $Z_{\ell}^{k}$ the output of layer $\ell$ after $k$ steps of GD. We make the following assumption on the pyramidal topology of the network, noting that this requirement is also common in prior work on the loss landscape (Nguyen & Hein, , ).

###### Assumption 4.1 .

(Pyramidal network topology)
Let $n_{1}\geq N$ and $n_{2}\geq n_{3}\geq\ldots\geq n_{L}.$

We make the following assumptions on the activation function $\sigma$ of the non-linear
layers.

###### Assumption 4.2 .

(Activation function)
Fix $\gamma\in(0,1)$ and $\beta\geq 1$.
Let $\sigma$ satisfy that:
(i) $\sigma^{\prime}(x)\in[\gamma,1]$, (ii) $|\sigma(x)|\leq|x|$ for every $x\in\mathbb{R}$, and (iii) $\sigma^{\prime}$ is $\beta$-Lipschitz.

In principle, $\sigma$ can change at all layers, as long as it satisfies the above assumption. As an example, one can consider a family of parameterized ReLU functions, smoothened by a Gaussian kernel:

$$ $\displaystyle\sigma(x)=-\frac{(1-\gamma)^{2}}{2\pi\beta}+\frac{\beta}{1-\gamma }\int_{-\infty}^{\infty}\mathop{\rm max}\nolimits(\gamma u,u)\,e^{-\frac{\pi \beta^{2}(x-u)^{2}}{(1-\gamma)^{2}}}du.$ (7) $$

One can readily verify that the activation in (7) satisfies Assumption 4.2 and it uniformly approximates the ReLU function over $\mathbb{R}$, see Lemma B.1 in Nguyen & Mondelli ().
Next, let us introduce some notation(^1^11To avoid confusion, we note that this notation is different from the one used in (Nguyen & Mondelli, ).) for the singular values of the weight matrices at initialization $\theta_{0}=(W_{\ell}^{0})_{\ell=1}^{L}$:

$$ $\begin{split}&\lambda_{\ell}=\sigma_{\textrm{min}}\left(W_{\ell}^{0}\right),\, \,\,\,\bar{\lambda}_{\ell}=\left\lVert W_{\ell}^{0}\right\rVert_{op}+\mathop{ \rm min}\nolimits_{\ell\in\{3,\ldots,L\}}\lambda_{\ell},\,\,\,\,\lambda_{i\to j }=\prod_{\ell=i}^{j}\lambda_{\ell},\,\,\,\,\bar{\lambda}_{i\to j}=\prod_{\ell= i}^{j}\bar{\lambda}_{\ell}.\\ \end{split}$ (8) $$

We also define $\lambda_{F}=\sigma_{\textrm{min}}\left(\sigma(W_{1}^{0}X)\right)$ as the smallest singular value of the output of the first hidden layer at initialization.
Finally, we make the following assumption on the initialization.

###### Assumption 4.3 .

(Initial conditions)

$$ $\displaystyle\lambda_{F}\lambda_{3\to L}\mathop{\rm min}\nolimits(\lambda_{F}, \mathop{\rm min}\nolimits_{\ell\in\{3,\ldots,L\}}\lambda_{\ell})\geq 8\gamma \sqrt{\left(\frac{2}{\gamma}\right)^{L}C_{0}(\theta_{0})}.$ (9) $$

We note that (9) can be satisfied by choosing a sufficiently small initialization for the second layer and a sufficiently large one for the remaining layers. In fact, the LHS of (9) depends on all the layer weights except the second, so this quantity can be made arbitrarily large. Next, by taking a sufficiently small second layer, the term $\sqrt{2C_{0}(\theta)}=\|Z_{L}-Y\|_{F}$ can be upper bounded by $2\|Y\|_{F}$. As a consequence, the RHS of (9) is at most $8\sqrt{2}\|Y\|_{F}\gamma\big{(}\frac{2}{\gamma}\big{)}^{L/2}$. As the LHS of (9) can be arbitrarily large, the inequality holds for a suitable initialization.

###### Theorem 4.4 .

Let the network satisfy Assumption 4.1,
$\sigma$ satisfy Assumption 4.2
and the initial conditions satisfy Assumption 4.3.
Fix $0<\epsilon_{1}\leq\frac{1}{2}\sqrt{\frac{(K-1)N}{K}}$, $\epsilon_{2}>0$, let $b\geq 1$ be s.t. $\|X_{:i}\|_{2}\leq b$ for all $i$, and run $k$ steps of $\lambda$-regularized GD with step size $\eta$, where

$$ $\begin{split}\lambda&\leq\mathop{\rm min}\nolimits\left(2\left(\frac{\gamma}{2 }\right)^{L-2}\lambda_{F}\lambda_{3\to L},\frac{2C_{0}(\theta_{0})}{\left \lVert\theta_{0}\right\rVert_{2}^{2}},\frac{\epsilon_{1}^{2}}{18(\left\lVert \theta_{0}\right\rVert_{2}+\lambda_{F}/2)^{2}}\right),\\ \eta&\leq\mathop{\rm min}\nolimits\left(\frac{1}{2\beta_{1}},\frac{1}{5N\beta b ^{3}\mathop{\rm max}\nolimits\left(1,\left(\frac{2\epsilon_{1}^{2}}{\lambda} \right)^{3L/2}\right)L^{5/2}},\frac{1}{2\lambda},\left(\frac{\lambda}{2 \epsilon_{1}^{2}}\right)^{L_{1}+L}\frac{\epsilon_{2}}{4\|X\|_{op}^{2}}\right), \\ k&\geq\left\lceil\frac{\log\frac{\lambda m_{\lambda}}{C_{\lambda}(\theta_{0})- \lambda m_{\lambda}}}{\log(1-\eta\frac{\alpha}{8})}\right\rceil+\left\lceil \frac{\log\frac{\lambda\epsilon_{2}}{4\epsilon_{1}^{2}}}{\log(1-\eta\lambda)} \right\rceil,\end{split}$ (10) $$

with $\beta_{1}=5N\beta b^{3}\left(\prod_{\ell=1}^{L}\mathop{\rm max}\nolimits(1,
\bar{\lambda}_{\ell})\right)^{3}L^{5/2}$, $m_{\lambda}=(1+\sqrt{4\lambda/\alpha})^{2}\left(\left\|\theta_{0}\right\|_{2}+
r_{0}\right)^{2}$, $r_{0}=\frac{1}{2}\mathop{\rm min}\nolimits(\lambda_{F},\mathop{\rm min}
\nolimits_{\ell\in\{3,\ldots,L\}}\lambda_{\ell})$, and $\alpha=2^{-(L-3)}\gamma^{L-2}\lambda_{F}\lambda_{3\to L}$.
Then, we have that

$$ $\text{NC1}(Z_{L-1}^{k})\leq\frac{r^{2}}{N}\frac{\Psi(\epsilon_{1}\sqrt{2}, \epsilon_{2},r)}{\left(\sqrt{\frac{K-1}{K}}-\frac{2\sqrt{2}}{\sqrt{N}}\epsilon _{1}\right)^{2}},$ (11) $$

with $\Psi(\epsilon_{1},\epsilon_{2},r)=r\left(\frac{\epsilon_{1}}{s_{K}(Y)-\epsilon
_{1}}+\sqrt{n_{L-1}\epsilon_{2}}\right)$ and

$$ $r=\mathop{\rm max}\nolimits\left(\epsilon_{1}\sqrt{\frac{2}{\lambda}},\left( \epsilon_{1}\sqrt{\frac{2}{\lambda}}\right)^{L-2}\|X\|_{op},\left(\epsilon_{1} \sqrt{\frac{2}{\lambda}}\right)^{L-1}\|X\|_{op}\right).$ (12) $$

In words, if regularization and learning rate are small enough and we run GD for sufficiently long (as in (10)), then the within-class variability vanishes (as in (11)). To interpret the result, note that the NC1 metric tends to 0 0 (i.e., the collapse is perfect) as $\epsilon_{1},\epsilon_{2}\to 0$. Since the terms $\beta_{1},m_{\lambda},r_{0},\alpha$ do not depend on $\epsilon_{1},\epsilon_{2}$ (but just on the network architecture and initialization), taking $\lambda$ of order $\epsilon_{1}^{2}$ and $\eta$ of order $\epsilon_{2}$ satisfies the first two requirements in (10), also giving that $r$ in (12) is of constant order. Finally, as $\log(1-x)\approx-x$ for small $x$, the quantity $\eta k$ – which quantifies the time of the dynamics, since it is the product of learning rate and number of GD steps – is of order $\log(1/\lambda)+\log(1/\epsilon_{2})/\lambda$. Below we provide a proof sketch deferring the full argument to Appendix B.

###### Assumption 4.1 .

###### Assumption 4.2 .

###### Assumption 4.3 .

###### Theorem 4.4 .

#### Proof sketch.

We show that the network trained via $\lambda$-regularized GD fulfills the three sufficient conditions for NC1 given by Theorem 3.1, i.e., approximate interpolation, approximate balancedness and bounded representations/weights. To do so, we distinguish two phases in the training dynamics.

The *first phase* lasts for logarithmic time in $1/\lambda$ (or, equivalently, $1/\epsilon_{1}$) and, here, the loss decreases exponentially fast to a value of at most $2\lambda m_{\lambda}\leq\epsilon_{1}^{2}$. As the learning rate is small enough, the loss cannot increase during the GD dynamics, which already gives approximate interpolation. To show the exponential convergence, we proceed in two steps. First, Lemma 4.1 in (Nguyen & Mondelli, ) gives that the unregularized loss $C_{0}(\theta)$ satisfies the Polyak-Lojasiewicz (PL) inequality

$$ $\left\|\nabla C_{0}(\theta)\right\|_{2}^{2}\geq\frac{\alpha}{2}C_{0}(\theta),$ (13) $$

for all $\theta$ in a ball centered at initialization $\theta_{0}$ and with sufficiently large radius (captured by $r_{0}$). Next, we show that, if $C_{0}(\theta)$ satisfies the $\alpha$-PL inequality in (13), then the regularized loss $C_{\lambda}(\theta)=C_{0}(\theta)+\frac{\lambda}{2}\left\|\theta\right\|_{2}^{2}$
satisfies a shifted $\alpha$-PL inequality, which implies exponential convergence. This second step is formalized by the proposition below proved in Appendix B.

###### Proposition 4.5 .

Let $C_{0}(\theta)$ satisfy the $\alpha$-PL inequality (13) in the ball $B(\theta_{0},r_{0})$.
Then, in the same ball, $C_{\lambda}(\theta)$ satisfies the inequality

$$ $\left\|\nabla C_{\lambda}(\theta))\right\|_{2}^{2}\geq\frac{\alpha}{4}\left(C_ {\lambda}(\theta)-\lambda m_{\lambda}\right),$ (14) $$

where $m_{\lambda}=(1+\sqrt{4\lambda/\alpha})^{2}\left(\left\|\theta_{0}\right\|_{2}+
r_{0}\right)^{2}$.
Furthermore, assume that $r_{0}\geq 8\sqrt{C_{\lambda}(\theta_{0})/\alpha}$ and $\nabla C_{0}(\theta)$ is $\beta_{1}$-Lipschitz in $B(\theta_{0},r_{0})$. Then, for any $\eta<1/(2\beta_{1})$, there exists

$$ $k_{1}\leq\left\lceil\frac{\log\frac{\lambda m_{\lambda}}{C_{\lambda}(\theta_{0 })-\lambda m_{\lambda}}}{\log(1-\eta\frac{\alpha}{8})}\right\rceil$ (15) $$

such that the $k_{1}$-th iterate of GD satisfies

$$ $C_{\lambda}(\theta_{k_{1}})\leq 2\lambda m_{\lambda},\qquad\qquad\left\|\theta _{k_{1}}-\theta_{0}\right\|_{2}\leq 8\sqrt{\frac{C_{\lambda}(\theta_{0})}{ \alpha}}\leq r_{0}.$ (16) $$

The *second phase* lasts for linear time in $1/\lambda$ (or, equivalently in $1/\epsilon_{1}^{2}$) and logarithmic time in $1/\epsilon_{2}$ and, here, the weight matrices in the linear part of the network become balanced. More precisely, we show that, if $W_{\ell}$ is a weight matrix of the linear part, $\|W_{\ell+1}^{\top}W_{\ell+1}-W_{\ell}W_{\ell}^{\top}\|_{op}$ decreases exponentially and the exponent scales with $1/\lambda$, which gives approximate balancedness.
Finally, as the regularization term in the loss is at most $\epsilon_{1}^{2}$, the operator norm of representations and weight matrices is bounded by $r$ as in (12), and the proof is completed by an application of Theorem 3.1.

###### Proposition 4.5 .

## 5 Orthogonality of Class Means (NC2) and Alignment with Last Weight Matrix (NC3)

To guarantee the orthogonality of class means and their alignment with the last weight matrix, the crux is to show that the condition number $\kappa(W_{L})$ of the last weight matrix $W_{L}$ is close to one. In fact, as $Z_{L-1}\approx W_{L}^{+}Y$, this implies that the last hidden representation $Z_{L-1}$ is approximately a rotation and rescaling of $Y$, which gives NC2, and a bound on NC3 of the form in (5) also follows.

The fact that $\kappa(W_{L})\approx 1$ is a consequence of the presence of many balanced linear layers. Indeed, balancedness implies
$W_{L}W_{L}^{\top}=(W_{L:L_{1}+1}W_{L:L_{1}+1}^{\top})^{\frac{1}{L_{2}}}$,
which gives that $\kappa(W_{L})=\kappa(W_{L:L_{1}+1})^{\frac{1}{L_{2}}}$.
Thus, if the conditioning of the product of the linear layers $W_{L:L_{1}+1}$ can be bounded independently of $L_{2}$, one can guarantee that the conditioning of $W_{L}$ approaches $1$ as $L_{2}\to\infty$.
As this is difficult to obtain in full generality (in particular the assumptions of Theorem 3.1 may not be sufficient), we show that the conditioning can be controlled *(i)* at any global minimizer, and *(ii)* when the parameters are ‘stable’ under large learning rates.

### 5.1 Global Minimizers

We first show that any set of parameters that approximately interpolate with small norm has bounded condition number. We will then show that with the right choice of widths and ridge, all global minimizers satisfy these two assumptions.

###### Proposition 5.1 .

Let $\sigma$ satisfy Assumption 4.2. Then, for any network that satisfies

- •
approximate interpolation, i.e., $\left\|Z_{L}-Y\right\|_{F}\leq\epsilon_{1}$,
- •
bounded parameters, i.e., $\left\|\theta\right\|_{2}^{2}\leq LK+c$,

the linear part $W_{L:L_{1}+1}$ satisfies

$$ $\kappa(W_{L:L_{1}+1})\leq\exp\left(\frac{1}{2}\left(c+L_{1}K\log K-2K\log\frac {s_{K}(Y)-\epsilon_{1}}{\left\|X\right\|_{op}}\right)\right).$ (17) $$

###### Theorem 5.2 .

Let $\sigma$ satisfy Assumption 4.2. Assume there exist parameters of the nonlinear part $\theta_{nonlin}=(W_{\ell})_{\ell=1}^{L_{1}}$ such that $Z_{L_{1}}=Y$ and $\|\theta_{nonlin}\|_{2}^{2}=c$. Then, at any global minimizer of the regularized loss $\mathcal{L}_{\lambda}(\theta)=\frac{1}{2}\left\|Y-Z_{L}\right\|_{F}^{2}+\frac{
\lambda}{2}\left\|\theta\right\|_{2}^{2}$
with $\lambda\leq\frac{\epsilon_{1}^{2}}{KL+c}$, we have

$$ $\begin{split}\kappa(W_{L:L_{1}+1})&\leq\left(\frac{\left\|X\right\|_{op}}{s_{K }(Y)-\epsilon_{1}}\right)^{K}\exp\left(\frac{1}{2}\left(c-L_{1}K+L_{1}K\log K \right)\right),\\ \kappa(W_{L})&\leq\left(\frac{\left\|X\right\|_{op}}{s_{K}(Y)-\epsilon_{1}} \right)^{\frac{K}{L_{1}}}\exp\left(\frac{1}{2L_{1}}\left(c-L_{1}K+L_{1}K\log K \right)\right).\end{split}$ (18) $$

This implies that the bounds on NC1, NC2 and NC3 in (2), (4) and (5), respectively, hold with $\kappa(W_{L})$ upper bounded as above and $\epsilon_{2}=0$.

The assumption that the parameters of the nonlinear part can be chosen to fit the labels $Y$ is guaranteed for large enough width *(i)* by relying on any traditional approximation result (Hornik et al., ; Leshno et al., ; Arora et al., ; He et al., ), or *(ii)* by taking the infinite time limit of any convergence results (Nguyen & Mondelli, ), or *(iii)* by taking the limit $\lambda\searrow 0$ in Proposition 4.5. A sketch of the arguments is below, with full proofs deferred to Appendix B.

###### Proposition 5.1 .

###### Theorem 5.2 .

#### Proof sketch.

To prove Proposition 5.1, we write the norm of the linear and nonlinear parts of the network in terms of the conditioning of $W_{L:L_{1}+1}$: for the linear part, this is a direct computation; for the nonlinear part, we use Theorem 1 of Dai et al. (), which lower bounds the norm of the parameters in terms of the product of the singular values, and then manipulate the latter quantity to obtain again the desired conditioning. Next, to prove Theorem 5.2, we pick the parameters of the nonlinear part $\theta_{nonlin}$ s.t. $Z_{L_{1}}=Y$ and $\|\theta_{nonlin}\|_{2}^{2}=c$, and set the linear layers to the identity. This leads to a total parameter norm of $KL_{2}+c$ and a regularized cost of $\frac{\lambda}{2}(KL_{2}+c)$, and it forces the global minimizer to satisfy the assumptions of Proposition 5.1, which gives the claim on $\kappa(W_{L:L_{1}+1})$. Then, as all local minimizers have balanced linear layers, $\kappa(W_{L})=\kappa(W_{L:L_{1}+1})^{\frac{1}{L_{1}}}$, which gives (18). Finally, the claim on NC1, NC2 and NC3 follows from an application of Theorem 3.1.

### 5.2 Large Learning Rates

Previous works have observed that the learning rates used in practice
are typically ‘too large’, i.e. the loss may not always
be strictly decreasing and GD diverges from
gradient flow (Cohen et al., ). Thankfully, instead of simply diverging, for large $\eta$ (but
not too large) the parameters naturally end up at the ‘edge of stability’: the top eigenvalue of the Hessian $\mathcal{H}C_{\lambda}$
is close to $\frac{2}{\eta}$, i.e., the threshold below which
GD is stable (Cohen et al., ; Lewkowycz et al., ). One can thus interpret GD with learning rate
$\eta$ as minimizing the cost $C_{\lambda}$ amongst parameters $\theta$
such that $\left\|\mathcal{H}C_{\lambda}\right\|_{op}<\frac{2}{\eta}$. These observations are supported by strong empirical evidence and have been also proved theoretically for
simple models (Damian et al., ), although a general result remains difficult to prove due to the chaotic behavior of GD for
large $\eta$. Specifically, the Hessian has the form

$$ $\mathcal{H}C_{\lambda}(\theta)=\left(\nabla_{\theta}Z_{L}\right)^{\top} \partial_{\theta}Z_{L}+\mathrm{Tr}\left[(Y-Z_{L})\nabla_{\theta}^{2}Z_{L} \right]+\lambda I_{P},$ (19) $$

with $\nabla_{\theta}Z_{L}\in\mathbb{R}^{P\times NK}$ and $P$ the number of parameters.
The first term is the Fisher information matrix: this is dual to
the Neural Tangent Kernel (NTK) $\Theta=\nabla_{\theta}Z_{L}\left(\nabla_{\theta}Z_{L}\right)^{\top}\in\mathbb{
R}^{NK\times NK}$ (Jacot et al., ). Therefore, at approximately
interpolating points, we have $\left\|\mathcal{H}C_{\lambda}(\theta)\right\|_{op}=\left\|\Theta\right\|_{op}+
O(\epsilon_{1})+O(\lambda)$, where $\epsilon_{1}$ is the interpolation error and $\lambda$ the regularization parameter.
We can thus interpret large learning as forcing a bound on the operator norm of the NTK. For networks that approximately interpolate the data with bounded NTK and bounded weights, the following proposition guarantees good conditioning of the weights in the linear part (and therefore NC2-3).

###### Proposition 5.3 .

For any network that satisfies

- •
bounded NTK, i.e., $\left\|\Theta\right\|_{op}=\mathop{\rm max}\nolimits_{A}\frac{\left\|\nabla_{
\theta}\mathrm{Tr}\left[Z_{L}A^{T}\right]\right\|_{2}^{2}}{\left\|A\right\|_{F
}^{2}}\leq CL_{2}$,
- •
approximate interpolation, i.e., $\left\|Z_{L}-Y\right\|_{F}\leq\epsilon_{1}$,
- •
bounded weights, i.e., $\left\|W_{\ell}\right\|_{op}\leq r$,

for any $M\leq L_{2}$, there is $\ell\in\{L_{1}+1,\dots,L_{1}+M\}$
such that $\kappa(W_{L:\ell})\leq\frac{\sqrt{CL_{2}}Kr}{\sqrt{M}\left(s_{K}(Y)-\epsilon_{
1}\right)}$.

Furthermore, any network that satisfies approximate interpolation and bounded weights is such that

$$ $\left\|\Theta\right\|_{op}\geq\frac{\left(s_{K}(Y)-\epsilon_{1}\right)^{2}}{K^ {2}r^{2}}L_{2}.$ (20) $$

As an example, by choosing $M=\frac{L_{2}}{2}$, we guarantee that there is at least one layer $\ell$ in the first half of the linear layers s.t.  $\kappa(W_{L:\ell})\leq\frac{\sqrt{2C}Kr}{s_{K}(Y)-\epsilon_{1}}$. Now, assuming the ‘edge of stability’ phenomenon, a learning rate of $\eta=\frac{\eta_{0}}{L_{2}}$ implies a bound
$\left\|\Theta\right\|_{op}\leq\left\|\mathcal{H}C_{\lambda}\right\|_{op}+O(
\epsilon_{1})+O(\lambda)\leq\frac{2L_{2}}{\eta_{0}}+O(\epsilon_{1})+O(\lambda)$, and thus Proposition 5.3 implies that there is a linear layer with bounded conditioning. Proposition 5.3 also suggests that one cannot take any significantly larger learning rate, since for any parameters that $\epsilon_{1}$-interpolate
and have $r$-bounded weights, the NTK satisfies the lower bound in (20), which implies that the learning rate must be smaller than $\frac{2r^{2}K^{2}}{L_{2}\left(s_{K}(Y)-\epsilon_{1}\right)^{2}}$.

###### Proposition 5.3 .

#### Proof sketch.

The idea is that, by suitably choosing $A$ in the evaluation of $\frac{\left\|\nabla_{\theta}\mathrm{Tr}\left[Z_{L}A^{T}\right]\right\|_{2}^{2}
}{\left\|A\right\|_{F}^{2}}$, the operator norm of the NTK is lower bounded by some constant times $\sum_{\ell=L_{1}+1}^{L}\kappa(W_{L:\ell})^{2}$. This implies the desired upper bound on $\kappa(W_{L:\ell})$ for some $\ell\in\{L_{1}+1,\ldots,L_{1}+M\}$, as well as the lower bound on $\|\Theta\|_{op}$ in (20). The details are deferred to Appendix B.

## 6 Numerical Results

In all experiments, we consider MSE loss and standard weight decay regularization. We train an MLP and a ResNet20 with an added MLP head on standard datasets (MNIST, CIFAR10), considering as backbone the first two layers for the MLP and the whole architecture before the linear head for the ResNet.
We evaluate the following metrics related to neural collapse: for NC1, we compute $\text{tr}(\Sigma_{W})/\text{tr}(\Sigma_{B}),$ where $\Sigma_{W},\Sigma_{B}$ are the within- and between-class variability matrices of the feature matrices, respectively; for NC2, we display the conditioning number of the class-mean matrix; for NC3, we use the average cosine angle between the rows of a weight matrix and the columns of the preceding class-mean matrix. Finally, for balancedness we use $\frac{\left\|W_{\ell+1}^{\top}W_{\ell+1}-W_{\ell}W_{\ell}^{\top}\right\|_{op}}
{\mathop{\rm min}\nolimits\left\{\left\|W_{\ell+1}^{\top}W_{\ell+1}\right\|_{
op},\left\|W_{\ell}W_{\ell}^{\top}\right\|_{op}\right\}}$ and for negativity we use $\frac{\left\|Z_{\ell}-\sigma(Z_{\ell})\right\|_{op}}{\left\|Z_{\ell}\right\|_{
op}}.$ We measure such metrics (and also index the layers) starting from the output of the backbone. Our findings can be summarized as follows (see Appendix A for additional complementary experiments).

#### The deeper the linear head, the more clear NC occurs.

We first test whether the models with deep linear heads exhibit NC and if that’s the case, whether it gets better as we deepen the linear head. In Figure 1, we show the NC metrics and the gram matrices of the class-mean matrices of the last layers for the training on CIFAR10 of ResNet20 with 6 extra layers of which the first three have a ReLU activation. We use weight decay of $0.001$ and learning rate of $0.001$, training for 5000 epochs (the learning rate drops ten-fold after 80% of the epochs in all our experiments).

Figure: Figure 1: Last 7 layers of a 9-layer MLP trained on MNIST with weight decay $0.0018$ and learning rate $0.001$. Top: NC1s, NC2s, balancednesses and negativities, from left to right. Results are averaged over 5 runs, and the confidence band at 1 standard deviation is displayed. Bottom: Class-mean matrices of the last three layers (i.e., the linear head), the first before the last ReLU.
Refer to caption: extracted/5903742/figures/single_hyperparam/nc1s_big.png

The plot clearly shows that the collapse is reached throughout the training. We also see that the NC2 metric improves progressively with each layer of the linear head, as predicted by our theory. This effect is also clearly visible from the gram matrices of the class-means (bottom row of Figure 1), which rapidly converge towards the identity. We note that these findings are remarkably consistent across a wide variety of hyperparameter settings and architectures.

Figure: Figure 2: Upper/Lower row: MLP/ResNet20 with a deep linear head. Left to right: NC1 in the last layer; NC1 in the first layer of the linear head; NC2 in the last layer; NC2 in the first layer of the linear head. All plots are a function of the number of layers in the linear head. Results are averaged over 50 runs (5 runs for each of the 10 hyperparameter setups), and the confidence band at 1 standard deviation is displayed.
Refer to caption: extracted/5903742/figures/nc_vs_depth/nc1_vs_depth_last.png

In Figure 2, we plot the dependence of the NC metrics on the number of layers in the linear head. We train an MLP on MNIST with 5 non-linear layers and a number of linear layers ranging from 1 to 5. We average over 5 runs per each combination of weight decay ($0.001,0.004$) and with learning rate of $0.001$. We also train the ResNet20 on CIFAR10 with one non-linear layer head and 1 to 6 linear layers on top. We use the same weight decay and learning rate. The relatively high variance of the results is due to averaging over rather strongly different weight decay values used per each depth. The plots clearly show that the NC2 significantly improves in the last layer as the depth of the linear head increases, while it gets slightly worse in the input layer to the linear head. This is consistent with our theory. The NC1 does not have a strong dependence with the number of layers.

#### Linear layers are increasingly balanced throughout training.

Figure 1 shows that the metric capturing the balancedness exponentially decreases until it plateaus at a rather small value (due to the large learning rate) and then again it exponentially decreases at a smaller rate after reducing the learning rate. The balancedness of non-linear layers instead plateaus at a significantly larger value.

#### Non-linear layers are increasingly balanced and linear, as the depth of the non-linear part increases.

Figure: Figure 3: Left to right: Minimum balancedness; mean balancedness; minimum negativity; mean negativity across non-linear layers of the head as a function of the number of non-linear layers. Results are averaged over 10 runs (5 runs for each of the 2 hyperparameter setups), and the confidence band at 1 standard deviation is displayed.
Refer to caption: extracted/5903742/figures/balance_negat_vs_depth/balancedness_vs_depth_min.png

In Figure 3, we show the dependence of balancedness and non-linearity of the non-linear layers as a function of the depth of the non-linear part. In particular, we plot minimum balancedness and non-linearity across all layers, as well as mean balancedness and non-linearity. The depth of the non-linear part ranges from 4 to 12 (the first two layers are considered as backbone and not measured), the learning rate is either $0.001$ or $0.002$ (5 runs each), and the weight decay is $0.016$ divided by the total number of layers. Balancedness clearly improves with depth, both on average per layer and in the most balanced layer. Similarly, the non-negativity of the layer that least uses the ReLU clearly decreases, and a decrease in the mean negativity is also reported (although less pronounced). Thus, these results suggest that the network tends to use non-linearities less as the depth increases and, in addition, it becomes more balanced. We also note that, in order to fit the data, regardless of the depth, the last non-linear layer(s) always exhibit significant negativity, heavily relying on the ReLU.

## Acknowledgements

M. M. and P. S. are funded by the European Union (ERC, INF^2, project number 101161364). Views and opinions expressed are however
those of the author(s) only and do not necessarily reflect those of the European Union or the
European Research Council Executive Agency. Neither the European Union nor the granting
authority can be held responsible for them.

## Appendix A Additional Experiments

We complement the experiments from Section 6 with additional numerical findings. We start by showing an analog of Figure 1 for MLP trained on MNIST to show that the behavior is robust with respect to the backbone and dataset. The results are shown in Figure 4. The architecture is an MLP with 5 non-linear layers followed by 4 linear layers. We use weight decay of $0.0018$ and learning rate of $0.001$, training for 10000 epochs.

Figure: Figure 4: Last 7 layers of a 9-layer MLP trained on MNIST with weight decay $0.0018$ and learning rate $0.001$. Top: NC1s, NC2s, balancednesses and negativities, from left to right. Results are averaged over 5 runs, and the confidence band at 1 standard deviation is displayed. Bottom: Class-mean matrices of the last four layers (i.e., the linear head).
Refer to caption: extracted/5903742/figures/single_hyperparam/appx/nc1s_big.png

The plot fully agrees with the one in Figure 1 with ResNet20 training on CIFAR10 in every qualitative aspect, and it only differs in the numerical values attained by the NC metrics. Therefore, the conclusion is that the NC is attained across different architectures and NC2 progressively improves as we get closer to the last layer of the DNN.

Next, we extend Figure 2 with experiments on ResNet20 trained on MNIST, which are shown in Figure 5. As above, the results match the interpretations discussed in Section 6, which proves the robustness of our findings across different achitectures, datasets and hyperparameter settings.

Figure: Figure 5: ResNet20 trained on MNIST with a deep linear head. Left to right: NC1 in the last layer; NC1 in the first layer of the linear head; NC2 in the last layer; NC2 in the first layer of the linear head. All plots are a function of the number of layers in the linear head. Results are based on 50 runs (5 runs for each of the 10 hyperparameter setups), and the confidence band at 1 standard deviation is displayed.
Refer to caption: extracted/5903742/figures/nc_vs_depth/RM/nc1_vs_depth_last.png

## Appendix B Deferred Proofs

See 3.1

###### Proof.

NC1: We start by proving the claim (2) on NC1.
We have

$$ $s_{K}(Y)-\epsilon_{1}\leq s_{K}(Y+(Z_{L}-Y))=s_{K}(Z_{L})=s_{K}(W_{L}Z_{L-1}) \leq s_{K}(W_{L})r.$ (21) $$

Denote $W_{L}^{+}$ the pseudoinverse of $W_{L}.$ Then, we have that $W_{L}^{+}W_{L}$ equals the projection $P$ on the row space of $W_{L}.$ We can now write

$$ $Z_{L-1}=PZ_{L-1}+(I-P)Z_{L-1}=W_{L}^{+}Y+W_{L}^{+}(Z_{L}-Y)+(I-P)W_{L-1}Z_{L-2}.$ $$

Note that

$$ $\left\|W_{L}^{+}(Z_{L}-Y)\right\|_{F}\leq\frac{\epsilon_{1}r}{s_{K}(Y)- \epsilon_{1}},$ $$

since $s_{K}(W_{L})\geq(s_{K}(Y)-\epsilon_{1})/r$ by (21). Furthermore,

$$ $\displaystyle\left\|(I-P)W_{L-1}Z_{L-2}\right\|_{F}^{2}$ $\displaystyle\leq\left\|(I-P)W_{L-1}\right\|_{F}^{2}\left\|Z_{L-2}\right\|_{op }^{2}\leq r^{2}\text{tr}((I-P)W_{L-1}W_{L-1}^{\top})$ $\displaystyle=r^{2}\text{tr}((I-P)(W_{L-1}W_{L-1}^{\top}-W_{L}^{\top}W_{L})) \leq r^{2}n_{L-1}\epsilon_{2}.$ $$

Putting these together, we have:

$$ $\left\lVert Z_{L-1}-W_{L}^{+}Y\right\rVert_{F}\leq r\left(\frac{\epsilon_{1}}{ s_{K}(Y)-\epsilon_{1}}+\sqrt{n_{L-1}\epsilon_{2}}\right)=\Psi(\epsilon_{1}, \epsilon_{2},r).$ (22) $$

From now on, we will drop the index $L-1$ and treat everything without layer label as belonging to that layer (the membership to any other layer will be indexed). We have that

$$ $\displaystyle\text{tr}(\Sigma_{W})$ $\displaystyle=\text{tr}\left(\frac{1}{N}\sum_{c,i}(z_{ci}-\mu_{c})(z_{ci}-\mu_ {c})^{\top}\right)=\frac{1}{N}\sum_{c,i}\left\|z_{ci}-\mu_{c}\right\|_{2}^{2}$ $\displaystyle\leq\frac{1}{N}\sum_{c,i}\left\|z_{ci}-(W_{L}^{+})_{c:}\right\|_{ 2}^{2}=\frac{1}{N}\left\|Z_{L-1}-W_{L}^{+}Y\right\|^{2}_{F}.$ $$

Furthermore,

$$ $\displaystyle\text{tr}(\Sigma_{B})$ $\displaystyle=\text{tr}\left(\frac{1}{K}\sum_{c=1}^{K}(\mu_{c}-\mu_{G})(\mu_{c }-\mu_{G})^{\top}\right)=\frac{1}{K}\sum_{c=1}^{K}\left\|\mu_{c}-\mu_{G}\right \|_{2}^{2}$ $\displaystyle\geq\frac{\left\|W_{L}\right\|^{2}_{op}}{r^{2}}\frac{1}{K}\sum_{c =1}^{K}\left\|\mu_{c}-\mu_{G}\right\|_{2}^{2}\geq\frac{1}{Kr^{2}}\sum_{c=1}^{K }\left\|\mu_{c}^{L}-\mu_{G}^{L}\right\|_{2}^{2}.$ $$

Now we proceed by lower-bounding the last term:

$$ $\displaystyle\frac{1}{K}\sum_{c=1}^{K}\left\|\mu_{c}^{L}-\mu_{G}^{L}\right\|_{ 2}^{2}$ $\displaystyle\geq\left(\frac{1}{K}\sum_{c=1}^{K}\left\|\mu_{c}^{L}-\mu_{G}^{L} \right\|_{2}\right)^{2}$ $\displaystyle\geq\left(\frac{1}{K}\sum_{c=1}^{K}\left\|\mu_{c}^{Y}-\mu_{G}^{Y} \right\|_{2}-\frac{1}{K}\sum_{c=1}^{K}\left\|\mu_{c}^{Y}-\mu_{c}^{L}\right\|_{ 2}-\left\|\mu_{G}^{Y}-\mu_{G}^{L}\right\|_{2}\right)^{2},$ $$

where $\mu_{c}^{Y},\mu_{G}^{Y}$ are the class and global means of the label matrix $Y.$ A direct computation yields that $\frac{1}{K}\sum_{c=1}^{K}\left\|\mu_{c}^{Y}-\mu_{G}^{Y}\right\|=\sqrt{\frac{K-
1}{K}}.$ Next we have:

$$ $\displaystyle\left\|\mu_{G}^{Y}-\mu_{G}^{L}\right\|_{2}$ $\displaystyle=\left\|\frac{1}{N}\sum_{c,i}z_{ci}^{L}-\frac{1}{N}\sum_{c,i}z_{ ci}^{Y}\right\|_{2}\leq\frac{1}{N}\sum_{c,i}\left\|z_{ci}^{L}-z_{ci}^{Y}\right \|_{2}$ $\displaystyle\leq\sqrt{\frac{1}{N}\sum_{c,i}\left\|z_{ci}^{L}-z_{ci}^{Y}\right \|_{2}^{2}}=\frac{1}{\sqrt{N}}\left\|Z_{L}-Y\right\|_{F}\leq\frac{\epsilon_{1} }{\sqrt{N}}.$ $$

Finally, for any fixed $c$ we have

$$ $\displaystyle\left\|\mu_{c}^{Y}-\mu_{c}^{L}\right\|_{2}=\left\|\frac{1}{n}\sum _{i}z_{ci}^{Y}-\frac{1}{n}\sum_{i}z_{ci}^{L}\right\|_{2}\leq\frac{1}{n}\sum_{i }\left\|z_{ci}^{Y}-z_{ci}^{L}\right\|_{2}.$ $$

Therefore we get:

$$ $\displaystyle\frac{1}{K}\sum_{c=1}^{K}\left\|\mu_{c}^{Y}-\mu_{c}^{L}\right\|_{ 2}\leq\frac{1}{N}\sum_{c,i}\left\|z_{ci}^{Y}-z_{ci}^{L}\right\|_{2}\leq\frac{1 }{\sqrt{N}}\left\|Z_{L}-Y\right\|_{F}\leq\frac{\epsilon_{1}}{\sqrt{N}}.$ $$

Upper bounding these terms in the above computation and dividing $\text{tr}(\Sigma_{W})$ by $\text{tr}(\Sigma_{B})$ yields (2).

Conditioning of $W_{L}$:
We now prove the claim (3) on the conditioning of $W_{L}$. We rely first on Lemma C.2 to relate the singular values of $W_{L}$ to the $L_{2}$-th root of those of $W_{L:L_{1}+1}$. We therefore obtain that

$$ $\displaystyle\frac{s_{1}(W_{L})^{2L_{2}}}{s_{K}(W_{L})^{2L_{2}}}$ $\displaystyle\leq\frac{s_{1}(W_{L:L_{1}+1})^{2}+\frac{L_{2}^{2}}{2}r^{2(L_{2}- 1)}\epsilon_{2}}{s_{K}(W_{L:L_{1}+1})^{2}-\frac{L_{2}^{2}}{2}r^{2(L_{2}-1)} \epsilon_{2}}$ $\displaystyle=\frac{s_{1}(W_{L:L_{1}+1})^{2}}{s_{K}(W_{L:L_{1}+1})^{2}}+\frac{ s_{K}(W_{L:L_{1}+1})^{2}\frac{L_{2}^{2}}{2}r^{2(L_{2}-1)}\epsilon_{2}+s_{1}(W_ {L:L_{1}+1})^{2}\frac{L_{2}^{2}}{2}r^{2(L_{2}-1)}\epsilon_{2}}{\left(s_{K}(W_{ L:L_{1}+1})^{2}-\frac{L_{2}^{2}}{2}r^{2(L_{2}-1)}\epsilon_{2}\right)s_{K}(W_{L :L_{1}+1})^{2}}.$ $$

We have that

$$ $s_{K}(Z_{L})\leq s_{K}(W_{L:L_{1}+1})\left\|Z_{L_{1}}\right\|_{op}\leq s_{K}(W _{L:L_{1}+1})\left\|X\right\|_{op}r^{L_{1}},$ $$

and therefore

$$ $s_{K}(W_{L:L_{1}+1})\geq\frac{s_{K}(Z_{L})}{\left\|X\right\|_{op}r^{L_{1}}} \geq\frac{s_{K}(Y)-\epsilon_{1}}{\left\|X\right\|_{op}r^{L_{1}}}.$ $$

This gives that

$$ $\frac{s_{1}(W_{L})^{2L_{2}}}{s_{K}(W_{L})^{2L_{2}}}\leq\frac{s_{1}(W_{L:L_{1}+ 1})^{2}}{s_{K}(W_{L:L_{1}+1})^{2}}+\frac{1+\kappa(W_{L:L_{1}+1})^{2}}{\left( \frac{(s_{K}(Y)-\epsilon_{1})^{2}}{\left\|X\right\|_{op}^{2}r^{2L_{1}}}-\frac{ L_{2}^{2}}{2}r^{2(L_{2}-1)}\epsilon_{2}\right)}\frac{L_{2}^{2}}{2}r^{2(L_{2}-1 )}\epsilon_{2}.$ $$

Thus, the following chain of inequalities gives (3):

$$ $\displaystyle\kappa(W_{L})$ $\displaystyle\leq\left(\kappa(W_{L:L_{1}+1})^{2}(1+\epsilon)+\epsilon\right)^{ \frac{1}{2L_{2}}}$ $\displaystyle\leq\kappa(W_{L:L_{1}+1})^{\frac{1}{L_{2}}}(1+\epsilon)^{\frac{1} {2L_{2}}}+{\kappa(W_{L:L_{1}+1})^{\frac{1}{L_{2}}-1}}(1+\epsilon)^{\frac{1}{2L _{2}}-1}\epsilon$ $\displaystyle\leq\kappa(W_{L:L_{1}+1})^{\frac{1}{L_{2}}}(1+\epsilon)^{\frac{1} {2L_{2}}}+{\kappa(W_{L:L_{1}+1})^{\frac{1}{L_{2}}-1}}\epsilon,$ $$

using the concavity of the $2L_{2}$-th root.

NC(2+3):
We start with the proof of the claim (5) on NC3, as the derivations will be used in the proof of the claim (4) on NC2 later. Since the cosine similarity does not depend on the scale of the involved matrices, we perform the following rescaling: denoting $\alpha=\left\|W_{L}\right\|$, we define $W_{L}^{\prime}=\frac{W_{L}}{\alpha}$ and $Z_{L-1}^{\prime}=\alpha Z_{L-1}.$ Similarly to before we denote $(z^{\prime})^{(L-1)}_{ci}$ the $i$-th sample of the $c$-th class in the matrix $Z_{L-1}^{\prime}.$ Then, we write

$$ $\displaystyle\text{NC3}(Z_{L-1}^{\prime},W_{L}^{\prime})$ $\displaystyle=\frac{1}{N}\sum_{c,i}\cos((z^{\prime})^{(L-1)}_{ci},(W_{L}^{ \prime})_{c:})=\frac{1}{N}\sum_{c,i}\frac{\left\langle(z^{\prime})^{(L-1)}_{ci },(W^{\prime}_{L})_{c:}\right\rangle}{\left\|(z^{\prime})^{(L-1)}_{ci}\right\| _{2}\left\|(W^{\prime}_{L})_{c:}\right\|_{2}}$ $\displaystyle\geq\frac{\sum_{c,i}\left\langle(z^{\prime})^{(L-1)}_{ci},(W^{ \prime}_{L})_{c:}\right\rangle}{N\kappa(W_{L})(1+\epsilon_{1})}=\frac{2\left \langle Z_{L-1}^{\prime},(W^{\prime}_{L})^{\top}Y\right\rangle}{2N\kappa(W_{L} )(1+\epsilon_{1})}$ $\displaystyle=\frac{\left\lVert Z^{\prime}_{L-1}\right\rVert_{F}^{2}+\left \lVert(W^{\prime}_{L})^{\top}Y\right\rVert_{F}^{2}-\left\lVert Z^{\prime}_{L-1 }-(W^{\prime}_{L})^{\top}Y\right\rVert_{F}^{2}}{2N\kappa(W_{L})(1+\epsilon_{1} )}.$ $$

Here, the first inequality follows from upper bounding $\left\lVert(W^{\prime}_{L})_{c:}\right\rVert_{2}$ trivially by 1 and

$$ $\left\lVert(z^{\prime})^{(L-1)}_{ci}\right\rVert_{2}\leq\kappa(W_{L})\left \lVert z^{(L)}_{ci}\right\rVert_{2}\leq\kappa(W_{L})\left(\left\lVert y_{c}-z^ {(L)}_{ci}\right\rVert_{2}+\left\lVert y_{c}\right\rVert_{2}\right)\leq\kappa( W_{L})(1+\epsilon_{1}).$ $$

Now,

$$ $\left\lVert Z^{\prime}_{L-1}\right\rVert_{F}^{2}\geq\left\lVert Z_{L}\right \rVert_{F}^{2}\geq\left(\left\lVert Y\right\rVert_{F}-\left\lVert Z_{L}-Y \right\rVert_{F}\right)^{2}=(\sqrt{N}-\epsilon_{1})^{2}.$ $$

Furthermore, we readily have that $\left\lVert(W^{\prime}_{L})^{\top}Y\right\rVert_{F}\geq\frac{\sqrt{N}}{\kappa(
W_{L})}.$ Finally, we have:

$$ $\displaystyle\left\lVert Z^{\prime}_{L-1}-(W_{L}^{\prime})^{\top}Y\right\rVert _{F}\leq\left\lVert Z^{\prime}_{L-1}-(W^{\prime}_{L})^{+}Y\right\rVert_{F}+ \left\lVert(W^{\prime}_{L})^{+}Y-(W^{\prime}_{L})^{\top}Y\right\rVert_{F}.$ $$

From (22), we obtain that

$$ $\displaystyle\left\lVert Z^{\prime}_{L-1}-(W_{L}^{\prime})^{+}Y\right\rVert_{F }\leq r^{2}\left(\frac{\epsilon_{1}}{s_{K}(Y)-\epsilon_{1}}+\sqrt{n_{L-1} \epsilon_{2}}\right)=r\Psi(\epsilon_{1},\epsilon_{2},r).$ (23) $$

Finally, we proceed with upper bounding $\left\lVert(W^{\prime}_{L})^{+}Y-(W^{\prime}_{L})^{\top}Y\right\rVert_{F}$, which can be done via a sandwich bound on the singular values using the conditioning number. This gives

$$ $\left\lVert(W^{\prime}_{L})^{+}Y-(W^{\prime}_{L})^{\top}Y\right\rVert_{F}\leq \sqrt{K}\left(\kappa(W_{L})-\frac{1}{\kappa(W_{L})}\right)\leq\sqrt{K}(\kappa( W_{L})^{2}-1).$ $$

Putting all the obtained bounds together, we get the desired bound (5) on NC3.

Finally, to pass from the bound (3) on $\kappa(W_{L})$ to the bound (4) on NC2, we use the inequality in (23) obtained in the proof of NC3. Note that obtaining a bound on $\kappa(\bar{Z}^{\prime}_{L-1})$ is equivalent to obtaining a bound on $\kappa(\bar{Z}_{L-1})$ since multiplying by a scalar does not change the condition number. Using (23) we get:

$$ $\displaystyle\left\lVert\bar{Z}^{\prime}_{L-1}-(W^{\prime}_{L})^{+}\right \rVert_{F}$ $\displaystyle=\left\lVert Z^{\prime}_{L-1}Y^{+}-(W^{\prime}_{L})^{+}YY^{+} \right\rVert_{F}\leq\left\lVert Y^{+}\right\rVert_{op}\left\lVert Z^{\prime}_{ L-1}-(W^{\prime}_{L})^{+}Y\right\rVert_{F}$ $\displaystyle\leq\frac{r\Psi(\epsilon_{1},\epsilon_{2},r)}{s_{K}(Y)}.$ $$

As $\|(W^{\prime}_{L})^{+}\|_{op}=\kappa(W_{L})$ and $s_{K}((W^{\prime}_{L})^{+})=1$, we conclude that

$$ $\kappa(\bar{Z}^{\prime}_{L-1})\leq\frac{\kappa(W_{L})+r\Psi(\epsilon_{1}, \epsilon_{2},r)(s_{K}(Y))^{-1}}{1-r\Psi(\epsilon_{1},\epsilon_{2},r)(s_{K}(Y)) ^{-1}},$ $$

which gives the desired bound in (4).

∎

See 4.5

###### Proof.

Let $\theta\in B(\theta_{0},r_{0})$. Then, the following chain of inequalities holds:

$$ $\begin{split}\left\|\nabla C_{0}(\theta)+\lambda\theta\right\|_{2}^{2}&\geq \left(\left\|\nabla C_{0}(\theta)\right\|_{2}-\lambda\left\|\theta\right\|_{2} \right)^{2}\\ &\geq\left(\sqrt{\frac{\alpha}{2}C_{0}(\theta)}-\lambda\left\|\theta\right\|_{ 2}\right)^{2}\\ &=\left(\sqrt{\frac{\alpha}{2}C_{\lambda}(\theta)-\frac{\alpha\lambda}{4}\left \|\theta\right\|_{2}^{2}}-\lambda\left\|\theta\right\|_{2}\right)^{2}\\ &\geq\left(\sqrt{\frac{\alpha}{2}C_{\lambda}(\theta)}-\left(\lambda+\sqrt{ \frac{\alpha\lambda}{4}}\right)\left\|\theta\right\|_{2}\right)^{2}\\ &\geq\frac{\alpha}{4}C_{\lambda}(\theta)-\lambda\left(\sqrt{\frac{\alpha}{4}}+ \sqrt{\lambda}\right)^{2}\left\|\theta\right\|_{2}^{2}\\ &\geq\frac{\alpha}{4}\left(C_{\lambda}(\theta)-\lambda\left(1+\sqrt{\frac{4 \lambda}{\alpha}}\right)^{2}\left(\left\|\theta_{0}\right\|_{2}+r_{0}\right)^{ 2}\right).\end{split}$ (24) $$

Here,
in the second line we use that $\sqrt{\frac{\alpha}{2}C_{0}(\theta)}\geq\lambda\left\|\theta\right\|_{2}$, which follows from $C_{\lambda}(\theta)\geq\lambda m_{\lambda}$ (otherwise, the claim is trivial); in the fourth line we use that $\sqrt{a-b}\geq\sqrt{a}-\sqrt{b}$ for $a\geq b$; in the fifth line we use
that $(a-b)^{2}\geq\frac{a^{2}}{2}-b^{2}$ for all $a,b$; and in the sixth line we use that $\theta$ is in the ball $B(\theta_{0},r_{0})$. As the LHS of (24) equals $\left\|\nabla C_{\lambda}(\theta))\right\|_{2}^{2}$, this proves (14).

Next, let $(\theta_{k})_{k\in\mathbb{N}}$ be the GD trajectory. Pick $k$ s.t. $\theta_{k}\in B(\theta_{0},r_{0})$. Then,

$$ $\displaystyle C_{\lambda}(\theta_{k+1})-C_{\lambda}(\theta_{k})=-\eta\int_{0}^ {1}\langle\nabla C_{\lambda}(\theta_{k}-s\eta\nabla C_{\lambda}(\theta_{k})), \nabla C_{\lambda}(\theta_{k})\rangle ds$ $\displaystyle=-\eta\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}^{2}+\eta \int_{0}^{1}\langle\nabla C_{\lambda}(\theta_{k}-s\eta\nabla C_{\lambda}( \theta_{k}))-\nabla C_{\lambda}(\theta_{k}),\nabla C_{\lambda}(\theta_{k}) \rangle ds$ $\displaystyle\leq-\eta\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}^{2}+ \eta\int_{0}^{1}\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}\left\|\nabla C _{\lambda}(\theta_{k}-s\eta\nabla C_{\lambda}(\theta_{k}))-\nabla C_{\lambda}( \theta_{k})\right\|_{2}ds$ $\displaystyle\leq-\eta\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}^{2}+ \eta^{2}\beta_{1}\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}^{2}$ $\displaystyle=-\eta(1-\eta\beta_{1})\left\|\nabla C_{\lambda}(\theta_{k}) \right\|_{2}^{2}$ $\displaystyle\leq-\frac{\eta}{2}\left\|\nabla C_{\lambda}(\theta_{k})\right\|_ {2}^{2}$ $\displaystyle\leq-\eta\frac{\alpha}{8}\left(C_{\lambda}(\theta_{k})-\lambda m_ {\lambda}\right).$ $$

Here, in the fourth line we use that the gradient is $\beta_{1}$-Lipschitz in $B(\theta_{0},r_{0})$; in the sixth line we use that $\eta<1/(2\beta_{1})$; and in the last line we use (14). Thus, as long as $\theta_{j}\in B(\theta_{0},r_{0})$ for all $j\in[k]$, by iterating the argument above, we have

$$ $C_{\lambda}(\theta_{k})-\lambda m_{\lambda}\leq\left(C_{\lambda}(\theta_{0})- \lambda m_{\lambda}\right)\left(1-\eta\frac{\alpha}{8}\right)^{k}.$ $$

This readily implies that, by letting $k_{1}$ be the first index $k$ s.t. $C_{\lambda}(\theta_{k})\leq 2\lambda m_{\lambda}$, (15) holds.
Finally, the distance $\left\|\theta_{k_{1}}-\theta_{0}\right\|_{2}$
is upper bounded by

$$ $\displaystyle\sum_{k=0}^{k_{1}-1}\eta$ $\displaystyle\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}=\eta\sum_{k=0}^ {k_{1}-1}\frac{\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}^{2}}{\left\| \nabla C_{\lambda}(\theta_{k})\right\|_{2}}$ $\displaystyle\leq\frac{4}{\sqrt{\alpha}}\sum_{k=0}^{k_{1}-1}\frac{C_{\lambda}( \theta_{k})-C_{\lambda}(\theta_{k+1})}{\sqrt{C_{\lambda}(\theta_{k})-\lambda m _{\lambda}}}$ $\displaystyle=\frac{4}{\sqrt{\alpha}}\sum_{k=0}^{k_{1}-2}\frac{C_{\lambda}( \theta_{k})-C_{\lambda}(\theta_{k+1})}{\sqrt{C_{\lambda}(\theta_{k})-\lambda m _{\lambda}}}+\frac{4}{\sqrt{\alpha}}\frac{C_{\lambda}(\theta_{k_{1}-1})-C_{ \lambda}(\theta_{k_{1}})}{\sqrt{C_{\lambda}(\theta_{k_{1}-1})-\lambda m_{ \lambda}}}$ $\displaystyle\leq\frac{8}{\sqrt{\alpha}}\sum_{k=0}^{k_{1}-2}\frac{C_{\lambda}( \theta_{k})-C_{\lambda}(\theta_{k+1})}{\sqrt{C_{\lambda}(\theta_{k+1})-\lambda m _{\lambda}}+\sqrt{C_{\lambda}(\theta_{k})-\lambda m_{\lambda}}}+\frac{8}{\sqrt {\alpha}}\frac{C_{\lambda}(\theta_{k_{1}-1})-\lambda m_{\lambda}}{\sqrt{C_{ \lambda}(\theta_{k_{1}-1})-\lambda m_{\lambda}}}$ $\displaystyle=\frac{8}{\sqrt{\alpha}}\sum_{k=0}^{k_{1}-2}\left(\sqrt{C_{ \lambda}(\theta_{k})-\lambda m_{\lambda}}-\sqrt{C_{\lambda}(\theta_{k+1})- \lambda m_{\lambda}}\right)+\frac{8}{\sqrt{\alpha}}\sqrt{C_{\lambda}(\theta_{k _{1}-1})-\lambda m_{\lambda}}$ $\displaystyle=\frac{8}{\sqrt{\alpha}}\sqrt{C_{\lambda}(\theta_{0})-\lambda m_{ \lambda}}\leq\frac{8}{\sqrt{\alpha}}\sqrt{C_{\lambda}(\theta_{0})},$ $$

which concludes the proof.
∎

See 4.4

###### Proof.

By Lemma 4.1 in (Nguyen & Mondelli, ), the loss $C_{0}(\theta)$ satisfies the $\alpha$-PL inequality with

$$ $\alpha=4\gamma^{L-2}\sigma_{\textrm{min}}\left(Z_{1}\right)\prod_{p=3}^{L} \sigma_{\textrm{min}}\left(W_{p}\right),$ $$

where we have also used that $\sigma^{\prime}$ is lower bounded by $\gamma$ by Assumption 4.2. Thus, by taking $r_{0}=\frac{1}{2}\mathop{\rm min}\nolimits(\lambda_{F},\mathop{\rm min}
\nolimits_{\ell\in\{3,\ldots,L\}}\lambda_{\ell})$, we have that, for all $\theta\in B(\theta_{0},r_{0})$, $C_{0}(\theta)$ satisfies the $\alpha$-PL inequality with $\alpha=2^{-(L-3)}\gamma^{L-2}\lambda_{F}\lambda_{3\to L}$.

By using Assumption 4.3 and that $\lambda$ is upper bounded by $\frac{2C_{0}(\theta_{0})}{\left\lVert\theta_{0}\right\rVert_{2}^{2}}$ in (10), one can readily verify that $r_{0}\geq 8\sqrt{C_{\lambda}(\theta_{0})/\alpha}$. Furthermore, by Lemma C.1, we have that $\nabla C_{0}(\theta)$ is $\beta_{1}$-Lipschitz for all $\theta\in B(\theta_{0},r_{0})$. Hence, we can apply Proposition 4.5, which gives that, for some $k_{1}$ upper bounded in (15),

$$ $C_{\lambda}(\theta_{k_{1}})\leq 2\lambda m_{\lambda}\leq\epsilon_{1}^{2},$ $$

where the last inequality uses again the upper bound on $\lambda$ in (10).

With gradient flow, the regularized loss would only
decrease further after $k_{1}$ steps. Since we are working with gradient descent,
we simply need to assume that the learning rate is small enough to
guarantee a decreasing loss. If the gradient $\nabla C_{\lambda}(\theta)$
is $\beta_{2}$-Lipschitz, then

$$ $\displaystyle C_{\lambda}(\theta_{k+1})$ $\displaystyle-C_{\lambda}(\theta_{k})=C_{\lambda}(\theta_{k}-\eta\nabla C_{ \lambda}(\theta_{k}))-C_{\lambda}(\theta_{k})$ $\displaystyle=-\eta\int_{0}^{1}\left\langle\nabla C_{\lambda}(\theta_{k}-s\eta \nabla C_{\lambda}(\theta_{k})),\nabla C_{\lambda}(\theta_{k})\right\rangle ds$ $\displaystyle\leq-\eta\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}^{2}$ $\displaystyle\hskip 11.00008pt+\eta\left\|\nabla C_{\lambda}(\theta_{k})\right \|_{2}\mathop{\rm max}\nolimits_{s\in[0,1]}\left\|\nabla C_{\lambda}(\theta_{k })-\nabla C_{\lambda}(\theta_{k}-s\eta\nabla C(\theta_{k}))\right\|_{2}$ $\displaystyle\leq-\eta\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}^{2}+ \eta^{2}\beta_{2}\left\|\nabla C_{\lambda}(\theta_{k})\right\|_{2}^{2}$ $\displaystyle=-\eta(1-\eta\beta_{2})\left\|\nabla C_{\lambda}(\theta_{k}) \right\|^{2}_{2},$ $$

which is non-positive as long as $\eta\leq 1/\beta_{2}$. Furthermore,
as long as the regularized loss is decreasing, the parameter norm is bounded by $\epsilon_{1}\sqrt{2/\lambda}$. Lemma C.1 then implies that the gradient is $\beta_{2}=5N\beta b^{3}\mathop{\rm max}\nolimits\left(1,\epsilon_{1}^{3L}\left
(\frac{2}{\lambda}\right)^{3L/2}\right)L^{5/2}$-Lipschitz and, therefore, $\eta\leq 1/\beta_{2}$ holds by (10). This allows us to conclude that, for all $k$ satisfying the lower bound in (10), $C_{\lambda}(\theta_{k})\leq\epsilon_{1}^{2}$, hence the network achieves approximate interpolation, i.e.,

$$ $\|Z_{L}^{k}-Y\|_{F}\leq\epsilon_{1}\sqrt{2}.$ (25) $$

Next, we show approximate balancedness. To do so, for $\ell\in\{L_{1}+2,\ldots,L-1\}$, we define

$$ $T_{\ell}^{k}:=(W_{\ell+1}^{k})^{\top}\cdots(W_{L}^{k})^{\top}(Z_{L}^{k}-Y)(Z_{ L_{1}}^{k})^{\top}(W_{L_{1}+1}^{k})^{\top}\cdots(W_{\ell-1}^{k})^{\top}.$ $$

Then, we have

$$ $\displaystyle W_{\ell}^{k+1}$ $\displaystyle(W_{\ell}^{k+1})^{\top}=((1-\eta\lambda)W_{\ell}^{k}-\eta T_{\ell }^{k})((1-\eta\lambda)W_{\ell}^{k}-\eta T_{\ell}^{k})^{\top}$ $\displaystyle=(1-\eta\lambda)^{2}W_{\ell}^{k}(W_{\ell}^{k})^{\top}-(1-\eta \lambda)\eta(W_{\ell}^{k}(T_{\ell}^{k})^{\top}+T_{\ell}^{k}(W_{\ell}^{k})^{ \top})+\eta^{2}T_{\ell}^{k}(T_{\ell}^{k})^{\top}.$ $$

Similarly,

$$ $\displaystyle(W_{\ell+1}^{k+1})^{\top}W_{\ell+1}^{k+1}$ $\displaystyle=(1-\eta\lambda)^{2}(W_{\ell+1}^{k})^{\top}W_{\ell+1}^{k}-(1-\eta \lambda)\eta((W_{\ell+1}^{k})^{\top}T_{\ell+1}^{k}+(T_{\ell+1}^{k})^{\top}W_{ \ell+1}^{k})+\eta^{2}(T_{\ell+1}^{k})^{\top}T_{\ell+1}^{k}.$ $$

Let us define

$$ $D_{\ell}^{k}=(W_{\ell+1}^{k})^{\top}W_{\ell+1}^{k}-W_{\ell}^{k}(W_{\ell}^{k})^ {\top}.$ $$

Since $T_{\ell}^{k}(W_{\ell}^{k})^{\top}=(W_{\ell+1}^{k})^{\top}T_{\ell+1}^{k}$ and $W_{\ell}^{k}(T_{\ell}^{k})^{\top}=(T_{\ell+1}^{k})^{\top}W_{\ell+1}^{k}$, we have

$$ $\displaystyle D_{\ell}^{k+1}=(1-\eta\lambda)^{2}D_{\ell}^{k}+\eta^{2}((T_{\ell +1}^{k})^{\top}T_{\ell+1}^{k}-T_{\ell}^{k}(T_{\ell}^{k})^{\top}).$ $$

Recall that, for all $k$ lower bounded in (10), $\|Z_{L}^{k}-Y\|_{F}\leq\epsilon_{1}\sqrt{2}$ and $\|W_{\ell}^{k}\|_{F}\leq\|\theta^{k}\|_{2}\leq\epsilon_{1}\sqrt{2/\lambda}$, which also implies that $\|Z_{L_{1}}\|_{F}\leq\left(\epsilon_{1}\sqrt{2/\lambda}\right)^{L_{1}}\|X\|_{op}$.
Thus,

$$ $\displaystyle\|D_{\ell}^{k+1}\|_{op}$ $\displaystyle\leq(1-\eta\lambda)^{2}\|D_{\ell}^{k}\|_{op}+\eta^{2}(\|T_{\ell}^ {k}\|_{op}^{2}+\|T_{\ell+1}^{k}\|_{op}^{2})$ $\displaystyle\leq(1-\eta\lambda)^{2}\|D_{\ell}^{k}\|_{op}+\eta^{2}\|Z_{L_{1}}^ {k}\|_{op}^{2}\|Z_{L}^{k}-Y\|_{op}^{2}(\prod_{j\neq\ell}\|W_{j}^{k}\|_{op}^{2} +\prod_{j\neq\ell+1}\|W_{j}^{k}\|_{op}^{2})$ $\displaystyle\leq(1-\eta\lambda)^{2}\|D_{\ell}^{k}\|_{op}+4\eta^{2}\epsilon_{1 }^{2}\left(\frac{2\epsilon_{1}^{2}}{\lambda}\right)^{L_{1}+L-1}\|X\|_{op}^{2}.$ $$

By using the upper bounds $\eta\leq\frac{1}{2\lambda}$ and $\eta\leq\left(\frac{\lambda}{2\epsilon_{1}^{2}}\right)^{L_{1}+L}\frac{\epsilon
_{2}}{4\|X\|_{op}^{2}}$ in (10), we have that:

- •
if $\|D_{\ell}^{k}\|_{op}\geq\epsilon_{2}$, then $\|D_{\ell}^{k+1}\|_{op}\leq(1-\eta\lambda)\|D_{\ell}^{k}\|_{op}$;
- •
if $\|D_{\ell}^{k}\|_{op}\leq\epsilon_{2}$, then $\|D_{\ell}^{k+1}\|_{op}\leq(1-\eta\lambda)\epsilon_{2}\leq\epsilon_{2}$.

This implies that, for all $\bar{k}\geq 0$ and $k_{1}\geq\left\lceil\frac{\log\frac{\lambda m_{\lambda}}{C_{\lambda}(\theta_{0
})-\lambda m_{\lambda}}}{\log(1-\eta\frac{\alpha}{8})}\right\rceil$,

$$ $\|D_{\ell}^{k_{1}+\bar{k}}\|_{op}\leq\mathop{\rm max}\nolimits((1-\eta\lambda) ^{\bar{k}}\|D_{\ell}^{k_{1}}\|_{op},\epsilon_{2}).$ $$

Note that

$$ $\|D_{\ell}^{k_{1}}\|_{op}\leq\|W_{\ell+1}^{k_{1}}\|_{op}^{2}+\|W_{\ell}^{k_{1} }\|_{op}^{2}\leq\frac{4\epsilon_{1}^{2}}{\lambda},$ $$

which allows us to conclude that, for all $k$ lower bounded in (10),

$$ $\|D_{\ell}^{k}\|_{2}\leq\epsilon_{2}.$ $$

Finally, we have the following bounds on the representations and weights at step $k$:

$$ $\displaystyle\left\|Z_{L-2}^{k}\right\|_{op}$ $\displaystyle\leq\left(\epsilon_{1}\sqrt{\frac{2}{\lambda}}\right)^{L-2}\|X\|_ {op},$ $\displaystyle\left\|Z_{L-1}^{k}\right\|_{op}$ $\displaystyle\leq\left(\epsilon_{1}\sqrt{\frac{2}{\lambda}}\right)^{L-1}\|X\|_ {op},$ $\displaystyle\left\|W_{\ell}^{k}\right\|_{op}$ $\displaystyle\leq\epsilon_{1}\sqrt{\frac{2}{\lambda}},\,\,\,\mbox{ for }\ell \in\{L_{1}+1,\ldots,L\}.$ $$

Hence, an application of Theorem 3.1 proves the desired result. ∎

See 5.1

###### Proof.

Let us split the parameter norm into the contribution from the nonlinear
and linear layers:

$$ $\left\|\theta\right\|_{2}^{2}=\left\|\theta_{nonlin}\right\|_{2}^{2}+\left\| \theta_{lin}\right\|_{2}^{2}.$ $$

We first lower bound both parts in terms of the product of
the linear part $W_{L:L_{1}+1}$. This then allows
us to bound the conditioning of $W_{L:L_{1}+1}$. We start with the nonlinear part:

$$ $\left\|Z_{L_{1}}\right\|_{F}\leq\left\|X\right\|_{op}\prod_{\ell=1}^{L_{1}} \left\|W_{\ell}\right\|_{F}\leq\left\|X\right\|_{op}\left(\frac{1}{L_{1}}\left \|\theta_{nonlin}\right\|_{2}^{2}\right)^{\frac{L_{1}}{2}},$ $$

which implies that

$$ $\left\|\theta_{nonlin}\right\|_{2}^{2}\geq L_{1}\left(\frac{\left\|Z_{L_{1}} \right\|_{F}}{\left\|X\right\|_{op}}\right)^{\frac{2}{L_{1}}}\geq L_{1}\left( \frac{\left\|Z_{L_{1}}\right\|_{op}}{\left\|X\right\|_{op}}\right)^{\frac{2}{L _{1}}}.$ $$

Since $\left\|Y-W_{L:L_{1}+1}Z_{L_{1}}\right\|_{F}\leq\epsilon_{1}$,
we have

$$ $s_{K}(Y)\leq s_{K}(W_{L:L_{1}+1}Z_{L_{1}})+\epsilon_{1}\leq s_{K}(W_{L:L_{1}+1 })\left\|Z_{L_{1}}\right\|_{op}+\epsilon_{1},$ $$

so that $\left\|Z_{L_{1}}\right\|_{op}\geq\frac{s_{K}(Y)-\epsilon_{1}}{s_{K}(W_{L:L_{1}
+1})}$
and thus

$$ $\left\|\theta_{nonlin}\right\|_{2}^{2}\geq L_{1}\left(\frac{s_{K}(Y)-\epsilon_ {1}}{s_{K}(W_{L:L_{1}+1})\left\|X\right\|_{op}}\right)^{\frac{2}{L_{1}}}.$ (26) $$

Next, for the linear part, we know from Theorem 1 of Dai et al. () that

$$ $\mathop{\rm min}\nolimits_{(W_{\ell})_{\ell=L_{1}+1}^{L}:A=W_{L:L_{1}+1}}\sum_ {\ell=L_{1}+1}^{L}\|W_{\ell}\|_{F}^{2}=(L-L_{1})\sum_{i=1}^{\mathrm{Rank}(A)}s _{i}(A)^{\frac{2}{L-L_{1}}},$ $$

so that

$$ $\left\|\theta_{lin}\right\|_{2}^{2}\geq(L-L_{1})\sum_{i=1}^{\mathrm{Rank}(W_{L :L_{1}+1})}s_{i}(W_{L:L_{1}+1})^{\frac{2}{L-L_{1}}}\geq(L-L_{1})K+2\log\left|W _{L:L_{1}+1}\right|_{+},$ (27) $$

where we used the fact that $x^{\frac{2}{L-L_{1}}}\geq 1+\frac{2}{L-L_{1}}\log x$.

We also have following bound on the condition number $\kappa(W_{L:L_{1}+1})$:

$$ $\kappa(W_{L:L_{1}+1})=\frac{s_{1}(W_{L:L_{1}+1})}{s_{K}(W_{L:L_{1}+1})}\leq \prod_{i=1}^{K}\frac{s_{i}(W_{L:L_{1}+1})}{s_{K}(W_{L:L_{1}+1})}=\frac{\left|W _{L:L_{1}+1}\right|_{+}}{s_{K}(W_{L:L_{1}+1})^{K}}.$ (28) $$

By combining (26), (27), and (28) (after applying the log on both sides of (28)), we obtain

$$ $\displaystyle\left\|\theta\right\|_{2}^{2}$ $\displaystyle\geq L_{1}\left(\frac{s_{K}(Y)-\epsilon_{1}}{s_{K}(W_{L:L_{1}+1}) \left\|X\right\|_{op}}\right)^{\frac{2}{L_{1}}}+(L-L_{1})K+2\log\left|W_{L:L_{ 1}+1}\right|_{+}$ $\displaystyle\geq L_{1}\left(\frac{s_{K}(Y)-\epsilon_{1}}{s_{K}(W_{L:L_{1}+1}) \left\|X\right\|_{op}}\right)^{\frac{2}{L_{1}}}+(L-L_{1})K+2\log\kappa(W_{L:L_ {1}+1})+2K\log s_{K}(W_{L:L_{1}+1}).$ $$

The above is lower bounded by the minimum over all possible choices of $s_{K}(W_{L:L_{1}+1})$.
This minimum would be attained at $s_{K}(W_{L:L_{1}+1})=K^{-\frac{L_{1}}{2}}\frac{s_{K}(Y)-\epsilon_{1}}{\left\|X
\right\|_{op}}$,
thus leading to the lower bound

$$ $\left\|\theta\right\|_{2}^{2}\geq LK+2\log\kappa(W_{L:L_{1}+1})-L_{1}K\log K+2 K\log\frac{s_{K}(Y)-\epsilon_{1}}{\left\|X\right\|_{op}}.$ $$

This implies

$$ $\displaystyle 2\log\kappa(W_{L:L_{1}+1})$ $\displaystyle\leq\left\|\theta\right\|_{2}^{2}-LK+L_{1}K\log K-2K\log\frac{s_{ K}(Y)-\epsilon_{1}}{\left\|X\right\|_{op}}$ $\displaystyle\leq c+L_{1}K\log K-2K\log\frac{s_{K}(Y)-\epsilon_{1}}{\left\|X \right\|_{op}}.$ $$

∎

See 5.2

###### Proof.

By assumption, there are parameters of the nonlinear part $\theta_{nonlin}$ such
that the representation at the end of the nonlinear layers already matches
the outputs $Z_{L_{1}}=Y$ with finite parameter norm $\left\|\theta_{nonlin}\right\|_{2}^{2}=c$.
We can now build a deeper network by simply setting all the linear
layers to be the $K$-dimensional identity, so that $\left\|\theta_{lin}\right\|_{2}^{2}=KL_{2}$,
leading to a total parameter norm of $\left\|\theta\right\|_{2}^{2}=c+KL_{2}$.
Since the outputs matches the labels exactly, we have that the
regularized loss is bounded by $\lambda(KL_{2}+c)/2$. This implies that any global minimizer $\theta^{*}$ satisfies

$$ $\displaystyle\left\|Y-Z_{L}\right\|^{2}_{F}$ $\displaystyle\leq\lambda(KL_{2}+c),$ $\displaystyle\left\|\theta^{*}\right\|_{2}^{2}$ $\displaystyle\leq KL_{2}+c.$ $$

Choosing $\lambda\leq\frac{\epsilon_{1}^{2}}{KL_{2}+c}$, we obtain
that the global minimizer $\theta^{*}$ must satisfy the assumptions
of Proposition 5.1, which readily gives the desired upper bound on $\kappa(W_{L:L_{1}+1})$.

To finish the proof, we show that all critical points of the regularized loss have balanced linear layers, so that we may relate the conditioning of $W_{L}$ to that of the product $W_{L:L_{1}+1}$. At any critical point, the gradient w.r.t. to any weight matrix amongst the linear layers must be zero, that is

$$ $0=\nabla_{W_{\ell}}C_{\lambda}(\theta)=W_{L:\ell+1}^{\top}(Z_{L}-Y)Z_{L_{1}}^{ \top}W_{\ell-1:L_{1}+1}^{\top}+\lambda W_{\ell}.$ $$

This implies that $W_{\ell}=-\frac{1}{\lambda}W_{L:\ell+1}^{\top}(Z_{L}-Y)Z_{L_{1}}^{\top}W_{\ell
-1:L_{1}+1}^{\top}$
for all linear layers $\ell$. Balancedness then follows directly

$$ $W_{\ell}W_{\ell}^{\top}=-\frac{1}{\lambda}W_{L:\ell+1}^{\top}(Z_{L}-Y)Z_{L_{1} }^{\top}W_{\ell:L_{1}+1}^{\top}=W_{\ell+1}^{\top}W_{\ell+1}.$ $$

Finally, the balancedness implies that

$$ $\kappa(W_{L})=\kappa(W_{L:L_{1}+1})^{\frac{1}{L_{1}}}\leq\left(\frac{\left\|X \right\|_{op}}{s_{K}(Y)-\epsilon_{1}}\right)^{\frac{K}{L_{1}}}\exp\left(\frac{ 1}{2L_{1}}\left(c-L_{1}K+L_{1}K\log K\right)\right).$ $$

∎

See 5.3

###### Proof.

As the NTK is a sum over all layers, we can lower bound it by the contribution of the linear layers only. Formally, we have

$$ $\displaystyle\left\|\nabla_{\theta}\mathrm{Tr}[Z_{L}A^{\top}]\right\|_{2}^{2}$ $\displaystyle\geq\left\|\nabla_{\theta_{lin}}\mathrm{Tr}[Z_{L}A^{\top}]\right \|_{2}^{2}$ $\displaystyle=\sum_{\ell=L_{1}+1}^{L}\left\|W_{\ell-1:L_{1}+1}Z_{L_{1}}A^{\top }W_{L:\ell+1}\right\|_{F}^{2}$ $$

Let $v_{1},\dots,v_{K}\in\mathbb{R}^{N}$ be $K$ orthonormal vectors
that span the preimage of $Z_{L}$ ($Z_{L}$ is rank $K$ as long
as $\epsilon_{1}\leq s_{K}(Y)$), and $e_{1},\dots,e_{K}\in\mathbb{R}^{K}$
be the standard basis of $\mathbb{R}^{K}$. We can then sum over $K^{2}$ possible choices of matrices $A=e_{i}v_{j}^{T}$ to obtain

$$ $\displaystyle K^{2}\|\Theta\|_{op}\geq\sum_{i,j=1}^{K}\left\|\nabla_{\theta}(e _{i}^{T}Z_{L}v_{j})\right\|_{2}^{2}$ $\displaystyle\geq\sum_{i,j=1}^{K}\sum_{\ell=L_{1}+1}^{L}\left\|W_{\ell-1:L_{1} +1}Z_{L_{1}}v_{j}\right\|_{2}^{2}\left\|e_{i}^{T}W_{L:\ell+1}\right\|_{2}^{2}.$ $$

We know use the fact that $P_{\mathrm{Im}W_{L:\ell}^{T}}W_{\ell-1:L_{1}+1}Z_{L_{1}}=W_{L:\ell}^{+}W_{L:
\ell}W_{\ell-1:L_{1}+1}Z_{L_{1}}=W_{L:\ell}^{+}Z_{L}$ to obtain the lower bound

$$ $\displaystyle K^{2}\|\Theta\|_{op}$ $\displaystyle\geq\sum_{i,j=1}^{K}\sum_{\ell=L_{1}+1}^{L}\left\|W_{L:\ell}^{+}Z _{L}v_{j}\right\|_{2}^{2}\left\|e_{i}^{T}W_{L:\ell+1}\right\|_{2}^{2}$ $\displaystyle=\sum_{\ell=L_{1}+1}^{L}\left\|W_{L:\ell}^{+}Z_{L}\right\|_{F}^{2 }\left\|W_{L:\ell+1}\right\|_{F}^{2}$ $\displaystyle\geq s_{K}(Z_{L})^{2}\sum_{\ell=L_{1}+1}^{L}\left\|W_{L:\ell}^{+} \right\|_{F}^{2}\frac{\left\|W_{L:\ell}\right\|_{F}^{2}}{\left\|W_{\ell}\right \|_{op}}$ $\displaystyle\geq\frac{\left(s_{K}(Y)-\epsilon_{1}\right)^{2}}{r^{2}}\sum_{ \ell=L_{1}+1}^{L}\kappa(W_{L:\ell})^{2}.$ $$

This then implies that

$$ $CL_{2}\geq\left\|\Theta\right\|_{op}\geq\frac{\left(s_{K}(Y)-\epsilon_{1} \right)^{2}}{K^{2}r^{2}}\sum_{\ell=L_{1}+1}^{L}\kappa(W_{L:\ell})^{2}.$ $$

Let us now assume by contradiction that for all the layers $\ell\in\{L_{1}+1,\dots,L_{1}+M\}$
we have $\kappa(W_{L:\ell})^{2}>\frac{CL_{2}K^{2}r^{2}}{M\left(s_{K}(Y)-\epsilon_{1}
\right)^{2}}$,
then

$$ $\frac{\left(s_{K}(Y)-\epsilon_{1}\right)^{2}}{K^{2}r^{2}}\sum_{\ell=L_{1}+1}^{ L}\kappa(W_{L:\ell})^{2}>CL_{2},$ $$

which yields a contradiction. Therefore, there must be
a layer within the layers $\ell\in\{L_{1}+1,\dots,L_{1}+M\}$ such that
$\kappa(W_{L:\ell})^{2}\leq\frac{CL_{2}K^{2}r^{2}}{M\left(s_{K}(Y)-\epsilon_{1}
\right)^{2}}$.

Furthermore, since $\kappa(W_{L:\ell})\geq 1$, we know that

$$ $\left\|\Theta\right\|_{op}\geq\frac{\left(s_{K}(Y)-\epsilon_{1}\right)^{2}}{K^ {2}r^{2}}L_{2}.$ $$

∎

###### Proof.

###### Proof.

###### Proof.

###### Proof.

###### Proof.

###### Proof.

## Appendix C Technical Results

###### Lemma C.1 .

Let $b\geq 1$ be s.t. $\|X_{:i}\|_{2}\leq b$ for all $i$. Then, inside the set of parameters with bounded
weights $\left\|W_{\ell}\right\|_{op}\leq r_{\ell}$ for all $\ell\in[L]$ and $r_{\ell}\geq 1$, the gradient
of the loss $\nabla C_{0}(\theta)$ is $5N\beta b^{3}\left(\prod_{j=1}^{L}r_{j}\right)^{3}L^{5/2}$-Lipschitz.

###### Proof.

Consider two parameters $\theta=(W_{\ell})_{\ell=1}^{L},\theta^{\prime}=(W_{\ell}^{\prime})_{\ell=1}^{L}$, and let $Z_{m},Z_{m}^{\prime}$ be the corresponding outputs at layer $m$. Then, we obtain the following telescopic sum:

$$ $\begin{split}(Z_{m})_{:i}-(Z_{m}^{\prime})_{:i}&=\sum_{\ell=1}^{L}(W_{m}\cdots W _{\ell+1}\circ\sigma\circ W_{\ell}\circ\sigma\circ W^{\prime}_{\ell-1}\cdots \circ W^{\prime}_{1})(X_{:i})\\ &-(W_{m}\cdots{}W_{\ell+1}\circ\sigma\circ W^{\prime}_{\ell}\circ\sigma\circ W ^{\prime}_{\ell-1}\circ\cdots\circ W^{\prime}_{1})(X_{:i}),\end{split}$ (29) $$

so that

$$ $\begin{split}\left\|(Z_{m})_{i:}-(Z_{m}^{\prime})_{i:}\right\|_{2}&\leq\sum_{ \ell=1}^{m}\prod_{j=1}^{m-1}r_{j}\left\|W_{\ell}-W^{\prime}_{\ell}\right\|_{F} \left\|X_{:i}\right\|_{2}\\ &\leq b\prod_{j=1}^{L-1}r_{j}\sqrt{m}\sqrt{\sum_{\ell=1}^{m}\left\|W_{\ell}-W^ {\prime}_{\ell}\right\|_{F}^{2}}\leq b\prod_{j=1}^{L-1}r_{j}\sqrt{m}\left\| \theta-\theta^{\prime}\right\|_{2}.\end{split}$ (30) $$

Now, the gradient equals

$$ $\nabla C_{0}(\theta)=\left(\sum_{i=1}^{N}\prod_{j=1}^{\ell-1}D_{j,i}W_{j}X_{:i }(Y_{:i}-(Z_{L})_{:i})^{\top}\prod_{j=\ell+1}^{L}W_{j}D_{j,i}\right)^{\top}_{ \ell=1,\dots,L}$ $$

where $D_{j,i}$ is a diagonal matrix with diagonal entries
given equal to the vector $\sigma^{\prime}((Z_{j})_{:i})$ for $j\in[L_{1}]$ and it is equal to the identity otherwise. Similarly, we define $D_{j,i}^{\prime}$ as a diagonal matrix with diagonal entries
given equal to the vector $\sigma^{\prime}((Z_{j}^{\prime})_{:i})$ for $j\in[L_{1}]$ and equal to the identity otherwise. Since $\sigma^{\prime}$
is $\beta$-Lipschitz, we have

$$ $\left\|D_{j,i}-D_{j,i}^{\prime}\right\|_{F}\leq\beta b\prod_{j=1}^{L-1}r_{j} \sqrt{m}\left\|\theta-\theta^{\prime}\right\|_{2}.$ $$

Furthermore since $\sigma^{\prime}\leq 1$ ,we have that $\left\|D_{j,i}\right\|_{op}\leq 1$.
By summing over all terms that need to be changed, we obtain

$$ $\displaystyle\left\|\nabla C_{0}(\theta)-\nabla C(\theta^{\prime})\right\|_{2}$ $\displaystyle\leq\sum_{\ell=1}^{L}\sum_{i=1}^{N}\left\|X_{:i}\right\|_{2}\left \|Y_{:i}-(Z_{L})_{:i}\right\|_{2}\sum_{m\neq\ell}\prod_{j\not\in\{m,L\}}r_{j} \left\|W_{m}-W^{\prime}_{m}\right\|_{F}$ $\displaystyle+\sum_{\ell=1}^{L}\sum_{i=1}^{N}\left\|X_{:i}\right\|_{2}\left\|Y _{:i}-(Z_{L})_{:i}\right\|_{2}\sum_{m=1}^{L_{1}}\left(\prod_{j=1}^{L-1}r_{j} \right)^{2}\beta\sqrt{m}\left\|\theta-\theta^{\prime}\right\|_{2}\left\|X_{:i} \right\|_{2}$ $\displaystyle+\sum_{\ell=1}^{L}\sum_{i=1}^{N}\left\|X_{:i}\right\|_{2}\left( \prod_{j=1}^{L-1}r_{j}\right)^{2}\sqrt{L}\left\|\theta-\theta^{\prime}\right\| _{2}\left\|X_{:i}\right\|_{2},$ $$

where the three terms correspond to the effect of changing the $W_{\ell}$’s,
the $D_{\ell}$’s and the $Z_{\ell}$’s respectively. This can
then be simplified to

$$ $\displaystyle\left\|\nabla C_{0}(\theta)-\nabla C_{0}(\theta)\right\|_{2}$ $\displaystyle\leq Nb^{2}\left(1+\prod_{j=1}^{L}r_{j}\right)\prod_{j=1}^{L-1}r_ {j}L^{3/2}\left\|\theta-\theta^{\prime}\right\|_{2}$ $\displaystyle+Nb^{3}\left(1+\prod_{j=1}^{L}r_{j}\right)\left(\prod_{j=1}^{L-1} r_{j}\right)^{2}LL_{1}^{3/2}\beta\left\|\theta-\theta^{\prime}\right\|_{2}$ $\displaystyle+Nb^{2}\left(\prod_{j=1}^{L-1}r_{j}\right)^{2}L^{3/2}\left\| \theta-\theta^{\prime}\right\|_{2}$ $\displaystyle\leq 5N\beta b^{3}\left(\prod_{j=1}^{L}r_{j}\right)^{3}L^{5/2} \left\|\theta-\theta^{\prime}\right\|_{2},$ $$

where we use the fact that $\left\|(Z_{L})_{:i}\right\|_{2}\leq b\prod_{j=1}^{L}r_{j}$.
∎

###### Lemma C.2 .

If the network satisfies

- •
approximate balancedness $\left\|W_{\ell+1}^{T}W_{\ell+1}-W_{\ell}W_{\ell}^{T}\right\|_{op}\leq\epsilon_
{2}$ for $\ell\in\{L_{1}+1,\cdots L-1\}$,
- •
bounded weights $\left\|W_{\ell}\right\|_{op}\leq r$ for $\ell\in\{L_{1}+1,\cdots,L\}$,

then we have $\left\|\left(W_{L}W_{L}^{\top}\right)^{L_{2}}-W_{L:L_{1}+1}W_{L:L_{1}+1}^{\top
}\right\|_{op}\leq\frac{L_{2}^{2}}{2}\epsilon_{2}r^{2(L_{2}-1)}$, where $L_{2}=L-L_{1}$.

###### Proof.

We denote $D_{\ell}:=W_{\ell}W_{\ell}^{\top}-W_{\ell+1}^{\top}W_{\ell+1}$ and

$$ $E_{\ell}=\sum_{i=L_{1}+1}^{\ell}\left(W_{\ell+1}^{\top}W_{\ell+1}\right)^{i-L_ {1}-1}D_{\ell}\left(W_{\ell}W_{\ell}^{\top}\right)^{\ell-i}.$ $$

We have that

$$ $\displaystyle\left(W_{\ell}W_{\ell}^{\top}\right)^{\ell-L_{1}}$ $\displaystyle=D_{\ell}\left(W_{\ell}W_{\ell}^{\top}\right)^{\ell-L_{1}-1}+W_{ \ell+1}^{\top}W_{\ell+1}\left(W_{\ell}W_{\ell}^{\top}\right)^{\ell-L_{1}-1}$ $\displaystyle=D_{\ell}\left(W_{\ell}W_{\ell}^{\top}\right)^{\ell-L_{1}-1}+W_{ \ell+1}^{\top}W_{\ell+1}D_{\ell}\left(W_{\ell}W_{\ell}^{\top}\right)^{\ell-L_{ 1}-2}+\left(W_{\ell+1}^{\top}W_{\ell+1}\right)^{2}\left(W_{\ell}W_{\ell}^{\top }\right)^{\ell-L_{1}-2}$ $\displaystyle=\dots$ $\displaystyle=\left(W_{\ell+1}^{\top}W_{\ell+1}\right)^{\ell-L_{1}}+\sum_{i=L_ {1}+1}^{\ell}\left(W_{\ell+1}^{\top}W_{\ell+1}\right)^{i-L_{1}-1}D_{\ell}\left (W_{\ell}W_{\ell}^{\top}\right)^{\ell-i}$ $\displaystyle=\left(W_{\ell+1}^{\top}W_{\ell+1}\right)^{\ell-L_{1}}+E_{\ell}.$ $$

Then,

$$ $\displaystyle W_{L:L_{1}+1}W_{L:L_{1}+1}^{\top}$ $\displaystyle=W_{L}\cdots W_{L_{1}+1}W_{L_{1}+1}^{\top}\cdots W_{L}^{\top}$ $\displaystyle=W_{L}\cdots W_{L_{1}+2}\left(W_{L_{1}+2}^{\top}W_{L_{1}+2}+E_{L_ {1}+1}\right)W_{L_{1}+2}^{\top}\cdots W_{L}^{\top}$ $\displaystyle=W_{L}\cdots W_{L_{1}+3}\left(W_{L_{1}+2}^{\top}W_{L_{1}+2}\right )^{2}W_{L_{1}+3}^{\top}\cdots W_{L}^{\top}+\mathcal{E}_{L_{1}+1}$ $\displaystyle=\dots$ $\displaystyle=\left(W_{L}W_{L}^{\top}\right)^{L_{2}}+\sum_{i=L_{1}+1}^{L-1} \mathcal{E}_{i},$ $$

where $\mathcal{E}_{i}=W_{L}\cdots W_{i+1}E_{i}W_{i+1}^{\top}\cdots W_{L}^{\top}$. As a result, we have

$$ $\left\|\left(W_{L}W_{L}^{\top}\right)^{L_{2}}-W_{L:L_{1}+1}W_{L:L_{1}+1}^{\top }\right\|_{op}=\left\|\sum_{i=L_{1}+1}^{L-1}\mathcal{E}_{i}\right\|_{op}\leq \sum_{i=L_{1}+1}^{L-1}\left\|\mathcal{E}_{i}\right\|_{op}.$ $$

Since $\|W_{k}\|_{op}\leq r$ for all $k=\{L_{1}+1,\cdots L\}$, we have $\|E_{\ell}\|_{op}\leq(\ell-L_{1})\epsilon_{2}r^{2(\ell-L_{1}-1)}$ and $\|\mathcal{E}_{\ell}\|_{op}\leq r^{2(L-\ell)}(\ell-L_{1})\epsilon_{2}r^{2(\ell
-L_{1}-1)}=(\ell-L_{1})\epsilon_{2}r^{2(L_{2}-1)}.$ Thus, we have

$$ $\left\|\left(W_{L}W_{L}^{\top}\right)^{L_{2}}-W_{L:L_{1}+1}W_{L:L_{1}+1}^{\top }\right\|_{op}\leq\frac{L_{2}^{2}}{2}\epsilon_{2}r^{2(L_{2}-1)}.$ $$

∎

###### Lemma C.1 .

###### Proof.

###### Lemma C.2 .

###### Proof.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: See Weaknesses for the most concerning questions. 
Additionaly:

Q1: Line 188: "... aproaches 1 ..." Shouldn't it be "... aproaches 2" based on (3)? Is it still sufficient for orthogonality claims (can the proof be adjusted to account for it)?

Q2: Proof sketch of Theorem 4.4., lines "306". Why and how is are "two phases" guaranteed to happen during GD training?
soundness: 3
strengths: S1: Novelty, pushing for (more) general results on timely and attractive topic of 'neural collapse' in DNN training
S2: Striving for a theoretically solid argumentation stating necessary assumptions (in Theorems and Propositions) and as Assumtions 4.1, 4.2, 4.3 ...
S3: Well done Introduction positioning paper within existing works (UFM, NTK and other specific results)
summary: Dear authors, thank you for submitting your work to ICLR 2025. The paper considers the phenomenon of 'neural collapse' and attempt to extend current rather special case results to brader class of networks, specifically, a deep neural (non-linear) networks with a wide first layer, funnel architecture and several, i.e. 2+, linear layers (head) before the output. After taking several assumptions, paper shows in series of Theorems (3.1, 4.4,5.2) and Propositions (4.5, 5.1,5.3) that GD training with weight decay (under further assumptions) leads to within class variability collapse (NC1). Results are supported by experiments on MNIST and CIFAR and MLP and ResNet + MLP head.
weaknesses: W1: Accessibility of the paper for a '15 mins' reader. The main results are formulated using (Abstract, l19) 'balancedness', (Abstract, l20) 'bounded conditioning' and other only later defined phrases and makes it hard to asses the attractivity/usefullness of paper unless reading it whole. It is recommended to rework (the Abstract at least, if no Conclusions are present) to make it clear and self-sustained.

W2: Quite a few assumptions are required in general. More over thay are added along the way, e.g., Assumption 4.2, $|\sigma(x)| \leq |x|$, etc., (which is ok, but also narrows the applicability of the final result). Some of those are very technical (especialy those in Theorems, such as in Theorem 3.1, (2),(3), (4)) but as well Assumption 4.3, Theorem 4.4. (10) and more) and with paper specific notation to learn, e.g., $\lambda_{3 \rightarrow L}$. It would help paper to have an thorough discussion on applicability/limitations of these assumptions. Perhaps at expense of shortening some 'proof sketches' refering them to SM?

W3: 'Discussion and Conclusion' sections are not presented. This is most likely due to space constrained and will be added in case of acceptance (?). Yet, I find it impacts the paper quality negatively. Especially in a light of the previous (W2) comments, the discussion and conclusions could have brought a critical 'back to reality' summary. Idealy it would bring more intuition for results and their application. Yet, I find it very important to have such Discussion reviewed before publishing ...

W4: Some references are rather inaccurately interpreted/generalized too much perhaps. For instance the lines 399-400 "...Thankfully, instead of simply diverging, for large η (but not too large) the parameters naturally end up at the ‘edge of stability’: the top eigenvalue of the Hessian is close to $2/\eta$ , i.e., the threshold below which GD is stable..." from (Cohen et all. 2021). Referenced work provides only experimental evidence for a phonomenon and only approximately, i.e., for certain settings operator norm of Hessian exceeds 'stability threshold' $2/\eta$, etc. Than approximation used on l410, especially $O(\epsilon_1)$ is only valid if $\nabla^2_{\theta} Z_L$ norm is bounded, which is ok for NTK regime, but not necessarily for large learning rate regime. Or is it?

W5: Following up on W4, Proposition 5.3, and other Theorem combine NTK with large learning rate regime, which sounds dangerous. Also requirement on wide first layer, suggest a NTK limit case is required. Could authors clarify a bit more on this relation?

Overall, I find it to be a solid attractive paper with technically legit reasoning, taking few shortcuts (some noted above) and with missing discussion and conclusions. I suggest authors to work on alleviating weaknesses and discussing limitations to improve contributions of this interesting work significantly.

---

confidence: 2
contribution: 4
initial_rating: 8
presentation: 3
questions: In Figure 3 the authors' numerical results show that non-linear layers are increasingly linear, as the depth of the non-linear part increases. Could the authors provide more insights into how this observation relates to their theoretical results and the mechanisms driving this increased linearity?
soundness: 4
strengths: This work provides an interesting theoretical insight on the role of training algorithm in the emergence of neural collapse, which I found especially exciting, and I think it opens up new directions for understanding the generalization properties of deep learning models.
summary: This paper presents an interesting theoretical advancement in understanding neural collapse in practical, end-to-end training scenarios, moving beyond the unconstrained features model. The authors provide a rigorous demonstration that neural collapse arises in networks with linear layers appended to a nonlinear backbone, given conditions of interpolation and balancedness. They show that these conditions hold for sufficiently wide networks trained with gradient descent and L2 regularization. The empirical results further support the theoretical findings, showcasing the robustness of neural collapse across various architectures and datasets.
weaknesses: In my opinion, this is a solid paper and I can not think of a weakness.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 1
questions: - What exactly is the role of these additional linear layers?  Are they only required for proving NC2 and NC3?
- Do these results suggest that neural networks with 1  non-layer and many linear layers can also exhibit neural collapse?
soundness: 3
strengths: - The authors provide the first theoretical analysis of neural collapse that does not rely on the unconstrained features model.
- The use of additional linear layers is novel and interesting technique. To the best of my knowledge such an architecture has not been studied.
- The results apply to deep neural networks trained with gradient descent (under a particular architecture/nonstandard activation function) as well as networks which are globally optimal for the weight decay regularized objective.
summary: This paper provides theoretical evidence for end-to-end training of true deep neural networks. This is contrary to previous works which primarily rely on the unconstrained features model. The paper provides explicit bounds on NC1, NC2 and NC3 for both globally optimal ($\ell_2$) regularized deep neural networks as well as neural networks trained with gradient descent. The results provide new insights on the role of additional linear layers, weight decay regularization and large learning rates with respect to neural collapse.
weaknesses: - My main concern with the paper is the writing. The theoretical statements are very difficult to parse this is especially the case for eqs 2-5 in Theorem 3.1. Can the authors provide a more interpretable version of this theorem (hide constants, big-O etc.) and delay the details to the appendix?
- I have the same concern about Theorem 4.4.

### Minor
- The proof of Theorem 3.1 in the Appendix could also be more clear. For example mentioning Weyl's inequality in the first inequality of (21)
- In the proof of Theorem 5.2 at the Appendix (line 1264) shouldn't it be
$$\kappa(W_L) = \kappa(W_{L:L_1+1})^{\frac{1}{L_2}}$$
not 
$$\kappa(W_L) = \kappa(W_{L:L_1+1})^{\frac{1}{L_1}}$$
- Same concern about $L_1$ vs $L_2$ in the statement of Theorem 5.2.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: 1. Are there no assumptions on the activation function $\sigma$ needed for Theorem 3.1?  None are stated in Section 3, although assumptino 4.2 appears later.  It seems odd that a potentially discontinuous and horribly behaved activation function would be permitted (I imagine Lipschitz is required?)

2. Assumption 4.2 seems more like a smooth leaky relu, which has appeared in prior work [Chatterjee arXiv:2203.16462, Frei et al COLT 2022]

3. The authors talk about (NC1)-(NC3), what about (NC4) from the original Papyan paper?

4. Are there any important differences between the proof of Theorem 4.4 and the proofs/results in Nguyen and Mondelli?   I'm not familiar with that work, but it seems like it's not super necessary to have details about optimization (e.g. the PL inequality) in the main section of the paper, it distracts from what I think are the more interesting and stronger results elswhere in it (it would also give additional space to have a proper conclusion and outline what the authors think are future directions).   I am assuming balancedness didn't appear in the prior analysis but does here.  A few sentences describing high-level differences between proofs and ideas would be helpful.  

5. Also, to be clear, the optimization analysis in Sec 4 is in the NTK regime right?  I didn't see this explicitly mentioned but it should be if it is known.
soundness: 4
strengths: The authors provide a novel characterization of conditions under which neural collapse can provably occur.  I am not an expert on the NC theory literature, but to my knowledge, no prior work identified balancedness and interpolation as being key factors which can allow for this to occur, and this is a pretty strong finding.   Since the balancedness of post-fixed linear layers is the strongest condition (as most common neural nets do not have multiple linear layers at the end, only a single one), the findings in Section 5 about how boundedness and interpolation can suffice for NC2+NC3 are also a nice addition.  The numerical findings nicely complement their theoretical ones.
summary: The authors identify a set of conditions under which various neural collapse phenomena provably occur in deep neural nets.  They consider neural nets which have a sequence of nonlinear layers and then a sequence of linear layers.  They find that approximate interpolation, weight balanced-ness, and boundedness suffice for deriving various neural collapse phenomena.  They then show that GD on networks with a "pyramidal" overparameterized topology (i.e., first width is >= number samples, remaining widths are decreasing), under suitable initialization and regularization, allow for one of the neural collapse phenomena to hold.  They then identify conditions (near-interpolation and small-norm) which ensure that global minimizers of the loss can satisfy all of the neural collapse phenomena.   Finally, they look at the neural collapse phenomena  from the perspective of the edge of stability via an analysis of the properties of the hessian under EoS assumptions.
weaknesses: There aren't any serious weaknesses to me.  The strongest assumptions--namely the need for multiple linear layers in order for the NC phenomenon to occur via the main theorem--seem necessary per experiments in Figure 1.  So it seems that these assumptions are strong for a fundamental reason.

The pyramidal structure assumption is a bit odd/strong, but only seems needed for the optimization result, which I don't think is the central point of the paper.  

The authors don't provide a conclusion or discussion section, and it would have been useful to have comments from the authors about what they think are the weakest parts of their work and what work should be prioritized in the future.  I think they could get some additional space by removing some of the details about optimization from Sec 4 since I don't think they're super important.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: I am wondering whether the layer balanced-ness property after training can be proved as a direct consequence of [1] in the author's setting.

[1] Du, Simon S., Wei Hu, and Jason D. Lee. "Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced." Advances in neural information processing systems 31 (2018).
soundness: 3
strengths: This work goes beyond the conventional unconstrained feature model (which many previous works have worked on) and proves that neural collapse can happen after training for full-connected neural networks satisfying some architectural conditions such as pyramidal topology and smooth activation. The conditions established in Theorem 3.1 of this work under which neural collapse can happen seem very reasonable. Indeed, later, the authors proved that those conditions can be satisfied by training a neural network via gradient descent which automatically implies that neural collapse can happen after training. 

I am not very familiar with the prior works along this line of research such as (Kothapalli & Tirer, 2024), (Hong & Ling, 2024) and many others mentioned in the introduction and related work. Based on my knowledge on neural collapse, I think this work made some interesting contributions towards understanding this phenomenon.
summary: This paper studies neural collapse phenomenon in training wide neural networks. The first result in this work establishes some general conditions (interpolation, balancedness between layers, well-conditioned-ness of weight matrices) such that NC1, NC2, and NC3 can hold. The second result considers training a wide neural network with gradient descent such that the aforementioned conditions hold after training, which implies neural collapse can happen after training.
weaknesses: 1. The analysis critically relies on the fact that the last two layers of the neural network are linear. I can definitely see this condition makes the problem a lot easier to analyze. I am wondering how hard it is to remove such restrictions. 
2. It seems the analysis of Theorem 4.4 relies on the neural network in the NTK regime, as the pyramidal topology assumption has appeared in previous works such as (Nguyen & Mondelli, 2020). I don't regard this as a major weakness even if it turned out to be true that the networks are in the NTK regime given the contribution of this work, however, I do appreciate clarification on this.

Meta-review: This paper explores neural collapse, a geometric structure observed in the last layer of deep neural networks (DNNs) at convergence, where training data is consistently represented. Moving beyond the unconstrained features model, the authors study DNNs with at least two linear layers and establish conditions for neural collapse, including low training error, balancedness of linear layers, and bounded conditioning of pre-linear features. They prove these conditions hold during gradient descent training with weight decay, particularly in networks with wide first layers and stable solutions. The authors claim that this work provides the first theoretical demonstration of neural collapse in end-to-end DNN training.

The reviewers raised the following strengths and weaknesses

Pros:

+ First to prove neural collapse in end-to-end DNN training without unconstrained features this paper introduces key conditions like balancedness and bounded conditioning.

+ This paper validates results with experiments on MNIST and CIFAR; aligns theory with training practices like gradient descent and weight decay.

+ The paper highlights the role of linear layers, large learning rates, and weight decay in achieving neural collapse.


Cons:

- Accessibility Issues: Dense initial presentation; abstract lacks clarity for general readers.

- Strong Assumptions: Relies on two linear layers, specific initialization, and pyramidal topology, limiting general applicability.

- Missing Aspects: No discussion of NC4; broader implications and applications were underexplored.

Some of these concerns were assuaged by the rebuttal of the authors and all reviewers are in favor of acceptance. I concur

Decision: Accept (Oral)