SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mariia Seleznova $^{1*}$  Dana Weitzner $^{2}$  Raja Giryes $^{2}$  Gitta Kutyniok $^{1}$  Hung-Hsu Chou $^{1}$
$^{1}$ Ludwig-Maximilians-Universität München  $^{2}$ Tel Aviv University
This work bridges two important concepts: the Neural Tangent Kernel (NTK), which captures the evolution of deep neural networks (DNNs) during training, and the Neural Collapse (NC) phenomenon, which refers to the emergence of symmetry and structure in the last-layer features of well-trained classification DNNs. We adopt the natural assumption that the empirical NTK develops a block structure aligned with the class labels, i.e., samples within the same class have stronger correlations than samples from different classes. Under this assumption, we derive the dynamics of DNNs trained with mean squared (MSE) loss and break them into interpretable phases. Moreover, we identify an invariant that captures the essence of the dynamics, and use it to prove the emergence of NC in DNNs with block-structured NTK. We provide large-scale numerical experiments on three common DNN architectures and three benchmark datasets to support our theory.
Deep Neural Networks (DNNs) are advancing the state of the art in many real-life applications, ranging from image classification to machine translation. Yet, there is no comprehensive theory that can explain a multitude of empirical phenomena observed in DNNs. In this work, we provide a theoretical connection between two such empirical phenomena, prominent in modern DNNs: Neural Collapse (NC) and Neural Tangent Kernel (NTK) alignment.
Neural Collapse. NC [39] emerges while training modern classification DNNs past zero error to further minimize the loss. During NC, the class means of the DNN's last-layer features form a symmetric structure with maximal separation angle, while the features of each individual sample collapse to their class means. This simple structure of the feature vectors appears favourable for generalization and robustness in the literature [12, 31, 40, 47]. Though NC is common in modern DNNs, explaining the mechanisms behind its emergence is challenging, since the complex non-linear training dynamics of DNNs evade analytical treatment.
Neural Tangent Kernel. The NTK [30] describes the gradient descent dynamics of DNNs in the function space, which provides a dual perspective to DNNs' evolution in the parameters space. This perspective allows to study the dynamics of DNNs analytically in the infinite-width limit, where the NTK is constant during training [30]. Hence, theoretical works often rely on the infinite-width NTK to analyze generalization of DNNs [1, 20, 28, 49]. However, multiple authors have argued that the infinite-width limit does not fully reflect the behaviour of realistic DNNs [2, 10, 22, 27, 36, 43], since constant NTK implies that no feature learning occurs during DNNs training.
NTK Alignment. While the infinite-width NTK is label-agnostic and does not change during training, the empirical NTK rapidly aligns with the target function in the early stages of training [5, 7, 44, 45]. In the context of classification, this manifests itself as the emergence of a block structure
*Correspondence to: Mariia Seleznova (selez@math.lmu.de).
in the kernel matrix, where the correlations between samples from the same class are stronger than between samples from different classes. The NTK alignment implies the so-called local elasticity of DNNs' training dynamics, i.e., samples from one class have little impact on samples from other classes in Stochastic Gradient Descent (SGD) updates [23]. Several recent works have also linked the local elasticity of training dynamics to the emergence of NC [33, 53]. This brings us to the main question of this paper: Is there a connection between NTK alignment and neural collapse?
Contribution. In this work, we consider a model of NTK alignment, where the kernel has a block structure, i.e., it takes only three distinct values: an inter-class value, an intra-class value and a diagonal value. We describe this model in Section 3. Within the model, we establish the connection between NTK alignment and NC, and identify the conditions under which NC occurs. Our main contributions are as follows:
- We derive and analyze the training dynamics of DNNs with MSE loss and block-structured NTK in Section 4. We identify three distinct convergence rates in the dynamics, which correspond to three components of the training error: error of the global mean, of the class means, and of each individual sample. These components play a key role in the dynamics.
- We show that NC emerges in DNNs with block-structured NTK under additional assumptions in Section 5.3. To the best of our knowledge, this is the first work to connect NTK alignment and NC. While previous contributions rely on the unconstrained features models [21, 38, 48] or other imitations of DNNs' training dynamics [53] to derive NC (see Appendix A for a detailed discussion of related works), we consider standard gradient flow dynamics of DNNs simplified by our assumption on the NTK structure.
- We analyze when NC does or does not occur in DNNs with NTK alignment, both theoretically and empirically. In particular, we identify an invariant of the training dynamics that provides a necessary condition for the emergence of NC in Section 5.2. Since DNNs with block-structured NTK do not always converge to NC, we conclude that NTK alignment is a more widespread phenomenon than NC.
- We support our theory with large-scale numerical experiments in Section 6. Source code to reproduce the results is available in the project's GitHub repository.
We consider the classification problem with  $C \in \mathbb{N}$  classes, where the goal is to build a classifier that returns a class label for any input  $x \in \mathcal{X}$ . In this work, the classifier is a DNN trained on a dataset  $\{(x_i, y_i)\}_{i=1}^N$ , where  $x_i \in \mathcal{X}$  are the inputs and  $y_i \in \mathbb{R}^C$  are the one-hot encodings of the class labels. We view the output function of the DNN  $f: \mathcal{X} \to \mathbb{R}^C$  as a composition of parametrized last-layer features  $h: \mathcal{X} \to \mathbb{R}^n$  and a linear classification layer parametrized by weights  $\mathbf{W} \in \mathbb{R}^{C \times n}$  and biases  $\mathbf{b} \in \mathbb{R}^C$ . Then the logits of the training data  $X = \{x_i\}_{i=1}^N$  can be expressed as follows:
$$
f (X) = \mathbf {W} \mathbf {H} + \mathbf {b} \mathbf {1} _ {N} ^ {\top}, \tag {1}
$$
where  $\mathbf{H} \in \mathbb{R}^{n \times N}$  are the features of the entire dataset stacked as columns and  $\mathbf{1}_N \in \mathbb{R}^N$  is a vector of ones. Though we omit the notion of the data dependence in the text to follow, i.e. we write  $\mathbf{H}$  without the explicit dependence on  $X$ , we emphasize that the features  $\mathbf{H}$  are a function of the data and the DNN's parameters, unlike in the previously studied unconstrained feature models [21, 38, 48].
We assume that the dataset is balanced, i.e. there are  $m \coloneqq N / C$  training samples for each class. Without loss of generality, we further assume that the inputs are reordered so that  $x_{(c - 1)m + 1},\ldots ,x_{cm}$  belong to class  $c$  for all  $c \in [C]$ . This will make the notation much easier later on. Since the dimension of features  $n$  is typically much larger than the number of classes, we also assume  $n > C$  in this work.
2.1 Neural Collapse
Neural Collapse (NC) is an empirical behaviour of classifier DNNs trained past zero error [39]. Let  $\langle h\rangle \coloneqq N^{-1}\sum_{i = 1}^{N}h(x_i)$  denote the global features mean and  $\langle h\rangle_c\coloneqq m^{-1}\sum_{x_i\in \mathrm{class}c}h(x_i)$ ,  $c\in [C]$  be the class means. Furthermore, define the matrix of normalized centered class means as  $\mathbf{M}\coloneqq [\langle \overline{h}\rangle_1 / \| \langle \overline{h}\rangle_1\| _2,\dots ,\langle \overline{h}\rangle_C / \| \langle \overline{h}\rangle_C\| _2]^\top \in \mathbb{R}^{n\times C}$ , where  $\langle \overline{h}\rangle_c = \langle h\rangle_c - \langle h\rangle_c\in [C]$ . We say that a DNN exhibits NC if the following four behaviours emerge as the training time  $t$  increases:
(NC1) Variability collapse: for all samples  $x_{i}^{c}$  from class  $c \in [C]$ , where  $i \in [m]$ , the penultimate layer features converge to their class means, i.e.  $\| h(x_{i}^{c}) - \langle h \rangle_{c} \|_{2} \to 0$ .
(NC2) Convergence to Simplex Equiangular Tight Frame (ETF): for all  $c, c' \in [C]$ , the class means converge to the following configuration:
$$
\| \langle h \rangle_ {c} - \langle h \rangle \| _ {2} - \| \langle h \rangle_ {c ^ {\prime}} - \langle h \rangle \| _ {2} \rightarrow 0, \quad \mathbf {M} ^ {\top} \mathbf {M} \rightarrow \frac {C}{C - 1} \left(\mathbb {I} _ {C} - \frac {1}{C} \mathbf {1} _ {C} \mathbf {1} _ {C} ^ {\top}\right).
$$
(NC3) Convergence to self-duality: the class means  $\mathbf{M}$  and the final weights  $\mathbf{W}^{\top}$  converge to each other:
$$
\left|\left| \mathbf {M} / \left|\left| \mathbf {M} \right|\right| _ {F} - \mathbf {W} ^ {\top} / \left|\left| \mathbf {W} ^ {\top} \right|\right| _ {F} \right|\right| _ {F} \rightarrow 0.
$$
(NC4) Simplification to Nearest Class Center (NCC): the classifier converges to the NCC decision rule behaviour:
$$
\operatorname * {a r g m a x} _ {c} (\mathbf {W} h (x) + \mathbf {b}) _ {c} \to \operatorname * {a r g m i n} _ {c} \| h (x) - \langle h \rangle_ {c} \| _ {2}.
$$
Though NC is observed in practice, there is currently no conclusive theory on the mechanisms of its emergence during DNN training. Most theoretical works on NC adopt the unconstrained features model, where features  $\mathbf{H}$  are free variables that can be directly optimized [21, 38, 48]. Training dynamics of such models do not accurately reflect the dynamics of real DNNs, since they ignore the dependence of the features on the input data and the DNN's trainable parameters. In this work, we make a step towards realistic DNN dynamics by means of the Neural Tangent Kernel (NTK).
2.2 Neural Tangent Kernel
The NTK  $\Theta$  of a DNN with the output function  $f:\mathcal{X}\to \mathbb{R}^C$  and trainable parameters  $\mathbf{w}\in \mathbb{R}^{P}$  (stretched into a single vector) is given by
$$
\Theta_ {k, s} \left(x _ {i}, x _ {j}\right) := \left\langle \nabla_ {\mathbf {w}} f _ {k} \left(x _ {i}\right), \nabla_ {\mathbf {w}} f _ {s} \left(x _ {j}\right) \right\rangle , \quad x _ {i}, x _ {j} \in \mathcal {X}, \quad k, s \in [ C ]. \tag {2}
$$
We also define the last-layer features kernel  $\Theta^h$ , which is a component of the NTK corresponding to the parameters up to the penultimate layer, as follows:
$$
\Theta_ {k, s} ^ {h} \left(x _ {i}, x _ {j}\right) := \left\langle \nabla_ {\mathbf {w}} h _ {k} \left(x _ {i}\right), \nabla_ {\mathbf {w}} h _ {s} \left(x _ {j}\right) \right\rangle , \quad x _ {i}, x _ {j} \in \mathcal {X}, \quad k, s \in [ n ]. \tag {3}
$$
Intuitively, the NTK captures the correlations between the training samples in the DNN dynamics. While most theoretical works consider the infinite-width limit of DNNs [30, 52], where the NTK can be computed theoretically, empirical studies have also extensively explored the NTK of finite-width networks [19, 36, 45, 49]. Unlike the label-agnostic infinite-width NTK, the empirical NTK aligns with the labels during training. We use this observation in our main assumption (Section 3).
2.3 Classification with MSE Loss
We study NC for DNNs with the mean squared error (MSE) loss given by
$$
\mathcal {L} (\mathbf {W}, \mathbf {H}, \mathbf {b}) = \frac {1}{2} \| f (X) - \mathbf {Y} \| _ {F} ^ {2}, \tag {4}
$$
where  $\mathbf{Y} \in \mathbb{R}^{C \times N}$  is a matrix of stacked labels  $y_{i}$ . While NC was originally introduced for the cross-entropy (CE) loss [39], which is more common in classification problems, the MSE loss is much easier to analyze theoretically. Moreover, empirical observations suggest that DNNs with MSE loss achieve comparable performance to using CE [14, 29, 41], which motivates the recent line of research on MSE-NC [21, 38, 48].
3 Block Structure of the NTK
Numerous empirical studies have demonstrated that the NTK becomes aligned with the labels  $\mathbf{Y}^{\top}\mathbf{Y}$  during the training process [7, 32, 45]. This alignment constitutes feature learning and is associated with better performance of DNNs [9, 13]. For classification problems, this means that the empirical
Figure 1: The NTK block structure of ResNet20 trained on MNIST. a) Traced kernel  $\sum_{k=1}^{C} \Theta_{k,k}(X)$  computed on a random data subset with 12 samples from each class. The samples are ordered as described in Section 2, so that the diagonal blocks correspond to pairs of inputs from the same class. b) Traced kernel  $\sum_{k=1}^{n} \Theta_{k,k}^{h}(X)$  computed on the same subset. c) Norms of the kernels  $\Theta_{k,s}(X)$  for all  $k, s \in [C]$ . d) Norms of the kernels  $\Theta_{k,s}^{h}(X)$  for all  $k, s \in [n]$ . The color bars show the values in each heatmap as a fraction of the maximal value in the heatmap. e) The alignment of the traced kernels from panes a and b with the class labels.
NTK develops an approximate block structure with larger kernel values corresponding to pairs of samples  $(x_{i}^{c}, x_{j}^{c})$  from the same class [44]. Figure 1 shows an example of such a structure emergent in the empirical NTK of ResNet20 trained on MNIST. Motivated by these observations, we assume that the NTK and the last-layer features kernel exhibit a block structure, defined as follows:
Definition 3.1 (Block structure of a kernel). We say a kernel  $\Theta : \mathcal{X} \times \mathcal{X} \to \mathbb{R}^{K \times K}$  has a block structure associated with  $(\lambda_1, \lambda_2, \lambda_3)$ , if  $\lambda_1 > \lambda_2 > \lambda_3 \geq 0$  and
$$
\Theta (x, x) = \lambda_ {1} \mathbb {I} _ {K}, \quad \Theta \left(x _ {i} ^ {c}, x _ {j} ^ {c}\right) = \lambda_ {2} \mathbb {I} _ {K}, \quad \Theta \left(x _ {i} ^ {c}, x _ {j} ^ {c ^ {\prime}}\right) = \lambda_ {3} \mathbb {I} _ {K}, \tag {5}
$$
where  $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$ .
Assumption 3.2. The  $NTK\Theta : \mathcal{X} \times \mathcal{X} \to \mathbb{R}^{C \times C}$  has a block structure associated with  $(\gamma_d, \gamma_c, \gamma_n)$ , and the penultimate kernel  $\Theta^h : \mathcal{X} \times \mathcal{X} \to \mathbb{R}^{n \times n}$  has a block structure associated with  $(\kappa_d, \kappa_c, \kappa_n)$ .
This assumption means that every kernel  $\Theta_{k,k}(X) := [\Theta_{k,k}(x_i, x_j)]_{i,j \in [N]}$  corresponding to an output neuron  $f_k, k \in [C]$  and every kernel  $\Theta_{p,p}^h(X)$  corresponding to a last-layer neuron  $h_p, p \in [n]$  is aligned with  $\mathbf{Y}^\top \mathbf{Y}$  (see Figure 1, panes a-b). Additionally, the "non-diagonal" kernels  $\Theta_{k,s}(X)$  and  $\Theta_{k,s}^h(X), k \neq s$  are equal to zero (see Figure 1, panes c-d). Moreover, if  $\gamma_c \gg \gamma_n$  and  $\kappa_c \gg \kappa_n$ , Assumption 3.2 can be interpreted as local elasticity of DNNs, defined below.
Definition 3.3 (Local elasticity [23]). A classifier is said to be locally elastic (LE) if its prediction or feature representation on point  $x_{i}^{c}$  from class  $c \in [C]$  is not significantly affected by performing SGD updates on data points from classes  $c' \neq c$ .
To see the relation between Assumption 3.2 and this definition, consider a Gradient Descent (GD) step of the output neuron  $f_{k}, k \in [C]$  with step size  $\eta$  performed on a single input  $x_{j}^{c^{\prime}}$  from class  $c^{\prime} \neq c$ . By the chain rule, block-structured  $\Theta$  implies locally-elastic predictions since
$$
f ^ {t + 1} \left(x _ {i} ^ {c}\right) - f ^ {t} \left(x _ {i} ^ {c}\right) = - \eta \Theta \left(x _ {i} ^ {c}, x _ {j} ^ {c ^ {\prime}}\right) \frac {\partial \mathcal {L} \left(x _ {j} ^ {c ^ {\prime}}\right)}{\partial f \left(x _ {j} ^ {c ^ {\prime}}\right)} + O \left(\eta^ {2}\right), \tag {6}
$$
i.e., the magnitude of the GD step of  $f(x_{i}^{c})$  is determined by the value of  $\Theta (x_i^c,x_j^{c'})$ . Similarly, block-structured kernel  $\Theta^h$  implies locally-elastic penultimate layer features because
$$
h ^ {t + 1} \left(x _ {i} ^ {c}\right) - h ^ {t} \left(x _ {i} ^ {c}\right) = - \eta \Theta^ {h} \left(x _ {i} ^ {c}, x _ {j} ^ {c ^ {\prime}}\right) \mathbf {W} ^ {\top} \frac {\partial \mathcal {L} \left(x _ {j} ^ {c ^ {\prime}}\right)}{\partial f \left(x _ {j} ^ {c ^ {\prime}}\right)} + O \left(\eta^ {2}\right). \tag {7}
$$
This observation provides a connection between our work and recent contributions suggesting a connection between NC and local elasticity [33, 53].
<sup>2</sup>We provide figures illustrating the NTK block structure on other architectures and datasets in Appendix C. <sup>3</sup>We discuss possible relaxations to our main assumption, where the "non-diagonal" components of the last-layer kernel  $\Theta_{k,s}^{h}$  are allowed to be non-zero, in Appendix D.
Eigenvalue
$$
\lambda_ {\text {s i n g l e}} = \gamma_ {d} - \gamma_ {c}
$$
$$
\mathbf{v}_{i}^{c} = \frac{1}{m - 1}\big(\underbrace{m - 1, - \mathbf{1}_{m - 1}^{\top}}_{\text{index $i > 0$, class $c <   0$}},\underbrace{\mathbf{0}_{N - m}^{\top}}_{\text{others} = 0}\big)^{\top}
$$
$$
\lambda_ {\text {c l a s s}} = \lambda_ {\text {s i n g l e}} + m \left(\gamma_ {c} - \gamma_ {n}\right)
$$
$$
\mathbf{v}_{c} = \frac{1}{C - 1}\bigl(\underbrace{(C - 1)\mathbf{1}_{m}^{\top}}_{\text{class} c > 0}, - \underbrace{\mathbf{1}_{N - m}^{\top}}_{\text{others} <   0}\bigr)^{\top}
$$
$$
\mathbf {v} _ {0} = \mathbf {1} _ {N}
$$
$$
N - C
$$
$$
\lambda_ {\text {g l o b a l}} = \lambda_ {\text {c l a s s}} + N \gamma_ {n}
$$
Table 1: Eigendecomposition of the block-structured NTK.
4 Dynamics of DNNs with NTK Alignment
4.1 Convergence
As a warm up for our main results, we analyze the effects of the NTK block structure on the convergence of DNNs. Consider a GD update of an output neuron  $f_{k}, k \in [C]$  with the step size  $\eta$ :
$$
f _ {k} ^ {t + 1} (X) = f _ {k} ^ {t} (X) - \eta \Theta_ {k, k} (X) (f _ {k} ^ {t} (X) - \mathbf {Y} _ {k}) + O (\eta^ {2}), \quad k = 1, \dots , C. \tag {8}
$$
Note that we have taken into account that  $\Theta_{k,s}$  is zero for  $k\neq s$  by our assumption. Denote the residuals corresponding to  $f_{k}$  as  $\mathbf{r}_k^\top \coloneqq f_k^\top (X) - \mathbf{Y}_k\in \mathbb{R}^N$ . Then we have the following dynamics for the residuals vector:
$$
\mathbf {r} _ {k} ^ {t + 1} = \left(1 - \eta \Theta_ {k, k} (X)\right) \mathbf {r} _ {k} ^ {t} + O \left(\eta^ {2}\right). \tag {9}
$$
The eigendecomposition of the block-structured kernel  $\Theta_{k,k}(X)$  provides important insights into this dynamics and is summarized in Table 1. We notice that the NTK has three distinct eigenvalues  $\lambda_{\mathrm{global}} \geq \lambda_{\mathrm{class}} \geq \lambda_{\mathrm{single}}$ , which imply different convergence rates for certain components of the error. Moreover, the eigenvectors associated with each of these eigenvalues reveal the meaning of the error components corresponding to each convergence rate. Indeed, consider the projected dynamics with respect to eigenvector  $\mathbf{v}_0$  and eigenvalue  $\lambda_{\mathrm{global}}$  from Table 1:
$$
\left\langle \mathbf {r} _ {k} ^ {t + 1}, \mathbf {v} _ {0} \right\rangle = \left(1 - \eta \lambda_ {\text {g l o b a l}}\right) \left\langle \mathbf {r} _ {k} ^ {t}, \mathbf {v} _ {0} \right\rangle , \tag {10}
$$
where we omitted  $O(\eta^2)$  for clarity. Now notice that the projection of  $\mathbf{r}_k^t$  onto the vector  $\mathbf{v}_0$  is in fact proportional to the average residual over the training set:
$$
\langle \mathbf {r} _ {k} ^ {t}, \mathbf {v} _ {0} \rangle = \langle \mathbf {r} _ {k} ^ {t}, \mathbf {1} _ {N} \rangle = N \langle \mathbf {r} _ {k} ^ {t} \rangle \tag {11}
$$
where  $\langle \cdot \rangle$  denotes the average over all the training samples  $x_{i}\in X$ . By a similar calculation, for all  $c\in [C]$  and  $i\in [m]$  we get interpretations of the remaining projections of the residual:
$$
\langle \mathbf {r} _ {k} ^ {t}, \mathbf {v} _ {c} \rangle = \frac {N}{C - 1} \left(\left\langle \mathbf {r} _ {k} ^ {t} \right\rangle_ {c} - \left\langle \mathbf {r} _ {k} ^ {t} \right\rangle\right), \quad \left\langle \mathbf {r} _ {k} ^ {t}, \mathbf {v} _ {i} ^ {c} \right\rangle = \frac {m}{m - 1} \left(\mathbf {r} _ {k} ^ {t} \left(x _ {i} ^ {c}\right) - \left\langle \mathbf {r} _ {k} ^ {t} \right\rangle_ {c}\right), \tag {12}
$$
We where  $\langle \cdot \rangle_{c}$  denotes the average over samples  $x_{i}^{c}$  from class  $c$ , and  $\mathbf{r}_k^\top (x_i^c)$  is the  $k$ th component of  $f^{\top}(x_i^c) - y_i^c$ . Combining (10), (11) and (12), we have the following convergence rates:
$$
\left\langle \mathbf {r} _ {k} ^ {t + 1} \right\rangle = (1 - \eta \lambda_ {\text {g l o b a l}}) \left\langle \mathbf {r} _ {k} ^ {t} \right\rangle , \tag {13}
$$
$$
\langle \mathbf {r} _ {k} ^ {t + 1} \rangle_ {c} - \langle \mathbf {r} _ {k} ^ {t + 1} \rangle = (1 - \eta \lambda_ {\text {c l a s s}}) (\langle \mathbf {r} _ {k} ^ {t} \rangle_ {c} - \langle \mathbf {r} _ {k} ^ {t} \rangle), \tag {14}
$$
$$
\mathbf {r} _ {k} ^ {t + 1} \left(x _ {i} ^ {c}\right) - \left\langle \mathbf {r} _ {k} ^ {t + 1} \right\rangle_ {c} = \left(1 - \eta \lambda_ {\text {s i n g l e}}\right) \left(\mathbf {r} _ {k} ^ {t} \left(x _ {i} ^ {c}\right) - \left\langle \mathbf {r} _ {k} ^ {t} \right\rangle_ {c}\right). \tag {15}
$$
Overall, this means that the global mean  $\langle \mathbf{r}\rangle$  of the residual converges first, then the class means, and finally the residual of each sample  $\mathbf{r}(x_i^c)$ . To simplify the notation, we define the following quantities:
$$
\mathbf {R} = f (X) - \mathbf {Y} = [ \mathbf {r} \left(x _ {1}\right), \dots , \mathbf {r} \left(x _ {N}\right) ], \tag {16}
$$
$$
\mathbf {R} _ {\text {c l a s s}} = \frac {1}{m} \mathbf {R Y} ^ {\top} \mathbf {Y} = \underbrace {[ \langle \mathbf {r} \rangle_ {1} , \dots , \langle \mathbf {r} \rangle_ {C} ]} _ {:= \mathbf {R} _ {1}} \otimes \mathbf {1} _ {m} ^ {\top}, \tag {17}
$$
$$
\mathbf {R} _ {\text {g l o b a l}} = \frac {1}{N} \mathbf {R} \mathbf {1} _ {N} \mathbf {1} _ {N} ^ {\top} = \langle \mathbf {r} \rangle \otimes \mathbf {1} _ {N} ^ {\top}, \tag {18}
$$
where  $\mathbf{R} \in \mathbb{R}^{C \times N}$  is the matrix of residuals,  $\mathbf{R}_{\mathrm{class}} \in \mathbb{R}^{C \times N}$  are the residuals averaged over each class and stacked  $m$  times, and  $\mathbf{R}_{\mathrm{global}} \in \mathbb{R}^{C \times N}$  are the residuals averaged over the whole training set stacked  $N$  times. According to the previous discussion,  $\mathbf{R}_{\mathrm{global}}$  converges to zero at the fastest rate, while  $\mathbf{R}$  converges at the slowest rate. The last phase, which we call the end of training, is when  $\mathbf{R}_{\mathrm{class}}$  and  $\mathbf{R}_{\mathrm{global}}$  have nearly vanished and can be treated as zero for the remaining training time. We will use this notion in several remarks, as well as in the proof of Theorem 5.2.
4.2 Gradient Flow Dynamics with Block-Structured NTK
We derive the dynamics of  $\mathbf{H},\mathbf{W},\mathbf{b}$  under Assumption 3.2 in Theorem 4.1. One can see that the block-structured kernel greatly simplifies the complicated dynamics of DNNs and highlights the role of each of the residual components identified in Section 4.1. We consider gradient flow, which is close to gradient descent for sufficiently small step size [16], to reduce the complications caused by higher order terms. The proof is given in Appendix B.1.
Theorem 4.1. Suppose Assumption 3.2 holds. Then the gradient flow dynamics of a DNN can be written as