SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\left\{\begin{array}{l}\mathbf {H} _ {1} ^ {\top} \mathbf {H} _ {1} \rightarrow \sqrt {\frac {\tilde {\lambda} _ {c}}{m}} \left(\mathbb {I} _ {C} - \beta \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right)\\\mathbf {W} \mathbf {W} ^ {\top} \rightarrow \sqrt {\frac {m}{\tilde {\lambda} _ {c}}} \left(\mathbb {I} _ {C} - \beta \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right).\end{array}\right. \tag {58}
$$
To establish NC2, recall that  $\mathbf{H}_1 = \left[\langle h\rangle_1,\dots ,\langle h\rangle_C\right]$  and that  $\mathbf{M}$ , as a normalized version of  $\mathbf{H}_1$ , satisfies
$$
\mathbf {M} ^ {\top} \mathbf {M} \rightarrow \frac {1}{1 - \beta} (\mathbb {I} _ {C} - \beta \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}) = \frac {C}{C - 1} (\mathbb {I} _ {C} - \frac {1}{C} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}).
$$
To establish NC3, note that from (55) and (58) together, it follows that the limits of  $\mathbf{M}$  and  $\mathbf{W}^{\top}$  only differ by a constant multiplier.
To establish NC4, note that using NC3 we can write
$$
\begin{array}{l} \underset {c} {\operatorname {a r g m a x}} (\mathbf {W} h (x) + \mathbf {b}) _ {c} = \underset {c} {\operatorname {a r g m a x}} (\mathbf {W} h (x)) _ {c} \quad (\mathbf {b} = \beta \mathbf {1} _ {C}) \\ \rightarrow \underset {c} {\operatorname {a r g m a x}} (\mathbf {M} ^ {\top} h (x)) _ {c} \tag {NC3} \\ = \operatorname * {a r g m i n} _ {c} \| h (x) - \langle h \rangle_ {c} \| _ {2}. \\ \end{array}
$$
B.4 General biases case
Proof. As in the proof of Theorem 5.2, at the end of training we have  $\mathbf{W}\mathbf{H}_1 + \mathbf{b}\mathbf{1}_C^\top = \mathbb{I}_C$ . Moreover, since  $\mathbf{E} = \mathbb{O}$  and  $\mathbf{H}_2 \to \mathbb{O}$ , we have
$$
\frac {1}{m} \mathbf {W} ^ {\top} \mathbf {W} - \frac {1}{\mu_ {\text {c l a s s}}} \mathbf {H} _ {1} \left(\mathbb {I} - \alpha \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right) \mathbf {H} _ {1} ^ {\top} \rightarrow \mathbb {O}. \tag {59}
$$
Multiplying the above expression to the left by  $\mathbf{W}$  and to the right by  $\mathbf{W}^{\top}$ , we obtain the general expression (29) for the matrix  $(\mathbf{WW}^{\top})^{2}$  mentioned in the main text:
$$
\left.\left(\mathbf {W} \mathbf {W} ^ {\top}\right) ^ {2} \rightarrow \frac {m}{\mu_ {\text {c l a s s}}} \left(\mathbb {I} _ {C} - \alpha \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top} + (1 - \alpha C) \left(C \mathbf {b b} ^ {\top} - \mathbf {b 1} _ {C} ^ {\top} - \mathbf {1} _ {C} \mathbf {b} ^ {\top}\right)\right). \right. \tag {60}
$$
This expression implies that the rows of the weights matrix may have varying separation angles in the general biases case, i.e., there is no symmetric structure is general. However, for constant biases  $\mathbf{b} = \beta \mathbf{1}_C$ , the above expression simplifies to
$$
\left(\mathbf {W} \mathbf {W} ^ {\top}\right) ^ {2} \rightarrow \frac {m}{\mu_ {\text {c l a s s}}} \left(\mathbb {I} _ {C} - \frac {1}{C} (1 - (1 - \alpha C) (1 - \beta C) ^ {2}) \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right). \tag {61}
$$
Since  $\alpha < 1 / C$  and  $(1 - \beta C)^2\geq 0$ , we have that  $(1 - (1 - \alpha C)(1 - \beta C)^2) / C\leq 1 / C$ . Therefore, the RHS of (61) is always p.s.d. and has a unique p.s.d square root proportional to  $\mathbb{I}_C - \gamma \mathbf{1}_C\mathbf{1}_C^\top$  for some constant  $\gamma < 1 / C$ . Denote  $\rho \coloneqq (1 - (1 - \alpha C)(1 - \beta C)^2) / C$ , then we have  $\gamma = (1 - \sqrt{1 - C\rho}) / C$ . Note that  $\rho < 1 / C$  ensures that  $\gamma$  is well defined. Then the configuration of the final weights is given by
$$
\mathbf {W} \mathbf {W} ^ {\top} \rightarrow \sqrt {\frac {m}{\mu_ {\text {c l a s s}}}} \left(\mathbb {I} _ {C} - \gamma \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right). \tag {62}
$$
This means that the norms of all the weights rows are still equal, as in NC2. However, since  $\gamma < 1 / C$  if  $\beta \neq 1 / C$ , the angle between these rows is smaller than in the ETF structure.
We can derive the configuration of the class means similarly by multiplying (59) to the left by  $\mathbf{H}_1^\top$  and to the right by  $\mathbf{H}_1$ . In the general biases case, we get
$$
\mathbf {H} _ {1} ^ {\top} \mathbf {H} _ {1} \left(\mathbb {I} _ {C} - \alpha \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right) \mathbf {H} _ {1} ^ {\top} \mathbf {H} _ {1} \rightarrow \frac {\mu_ {\text {c l a s s}}}{m} \left(\mathbb {I} _ {C} - \mathbf {b} \mathbf {1} _ {C} ^ {\top} - \mathbf {1} _ {C} \mathbf {b} ^ {\top} + \| \mathbf {b} \| _ {2} ^ {2} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right). \tag {63}
$$
As with the weights, we see that this is not a symmetric structure in general. Thus, NC2 does not hold in the general biases case. However, for the constant biases  $\mathbf{b} = \beta \mathbf{1}_C$ , the above expression simplifies to
$$
\mathbf {H} _ {1} ^ {\top} \mathbf {H} _ {1} \left(\mathbb {I} _ {C} - \alpha \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right) \mathbf {H} _ {1} ^ {\top} \mathbf {H} _ {1} \rightarrow \frac {\mu_ {\text {c l a s s}}}{m} \left(\mathbb {I} _ {C} - \beta \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right) ^ {2}. \tag {64}
$$
Analogously to the previous derivations, we get that the unique p.s.d. square root of the RHS is given by  $\mathbb{I}_C - \tilde{\rho}\mathbf{1}_C\mathbf{1}_C^\top$ , where  $\tilde{\rho} \coloneqq (1 - |1 - \beta C|) / C < 1 / C$  for  $\beta \neq 1 / C$ . On the other hand, the unique p.s.d root of  $\mathbb{I} - \alpha \mathbf{1}_C\mathbf{1}_C^\top$  is given by  $\mathbb{I}_C - \phi \mathbf{1}_C\mathbf{1}_C^\top$ , where  $\phi \coloneqq (1 - \sqrt{1 - \alpha C}) / C$ . Thus, we have the following
$$
\sqrt {\frac {m}{\mu_ {\text {c l a s s}}}} \mathbf {H} _ {1} ^ {\top} \mathbf {H} _ {1} \left(\mathbb {I} _ {C} - \phi \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right)\rightarrow \mathbb {I} _ {C} - \tilde {\rho} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}. \tag {65}
$$
Therefore, the structure of the last-layer features class means is given by
$$
\mathbf {H} _ {1} ^ {\top} \mathbf {H} _ {1} \rightarrow \sqrt {\frac {\mu_ {\text {c l a s s}}}{m}} \left(\mathbb {I} _ {C} - \tilde {\rho} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right)\left(\mathbb {I} _ {C} - \frac {\phi}{1 + \phi C} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right) = \sqrt {\frac {\mu_ {\text {c l a s s}}}{m}} \left(\mathbb {I} _ {C} - \theta \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right), \tag {66}
$$
where  $\theta := \tilde{\rho} + \phi / (1 + \phi C) - C \tilde{\rho} \phi / (1 + \phi C) < 1 / C$  for  $\beta \neq 1 / C$ . Thus, similarly to the classifier weights  $\mathbf{W}$ , the last-layer features class means form a symmetric structure with equal lengths and a separation angle smaller than in the ETF. However, the centralized class means given by  $\mathbf{M} = \mathbf{H}_1(\mathbb{I}_C - \mathbf{1}_C \mathbf{1}_C^T / C)$  still form the ETF structure:
$$
\mathbf {M} ^ {\top} \mathbf {M} \rightarrow \sqrt {\frac {\mu_ {\text {c l a s s}}}{m}} \left(\mathbb {I} _ {C} - \frac {1}{C} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right). \tag {67}
$$
This holds since the component proportional to  $\mathbf{1}_C\mathbf{1}_C^\top$  on the RHS of equation (66) lies in the kernel of the ETF matrix  $(\mathbb{I}_C - \mathbf{1}_C\mathbf{1}_C^\top /C)$ . Thus, we conclude that NC2 holds in case of equal biases, while NC3 does not.
Remark on  $\alpha \to 0$  case: Simplifying the expressions for constants  $\gamma$  and  $\theta$ , which define the angles in the configurations of the weights and the class means above, we get the following:
$$
\gamma = \frac {1}{C} (1 - | 1 - \beta C | \sqrt {1 - \alpha C}), \quad \theta = \frac {1}{C} (1 - \frac {| 1 - \beta C |}{2 - \sqrt {1 - \alpha C}}). \tag {68}
$$
Analyzing these expressions, we find that they are equal only if  $1 - \alpha C = 1$ , i.e.  $\alpha = 0$ . However, this can only hold if  $\kappa_{n} = 0$  by definition of  $\alpha$ , i.e., when the kernel  $\Theta^{h}$  is zero on pairs of samples from different classes. While  $\alpha \neq 0$  in general, there are certain settings where  $\alpha$  approaches zero. Simplifying the expression for  $\alpha$ , we can get the following
$$
\alpha = \frac {1}{\frac {\kappa_ {c}}{\kappa_ {n}} \left(1 - \frac {1}{m}\right) + \frac {\kappa_ {d}}{\kappa_ {n}} \frac {1}{m} + (C - 1)}. \tag {69}
$$
One can see that  $\alpha \rightarrow 0$  if  $C\to \infty$  or when  $\kappa_c / \kappa_n\to \infty$ . Since the kernel  $\Theta^h$  is strongly aligned with the labels in our numerical experiments, the value of  $\kappa_{c} / \kappa_{n}$  is large in practice. Thus,  $\alpha$  is not zero but indeed significantly smaller than  $1 / C$ . Thus, in our numerical experiments the angles  $\theta$  and  $\gamma$  are close to each other. However, we note that the equality of these two angles does not imply NC3, since the value of  $\theta$  characterizes the angles between the non-centralized class means.
Remark on  $\alpha \to 1 / C$  case: If  $\alpha = 1 / C$ , the equation (63) for the structure of the features class means with general (not equal) biases simplifies to
$$
\mathbf {M} ^ {\top} \mathbf {M} \rightarrow \frac {\mu_ {\text {c l a s s}}}{m} \left(\mathbb {I} _ {C} - \frac {1}{C} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right), \tag {70}
$$
i.e., in this case the class means always exhibit the ETF structure, even without the assumption that all the biases are equal. Moreover, in this case  $\gamma = 1 / C$  as well. Thus, both NC2 and NC3 hold. While by definition  $\alpha < 1 / C$ , we can analyze the cases when it approaches  $1 / C$  using the expression (69) again. One can see that when  $m\to \infty$  and  $\kappa_c / \kappa_n\rightarrow 1$ , we have  $\alpha \rightarrow 1 / C$ . However, the requirement  $\kappa_{c} / \kappa_{n}\rightarrow 1$  implies that the kernel  $\Theta^h$  does not distinguish between pairs of samples from the same class and from different classes. Such a property of the kernel is associated with poor generalization performance and does not occur in our numerical experiments.
C Numerical experiments
Implementation details We use JAX [8] and Flax (neural network library for JAX) [25] to implement all the DNN architectures and the training routines. This choice of the software allows to compute the empirical NTK of any DNN architecture effortlessly and efficiently. We compute the values of kernels  $\Theta$  and  $\Theta^h$  on the whole training batch ( $m = 12$  samples per class, 120 samples in total) in case of ResNet20 and DenseNet40 to approximate the values  $(\gamma_d, \gamma_c, \gamma_n)$  and  $(\kappa_d, \kappa_c, \kappa_n)$ , as well as the NTK alignment metrics, and compute the invariant  $\mathbf{E}$  using these values. Since VGG11 and VGG16 architectures are much larger (over 10 million parameters) and computing their Jacobians is very memory-intensive, we use  $m = 4$  samples per class (i.e., 40 samples in total) to approximate the kernels of these models. We compute all the other training metrics displayed in panes a-e of Figures 3, 4, 5, 6, 7, 8, 9, 10, 11 on the whole last batch of every second training epoch for all the architectures. The test accuracy is computed on the whole test set. To produce panes f-h of the same figures, we only compute the NC metrics and the test accuracy one time after 400 epochs of training for every learning rate. We use 30 logarithmically spaced learning rates in the range  $\eta \in [10^{-4}, 10^{0.25}]$  for ResNet20 trained on MNIST and VGG11 trained on MNIST. For all the other architecture-dataset pairs we only compute the last 20 of these learning rates to reduce the computational costs, since the smallest learning rates do not yield models with acceptable performance.
Compute We executed the numerical experiments mainly on NVIDIA GeForce RTX 3090 Ti GPUs, each model was trained on a single GPU. In this setup, a single training run displayed in panes a-e of Figures 3, 4, 5, 6, 7, 8, 9, 10, 11 took approximately 3 hours for ResNet20, 6 hours for DenseNet40, 7 hours for VGG11, and 11 hours for VGG16. This adds up to a total of 312 hours to compute panes a-e of the figures. The computation time is mostly dedicated not to the training routine itself but to the large number of computationally-heavy metrics, which are computed every second epoch of a training run. Indeed, to approximate the values of  $\Theta$  and  $\Theta^h$ , one needs to compute  $C(C + 1) + n(n + 1)$  kernels on a sample of size  $mC$  from the dataset, and each of the kernels requires computing a
gradient with respect to numerous parameters of a DNN. Additionally, the graphs in panes f-h of the same figures take around 1.5 hours for each learning rate value for ResNet20, 3 hours for DenseNet40, and 4 hours for VGG11 and VGG16, which adds up to approximately 1350 computational hours.
Results We include experiments on the following architecture-dataset pairs:
Figure 3: VGG11 trained on MNIST
Figure 4: VGG11 trained on FashionMNIST
Figure 5: VGG16 trained on CIFAR10
Figure 6: ResNet20 trained on MNIST
Figure 7: ResNet20 trained on FashionMNIST
Figure 8: ResNet20 trained on CIFAR10
Figure 9: DenseNet40 trained on MNIST
Figure 10: DenseNet40 trained on FashionMNIST
Figure 11: DenseNet40 trained on CIFAR10
The experiments setup is described in Section 6. Panes a-h of Figures 3, 4, 5, 6, 7, 8, 9, 10, 11 are analogous to the same panes of Figure 2. We include additional pane i here, which displays the norms of the invariant terms corresponding to the feature matrix components  $\mathbf{H}_1$  and  $\mathbf{H}_2$ , and the global features mean  $\langle h\rangle$  at the end of training. One can see that the global features mean is relatively small in comparison with the class-means in every setup, and the "variance" term  $\mathbf{H}_2$  is small for models that exhibit NC. We also add pane j, which displays the alignment of kernels  $\Theta$  and  $\Theta^h$  for every model at the end of training. One can see that the kernel alignments is typically stronger in models that exhibit NC.
C.1 Additional examples of the NTK block structure
We include the following additional illustrative figures (analogous to Figure 1 in the main text) that show the NTK block structure in dataset-architecture pairs covered in our experiments:
Figure 13: VGG11 trained on MNIST
Figure 14: VGG11 trained on FashionMNIST
Figure 15: VGG16 trained on CIFAR10
Figure 16: ResNet20 trained on FashionMNIST
Figure 17: ResNet20 trained on CIFAR10
Figure 18: DenseNet40 trained on MNIST
Figure 19: DenseNet40 trained on FashionMNIST
Figure 11: DenseNet40 trained on CIFAR10
Overall, the block structure pattern is visible in the traced kernels in all the figures. As expected, the block structure is more pronounced in the kernels where the final alignment values are higher. While the norms of the "non-diagonal" components of the kernels are generally smaller than the "diagonal" components in panes c) and d), we notice that there is a large variability in the norms of the "diagonal" components in some settings. This means that different neurons of the penultimate layer and different classification heads may contribute to the kernel unequally in some settings. Moreover, certain "non-diagonal" components of the last-layer kernel may have non-negligible effect in some settings. We discuss how one could generalize our analysis to account for these properties of the NTK in Appendix D.
C.2 Preliminary experiments with CE loss
While CE loss is a common choice for training DNN classifiers, our theoretical analysis and the experimental results only cover DNNs trained with MSE loss. For completeness, we provide experimental results for ResNet20 trained on MNIST with CE loss in Figure 12. One can see that
smaller invariant norm and higher invariant alignment correlate with NC in the figure. However, DNNs trained with CE loss overall reach better NC metrics but have much larger norm of the invariant in comparison with DNNs trained with MSE loss.
D Relaxation of the NTK Block-Structure Assumption
In this section, we first derive the dynamics equations of DNNs with a general block structure assumption on the last-layer kernel  $\Theta^h$  (analogous to the equations presented in Theorem 4.1 and Theorem 5.1). Then we discuss a possible relaxation of Assumption 3.2, under which our main result regarding NC in Theorem 5.2 still holds.
D.1 Dynamics under General Block Structure Assumption
We first formulate the most general form of the block structure assumption on  $\Theta^h$  as follows:
Assumption D.1. Assume that  $\Theta^h:\mathcal{X}\times \mathcal{X}\to \mathbb{R}^{n\times n}$  has the following block structure
$$
\Theta^ {h} (x, x) = \mathbf {A} _ {d} + \mathbf {A} _ {c} + \mathbf {A} _ {n}, \quad \Theta^ {h} \left(x _ {i} ^ {c}, x _ {j} ^ {c}\right) = \mathbf {A} _ {c} + \mathbf {A} _ {n}, \quad \Theta^ {h} \left(x _ {i} ^ {c}, x _ {j} ^ {c ^ {\prime}}\right) = \mathbf {A} _ {n}, \tag {71}
$$
where  $\mathbf{A}_{d,c,n} \in \mathbb{R}^{n \times n}$  are arbitrary p.s.d. matrices. Here  $x_{i}^{c}$  and  $x_{j}^{c}$  are two distinct inputs from the same class, and  $x_{j}^{c'}$  is an input from class  $c' \neq c$ .
This assumption means that every kernel matrix  $\Theta_{k,s}^{h}(X), k, s \in [1,n]$  still has at most three distinct values, corresponding to the inter-class, intra-class, and the diagonal values of the kernel. However, these values are arbitrary and may depend on the choice of  $k, s \in [1,n]$ .
Under the general block structure assumption, the gradient flow dynamics of DNNs with MSE loss takes the following form:
$$
\left\{ \begin{array}{l l} \dot {\mathbf {H}} = & - \mathbf {A} _ {d} \mathbf {W} ^ {\top} \mathbf {R} + m \mathbf {A} _ {c} \mathbf {W} ^ {\top} \mathbf {R} _ {\text {c l a s s}} + N \mathbf {A} _ {n} \mathbf {W} ^ {\top} \mathbf {R} _ {\text {g l o b a l}} \\ \dot {\mathbf {W}} = & - \mathbf {R} \mathbf {H} ^ {\top} \\ \dot {\mathbf {b}} = & - \mathbf {R} _ {\text {g l o b a l}} \mathbf {1} _ {N}. \end{array} \right. \tag {72}
$$
This is the generalized version of the dynamics presented in Theorem 4.1. Consequently, the decomposed dynamics presented in Theorem 5.1 takes the following form under the general block structure assumption:
$$
\left\{ \begin{array}{l l} \dot {\mathbf {H}} _ {1} & = - \left(\mathbf {A} _ {d} + m \mathbf {A} _ {c}\right) \mathbf {W} ^ {\top} \mathbf {R} _ {1} - m \mathbf {A} _ {n} \mathbf {W} ^ {\top} \mathbf {R} _ {1} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top} \\ \dot {\mathbf {H}} _ {2} & = - \mathbf {A} _ {d} \mathbf {W} ^ {\top} \mathbf {W} \mathbf {H} _ {2} \\ \dot {\mathbf {W}} & = - m \left(\mathbf {R} _ {1} \mathbf {H} _ {1} ^ {\top} + \mathbf {W} \mathbf {H} _ {2} \mathbf {H} _ {2} ^ {\top}\right) \\ \dot {\mathbf {b}} & = - m \mathbf {R} _ {1} \mathbf {1} _ {C}. \end{array} \right. \tag {73}
$$
The derivation of the above dynamics equations are identical to the proofs of Theorem 4.1 and Theorem 5.1 presented in Appendix B.
Rotation invariance We notice that the dynamics of  $(\mathbf{W},\mathbf{H})$  in (72) has to be rotation invariant, i.e., the equations should not be affected by a change of variables  $\mathbf{W}\rightarrow \mathbf{WQ}$ ,  $\mathbf{H}\rightarrow \mathbf{Q}^{\top}\mathbf{H}$  for any orthogonal matrix  $\mathbf{Q}$ . This holds since the loss function only depends on the product  $\mathbf{WH}$ , which does not change under rotation. This requirement puts conditions on the behavior of  $\mathbf{A}_{d,c,n}$  under rotation. Indeed, assume that the rotation  $\mathbf{W}\rightarrow \mathbf{WQ}$ ,  $\mathbf{H}\rightarrow \mathbf{Q}^{\top}\mathbf{H}$  for some  $\mathbf{Q}$  corresponds to the following change of the kernel:
$$
\mathbf {A} _ {d, c, n} \rightarrow \tilde {\mathbf {A}} _ {d, c, n} (\mathbf {Q}), \tag {74}
$$
then the rotation invariance of the dynamics implies the following equality for any  $\mathbf{Q}$ :
$$
\begin{array}{l} \mathbf {Q} \tilde {\mathbf {A}} _ {d} (\mathbf {Q}) \mathbf {Q} ^ {\top} \mathbf {W} ^ {\top} \mathbf {R} + m \mathbf {Q} \tilde {\mathbf {A}} _ {c} (\mathbf {Q}) \mathbf {Q} ^ {\top} \mathbf {W} ^ {\top} \mathbf {R} _ {\text {c l a s s}} + N \mathbf {Q} \tilde {\mathbf {A}} _ {n} (\mathbf {Q}) \mathbf {Q} ^ {\top} \mathbf {W} ^ {\top} \mathbf {R} _ {\text {g l o b a l}} (75) \\ = \mathbf {A} _ {d} \mathbf {W} ^ {\top} \mathbf {R} + m \mathbf {A} _ {c} \mathbf {W} ^ {\top} \mathbf {R} _ {\text {c l a s s}} + N \mathbf {A} _ {n} \mathbf {W} ^ {\top} \mathbf {R} _ {\text {g l o b a l}}. (76) \\ \end{array}
$$
These equations are satisfied trivially with our initial assumption, where  $\mathbf{A}_{d,c,n} = \tilde{\mathbf{A}}_{d,c,n}(\mathbf{Q})\propto \mathbb{I}_n$  However, as we can see, any generalized assumption should specify the behavior of the kernel under rotation, and satisfy the above equation.
For general  $\mathbf{A}_{d,c,n}$ , the following behavior under rotation trivially satisfies the above condition:  $\tilde{\mathbf{A}}_{d,c,n}(\mathbf{Q}) = \mathbf{Q}^\top \mathbf{A}_{d,c,n}\mathbf{Q}$ . This behaviour of the kernel under rotation is intuitive, since it implies that the gradients of the last-layer features  $h$  are rotated in the same way as the features. However, we note that gradients of parametrized functions do not in general behave this way, since the rotation of the function has to be realized by a certain change of parameters. Consider, for instance, a one-hidden-layer linear network with weights  $\mathbf{V}$  in the first layer. Then we have  $\mathbf{H} = \mathbf{V}X$ , and a rotation  $\mathbf{H} \to \mathbf{Q}^\top \mathbf{H}$  corresponds to the change of parameters  $\mathbf{V} \to \mathbf{Q}^\top \mathbf{V}$ . In this case, the kernel does not change under rotation, i.e.,  $\tilde{\mathbf{A}}_{d,c,n}(\mathbf{Q}) = \mathbf{A}_{d,c,n}$ .
Dynamics invariant We note that the dynamics in 73 does not in general have an invariant analogous to the one we identified in Theorem 5.1. Indeed, if we define a quantity  $\mathbf{E} := \mathbf{W}^{\top}\mathbf{W} - c_{1}\mathbf{H}_{1}\mathbf{H}_{1}^{\top} - c_{2}\mathbf{H}_{2}\mathbf{H}_{2}^{\top}$  for some constants  $c_{1,2} \in \mathbb{R}$ , and additionally assume centered global means  $\mathbf{H}_1\mathbf{1}_C = 0$ , we get the following expression for the derivative of  $\mathbf{E}$ :
$$
\begin{array}{l} \dot {\mathbf {E}} = \left(c _ {1} \left(\mathbf {A} _ {d} + m \mathbf {A} _ {c}\right) - m \mathbb {I} _ {n}\right) \mathbf {W} ^ {\top} \mathbf {R} _ {1} \mathbf {H} _ {1} ^ {\top} - \mathbf {H} _ {1} \mathbf {R} _ {1} ^ {\top} \mathbf {W} \left(c _ {1} \left(\mathbf {A} _ {d} + m \mathbf {A} _ {c}\right) ^ {\top} - m \mathbb {I} _ {n}\right) (77) \\ + \left(c _ {2} \mathbf {A} _ {d} - m \mathbb {I} _ {n}\right) \mathbf {W} ^ {\top} \mathbf {W} \mathbf {H} _ {2} \mathbf {H} _ {2} ^ {\top} - \mathbf {H} _ {2} \mathbf {H} _ {2} ^ {\top} \mathbf {W} ^ {\top} \mathbf {W} \left(c _ {2} \mathbf {A} _ {d} ^ {\top} - m \mathbb {I} _ {n}\right), (78) \\ \end{array}
$$
which is not equal to zero with arbitrary matrices  $\mathbf{A}_{d,c}$ .
D.2 Neural Collapse under Relaxed Block Structure Assumption
We now propose a relaxation of our main assumption, under which our main result regarding NC in Theorem 5.2 still holds. In terms of Assumption D.1 on the general block structure of  $\Theta^h$ , our initial Assumption 3.2 in the main text is the special case with  $\mathbf{A}_n = \kappa_n\mathbb{I}_n$ ,  $\mathbf{A}_c = (\kappa_c - \kappa_n)\mathbb{I}_n$ ,  $\mathbf{A}_d = (\kappa_d - \kappa_c)\mathbb{I}_n$ . The relaxed assumption can be formulated as follows in terms of matrices  $\mathbf{A}_{d,c,n}$ :
Assumption D.2. Assume that  $\mathbf{A}_n$  is an arbitrary p.s.d. matrix and  $(\mathbf{A}_c, \mathbf{A}_d)$  satisfy the following conditions:
$$
\mathbf {A} _ {c} = \kappa_ {c} \mathbb {I} _ {n} + \mathbf {N} _ {c}, \mathbf {A} _ {d} = \kappa_ {d} \mathbb {I} _ {n} + \mathbf {N} _ {d}, \tag {79}
$$
where  $\mathbf{N}_{c,d}^{\top} \in \ker (\mathbf{R}^{\top}\mathbf{W})$ , i.e.,  $\mathbf{N}_{c,d}\mathbf{W}^{\top}\mathbf{R} = \mathbb{0}$ . Further, assume that the kernel changes under rotation with an orthogonal matrix  $\mathbf{Q}$  as follows: