SlowGuess/NeurIPS_2023_OCR huggingface dataset:
1. For gradient descent,  $f(\theta) = \theta - h\nabla l$ . If  $l \in C^2$ ,  $\nabla l$  is L-Lipschitz and  $h \leq 1 / L$ , then  $f$  is a diffeomorphism (Proposition 4.5 [44]), thus satisfying our assumption. In optimization literature, the above sufficient condition is mild and standard assumption to ensure the convergence of GD.
2. For gradient flow,  $\frac{\mathrm{d}}{\mathrm{d}t}\theta = -\nabla l$ . Our assumption immediately follows existence and uniqueness of ODE solution. For example, if  $l \in C^2$  and  $\nabla l$  is L-Lipschitz, then  $f_t(\theta)$  is defined for all  $t$  and  $\theta$  (Corollary 2.6 [67]) and is a diffeomorphism (Theorem 2.10 [67]).
3. Momentum SGD and Adagrad always satisfy our assumption on enriched state space (Appendix F.2).
4. Same assumption was also used in previous research [50]. Moreover, many previous works use assumptions that imply our assumption as discussed in 1.
In Section 7, we only analyze the infinitesimal step-size. The main reason for this choice is that it allows the training dynamics of neural networks to be described by an ordinary differential equation (ODE), which simplifies the analysis. While it's possible that our approach could still be viable with discrete step sizes, it would make the theoretical analysis much more challenging. Practically, only discrete step sizes are implementable in experiments. However, our findings show that the theoretical predictions based on infinitesimal steps align closely with the experimental outcomes, as demonstrated in Figure 2.
One strong assumption imposed by this paper is an infinite width for the neural network. This assumption is vital for the feasibility of our theoretical analysis in Section 7 and is a common stance in numerous theoretical studies of neural networks [55, 73, 42, 51, 18, 39, 43, 76, 3]. The infinite-width perspective simplifies the dynamics and analysis but, admittedly, deviates from the practical, finite-width architectures typically used in real-world applications. Several researchers [29, 20, 35, 46, 63] have tried to understand the impact of finite width, known as "finite-width corrections". Delving into this adds complexity to the theoretical analysis, and hence it's left as a potential topic for future research. Moreover, we find that the dynamics in function space, derived from the theory of infinitely wide networks, closely resembles the behavior of actual finite-width networks, as can be seen in Figure 2.
C Additional Proofs
Proof of Theorem 2. First we show that sampling efficiency is determined by the  $\chi^2$  divergence.
$$
\frac {1}{\operatorname {e f f} _ {\lambda}} = \frac {\mathbb {E} _ {\theta \sim p} \left[ w _ {\lambda} ^ {2} (\theta) \right]}{Z _ {\lambda} ^ {2}} = \frac {1}{Z _ {\lambda} ^ {2}} \mathbb {E} _ {\theta \sim q} \left[ \left(e ^ {- \lambda r (f (\theta))} \frac {p (f (\theta))}{q (f (\theta))}\right) ^ {2} \right] = \mathbb {E} _ {\theta \sim q} \left[ \left(\frac {p _ {\lambda} (\theta)}{q (\theta)}\right) ^ {2} \right] = D _ {\chi^ {2}} \left(p _ {\lambda} \| q\right) + 1 \tag {12}
$$
The inequality in Eq. (6) is from Theorem 5 in [25].
Proof of Theorem 3. Based on some algebra and calculus, we have
$$
(Z _ {\lambda} ^ {\beta}) ^ {2} \frac {\mathrm {d}}{\mathrm {d} \beta} D _ {\mathrm {K L}} (p _ {\lambda} ^ {\beta} \| p _ {\lambda}) = \left(\mathbb {E} _ {\theta \sim p} [ \log \frac {\beta}{w _ {\lambda} (\theta)} \mathbb {1} (w _ {\lambda} (\theta) \geq \beta) ]\right) \times (\mathbb {E} _ {\theta \sim p} [ w _ {\lambda} (\theta) \mathbb {1} (w _ {\lambda} (\theta) <   \beta) ]) \leq 0
$$
For sampling efficiency, we have
$$
\mathbb {E} _ {\theta \sim p} [ v _ {\beta} (w _ {\lambda} (\theta)) ] \frac {\mathrm {d}}{\mathrm {d} \beta} \log \frac {1}{\mathrm {e f f} _ {\lambda} ^ {\beta}} = 2 \mathbb {E} _ {\theta \sim p} [ \Im (w _ {\lambda} (\theta) \geq \beta) ] \times (\frac {\beta \mathbb {E} _ {\theta \sim p} [ v _ {\beta} (w _ {\lambda} (\theta)) ]}{\mathbb {E} _ {\theta \sim p} [ v _ {\beta} (w _ {\lambda} (\theta)) ^ {2} ]} - 1) \geq 0
$$
Proof of Eq. (27). We discard the parameter  $\theta$ ,  $s_a$  without loss of generality.
$$
\nabla_ {W _ {i}} y _ {d} = c _ {i} x _ {i - 1} \nabla_ {y _ {i}} y _ {d} \tag {13}
$$
$$
\begin{array}{l} \forall j \geq i, \nabla_ {y _ {j}} y _ {j + 1} y _ {j} ^ {(i)} = c _ {j + 1} W _ {j + 1} \nabla_ {y _ {j}} x _ {j} y _ {j} ^ {(i)} \\ = c _ {j + 1} W _ {j + 1} \nabla_ {y _ {j}} \operatorname {d i a g} \left(\sigma_ {j} ^ {\prime} \left(y _ {j}\right)\right) y _ {j} ^ {(i)} \tag {14} \\ = y _ {j + 1} ^ {(i)} \\ \end{array}
$$
We finish the proof by combining above results.
$$
\begin{array}{l} \operatorname {T r} \left(\nabla_ {W _ {i}} y _ {d} W _ {i}\right) = \operatorname {T r} \left(c _ {i} x _ {i - 1} \nabla_ {y _ {i}} y _ {d} W _ {i}\right) \\ = \operatorname {T r} \left(\nabla_ {y _ {i}} y _ {d} c _ {i} W _ {i} x _ {i - 1}\right) \\ = \operatorname {T r} \left(\nabla_ {y _ {i}} y _ {d} y _ {i} ^ {(i)}\right) \tag {15} \\ = \operatorname {T r} \left(\nabla_ {y _ {d - 1}} y _ {d} \dots \nabla_ {y _ {i}} y _ {i + 1} y _ {i} ^ {(i)}\right) \\ = y _ {d} ^ {(i)} \\ \end{array}
$$
$$
\begin{array}{l} \nabla_ {W _ {i}} y _ {d} = c _ {i} x _ {i - 1} \nabla_ {y _ {i}} y _ {d} \\ = \left(c _ {i} x _ {i - 1}\right) \left(\nabla_ {y _ {d - 1}} y _ {d} \dots \nabla_ {y _ {i}} y _ {i + 1}\right) \tag {16} \\ = \left(c _ {i} x _ {i - 1}\right) \left(\left(c _ {d} W _ {d}\right) \operatorname {d i a g} \left(\sigma_ {d - 1} ^ {\prime} \left(y _ {d - 1}\right)\right) \dots c _ {i + 1} W _ {i + 1} \operatorname {d i a g} \left(\sigma_ {i} ^ {\prime} \left(y _ {i}\right)\right)\right) \\ \end{array}
$$
Only  $\sigma_j^\prime (y_j)$  for  $i\leq j < d$  depends on  $W_{i}$  in the above formula.
$$
\nabla_ {\left(W _ {i}\right) _ {\alpha , \beta}} y _ {d} = \left(c _ {i} x _ {i - 1}\right) _ {\beta} \left(\nabla_ {x _ {j}} y _ {d} \operatorname {d i a g} \left(\sigma_ {j} ^ {\prime} \left(y _ {j}\right)\right) \nabla_ {\left(y _ {i}\right) _ {\alpha}} y _ {j}\right) \tag {17}
$$
$$
\begin{array}{l} \nabla_ {(W _ {i}) _ {\alpha , \beta}} \nabla_ {(W _ {i}) _ {\alpha , \beta}} y _ {d} = (c _ {i} x _ {i - 1}) _ {\beta} \left(\sum_ {j = i} ^ {d - 1} (\nabla_ {x _ {j}} y _ {d}) \left(\sigma_ {j} ^ {\prime \prime} \left(y _ {j}\right) \odot \nabla_ {(W _ {i}) _ {\alpha , \beta}} y _ {j} \odot \nabla_ {(y _ {i}) _ {\alpha}} y _ {j}\right)\right) \tag {18} \\ = (c _ {i} x _ {i - 1}) _ {\beta} ^ {2} \left(\sum_ {j = i} ^ {d - 1} \left(\nabla_ {x _ {j}} y _ {d}\right) \left(\sigma_ {j} ^ {\prime \prime} (y _ {j}) \odot \nabla_ {(y _ {i}) _ {\alpha}} y _ {j} \odot \nabla_ {(y _ {i}) _ {\alpha}} y _ {j}\right)\right) \\ \end{array}
$$
$$
\sum_ {\alpha = 1} ^ {n _ {i}} \sum_ {\beta = 1} ^ {n _ {i - 1}} \nabla_ {(W _ {i}) _ {\alpha , \beta}} \nabla_ {(W _ {i}) _ {\alpha , \beta}} y _ {d} = \sum_ {j = i} ^ {d - 1} \left(\nabla_ {x _ {j}} y _ {d}\right) \left(\sigma_ {j} ^ {\prime \prime} (y _ {j}) \odot \xi^ {(i, j)}\right) \tag {19}
$$
$$
\begin{array}{l} \operatorname {T r} H = \sum_ {i = 1} ^ {d} \sum_ {\alpha = 1} ^ {n _ {i}} \sum_ {\beta = 1} ^ {n _ {i - 1}} \nabla_ {(W _ {i}) _ {\alpha , \beta}} \nabla_ {(W _ {i}) _ {\alpha , \beta}} y _ {d} \\ = \sum_ {i = 1} ^ {d} \sum_ {j = i} ^ {d - 1} \left(\nabla_ {x _ {j}} y _ {d}\right) \left(\sigma_ {j} ^ {\prime \prime} \left(y _ {j}\right) \odot \xi^ {(i, j)}\right) \tag {20} \\ = \sum_ {j = 1} ^ {d - 1} \left(\nabla_ {x _ {j}} y _ {d}\right) \left(\sigma_ {j} ^ {\prime \prime} (y _ {j}) \odot \left(\sum_ {i = 1} ^ {j} \xi^ {(i, j)}\right)\right) \\ \end{array}
$$
D Analysis of Infinitely Wide Network
We first give the analytic form of limits in Theorem 5.
The covariant matrix  $\Sigma$  is defined recursively. We first define series of Gaussian random variables  $y_{k}^{(i)}(s_{a}),\gamma_{k}^{(i)}(s_{a})$ , where variables with different subscript  $k$  are independent. We also let  $y_{k}(s_{a}) = y_{k}^{(d)}(s_{a})$ . The covariance of  $y_{k}^{(i)}(s_{a})$  and  $y_{k}^{(j)}(s_{b})$  are denoted as  $\Sigma_k[y^{(i)}(s_a),y^{(j)}(s_b)]$ . Similarly,
we can define  $\Sigma_k[y^{(i)}(s_a),\gamma^{(j)}(s_b)]$  and  $\Sigma_k[\gamma^{(i)}(s_a),\gamma^{(j)}(s_b)]$ .
$$
\Sigma_ {1} \left[ y ^ {(i)} \left(s _ {a}\right), y ^ {(j)} \left(s _ {b}\right) \right] = \Sigma_ {1} \left[ y ^ {(i)} \left(s _ {a}\right), \gamma^ {(j)} \left(s _ {b}\right) \right] = \Sigma_ {1} \left[ \gamma^ {(i)} \left(s _ {a}\right), \gamma^ {(j)} \left(s _ {b}\right) \right] = s _ {a} ^ {T} s _ {b} / n _ {0}
$$
$$
\Sigma_ {k + 1} \left[ y ^ {(i)} \left(s _ {a}\right), y ^ {(j)} \left(s _ {b}\right) \right] = \mathbb {E} \left[ I _ {i, k} \left(y _ {k} \left(s _ {a}\right), y _ {k} ^ {(i)} \left(s _ {a}\right)\right) I _ {j, k} \left(y _ {k} \left(s _ {b}\right), y _ {k} ^ {(j)} \left(s _ {b}\right)\right) \right]
$$
$$
\Sigma_ {k + 1} \left[ y ^ {(i)} \left(s _ {a}\right), \gamma^ {(j)} \left(s _ {b}\right) \right] = \mathbb {E} \left[ I _ {i, k} \left(y _ {k} \left(s _ {a}\right), y _ {k} ^ {(i)} \left(s _ {a}\right)\right) J _ {j, k} \left(y _ {k} \left(s _ {b}\right), \gamma_ {k} ^ {(j)} \left(s _ {b}\right)\right) \right]
$$
$$
\Sigma_ {k + 1} \left[ \gamma^ {(i)} \left(s _ {a}\right), \gamma^ {(j)} \left(s _ {b}\right) \right] = \mathbb {E} \left[ J _ {i, k} \left(y _ {k} \left(s _ {a}\right), \gamma_ {k} ^ {(i)} \left(s _ {a}\right)\right) J _ {j, k} \left(y _ {k} \left(s _ {b}\right), \gamma_ {k} ^ {(j)} \left(s _ {b}\right)\right) \right]
$$
$$
I _ {i, k} \left(y _ {k}, y _ {k} ^ {(i)}\right) = \left\{ \begin{array}{l l} \sigma_ {k} \left(y _ {k}\right) & i > k \\ \sigma_ {k} ^ {\prime} \left(y _ {k}\right) y _ {k} ^ {(i)} & i \leq k \end{array} \right. \tag {21}
$$
$$
J _ {i, k} (y _ {k}, \gamma_ {k} ^ {(i)}) = \left\{ \begin{array}{l l} \sigma_ {k} (y _ {k}) & i > k \\ \sigma_ {k} ^ {\prime \prime} (y _ {k}) & i = k \\ \sigma_ {k} ^ {\prime} (y _ {k}) \gamma_ {k} ^ {(i)} & i <   k \end{array} \right.
$$
In the infinite width limit, the covariance at initialization of  $(Y^{(i)}(\theta (0)))_a$  and  $(Y^{(j)}(\theta (0)))_b$  is  $\Sigma_{d}[y^{(i)}(s_{a}),y^{(j)}(s_{b})]$ . Similarly, the covariance at initialization of  $(Y^{(i)}(\theta (0)))_a$  and  $(\Gamma^{(j)}(\theta (0)))_b$  is  $\Sigma_{d}[y^{(i)}(s_{a}),\gamma^{(j)}(s_{b})]$ , the covariance at initialization of  $(\Gamma^{(i)}(\theta (0)))_a$  and  $(\Gamma^{(j)}(\theta (0)))_b$  is  $\Sigma_{d}[\gamma^{(i)}(s_a),\gamma^{(j)}(s_b)]$ .
Other matrices limits are
$$
\begin{array}{l} \left(\Theta^ {(i)}\right) _ {a, b} = \sum_ {j = 1} ^ {d} \left(\prod_ {k = j} ^ {d - 1} \mathbb {E} \left[ \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {a}\right)\right) \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {b}\right)\right) \right]\right) \Sigma_ {j} \left[ y ^ {(i)} \left(s _ {a}\right), y \left(s _ {b}\right) \right] (22) \\ + \sum_ {j = 1} ^ {d} \sum_ {k = \max  (i, j)} ^ {d - 1} \left(\prod_ {\substack {l = j \\ l \neq k}} ^ {d - 1} \mathbb {E} \left[ \sigma_ {l} ^ {\prime} \left(y _ {l} \left(s _ {a}\right)\right) \sigma_ {l} ^ {\prime} \left(y _ {l} \left(s _ {b}\right)\right) \right]\right) \mathbb {E} \left[ y _ {k} ^ {(i)} \left(s _ {a}\right) \sigma_ {k} ^ {\prime \prime} \left(y _ {k} \left(s _ {a}\right)\right) \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {b}\right)\right) \right] \Sigma_ {j} [ y \left(s _ {a}\right), y \left(s _ {b}\right) ], (23) \\ \end{array}
$$
$$
\begin{array}{l} \left(\Phi^ {(i)}\right) _ {a, b} = \sum_ {j = i + 1} ^ {d} \left(\prod_ {k = j} ^ {d - 1} \mathbb {E} \left[ \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {a}\right)\right) \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {b}\right)\right) \right]\right) \Sigma_ {j} \left[ \gamma^ {(i)} \left(s _ {a}\right), y \left(s _ {b}\right) \right] \\ + \sum_ {j = 1} ^ {i} \left(\mathbb {E} \left[ \sigma_ {i} ^ {\prime \prime \prime} \left(y _ {i} \left(s _ {a}\right)\right) \sigma_ {i} ^ {\prime} \left(y _ {i} \left(s _ {b}\right)\right) \right] \prod_ {\substack {k = j \\ k \neq i}} ^ {d - 1} \mathbb {E} \left[ \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {a}\right)\right) \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {b}\right)\right) \right]\right) \Sigma_ {j} [ y (s _ {a}), y (s _ {b}) ] (24) \\ + \sum_ {j = 1} ^ {d} \sum_ {k = \max  (i + 1, j)} ^ {d - 1} \left(\prod_ {\substack {l = j \\ l \neq k}} ^ {d - 1} \mathbb {E} \left[ \sigma_ {l} ^ {\prime} \left(y _ {l} \left(s _ {a}\right)\right) \sigma_ {l} ^ {\prime} \left(y _ {l} \left(s _ {b}\right)\right) \right]\right) \mathbb {E} \left[ \gamma_ {k} ^ {(i)} \left(s _ {a}\right) \sigma_ {k} ^ {\prime \prime} \left(y _ {k} \left(s _ {a}\right)\right) \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {b}\right)\right) \right] \Sigma_ {j} [ y \left(s _ {a}\right), y \left(s _ {b}\right) ], (25) \\ \end{array}
$$
$$
\left. \left(\Xi^ {(i, j)}\right) _ {a} = \left(\prod_ {k = i} ^ {j - 1} \mathbb {E} \left[ \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {a}\right)\right) \sigma_ {k} ^ {\prime} \left(y _ {k} \left(s _ {a}\right)\right) \right]\right) \Sigma_ {i} [ y (s _ {a}), y (s _ {a}) ]. \right. \tag {26}
$$
Before embarking on the proof, we first give an intuition of above results. We first notice that  $y_{d}^{(i)}(\theta ,s_{a})$  can be defined recursively:
$$
\forall j \leq i \leq d, \quad y _ {j} ^ {(i)} (\theta , s _ {a}) = y _ {j} (\theta , s _ {a}), \tag {27}
$$
$$
\forall i \leq j <   d, y _ {j + 1} ^ {(i)} (\theta , s _ {a}) = c _ {j + 1} W _ {j + 1} \operatorname {d i a g} \left(\sigma_ {j} ^ {\prime} \left(y _ {j} (\theta , s _ {a})\right)\right) y _ {j} ^ {(i)} (\theta , s _ {a}).
$$
We can give a similar definition by forward propagation for  $\gamma_j^{(i)}(\theta, s_a) = (\nabla_{x_i} y_j(\theta, s_a)) \sigma_i''(y_i(\theta, s_a))$ , which covers the case when  $j \leq i$ :
$$
\begin{array}{l} \forall j \leq i \leq d, \gamma_ {j} ^ {(i)} (\theta , s _ {a}) = y _ {j} (\theta , s _ {a}), \\ \gamma_ {i + 1} ^ {(i)} (\theta , s _ {a}) = c _ {i + 1} W _ {i + 1} \sigma_ {i} ^ {\prime \prime} \left(\gamma_ {i} ^ {(i)} (\theta , s _ {a})\right), \tag {28} \\ \end{array}
$$
$$
\forall i <   j <   d, \gamma_ {j + 1} ^ {(i)} (\theta , s _ {a}) = c _ {j + 1} W _ {j + 1} \mathrm {d i a g} (\sigma_ {j} ^ {\prime} (y _ {j} (\theta , s _ {a}))) \gamma_ {j} ^ {(i)} (\theta , s _ {a}).
$$
Therefore, both  $y_{d}^{(i)}$  and  $\gamma_{d}^{(i)}$  can be regarded as output from some other "neural network". We compares these forward pass with original neural network as follows:
$$
\forall i \leq j <   d, y _ {j + 1} (\theta , s _ {a}) = c _ {j + 1} W _ {j + 1} \sigma_ {j} \left(y _ {j} (\theta , s _ {a})\right). \tag {29}
$$
Both  $y_{d}^{(i)}$  and  $\gamma_{d}^{(i)}$  propagate in a same way as original network until  $i$ -th layer. After that,  $y_{d}^{(i)}$  propagate by linear transformation, with matrix  $\mathrm{diag}(\sigma_j'(y_j(\theta ,s_a)))$  being controlled by original network.  $\gamma_{d}^{(i)}$  uses a different activation function  $\sigma_i^{\prime \prime}$  which is different than original network with  $\sigma_{i}$  in  $i$ -th layer, and propagate by linear transformation in a similar way as  $y_{d}^{(i)}$ .
When we consider the  $y_{d}, y_{d}^{(i)}$  and  $\gamma_{d}^{(i)}$  as output from a multiple branch neural network, the  $\Sigma$  in Eq. (21) is just conjugate kernel, and  $\Theta^{(i)}, \Phi^{(i)}, \Xi^{(i,j)}$  in Eq. (22) are fragment of NTK of the larger network.
Proof of Theorem 5. We first prove the convergence to Gaussian distribution. It follows the proof of convergence of  $\Theta^{(i)},\Phi^{(i)},\Xi^{(i,j)}$ . Finally we prove that these values don't change during training.
We note that in the forward propagation as shown in Eqs. (27) to (29), we only need to perform two kinds of operations: matrix vector multiplication and map vector into matrix by diag. The latter operation could be rewrite into element-wise product which is a non-linear element-wise vector operation, e.g.  $\mathrm{diag}(\sigma_j'(y_j(\theta ,s_a)))y_j^{(i)}(\theta ,s_a) = \sigma_j'(y_j(\theta ,s_a))\odot y_j^{(i)}(\theta ,s_a)$ .
Therefore, the calculation of  $y_{d}$ ,  $y_{d}^{(i)}$  and  $\gamma_{d}^{(i)}$  readily satisfies the standard of Tensor Program which is firstly introduced in [77] and further developed in [76]. Since we require the continuous polynomially bounded third order derivative for  $\sigma$ , the general result in [76] for Tensor Program could apply, and the convergence in law of outputs to Gaussian distribution and almost surely convergence of NTK at initialization is justified. What is left is the computation of the covariance and NTK at the infinite width limit.
The variance could be computed recursively. Given that the covariance matrix of each previous layer is established, and the convergence to Gaussian distribution of each element in previous layer, it is easy to verify the covariance matrix for next layer follows Eq. (21).
We next compute the NTK limit. Note that  $\Theta^{(d)}$ , as gradients inner product, is exactly the same NTK that has been studied in [39, 3, 18]. Moreover,  $(\Xi^{(i,j)})_{\alpha}$  is just the diagonal NTK if we consider a sub neural network which coincide with original network in previous  $j - 1$  layers, but use one single hidden unit  $(y_j)_\alpha$  at  $j$ -th layer as output. Therefore, we only need to take care about  $\Theta^{(i)}$  and  $\Phi^{(i)}$  for  $i < d$ .
$$
\begin{array}{l} \left(\Theta^ {(i)}\right) _ {a, b} = \left(\nabla_ {\theta} y _ {d} ^ {(i)} (s _ {a})\right) ^ {T} \nabla_ {\theta} y _ {d} (s _ {b}) \\ = \sum_ {j = 1} ^ {d} \operatorname {T r} \left(\nabla_ {W _ {j}} y _ {d} ^ {(i)} \left(s _ {a}\right) ^ {T} \nabla_ {W _ {j}} y _ {d} \left(s _ {b}\right)\right) \tag {30} \\ \end{array}
$$
$$
\nabla_ {W _ {j}} y _ {d} ^ {(i)} \left(s _ {a}\right) = c _ {j} x _ {j - 1} ^ {(i)} \left(s _ {a}\right) \nabla_ {y _ {j}} y _ {d} ^ {(i)} \left(s _ {a}\right) + \sum_ {k = \max  (i, j)} ^ {d - 1} c _ {j} x _ {j - 1} \left(s _ {a}\right) \left(\nabla_ {x _ {k}} y _ {d} \left(s _ {a}\right) ^ {T} \odot y _ {k} ^ {(i)} \left(s _ {a}\right) \odot \sigma_ {k} ^ {\prime \prime} \left(y _ {k} \left(s _ {a}\right)\right)\right) ^ {T} \nabla_ {y _ {j}} y _ {k} \left(s _ {a}\right) \tag {31}
$$
The first term is easy to calculate and shares similar form with classical NTK after computing the trace  $\mathrm{Tr}\Bigl (\nabla_{W_j}y_d^{(i)}(s_a)^T\nabla_{W_j}y_d(s_b)\Bigr)$ .
$$
\nabla_ {y _ {j}} y _ {d} ^ {(i)} (s _ {a}) = c _ {d} W _ {d} \operatorname {d i a g} \left(\sigma_ {i - 1} ^ {\prime} \left(y _ {i - 1} (s _ {a})\right)\right) \dots c _ {j + 1} W _ {j + 1} \operatorname {d i a g} \left(\sigma_ {j} ^ {\prime} \left(y _ {j} (s _ {a})\right)\right) = \nabla_ {y _ {j}} y _ {d} (s _ {a}) \tag {32}
$$
Combining above results, we have
$$
\left(\Theta^ {(i)}\right) _ {a, b} = \sum_ {j = 1} ^ {d} \nabla_ {y _ {j}} y _ {d} \left(s _ {a}\right) \nabla_ {y _ {j}} y _ {d} \left(s _ {b}\right) ^ {T} c _ {j} ^ {2} \left(x _ {j - 1} ^ {(i)} \left(s _ {a}\right) x _ {j - 1} \left(s _ {b}\right)\right) \tag {33}
$$
The limit of first term is  $\lim_{k\to j}\nabla_{y_j}y_d(s_a)\nabla_{y_j}y_d(s_b)^T = \prod_{k = j}^{d - 1}\mathbb{E}[\sigma_k'(y_k(s_a))\sigma_k'(y_k(s_b))]$ . In order to calculate the second term  $c_{j}^{2}(x_{j - 1}^{(i)}(s_{a})x_{j - 1}(s_{b}))$ , we note that  $\mathbb{E}_{w\sim \mathcal{N}(0,I)}[ww^T] = I$  and  $(W_{j})_{\alpha}$  is a standard normal random vector, then we have
$$
c _ {j} ^ {2} (x _ {j - 1} ^ {(i)} (s _ {a}) x _ {j - 1} (s _ {b})) = \mathbb {E} [ (c _ {j} (W _ {j}) _ {\alpha} x _ {j - 1} ^ {(i)} (s _ {a})) ^ {T} c _ {j} (W _ {j}) _ {\alpha} x _ {j - 1} (s _ {b}) ] = \mathbb {E} [ c _ {j} ^ {2} (y _ {j} ^ {(i)} (s _ {a}) y _ {j} (s _ {b})) ]
$$
so the limit is  $\Sigma_{j}[y^{(i)}(s_{a}),y(s_{b})]$
All other terms comes from second order derivative of non-linear unit  $\sigma_{k}$  and doesn't show up in classical NTK result. We note in  $\mathrm{Tr}\Big((c_jx_{j - 1}(s_a)(\nabla_{x_k}y_d(s_a)^T\odot y_k^{(i)}(s_a)\odot \sigma_k^{\prime \prime}(y_k(s_a)))^T\nabla_{y_j}y_k(s_a))^T\nabla_{W_j}y_d(s_b)\Big)$
$(c_{j}^{2}x_{j - 1}(s_{a})^{T}x_{j - 1}(s_{b}))(\nabla_{y_{j}}y_{d}(s_{b})\nabla_{y_{j}}y_{k}(s_{a})^{T}(\nabla_{x_{k}}y_{d}(s_{a})^{T}\odot y_{k}^{(i)}(s_{a})\odot \sigma_{k}^{\prime \prime}(y_{k}(s_{a})))$ $c_j^2 x_{j - 1}(s_a)^T x_{j - 1}(s_b) = \mathcal{O}(1)$  converge to a fixed value, in  $\nabla_{y_j}y_d(s_b)\nabla_{y_j}y_k(s_a)^T\odot \nabla_{x_k}y_d(s_a) =$ $\nabla_{x_k}y_d(s_b)\mathrm{diag}\sigma_k'(y_k(s_a))\nabla_{y_j}y_k(s_b)\nabla_{y_j}y_k(s_a)^T\odot \nabla_{x_k}y_d(s_a),\nabla_{y_j}y_k(s_b)\nabla_{y_j}y_k(s_a)^T$  converges to fixed values times identity matrix, and  $\nabla_{x_k}y_d(s_b)^T\nabla_{x_k}y_d(s_b)$  also converges to fixed values times identity matrix. The rest terms converge to its expectation which contains  $y_{k}^{(i)}(s_{a}),\sigma_{k}^{\prime \prime}(y_{k}(s_{a}))$  and  $\sigma_k^\prime (y_k(s_b))$
$$
\begin{array}{l} \left(\Phi^ {(i)}\right) _ {a, b} = \left(\nabla_ {\theta} \gamma_ {d} ^ {(i)} \left(s _ {a}\right)\right) ^ {T} \nabla_ {\theta} y _ {d} \left(s _ {b}\right) \\ = \sum_ {j = 1} ^ {i} \operatorname {T r} \left(\nabla_ {W _ {j}} \gamma_ {d} ^ {(i)} \left(s _ {a}\right) ^ {T} \nabla_ {W _ {j}} y _ {d} \left(s _ {b}\right)\right) + \sum_ {j = i + 1} ^ {d} \operatorname {T r} \left(\nabla_ {W _ {j}} \gamma_ {d} ^ {(i)} \left(s _ {a}\right) ^ {T} \nabla_ {W _ {j}} y _ {d} \left(s _ {b}\right)\right) \tag {34} \\ \end{array}
$$
Note that  $\gamma_d^{(i)}(\theta, s_a) = (\nabla_{x_i} y_d(\theta, s_a)) \sigma_i''(y_i(\theta, s_a))$  and when  $j \leq i$  the value  $y_i(\theta, s_a)$  depends on  $W_i$ . Therefore we conduct a separate discussion. When  $j > i$ , the limit value could be derived in a way similar to  $(\Theta^{(i)})_{a,b}$ . For  $j \leq i$ ,
$$
\begin{array}{l} \nabla_ {W _ {j}} \gamma_ {d} ^ {(i)} (s _ {a}) = c _ {j} x _ {j - 1} ^ {(i)} (s _ {a}) \nabla_ {y _ {j}} \gamma_ {d} ^ {(i)} (s _ {a}) \\ + \sum_ {k = j} ^ {d - 1} c _ {j} x _ {j - 1} \left(s _ {a}\right) \left(\nabla_ {x _ {k}} y _ {d} \left(s _ {a}\right) ^ {T} \odot \gamma_ {k} ^ {(i)} \left(s _ {a}\right) \odot \sigma_ {k} ^ {\prime \prime} \left(y _ {k} \left(s _ {a}\right)\right)\right) ^ {T} \nabla_ {y _ {j}} y _ {k} \left(s _ {a}\right) \tag {35} \\ \end{array}
$$
The summations are again from the second order derivative of non-linear unit and the limit could be derived in a similar way to  $\Phi^{(i)}$ . We focus on first term is somewhat different since it depends on third order derivative of non-linear unit.
$$
\nabla_ {y _ {j}} \gamma_ {d} ^ {(i)} \left(s _ {a}\right) = c _ {d} W _ {d} \operatorname {d i a g} \left(\sigma_ {i - 1} ^ {\prime} \left(y _ {i - 1} \left(s _ {a}\right)\right)\right) \dots c _ {i + 1} W _ {i + 1} \operatorname {d i a g} \left(\sigma_ {i} ^ {\prime \prime \prime} \left(y _ {i} \left(s _ {a}\right)\right)\right) \dots \dots c _ {j + 1} W _ {j + 1} \operatorname {d i a g} \left(\sigma_ {j} ^ {\prime} \left(y _ {j} \left(s _ {a}\right)\right)\right) \tag {36}
$$
$\nabla_{y_j}\gamma_d^{(i)}(s_a)$  differs from  $\nabla_{y_j}y_d^{(i)}(s_a)$  and  $\nabla_{y_j}y_d(s_a)$  in one term  $\sigma_i^{\prime \prime \prime}(y_i(s_a))$  and that fact is reflected in Eq. (25).
Finally we prove that  $\Theta^{(i)},\Phi^{(i)},\Xi^{(i,j)}$  are constant during training. Since these values are all inner product of gradients, we just need to show gradients  $\nabla_{W_j}y_d(\theta ,s_a)$ $\nabla_{W_j}y_d^{(i)}(\theta ,s_a)$  and  $\nabla_{W_j}\gamma_d^{(i)}(\theta ,s_a)$  doesn't change too much. We write down the formula of these gradient without the
summation terms which could be controlled similarly.