SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[1] Martin Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. TensorFlow: a system for large-scale machine learning. In OSDI, 2016.
[2] Francisco E Acosta, Sophia Sanborn, Khanh Dao Duc, Manu Madhav, and Nina Miolane. Quantifying local extrinsic curvature in neural manifolds. arXiv preprint arXiv:2212.10414, 2022.
[3] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-276, 1998.
[4] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020.
[5] Javier Antorán, David Janz, James U Allingham, Erik Daxberger, Riccardo Rb Barbano, Eric Nalisnick, and José Miguel Hernández-Lobato. Adapting the linearised Laplace model evidence for modern deep learning. In ICML, 2022.
[6] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In NeurIPS, 2019.
[7] Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. In ICLR, 2019.
[8] Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. arXiv preprint arXiv:2205.09745, 2022.
[9] Georgios Arvanitidis, Lars Kai Hansen, and Søren Hauberg. Latent space oddity: on the curvature of deep generative models. In ICLR, 2018.
[10] Alberto Bernacchia, Máté Lengyel, and Guillaume Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. In NIPS, 2018.
[11] Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
[61] Behnam Neyshabur, Russ R Salakhutdinov, and Nati Srebro. Path-SGD: Path-normalized optimization in deep neural networks. In NIPS, 2015.
[62] Kazuki Osawa, Shigang Li, and Torsten Hoefler. PipeFisher: Efficient training of large language models using pipelining and Fisher information matrices. arXiv preprint arXiv:2211.14133, 2022.
[63] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. JMLR, 22(57), 2021.
[64] Omiros Papaspiliopoulos, Gareth O Roberts, and Martin Sköld. A general framework for the parametrization of hierarchical models. Statistical Science, 2007.
[65] Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. In ICLR, 2014.
[66] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. In NeurIPS, 2019.
[67] Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. In NeurIPS, 2021.
[68] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICML, 2015.
[69] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. In NIPS, 2016.
[70] Jascha Sohl-Dickstein, Roman Novak, Samuel S Schoenholz, and Jaehoon Lee. On the infinite width limit of neural networks with a standard parameterization. arXiv preprint arXiv:2001.07301, 2020.
[71] Yang Song, Jiaming Song, and Stefano Ermon. Accelerating natural gradient with higher-order invariance. In ICML, 2018.
[72] Michalis K Titsias. Learning model reparametrizations: Implicit variational inference by fitting MCMC distributions. arXiv preprint arXiv:1708.01529, 2017.
[73] Yusuke Tsuzuku, Issei Sato, and Masashi Sugiyama. Normalized flat minima: Exploring scale invariant definition of flat minima for neural networks using PAC-Bayesian analysis. In ICML, 2020.
[74] Jesse van Oostrum, Johannes Müller, and Nihat Ay. Invariance properties of the natural gradient in overparametrised systems. Information Geometry, 2022.
[75] Greg Yang and Edward J Hu. Tensor programs IV: Feature learning in infinite-width neural networks. In ICML, 2021.
[76] Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs V: Tuning large neural networks via zero-shot hyperparameter transfer. In NeurIPS, 2022.
[77] Liu Yang and Rong Jin. Distance metric learning: A comprehensive survey. Michigan State University, 2 (2), 2006.
[78] Guodong Zhang, James Martens, and Roger B Grosse. Fast convergence of natural gradient descent for over-parameterized neural networks. In NeurIPS, 2019.
Appendix A A Simple Working Example
Let  $\Theta := \mathbb{R}^d$  and let  $\mathcal{L} : \mathbb{R}^d \to \mathbb{R}$  defined by  $\mathcal{L}(\pmb{\theta}) := \sum_{i=1}^{d} (\pmb{\theta}_i)^2$  be a twice-differentiable function on  $\Theta$ . Suppose  $\varphi : \Theta \to \Psi$  defined by  $\varphi(\pmb{\theta}) = 3\pmb{\theta} =: \psi$  be the reparametrization of choice. In this example, we shall see how this often-used reparametrization leads to pathologies in the values derived from the Hessian of  $\mathcal{L}$  computed by an autodiff system.
Notice that the Jacobian of  $\varphi$  is given by  $J(\pmb{\theta}) = \mathrm{diag}(3, \dots, 3) \in \mathbb{R}^{d \times d}$  and its inverse is  $J^{-1}(\psi) = \mathrm{diag}(1/3, \dots, 1/3)$ . The automodiff Hessian matrix is given by  $\pmb{H}(\pmb{\theta}) \equiv 2\pmb{I} \in \mathbb{R}^{d \times d}$ . When we transform  $\pmb{\theta} \rightarrow \pmb{\psi}$ , the function  $\mathcal{L}(\pmb{\theta})$  becomes  $\hat{\mathcal{L}}(\psi) := \sum_{i=1}^{d} (1/3\psi_i)^2$ . The automodiff will take into account this change and thus will compute the Hessian  $\hat{\pmb{H}}$  of  $\hat{\mathcal{L}}$  as
$$
\hat {\boldsymbol {H}} _ {i j} (\boldsymbol {\psi}) = \operatorname {d i a g} (1 / 3, \dots , 1 / 3) (2 \boldsymbol {I}) \operatorname {d i a g} (1 / 3, \dots , 1 / 3) = \operatorname {d i a g} (2 / 9, \dots , 2 / 9).
$$
Notice that,  $\operatorname*{det}\pmb {H}(\pmb {\theta})\equiv 2^d$  , meanwhile, we have det  $\hat{H} (\psi)\equiv (2 / 9)^d$  . Therefore, the Hessian determinant computed by an autodiff system is not invariant under reparametrization.
But if instead we explicitly take into account the metric into the Hessian, which by default is (implicitly) chosen to be  $\dot{G} (\theta)\equiv I$ , resulting in the seemingly redundant expression  $E(\theta) = G(\theta)^{-1}H(\theta) = I^{-1}(2I)$ , and perform the correct transformation on both identity matrices, we obtain
$$
\begin{array}{l} \hat {\boldsymbol {E}} (\boldsymbol {\psi}) = \left(\operatorname {d i a g} (1 / 3, \dots , ^ {1} / 3) \boldsymbol {I} \operatorname {d i a g} (1 / 3, \dots , ^ {1} / 3)\right) ^ {- 1} \operatorname {d i a g} (1 / 3, \dots , ^ {1} / 3) (2 \boldsymbol {I}) \operatorname {d i a g} (1 / 3, \dots , ^ {1} / 3) \\ = \operatorname {d i a g} (3, \dots , 3) \operatorname {d i a g} (3, \dots , 3) \operatorname {d i a g} (1 / 3, \dots , 1 / 3) (2 I) \operatorname {d i a g} (1 / 3, \dots , 1 / 3) \\ = \operatorname {d i a g} (3, \dots , 3) (2 I) \operatorname {d i a g} ^ {(1 / 3, \dots , ^ {1 / 3})}. \\ \end{array}
$$
Thus, the determinant of the transformed "metric-conditioned Hessian" equals
$$
\det \hat {\pmb {E}} (\psi) = (\det \mathrm {d i a g} (3, \ldots , 3)) (\det 2 \pmb {I}) (\det \mathrm {d i a g} (1 / 3, \ldots , 1 / 3)) = 3 ^ {d} 2 ^ {d} (1 / 3) ^ {d} = 2 ^ {d},
$$
which coincides with the determinant of the original  $E$ .
Appendix B Derivations
Note. Throughout this section, we use the Einstein summation convention: If the same index appears twice, once as an upper index and once as a lower index, we sum them over. For example:  $z = x^{i}y_{i}$  means  $z = \sum_{i}x^{i}y_{i}$  and  $B_{j}^{k} = A_{ij}^{k}x^{i}$  means  $B_{j}^{k} = \sum_{k}A_{ij}^{k}x^{i}$ , etc. Specifically for partial derivatives, the index of the denominator is always treated as a lower index.
B.1 The Riemannian Hessian Under Reparametrization
Let  $\mathcal{L}: M \to \mathbb{R}$  be a function on a Riemannian manifold  $M$  with metric  $G$ . The Riemannian Hessian  $\mathrm{Hess} \mathcal{L}$  of  $\mathcal{L}$  is defined in coordinates  $\theta$  by
$$
\boldsymbol {H} _ {i j} = \frac {\partial^ {2} \mathcal {L}}{\partial \boldsymbol {\theta} ^ {i} \partial \boldsymbol {\theta} ^ {j}} - \Gamma_ {i j} ^ {k} \frac {\partial \mathcal {L}}{\partial \boldsymbol {\theta} ^ {k}}, \tag {6}
$$
where  $\Gamma_{ij}^{k}$  is the connection coefficient.
Under a change of coordinates  $\varphi :\pmb {\theta}\mapsto \pmb{\psi}$  , we have  $\tilde{\mathcal{L}} = \mathcal{L}\circ \varphi^{-1}$  and
$$
\tilde {\Gamma} _ {i j} ^ {k} = \Gamma_ {m n} ^ {o} \frac {\partial \psi^ {k}}{\partial \theta^ {o}} \frac {\partial \boldsymbol {\theta} ^ {m}}{\partial \psi^ {i}} \frac {\partial \boldsymbol {\theta} ^ {n}}{\partial \psi^ {j}} + \frac {\partial^ {2} \boldsymbol {\theta} ^ {o}}{\partial \psi^ {i} \partial \psi^ {j}} \frac {\partial \psi^ {k}}{\partial \theta^ {o}}, \tag {7}
$$
where  $m, n, o$  are just dummy indices—present to express summations. Note that the transformation rule for  $\Gamma_{ij}^{k}$  implies that it is not a tensor—to be a tensor, there must not be the second term in the previous formula.
Using (7) and the standard chain & product rules to transform the partial derivatives in (6), we obtain the coordinate representation of the transformed Hessian  $\mathrm{Hess}\mathcal{L}$ :
In the matrix form, we can write the above as  $\pmb{H} = \pmb{J}^{-1}\pmb{H}\pmb{J}^{-1}$ , where  $\pmb{J}$  is the Jacobian of  $\varphi$ . Thus, the Riemannian Hessian at any  $\theta$  (not just at critical points) transforms just like the metric and thus invariant as discussed in Example 4. Note: this only holds when the term containing the connection coefficients  $\Gamma_{ij}^{k}$  is explicitly considered. In particular, the Euclidean Hessian does not follow this tensorial transformation under autodiff due to the fact that (i)  $\Gamma_{ij}^{k} = 0$  for any  $i,j,k$  and thus dropped from the equation, and (ii) autodiff is not designed to handle advanced geometric objects like  $\Gamma_{ij}^{k}$ .
B.2 Hessian-Trace Under Reparametrization
Let  $\mathcal{L}:\mathbb{R}^d\to \mathbb{R}$  be a function of  $\mathbb{R}^d$  under the Cartesian coordinates and  $G$  a Riemannian metric. The Riemannian trace of the Hessian matrix  $\pmb{H}$  of  $\mathcal{L}$  is defined by [47]:
$$
\left(\operatorname {t r} _ {\boldsymbol {G}} \boldsymbol {H}\right) (\boldsymbol {\theta}) = \operatorname {t r} \left(\boldsymbol {G} \left(\boldsymbol {\theta}\right) ^ {- 1} \boldsymbol {H} (\boldsymbol {\theta})\right). \tag {9}
$$
That is, it is defined as the standard trace of the Hessian operator  $E$ .
Let  $\varphi : \pmb{\theta} \mapsto \psi$  be a reparametrization on  $\mathbb{R}^d$ . Then, using (2) and the property  $\operatorname{tr}(\boldsymbol{A}\boldsymbol{B}) = \operatorname{tr}(\boldsymbol{B}\boldsymbol{A})$  twice, the Riemannian trace of the Hessian transforms into
$$
\begin{array}{l} (\operatorname {t r} _ {\tilde {G}} \tilde {H}) (\psi) = \operatorname {t r} (\tilde {E} (\psi)) \\ = \operatorname {t r} \left(\left(\boldsymbol {J} ^ {- 1} (\boldsymbol {\psi})\right) ^ {- 1} \boldsymbol {G} \left(\varphi^ {- 1} (\boldsymbol {\psi})\right) \boldsymbol {H} \left(\varphi^ {- 1} (\boldsymbol {\psi})\right) \boldsymbol {J} ^ {- 1} (\boldsymbol {\psi})\right) \tag {10} \\ = \operatorname {t r} (\boldsymbol {G} (\varphi^ {- 1} (\boldsymbol {\psi})) \boldsymbol {H} (\varphi^ {- 1} (\boldsymbol {\psi}))) \\ = \left(\operatorname {t r} _ {\boldsymbol {G}} \boldsymbol {H}\right) \left(\varphi^ {- 1} (\psi)\right). \\ \end{array}
$$
Since  $\psi = \varphi (\pmb {\theta})$ , we have that  $(\mathrm{tr}_{\tilde{G}}\tilde{H})(\psi) = (\mathrm{tr}_G H)(\theta)$  for any given  $\pmb{\theta}$ . Therefore, the trace of the Hessian operator (or the Riemannian trace of the Hessian) is invariant.
B.3 Hessian-Eigenvalues Under Reparametrization
We use the setting from the preceding section. Recall that  $\lambda$  is an eigenvalue of the linear map  $E(\pmb{\theta}) = G(\pmb{\theta})^{-1}H(\pmb{\theta})$  on the tangent space at  $z \in \mathbb{R}^d$  that is represented  $\pmb{\theta}$  if  $E(\pmb{\theta})\pmb{v} = \lambda \pmb{v}$  for an eigenvector  $\pmb{v} \in T_z\mathbb{R}^d$ . We shall show that  $\tilde{\lambda}$ , the eigenvalue under the reparametrization  $\varphi: \pmb{\theta} \mapsto \psi$ , equals the original eigenvalue  $\lambda$ .
Using the transformation rule of  $\pmb{E}(\pmb{\theta})$  in (2) and the transformation rule of tangent vectors in (4), along with the relation  $(\pmb{J}^{-1}(\psi))^{-1} = \pmb{J}(\varphi^{-1}(\psi))$ , we get
$$
\tilde {\boldsymbol {E}} (\psi) \tilde {\boldsymbol {v}} = \bar {\lambda} \tilde {\boldsymbol {v}}
$$
$$
\begin{array}{l} (\boldsymbol {J} ^ {- 1} (\boldsymbol {\psi})) ^ {- 1} \boldsymbol {G} (\varphi^ {- 1} (\boldsymbol {\psi})) \boldsymbol {H} (\varphi^ {- 1} (\boldsymbol {\psi})) \boldsymbol {J} ^ {- 1} (\boldsymbol {\phi}) \boldsymbol {J} (\varphi^ {- 1} (\boldsymbol {\psi})) \boldsymbol {v} = \tilde {\lambda} \boldsymbol {J} (\varphi^ {- 1} (\boldsymbol {\psi})) \boldsymbol {v} \\ \boldsymbol {J} \left(\varphi^ {- 1} (\psi)\right) \boldsymbol {G} \left(\varphi^ {- 1} (\psi)\right) \boldsymbol {H} \left(\varphi^ {- 1} (\psi)\right) \boldsymbol {v} = \tilde {\lambda} \boldsymbol {J} \left(\varphi^ {- 1} (\psi)\right) \boldsymbol {v} \tag {11} \\ \left. \left(\boldsymbol {J} (\varphi^ {- 1} (\boldsymbol {\psi}))\right) ^ {- 1} \boldsymbol {J} (\varphi^ {- 1} (\boldsymbol {\psi})) \boldsymbol {G} (\varphi^ {- 1} (\boldsymbol {\psi})) \boldsymbol {H} (\varphi^ {- 1} (\boldsymbol {\psi})) \boldsymbol {v} = \tilde {\lambda} \boldsymbol {v} \right. \\ \boldsymbol {E} (\varphi^ {- 1} (\psi)) \boldsymbol {v} = \bar {\lambda} \boldsymbol {v}. \\ \end{array}
$$
Since  $\varphi^{-1}(\psi) = \theta$ , we arrive back at the original equation before the reparametrization and thus we conclude that  $\tilde{\lambda} = \lambda$ .
Appendix C The Invariance of the Fisher Metric
We have seen that any metric on the parameter space yields inv(equi)variance when all transformation rules of geometric objects are followed. However, in practical settings, one uses autodiff libraries [1, 66], whose job is to compute only the elementary rules of derivatives such as the chain and product rules. Notice that derivatives and Hessians (at least at critical points) are transformed correctly under autodiff, while the metric is not in general. It is thus practically interesting to find a family of metrics that transform correctly and automatically under reparametrization, given only an autodiff library. Under those metrics, the aforementioned inv(equi)variances can thus be obtained effortlessly.
Martens [55, Sec. 12] mentions the following family of curvature matrices satisfying this transformation behavior
$$
\boldsymbol {B} (\boldsymbol {\theta}) \propto \underset {\boldsymbol {x}, \boldsymbol {y} \sim D} {\mathbb {E}} \left[ \boldsymbol {J} (\boldsymbol {\theta}; \boldsymbol {x}) ^ {\top} \boldsymbol {A} (\boldsymbol {\theta}, \boldsymbol {x}, \boldsymbol {y}) \boldsymbol {J} (\boldsymbol {\theta}; \boldsymbol {x}) \right], \tag {12}
$$
where  $J(\pmb{\theta}; \cdot)$  is the network's Jacobian  $\partial f(\cdot; \pmb{\theta}) / \partial \pmb{\theta}$ , with an arbitrary data distribution  $D$  and an invertible matrix  $A$  that transforms like  $A(\varphi^{-1}(\psi))$  under  $\varphi: \pmb{\theta} \mapsto \psi$ . Under a reparametrization  $\varphi$ , an autodiff library will compute
$$
\hat {\boldsymbol {B}} (\psi) = \boldsymbol {J} ^ {- 1} (\psi) ^ {\top} \boldsymbol {B} \left(\varphi^ {- 1} (\psi)\right) \boldsymbol {J} ^ {- 1} (\psi), \tag {13}
$$
given  $\mathcal{L} \circ \varphi^{-1}$ , due to the elementary transformation rule of  $J(\theta; \cdot)$ , just like Example 4(c). This family includes the Fisher and the generalized Gauss-Newton matrices, as well as the empirical Fisher matrix [44, 55].
However, note that any  $B$  as above is sufficient. Thus, this family is much larger than the Fisher metric, indicating that automatic invariance is not unique to that metric and its denominations. Moreover, in practice, these metrics are often substituted by structural approximations, such as their diagonal and Kronecker factorization [56]. This restricts the kinds of reparametrizations under which these approximate metrics transform automatically: Take the diagonal of (12) as an example. Its  $\psi$ -representation computed by the autodiff library will only match the correct transformation for element-wise reparametrizations, whose Jacobian  $J(\theta)$  is diagonal. On top of that, in practical algorithms, such approximations are usually combined with additional techniques such as damping or momentum, which further break their automaticness.
Due to the above and because any metric is theoretically in(equi)variant if manual intervention is performed, the inv(equi)variance property of the Fisher metric should not be the determining factor of using the Fisher metric. Instead, one should look into its more unique properties such as its statistical efficiency [3] and guarantees in optimization [34, 78]. Automatic transformation is but a cherry on top.
It is interesting for future work to extend autodiff libraries to take into account the invariant transformation rules of geometric objects. By doing so, any metric—not just the Fisher metric—will yield invariance automatically.
Appendix D General Manifold: Local Coordinate Charts
In the main text, we have exclusively used global coordinate charts  $(\mathbb{R}^d,\theta)$  and  $(\mathbb{R}^d,\psi)$  these charts cover the entire  $\mathbb{R}^d$  and  $\theta ,\psi$  are homeomorphisms on  $\mathbb{R}^d$ . However, there exist other coordinate systems that are not global, i.e. they are constructed using multiple charts  $\{(U_i\subseteq \mathbb{R}^d,\theta_i:U_i\to$ $\Theta_{i}\subseteq \mathbb{R}^{d})\}_{i}$  s.t.  $\mathbb{R}^d = \cup_i U_i$  and each  $\theta_{i}$  is a diffeomorphism on  $U_{i}$ .
If  $\{(U_i,\theta_i)\}_i$  and  $\{(V_j,\psi_j)\}_j$  be two local coordinate systems of  $\mathbb{R}^d$ . Let  $(U_i,\theta_i)$  and  $(V_j,\psi_j)$  be an arbitrary pair of charts from the above collections where  $U_i\cap V_j\neq \emptyset$ . In this setting, reparametrization between these two charts amounts to the condition that the transition map  $\varphi$  between  $\theta (U_i\cap V_j)$  and  $\psi (U_i\cap V_j)$  is a diffeomorphism, see Fig. 6. Reparametrization on the whole manifold can then be defined if for all pairs of two charts from two coordinate systems, their transition maps are all diffeomorphism whenever they overlap with each other.
Note that this definition is a generalization of global coordinate charts. In this case, there is only a single chart for each coordinate system, i.e.  $(U,\theta)$  and  $V,\psi$ , and they trivially overlap since  $U = \mathbb{R}^d$  and  $V = \mathbb{R}^d$ . Moreover, there is only a single diffeomorphism of concern, as shown in Fig. 3.
Our results in the main text hold in this general case by simply replacing the domain for the discussion to be the overlapping regions of a pair of two charts, instead of the whole  $\mathbb{R}^d$ .
Appendix E Details on Applications
E.1 Infinite-Width Neural Networks
Figure 6: The diffeomorphism  $\varphi$  is defined from  $\theta(U \cap V)$  to  $\psi(U \cap V)$ .
Note We use a 2-layer NN without bias for clarity. The extension to deep networks follows directly by induction—see e.g. [6, 45]. We also use the same prior variance  $\sigma^2$  without loss of generalization for simplicity. See also Kristiadi et al. [41, Appendix A] for further intuition. Recall the property of Gaussians under linear transformation:  $z \sim \mathcal{N}(\mu, \Sigma) \Rightarrow Az \sim \mathcal{N}(A\mu, A\Sigma A^\top)$ .
Neural-network Gaussian process (NNGP)
Let  $f(\pmb{x}): \mathbb{R}^d \to \mathbb{R}$  defined by  $f(\pmb{x}) \coloneqq \pmb{w}^\top \phi(\pmb{W}\pmb{x})$  be a real-valued, 2-layer NN with parameter  $\pmb{\theta} \coloneqq \{\pmb{W} \in \mathbb{R}^{h \times n}, \pmb{w} \in \mathbb{R}^h\}$  and component-wise nonlinearity  $\phi$ . Note that we assume  $\pmb{x}$  is i.i.d. Let  $\operatorname{vec}(\pmb{W}) \sim \mathcal{N}(\pmb{0}, \sigma^2\pmb{I})$  and  $\pmb{w} \sim \mathcal{N}(\pmb{0}, \frac{\sigma^2}{h}\pmb{I})$  be priors over the weights. This parametrization of the weights and the priors is called the standard parametrization (SP) [45, 60].
Step (a) Given a particular preactivation value  $z$ , we have a linear model  $f(\boldsymbol{x}) = \boldsymbol{w}^\top \phi(\boldsymbol{z})$ . We can view this as a Gaussian process with mean and covariance
$$
\mathbb {E} [ f (\boldsymbol {x}) ] = \boldsymbol {0} ^ {\top} \phi (\boldsymbol {z}) = 0,
$$
$$
\operatorname {C o v} [ f (\boldsymbol {x}), f (\boldsymbol {x} ^ {\prime}) ] = \frac {\sigma^ {2}}{h} \phi (\boldsymbol {z}) ^ {\top} \phi (\boldsymbol {z} ^ {\prime}) = \frac {\sigma^ {2}}{h} \sum_ {i = 1} ^ {h} \phi (\boldsymbol {z} _ {i}) \phi (\boldsymbol {z} _ {i} ^ {\prime}).
$$
Taking the limit as  $h \to \infty$ , the mean stays trivially zero and we have by the law of large numbers:
$$
K(\boldsymbol {x},\boldsymbol{x}^{\prime}):= \lim_{h\to \infty}\operatorname {Cov}[f(\boldsymbol {x}),f(\boldsymbol{x}^{\prime})] = \lim_{h\to \infty}\frac{\sigma^{2}}{h}\sum_{i = 1}^{h}\phi (\boldsymbol{z}_{i})  \phi (\boldsymbol{z}_{i}^{\prime}) = \sigma^{2}  \mathbb{E}_{\boldsymbol{z}_{i},\boldsymbol{z}_{i}^{\prime}}[\phi (\boldsymbol{z}_{i})\phi (\boldsymbol{z}_{i}^{\prime})].
$$
In particular, both the mean and covariance over the output does not depend on the particular realization of the hidden units  $\phi(z)$ . That is, they only depend on the distribution of  $z$  induced by the prior, which we will obtain now.
Step (b) Notice that each  $z_{i} = W_{i}^{\top}\pmb{x}$  where  $\pmb{W}_{i}$  is the  $i$ -th row of  $\pmb{W}$ . This is a linear model and thus,  $\pmb{z}$  is distributed as a GP with mean and covariance
$$
\mathbb {E} [ \boldsymbol {z} _ {i} ] = \boldsymbol {0} ^ {\top} \boldsymbol {x} = 0,
$$
$$
\operatorname {C o v} \left[ \boldsymbol {z} _ {i}, \boldsymbol {z} _ {i} ^ {\prime} \right] = \sigma^ {2} \boldsymbol {x} ^ {\top} \boldsymbol {x} ^ {\prime} =: K _ {z} \left(\boldsymbol {z} _ {i}, \boldsymbol {z} _ {i} ^ {\prime}\right).
$$
So,  $\mathbf{z}_i \sim \mathcal{GP}(0, K_z)$  Since the prior over  $\mathbf{W}$  is i.i.d., this holds for all  $i = 1, \dots, h$ . We can thus now compute the expectation in  $K(\mathbf{x}, \mathbf{x}')$ : it is done w.r.t. this GP over  $\mathbf{z}_i$ .
To obtain the GP over the function output of a deep network, simply apply steps (a) and (b) above recursively. The crucial message from this derivation is that as the width of each layer of a deep net goes to infinity, the network loses representation power—the output of each layer only depends on the prior, and not on particular values (e.g. learned) of the previous hidden units. In this sense, an infinite-width  $L$ -layer NN is simply a linear model with a constant feature extractor induced by the network's first  $L - 1$  layers that are fixed at initialization. Note that the kernel  $K$  over the function output is called the NNGP kernel [45].
Neural tangent kernel (NTK)
Let us transform  $\pmb{w}$  into  $\pmb{v} \coloneqq \frac{\sigma}{\sqrt{h}}\pmb{w}$  and  $\pmb{W}$  into  $\pmb{V} \coloneqq \frac{\sigma}{\sqrt{h}}\pmb{W}$  and define the prior to be  $\pmb{w} \sim \mathcal{N}(\pmb{0},\pmb{I})$  and  $\operatorname{vec}(\pmb{W}) \sim \mathcal{N}(\pmb{0},\pmb{I})$ . Then, we define the transformed network as  $\hat{f}(\pmb{x}) \coloneqq \pmb{v}^{\top}\phi(\pmb{V}\pmb{x}) = \frac{\sigma}{\sqrt{h}}\pmb{w}^{\top}\phi\left(\frac{\sigma}{\sqrt{h}}\pmb{W}\pmb{x}\right)$  with parameter  $\psi \coloneqq \{\pmb{V} \in \mathbb{R}^{h\times n}, \pmb{v} \in \mathbb{R}^{h}\}$ . This is called the NTK parametrization (NTP) [29]. We will see below that even though  $\pmb{v}$ ,  $\pmb{V}$  have the same prior distributions as  $\pmb{w}$ ,  $\pmb{W}$  in the SP, they have different behavior in terms of the NTK.