SlowGuess/NeurIPS_2023_OCR huggingface dataset:
As before, let us assume a particular preactivation value  $z$ . The empirical NTK (i.e. finite-width NTK) on the last layer is defined by:
$$
\hat {\mathcal {K}} (\boldsymbol {x}, \boldsymbol {x} ^ {\prime}) := \langle \nabla_ {\boldsymbol {w}} \hat {f} (\boldsymbol {x}), \nabla_ {\boldsymbol {w}} \hat {f} (\boldsymbol {x} ^ {\prime}) \rangle = \frac {\sigma^ {2}}{h} \sum_ {i = 1} ^ {h} \phi (\boldsymbol {z} _ {i}) \phi (\boldsymbol {z} _ {i} ^ {\prime}).
$$
The (asymptotic)  $NTK$  is obtained by taking the limit of  $h\to \infty$ :
$$
\mathcal {K} \left(\boldsymbol {x}, \boldsymbol {x} ^ {\prime}\right) := \lim  _ {h \rightarrow \infty} \hat {\mathcal {K}} \left(\boldsymbol {x}, \boldsymbol {x} ^ {\prime}\right) = \frac {\sigma^ {2}}{h} \sum_ {i = 1} ^ {h} \phi \left(\boldsymbol {z} _ {i}\right) \phi \left(\boldsymbol {z} _ {i} ^ {\prime}\right) = \underset {\boldsymbol {z} _ {i}, \boldsymbol {z} _ {i} ^ {\prime}} {\mathbb {E}} \left[ \phi \left(\boldsymbol {z} _ {i}\right) \phi \left(\boldsymbol {z} _ {i} ^ {\prime}\right)\right], \tag {14}
$$
which coincides the NNGP kernel  $K$ . Crucially, this is obtained via a backward propagation from the output of the network and thus the linear-Gaussian property we have used to derive the NNGP via forward propagation does not apply. This is why the scaling of  $\frac{\sigma}{h}$  is required in the NTP. That is, using the SP, the empirical NTK is not scaled by  $\frac{\sigma^2}{h}$  and thus when taking the limit to obtain  $\mathcal{K}$ , the sum diverges and the limit does not exist.
Is the NTP a reparametrization of the SP?
It is tempting to treat the NTP as a reparametrization of the SP—in fact, it is standard in the literature to treat them as two different parametrizations of the same network. However, we show that geometrically, this is inaccurate. Indeed from the geometric perspective, if two functions are reparametrization of each other, they should be invariant, as we have discussed in the main text. Instead, we show that the different limiting behaviors are present because the NTP and SP assume two different functions and two different priors—they are not connected by a reparametrization. This clears up confusion and provides a foundation for future work in this field: To obtain a desired limiting behavior, study the network architecture and its prior, instead of the parametrization.
Suppose  $\psi$  in the NTP is a reparametrization of  $\theta$  in the SP. Then the function  $\varphi : \theta \mapsto \psi$  defined by  $\theta \mapsto \frac{\sigma}{\sqrt{h}}\theta$  is obviously the smooth reparametrization with an invertible (diagonal) Jacobian  $J(\theta) = \frac{\sigma}{\sqrt{h}}I$ . In this case, the network in the NTP must be defined by  $\tilde{f} = f \circ \varphi^{-1}$ , where  $f$  is the SP-network, by Example 4. That is, with some abuse of notation,
$$
\tilde {f} (\boldsymbol {x}) = \varphi^ {- 1} (\boldsymbol {v}) ^ {\top} \phi (\varphi^ {- 1} (\boldsymbol {V}) \boldsymbol {x}) = \boldsymbol {w} ^ {\top} \phi (\boldsymbol {W} \boldsymbol {x}) = f (\boldsymbol {x}).
$$
This is different from the definition of the NTP-network  $\hat{f}(\boldsymbol{x}) = \boldsymbol{v}^\top \phi(\boldsymbol{V}\boldsymbol{x})$ . So, obviously, the NTP is not the reparametrization of the SP. Therefore, a clearer way of thinking about the NTP and SP is to treat them as two separate network functions (i.e. two separate architectures)—the scaling factor  $\frac{\sigma}{\sqrt{h}}$  should be thought of as part of the layer's functional form instead of as part of the parameter. In particular, they are not two representations of a single abstract function.
To verify this, let us compute the NTK of  $\tilde{f}(\pmb{x})$  (i.e. treating the scaling as a reparametrization) at its last layer. The derivation is based on Section 3.2. First, notice that the differential  $\nabla_{\pmb{w}} f(\pmb{x})$  transforms into  $\pmb{J}^{-1}(\pmb{v})^{\top} \pmb{\nabla} \tilde{f}(\pmb{x})|_{\varphi^{-1}(\pmb{v})}$  for any  $\pmb{x} \in \mathbb{R}^n$ . Next, notice that the Euclidean metric transforms into  $\tilde{\pmb{G}}(\pmb{v}) := \pmb{J}^{-1}(\pmb{v})^{\top} \pmb{J}^{-1}(\pmb{v})$ . So the gradient transforms into  $\pmb{J}(\varphi^{-1}(\pmb{v})) \pmb{\nabla} f(\pmb{x})|_{\varphi^{-1}(\pmb{v})}$ . Therefore,
<sup>3</sup>It still applies for obtaining the distribution of  $z_{i}$ . The NTK can thus be thought of as a kernel that arises from performing forward and backward propagations once at initialization [6, 75]. This can be seen in the expression of the NTKs on lower layers which decompose into the NNGP and an expression similar to (14), but involving the derivative of  $\phi$  [29].
<sup>4</sup>We use the gradient to get the NTK since otherwise it does not make sense to take the inner product of differentials w.r.t. the metric.
Table 2: Test accuracies, averaged over 5 random seeds.
Methods	MNIST	FMNIST	CIFAR10	CIFAR100
SGD	99.3	92.9	94.9	76.8
ADAM	99.2	92.6	92.4	71.9
the empirical NTK kernel  $\hat{\mathcal{K}}_{\Psi}$  for  $\tilde{f}$  is given by
$$
\begin{array}{l} \hat {\mathcal {K}} _ {\Psi} (\boldsymbol {x}, \boldsymbol {x} ^ {\prime}) = \left\langle J \left(\varphi^ {- 1} (\boldsymbol {v})\right) \nabla f (\boldsymbol {x}) \right| _ {\varphi^ {- 1} (\boldsymbol {v})}, J \left(\varphi^ {- 1} (\boldsymbol {v})\right) \nabla f \left(\boldsymbol {x} ^ {\prime}\right) \mid_ {\varphi^ {- 1} (\boldsymbol {v})} \rangle \tilde {\boldsymbol {G}} (\boldsymbol {v}) \\ = \left(\boldsymbol {J} \left(\varphi^ {- 1} (\boldsymbol {v})\right) \nabla f (\boldsymbol {x}) \mid_ {\varphi^ {- 1} (\boldsymbol {v})}\right) ^ {\top} \tilde {\boldsymbol {G}} (\boldsymbol {v}) \boldsymbol {J} \left(\varphi^ {- 1} (\boldsymbol {v})\right) \nabla f \left(\boldsymbol {x} ^ {\prime}\right) \mid_ {\varphi^ {- 1} (\boldsymbol {v})} \\ = \left(\nabla f (\boldsymbol {x}) | _ {\varphi^ {- 1} (\boldsymbol {v})}\right) ^ {\top} \underbrace {\boldsymbol {J} \left(\varphi^ {- 1} (\boldsymbol {v})\right) ^ {\top} \boldsymbol {J} ^ {- 1} (\boldsymbol {v}) ^ {\top} \boldsymbol {J} ^ {- 1} (\boldsymbol {v}) \boldsymbol {J} \left(\varphi^ {- 1} (\boldsymbol {v})\right)} \nabla f \left(\boldsymbol {x} ^ {\prime}\right) | _ {\varphi^ {- 1} (\boldsymbol {v})} \\ = \left\langle \nabla f (\boldsymbol {x}) \right| _ {\varphi^ {- 1} (\boldsymbol {v})}, \left. \nabla f (\boldsymbol {x} ^ {\prime}) \right| _ {\varphi^ {- 1} (\boldsymbol {v})} \rangle . \\ \end{array}
$$
Thus, the empirical NTK is invariant and the asymptotic NTK also is. Therefore, we still have a problem with the NTK blow-up in this parametrization. This reinforces the fact that the difference between the SP and NTP is not because of parametrization.
Additionally, let us now inspect the priors in the SP and NTP. In the SP, the prior is  $\mathcal{N}(\pmb{\theta} \mid \mathbf{0}, \sigma^2 / h\mathbf{I})$ . Therefore, so that we have the same prior in both  $\Theta$  and  $\Psi$ , the prior of  $\psi = \varphi(\pmb{\theta})$  must be  $\mathcal{N}(\psi \mid \mathbf{0}, \mathbf{I})$ . This is obviously not the case since we have  $\mathcal{N}(\psi \mid \mathbf{0}, \sigma^2 / h\mathbf{I})$  because the NTP explicitly defines  $\mathcal{N}(\pmb{\theta} \mid \mathbf{0}, \mathbf{I})$  as the prior of  $\pmb{\theta}$ . Thus, not only that the SP and NTP assume two different architectures, but they also assume two different priors altogether. It is thus not surprising that the distribution over their network outputs  $f(x)$ ,  $\hat{f}(x)$  are different, both in the finite- and infinite-width regimes.
**Implication** In his seminal work, Neal [60] concluded that the fact that infinite-width NNs are Gaussian processes disappointing. However, as we have seen in the discussion above, different functional forms, architectures, and priors of NNs yield different limiting behaviors. Therefore, this gives us hope that meaningful, non-GP infinite-width NNs can be obtained. Indeed, Yang and Hu [75], Yang et al. [76] have recently shown us a way to do so. However, they argue that their feature-learning limiting behavior is due to a different parametrization, contrary to the present work. Our work thus complements theirs and opens up the avenue for constructing non-trivial infinite-width NNs in a "Bayesian" way, in the sense that we achieve the desired limiting behaviors by varying the model and the prior.
E.2 Biases of Preconditioned Optimizers
For MNIST and FMNIST, the network is LeNet. Meanwhile, we use the WiderResNet-16-4 model for CIFAR-10 and -100. For ADAM, we use the default setting suggested by Kingma and Ba [39]. For SGD, we use the commonly-used learning rate of 0.1 with Nesterov momentum 0.9 [26]. The cosine annealing method is used to schedule the learning rate for 100 epochs. The test accuracies are in Table 2. Additionally, in Table 3, we discuss the effect of reparametrization to sharpness on ADAM and SGD.
E.3 Laplace Marginal Likelihood
Let  $\theta_{\mathrm{MAP}}$  be a MAP estimate in an arbitrary  $\theta$ -coordinates of  $\mathbb{R}^d$ , obtained by minimizing the MAP loss  $\mathcal{L}_{\mathrm{MAP}}$ . Let  $\log h = -\mathcal{L}_{\mathrm{MAP}}$ -note that  $\mathcal{L}_{\mathrm{MAP}}$  itself is a log-density function. The Laplace marginal likelihood [18, 28, 54] is obtained by performing a second-order Taylor's expansion:
$$
\log h (\boldsymbol {\theta}) \approx \log h \left(\boldsymbol {\theta} _ {\mathrm {M A P}}\right) - \frac {1}{2} \left(\boldsymbol {\theta} - \boldsymbol {\theta} _ {\mathrm {M A P}}\right) ^ {\top} \boldsymbol {H} \left(\boldsymbol {\theta} _ {\mathrm {M A P}}\right) \left(\boldsymbol {\theta} - \boldsymbol {\theta} _ {\mathrm {M A P}}\right),
$$
where  $H(\theta_{\mathrm{MAP}})$  is the Hessian matrix of  $\mathcal{L}_{\mathrm{MAP}}$  at  $\theta_{\mathrm{MAP}}$ . Then, by exponentiating and taking the integral over  $\theta$ , we have
$$
Z \left(\boldsymbol {\theta} _ {\mathrm {M A P}}\right) \approx h \left(\boldsymbol {\theta} _ {\mathrm {M A P}}\right) \int_ {\mathbb {R} ^ {d}} \exp \left(- \frac {1}{2} \left(\boldsymbol {\theta} - \boldsymbol {\theta} _ {\mathrm {M A P}}\right) ^ {\top} \boldsymbol {H} \left(\boldsymbol {\theta} _ {\mathrm {M A P}}\right) \left(\boldsymbol {\theta} - \boldsymbol {\theta} _ {\mathrm {M A P}}\right)\right) d \boldsymbol {\theta}. \tag {15}
$$
Table 3: Hessian-based sharpness measures can change under reparametrization without affecting the model's generalization (results on CIFAR-10). The generalization gap is the test accuracy, subtracted from the train accuracy—lower is better. Under the default parametrization, SGD achieves lower sharpness and generalizes better than ADAM which achieves higher sharpness. However, one can reparametrize SGD's minimum s.t. it achieves much higher (or lower) sharpness than ADAM while retaining the same generalization performance. Hence, it is hard to study the correlation between sharpness and generalization. This highlights the need for invariance.
Optimizer	Reparametrization ψMAP = φ(θMAP)	Generalization gap [%]	Sharpness tr(H(ψMAP))
ADAM	ψMAP = θMAP	7.2 ± 0.2	1929.8 ± 61.2
SGD	ψMAP = θMAP	5.2 ± 0.2	1531.8 ± 14.2
ψMAP = 1/2θMAP	6143.7 ± 60.8
ψMAP = 2θMAP	383.6 ± 3.3
Since the integral is the normalization constant of the Gaussian  $\mathcal{N}(\pmb{\theta} \mid \pmb{\theta}_{\mathrm{MAP}}, \pmb{H}(\pmb{\theta}_{\mathrm{MAP}}))$ , we obtain the Laplace log-marginal likelihood (LML):
$$
\log Z (\boldsymbol {\theta} _ {\mathrm {M A P}}) = - \mathcal {L} _ {\mathrm {M A P}} (\boldsymbol {\theta} _ {\mathrm {M A P}}) - \frac {d}{2} \log (2 \pi) + \log \det  \boldsymbol {H} (\boldsymbol {\theta} _ {\mathrm {M A P}}).
$$
Notice that  $H(\theta)$  is a bilinear form, acting on the tangent vector  $d(\theta_{\mathrm{MAP}}) \coloneqq (\theta - \theta_{\mathrm{MAP}})$ . Under a reparametrization  $\varphi : \theta \mapsto \psi$  with  $\psi_{\mathrm{MAP}} = \varphi(\theta_{\mathrm{MAP}})$ , the term inside the exponent in (15) transforms into
$$
\begin{array}{l} - \frac {1}{2} \left(\boldsymbol {J} \left(\varphi^ {- 1} \left(\psi_ {\mathrm {M A P}}\right)\right) \boldsymbol {d} \left(\varphi^ {- 1} \left(\psi_ {\mathrm {M A P}}\right)\right)\right) ^ {\top} \left(\boldsymbol {J} ^ {- 1} \left(\psi_ {\mathrm {M A P}}\right) ^ {\top} \boldsymbol {H} \left(\varphi^ {- 1} \left(\psi_ {\mathrm {M A P}}\right)\right) \boldsymbol {J} ^ {- 1} \left(\psi_ {\mathrm {M A P}}\right)\right) \\ J \left(\varphi^ {- 1} \left(\psi_ {\mathrm {M A P}}\right)\right) d \left(\varphi^ {- 1} \left(\psi_ {\mathrm {M A P}}\right)\right), \\ \end{array}
$$
due to the transformations of the tangent vector and the bilinear-Hessian. This simplifies into
$$
\exp \left(- \frac {1}{2} \boldsymbol {d} \left(\varphi^ {- 1} \left(\psi_ {\mathrm {M A P}}\right)\right) ^ {\top} \boldsymbol {H} \left(\varphi^ {- 1} \left(\psi_ {\mathrm {M A P}}\right)\right) \boldsymbol {d} \left(\varphi^ {- 1} \left(\psi_ {\mathrm {M A P}}\right)\right)\right)
$$
which always equals the original integrand in (15). Thus, the integral evaluates to the same value. Hence, the last two terms of  $\log Z(\theta_{\mathrm{MAP}})$  transform into  $-\frac{d}{2}\log (2\pi) + \log \det H(\varphi^{-1}(\psi_{\mathrm{MAP}}))$  in  $\psi$ -coordinates. This quantity is thus invariant under reparametrization since it behaves like standard functions.
E.3.1 Experiment Setup
We use the toy regression dataset of size 150. Training inputs are sampled uniformly from  $[0,8]$ , while training targets are obtained via  $y = \sin x + \epsilon$ , where  $\epsilon \sim \mathcal{N}(0,0.3^2)$ . The network is a 1-hidden layer TanH network trained for 1000 epochs.