SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The asymptotic normality of kernel-based regression roughly states that under some mild conditions,  $\sqrt{n} (g_{\pi_n,\lambda_n} - g_{\pi ,\lambda_0})$  is asymptotically normal as  $n\to \infty$ . This result provides the theoretical foundations to conduct asymptotic statistical inference for kernel-based regression, in particular, building asymptotic confidence interval for the ground-truth  $g_{\pi ,\lambda_0}$ .
To be rigorous, we first list the following assumptions introduced by [50]:
Assumption B.3. Let  $(\Omega, \mathcal{A}, \mathcal{Q})$  be a probability space. Suppose that  $\mathcal{X} \subset \mathbb{R}^d$  is closed and bounded with Borel- $\sigma$ -algebra  $\mathcal{B}(\mathcal{X})$ , and  $\mathcal{Y} \subset \mathbb{R}$  is closed with Borel- $\sigma$ -algebra  $\mathcal{B}(\mathcal{Y})$ . The Borel- $\sigma$ -algebra of  $\mathcal{X} \times \mathcal{Y}$  is denoted by  $\mathcal{B}(\mathcal{X} \times \mathcal{Y})$ . Let  $\mathcal{H}$  be an RKHS with kernel  $k$  and let  $\pi$  be a probability measure on  $(\mathcal{X} \times \mathcal{Y}, \mathcal{B}(\mathcal{X} \times \mathcal{Y}))$ . Assume that the kernel of  $\mathcal{H}$ ,  $k: \mathcal{X} \times \mathcal{X} \to \mathbb{R}$  is the restriction of an  $m$ -times continuously differentiable kernel  $\tilde{k}: \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}$  such that  $m > d/2$  and  $k \neq 0$ . Let  $\lambda_0 \in (0, +\infty)$  be any positive number. Suppose that the sequence  $\lambda_n$  satisfies that  $\lambda_n - \lambda_0 = o\left(\frac{1}{\sqrt{n}}\right)$ .
Assumption B.4. Let  $\mathcal{L}:\mathcal{V}\times \mathbb{R}\to [0, + \infty)$  be a loss function satisfying the following conditions:
-  $\mathcal{L}$  is a convex loss, i.e.,  $z \mapsto L(y,z)$  is convex in  $z$  for every  $y \in \mathcal{V}$ .
- The partial derivatives
$$
\mathcal {L} ^ {\prime} (y, z) = \frac {\partial}{\partial z} \mathcal {L} (y, z), \quad \mathcal {L} ^ {\prime \prime} (y, z) = \frac {\partial^ {2}}{\partial^ {2} z} \mathcal {L} (y, z)
$$
exist for every  $(y,z)\in \mathcal{V}\times \mathbb{R}$
The maps
$$
(y, z) \mapsto \mathcal {L} ^ {\prime} (y, z), \quad (y, z) \mapsto \mathcal {L} ^ {\prime \prime} (y, z)
$$
are continuous.
- There is a  $b \in L^{2}(\pi_{Y})$ , and for every  $a \in (0, +\infty)$ , there is a  $b_{a}' \in L^{2}(\pi_{Y})$  and a constant  $b_{a}'' \in (0, +\infty)$  such that, for every  $y \in \mathcal{V}$ ,
$$
| \mathcal {L} (y, z) | \leq b (y) + | z | ^ {p} \forall z, \sup  _ {z \in [ - a, a ]} | \mathcal {L} ^ {\prime} (y, z) | \leq b _ {a} ^ {\prime} (y), \sup  _ {z \in [ - a, a ]} | \mathcal {L} ^ {\prime \prime} (y, z) | \leq b _ {a} ^ {\prime \prime}
$$
where  $p \geq 1$  is a constant.
Note that the conditions on the loss function  $\mathcal{L}$  in Assumption B.4 are satisfied, e.g., for the logistic loss for classification or least-square loss for regression with  $\pi$  such that  $\mathbb{E}_{\pi_Y}[Y^4] < \infty$ . Therefore, Assumption B.4 is reasonable for the kernel ridge regression problem we consider in this work.
Theorem B.5 (Theorem 3.1 in [50]). Suppose that Assumptions B.3 and B.4 hold. Then there is a tight, Borel-measurable Gaussian process  $\mathbb{G}:\Omega \to \mathcal{H},\omega \mapsto \mathbb{G}(\omega)$  such that
$$
\sqrt {n} \left(g _ {\pi_ {n}, \lambda_ {n}} - g _ {\pi , \lambda_ {0}}\right) \Rightarrow \mathbb {G} i n \mathcal {H}
$$
where  $\Rightarrow$  represents "converges weakly". The Gaussian process  $\mathbb{G}$  is zero-mean; i.e.,  $\mathbb{E}[\langle g,\mathbb{G}\rangle_{\mathcal{H}}] = 0$  for every  $g\in \mathcal{H}$ .
Note that Theorem B.5 implies that for any  $g \in \mathcal{H}$ , the random variable
$$
\Omega \to \mathbb {R}, \quad \omega \mapsto \langle g, \mathbb {G} (\omega) \rangle_ {\mathcal {H}}
$$
has a zero-mean normal distribution. In particular, letting  $g = k_x$ , the reproducing property of  $k$  implies that,
$$
\sqrt {n} \left(g _ {\pi_ {n}, \lambda_ {n}} (x) - g _ {\pi , \lambda_ {0}} (x)\right) \Rightarrow \mathbb {G} (x)
$$
To obtain the variance of this limiting distribution, denoted as  $\xi^2 (x)$ , we review the tool of influence functions as follows.
Let  $P$  be a general distribution on the domain  $\mathcal{X}\times \mathcal{Y}$ . Let  $z = (z_{x},z_{y})\in \mathcal{X}\times \mathcal{Y}$  be a general point. Let  $T$  be a general statistical functional that maps a distribution  $P$  to a real value, that is,  $T:P\to T(P)\in \mathbb{R}$ . The Gâteaux derivative of  $T(P),T_P'(\cdot)$ , is defined as
$$
T _ {P} ^ {\prime} (Q) := \lim  _ {\varepsilon \rightarrow 0} \frac {T (P + \varepsilon Q) - T (P)}{\varepsilon}
$$
$$
P _ {\varepsilon , z} := (1 - \varepsilon) P + \varepsilon \delta_ {z} = P + \varepsilon (\delta_ {z} - P)
$$
where  $\delta_z$  denotes the point mass distribution in  $z$ . Then the influence function of  $T(P)$  at the point  $z$  is a special Gâteaux derivative defined as
$$
I F (z; T, P) := \lim  _ {\varepsilon \rightarrow 0} \frac {T \left(P _ {\varepsilon , z}\right) - T (P)}{\varepsilon} = T _ {P} ^ {\prime} (\delta_ {z} - P).
$$
Under some mild conditions (e.g.,  $T$  is Hadamard differentiable [38, 110]), the central limit theorem for the statistical functional  $T$  holds:
$$
\sqrt {n} (T (\pi_ {n}) - T (\pi)) \Rightarrow \mathcal {N} \left(0, \int_ {z} I F ^ {2} (z; T, \pi) d \pi (z)\right)
$$
as  $n\to \infty$  where  $\pi_{n}$  is the empirical distribution associated with  $n$  samples i.i.d. drawn from  $\pi$ .
Applying the above concepts to the kernel-based regression, we let  $T(P)$  be the solution to the following problem:
$$
\min _ {g \in \mathcal {H}} \mathbb {E} _ {P} \left[ (Y - g (X)) ^ {2} \right] + \lambda_ {0} \| g \| _ {\mathcal {H}} ^ {2}.
$$
Note that [50] has shown that  $\sqrt{n} (g_{\pi_n,\lambda_n}(x) - g_{\pi ,\lambda_0}(x))$  has the same limiting behavior as  $\sqrt{n} (g_{\pi_n,\lambda_0}(x) - g_{\pi ,\lambda_0}(x))$  so only the limiting regularization hyper-parameter  $\lambda_0$  is used in  $T(P)$ . The proof of Theorem 3.1 in [50] provides the closed-formed expression of the Gâteaux derivative  $T_P'(Q)$  as follows:
$$
T _ {P} ^ {\prime} (Q) = - S _ {P} ^ {- 1} (\mathbb {E} _ {Q} [ \mathcal {L} ^ {\prime} (Y, g _ {P, \lambda_ {0}} (X)) \Phi (X) ])
$$
for  $Q \in \mathrm{lin}(B_S)$  where  $\mathrm{lin}(B_S)$  corresponds to a subset of finite measures on  $(\mathcal{X} \times \mathcal{Y}, \mathcal{B}(\mathcal{X} \times \mathcal{Y}))$  (See Proposition A.8. in [50]) and  $S_P: \mathcal{H} \to \mathcal{H}$  is defined by
$$
S _ {P} (f) = 2 \lambda_ {0} f + \mathbb {E} _ {P} \left[ \mathcal {L} ^ {\prime \prime} \left(Y, g _ {P, \lambda_ {0}} (X)\right) f (X) \Phi (X) \right]
$$
where  $\Phi$  is the feature map of  $\mathcal{H}$  (e.g., we can simply take  $\Phi(x)(\cdot) = k(x, \cdot)$ ). Note that  $S_P$  is a continuous linear operator that is invertible (See Proposition A.5. in [50]). In particular, letting  $Q = \delta_z - P$ , we obtain that
$$
T _ {P} ^ {\prime} (\delta_ {z} - P) = S _ {P} ^ {- 1} \left(\mathbb {E} _ {P} \left[ \mathcal {L} ^ {\prime} \left(Y, g _ {P, \lambda_ {0}} (X)\right) \Phi (X) \right]\right) - \mathcal {L} ^ {\prime} \left(z _ {y}, g _ {P, \lambda_ {0}} (z _ {x})\right) S _ {P} ^ {- 1} \Phi (z _ {x}). \tag {14}
$$
Note that the above special Gâteaux derivative (influence function) of  $T(P), T_P'(\delta_z - P)$ , had been derived initially in [25].
The proof of Theorem 3.1 in [50] shows that  $T(P)$  is Hadamard differentiable, and the asymptotic normality holds. Using the above expression of the influence function, we conclude that:
$$
\sqrt {n} \left(g _ {\pi_ {n}, \lambda_ {n}} (x) - g _ {\pi , \lambda_ {0}} (x)\right) \Rightarrow \mathcal {N} \left(0, \xi^ {2} (x)\right)
$$
$$
\xi^ {2} (x) = \int_ {z \in \mathcal {X} \times \mathcal {Y}} I F ^ {2} (z; T, \pi) (x) d \pi (z)
$$
and  $IF(z;T,\pi)$  is given by (14):
$$
I F (z; T, \pi) = T _ {\pi} ^ {\prime} (\delta_ {z} - \pi) = S _ {\pi} ^ {- 1} (\mathbb {E} _ {\pi} [ \mathcal {L} ^ {\prime} (Y, g _ {\pi , \lambda_ {0}} (X)) \Phi (X) ]) - \mathcal {L} ^ {\prime} (z _ {y}, g _ {\pi , \lambda_ {0}} (z _ {x})) S _ {\pi} ^ {- 1} \Phi (z _ {x}).
$$
Summarizing the above discussion, Proposition 3.4 follows.
B.3 Infinitesimal Jackknife for Kernel-Based Regression
In this section, we follow up on our discussion in Section 3.3 on the infinitesimal jackknife variance estimation. We derive the closed-formed expression of the infinitesimal jackknife variance estimation for  $\xi^2 (x)$  in Section B.2 in the kernel ridge regression. We also show the consistency of infinitesimal jackknife variance estimation in general kernel-based regression. Our consistency result appears new in the literature.
First, we estimate  $\xi^2 (x_0)$  as follows:
$$
\hat {\xi} ^ {2} (x _ {0}) = \frac {1}{n} \sum_ {z _ {i} \in \mathcal {D}} I F ^ {2} (z _ {i}; T, \pi_ {n}) (x _ {0})
$$
where  $\mathcal{D} = \{(x_1, y_1), (x_2, y_2), \ldots, (x_n, y_n)\}$  and  $\pi_n$  is the empirical distribution associated with the data  $\mathcal{D}$ . This estimator is known as the infinitesimal jackknife variance estimator [65]. In this section, we use  $x_0$  instead of  $x$  in the variance or influence function at a test point  $x_0$  to avoid confusion between  $x_0$  and  $z_x$ . The latter  $z_x$  is referred to the  $x$ -componenet of  $z$ .
Based on [25, 31, 50], we derive the closed-form expression of the variance estimation  $\hat{\xi}^2 (x_0)$  as follows.
Theorem B.6 (Expression of infinitesimal jackknife for kernel ridge regression). Suppose that Assumptions B.3 and B.4 hold. For the least squares loss  $\mathcal{L}(h(X),Y) = (h(X) - Y)^2$  in kernel ridge regression,  $IF(z;T,\pi_n)$  is given by
$$
I F (z; T, \pi_ {n}) (x _ {0}) = \boldsymbol {k} (x _ {0}, \boldsymbol {x}) ^ {T} (\boldsymbol {k} (\boldsymbol {x}, \boldsymbol {x}) + \lambda_ {0} n \boldsymbol {I}) ^ {- 1} M _ {z} (\boldsymbol {x}) - M _ {z} (x _ {0}) \tag {15}
$$
where  $M_z(\pmb{x}) \coloneqq (M_z(x_1), \dots, M_z(x_n))^T$  and
$$
M _ {z} (x _ {i}) := g _ {\pi_ {n}, \lambda_ {0}} (x _ {i}) - \frac {1}{\lambda_ {0}} (z _ {y} - g _ {\pi_ {n}, \lambda_ {0}} (z _ {x})) k (z _ {x}, x _ {i})
$$
for  $x_{i} = x_{0},x_{1},\dots ,x_{n}$
Note that although Hadamard differentiability is able to guarantee the central limit theorem of  $T$ , the consistency of variance estimation generally requires more than Hadamard differentiability. In general, we need some additional continuity condition (such as continuously Gâteaux differentiability) to guarantee that  $IF^2(z; T, \pi_n)$  is indeed "close" to  $IF^2(z; T, \pi)$  and  $\frac{1}{n} \sum_{z_i \in \mathcal{D}} IF^2(z_i; T, \pi_n)(x_0)$  is indeed "close" to  $\int_{z \in \mathcal{X} \times \mathcal{Y}} IF^2(z; T, \pi)(x_0)d\pi(z)$ . We show that this is achievable for the infinitesimal jackknife variance estimator for kernel ridge regression by only imposing a weak assumption.
Theorem B.7 (Consistency of infinitesimal jackknife for kernel-based regression). Suppose that Assumptions B.3 and B.4 hold. Moreover, assume that  $\mathcal{V} \subset \mathbb{R}$  is bounded, and  $b$  and  $b_{a}^{\prime}$  in Assumption B.4 are bounded on  $\mathcal{V}$ . Then we have
$$
\hat {\xi} ^ {2} (x _ {0}) = \frac {1}{n} \sum_ {z _ {i} \in \mathcal {D}} I F ^ {2} (z _ {i}; T, \pi_ {n}) (x _ {0}) \to \int_ {z \in \mathcal {X} \times \mathcal {Y}} I F ^ {2} (z; T, \pi) (x) d \pi (z) = \xi^ {2} (x _ {0}), a. s.
$$
as  $n\to \infty$ . Hence, an asymptotically exact  $(1 - \alpha)$ -level confidence interval of  $g_{\pi ,\lambda_0}(x_0)$  is
$$
\left[ g _ {\pi_ {n}, \lambda_ {n}} (x _ {0}) - \frac {\hat {\xi} (x _ {0})}{n} q _ {1 - \frac {\alpha}{2}}, g _ {\pi_ {n}, \lambda_ {n}} (x _ {0}) + \frac {\hat {\xi} (x _ {0})}{n} q _ {1 - \frac {\alpha}{2}} \right]
$$
where  $q_{\alpha}$  is the  $\alpha$ -quantile of the standard Gaussian distribution  $\mathcal{N}(0,1)$ .
The proof of Theorems B.6 and B.7 is given in Appendix E.
B.4 PNC-Enhanced Infinitesimal Jackknife
In this section, we return to our problem setting in Section 3.3. We apply the general results in Section B.3 to our PNC predictor and develop the PNC-enhanced infinitesimal jackknife confidence interval for over-parameterized neural networks.
Recall that the statistical functional  $T_{1}$  is associated with the following problem:
$$
\hat {h} _ {n, \theta^ {b}} (\cdot) - \hat {\phi} _ {n, \theta^ {b}} (\cdot) - \bar {s} (\cdot) = \min  _ {g \in \mathcal {H}} \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ \left(y _ {i} - \bar {s} (x _ {i}) - g (x _ {i})\right) ^ {2} \right] + \lambda_ {n} \| g \| _ {\mathcal {H}} ^ {2}
$$
Consider the PNC-enhanced infinitesimal jackknife variance estimator:
$$
\hat {\sigma} ^ {2} (x _ {0}) = \frac {1}{n} \sum_ {z _ {i} \in \mathcal {D}} I F ^ {2} (z _ {i}; T _ {1}, \pi_ {n}) (x _ {0}).
$$
Theorem B.8 (Expression of PNC-enhanced infinitesimal jackknife). Suppose that Assumption C.2 holds. Suppose that Assumptions B.3 and B.4 hold when  $Y$  is replaced by  $Y - \bar{s}(X)$ . Then  $IF(z; T_1, \pi_n)$  is given by
$$
I F (z; T _ {1}, \pi_ {n}) (x _ {0}) = \boldsymbol {K} (x _ {0}, \boldsymbol {x}) ^ {T} (\boldsymbol {K} (\boldsymbol {x}, \boldsymbol {x}) + \lambda_ {0} n \boldsymbol {I}) ^ {- 1} M _ {z} (\boldsymbol {x}) - M _ {z} (x _ {0}) \tag {16}
$$
where  $M_z(\pmb{x}) \coloneqq (M_z(x_1),\dots,M_z(x_n))^T$  and
$$
M _ {z} (x _ {i}) := \hat {h} _ {n, \theta^ {b}} (x _ {i}) - \hat {\phi} _ {n, \theta^ {b}} (x _ {i}) - \bar {s} (x _ {i}) - \frac {1}{\lambda_ {0}} \left(z _ {y} - \left(\hat {h} _ {n, \theta^ {b}} \left(z _ {x}\right) - \hat {\phi} _ {n, \theta^ {b}} \left(z _ {x}\right)\right)\right) K (z _ {x}, x _ {i})
$$
for  $x_{i} = x_{0},x_{1},\dots ,x_{n}$  .Hence
$$
\begin{array}{l} \hat {\sigma} ^ {2} (x _ {0}) = \frac {1}{n} \sum_ {z _ {i} \in \mathcal {D}} I F ^ {2} \left(z _ {i}; T _ {1}, \pi_ {n}\right) \left(x _ {0}\right) \\ = \frac {1}{n} \sum_ {z _ {i} \in \mathcal {D}} \left(\boldsymbol {K} \left(x _ {0}, \boldsymbol {x}\right) ^ {T} \left(\boldsymbol {K} (\boldsymbol {x}, \boldsymbol {x}) + \lambda_ {0} n \boldsymbol {I}\right) ^ {- 1} M _ {z _ {i}} (\boldsymbol {x}) - M _ {z _ {i}} \left(x _ {0}\right)\right) ^ {2} \\ \end{array}
$$
Theorem B.8 immediately follows from Theorem B.6.
Theorem B.9 (Exact coverage of PNC-enhanced infinitesimal jackknife confidence interval). Suppose that Assumption C.2 holds. Suppose that Assumptions B.3 and B.4 hold when  $Y$  is replaced by  $Y - \bar{s}(X)$ . Moreover, assume that  $\mathcal{Y} \subset \mathbb{R}$  is bounded. Then we have
$$
\hat {\sigma} ^ {2} (x _ {0}) = \frac {1}{n} \sum_ {z _ {i} \in \mathcal {D}} I F ^ {2} (z _ {i}; T _ {1}, \pi_ {n}) (x _ {0}) \to \int_ {z \in \mathcal {X} \times \mathcal {Y}} I F ^ {2} (z; T _ {1}, \pi) (x) d \pi (z) = \sigma (x _ {0}), a. s.
$$
as  $n\to \infty$ . Hence, an asymptotically exact  $(1 - \alpha)$ -level confidence interval of  $h^{*}(x_{0})$  is
$$
\left[ \hat {h} _ {n, \theta^ {b}} (x _ {0}) - \hat {\phi} _ {n, \theta^ {b}} (x _ {0}) - \frac {\hat {\sigma} (x _ {0})}{n} q _ {1 - \frac {\alpha}{2}}, \hat {h} _ {n, \theta^ {b}} (x _ {0}) - \hat {\phi} _ {n, \theta^ {b}} (x _ {0}) + \frac {\hat {\sigma} (x _ {0})}{n} q _ {1 - \frac {\alpha}{2}} \right]
$$
where  $q_{\alpha}$  is the  $\alpha$ -quantile of the standard Gaussian distribution  $\mathcal{N}(0,1)$  and the computation of  $\hat{\sigma} (x_0)$  is given by Theorem B.8.
Note that in our problem setting, we can set  $b(y) = y^2$  and  $b_{a}^{\prime} = a + |y|$  for kernel ridge regression. Since  $\mathcal{V}$  is bounded,  $b$  and  $b_{a}^{\prime}$  are naturally bounded on  $\mathcal{V}$ . Therefore, Theorem B.9 follows from Theorem B.7.
Appendix C Theory of Neural Tangent Kernel
In this section, we provide a brief review of the theory of neural tangent kernel (NTK). Then we discuss particularly one result employed by this paper, i.e.,
The shifted kernel ridge regressor using NTK with a shift from an initial function  $s_{\theta^b}$  is exactly the linearized neural network regressor that starts from the initial network  $s_{\theta^b}$ .
NTK [64] has attracted a lot of attention since it provides a new perspective on training dynamics, generalization, and expressibility of over-parameterized neural networks. Recent papers show that when training an over-parametrized neural network, the weight matrix at each layer is close to its
initialization [85, 35]. An over-parameterized neural network can rapidly reduce training error to zero via gradient descent, and thus finds a global minimum despite the objective function being non-convex [34, 4-6, 122]. Moreover, the trained network also exhibits good generalization property [19, 20]. These observations are implicitly described by a notion, NTK, suggested by [64]. This kernel is able to characterize the training behavior of sufficiently wide fully-connected neural networks and build a new connection between neural networks and kernel methods. Another line of work is to extend NTK for fully-connected neural networks to CNTK for convolutional neural networks [7, 114, 86]. NTK is a fruitful and rapidly growing area.
We focus on the result that is used as the foundation of our epistemic uncertainty quantification: Since the gradient of over-parameterized neural networks is nearly constant and close to its initialization, training networks is very similar to a shifted kernel ridge regression where the feature map of the kernel is given by the gradient of networks [80]. Multiple previous works [7, 80, 54, 59, 118, 15] have established some kind of equivalences between the neural network regressor and the kernel ridge regressor based on a wide variety of assumptions or settings, such as the full-rank assumption of NTK, or zero-valued initialization, or introducing a small multiplier  $\kappa$ , or normalized training inputs  $\| x_{i}\|_{2} = 1$ . These results are also of different forms. Since a uniform statement is not available, we do not intend to dig into those detailed assumptions or theorems as they are not the focus of this paper (Interested readers may refer to the above papers).
In the following, we provide some discussions on how to obtain the equivalence in Proposition 3.1 adopted in the main paper, based on a less (technically) rigorous assumption. This assumption is borrowed from the linearized neural network property introduced in [80], where they replace the outputs of the neural network with their first-order Taylor expansion, called the linearized neural network. They show that in the infinite width limit, the outputs of the neural network are the same as the linearized neural network. [80] does not add regularization in the loss. Instead, we introduce a regularization  $\lambda_{n} > 0$  in the loss since regularization is common and useful in practice. Moreover, it can guarantee the stable computation of the inversion of the NTK Gram matrix and can be naturally linked to kernel ridge regression. In the following, we will review the linearized neural network assumption (Assumption C.1) as well as other network specifications. Based on them, we show that Proposition C.3 holds, which provides the starting point for our uncertainty quantification framework. We also provide some additional remarks in Section D.
NTK parameterization. We consider a fully-connected neural network with any depth defined formally as follows. Suppose the network consists of  $L$  hidden layers. Denote  $g^{(0)}(x) = x$  and  $d_0 = d$ . Let
$$
f ^ {(l)} (x) = W ^ {(l)} g ^ {(l - 1)} (x), g ^ {(l)} (x) = \sqrt {\frac {c _ {\sigma}}{d _ {l}}} \sigma (f ^ {(l)} (x))
$$
where  $W^{(l)} \in \mathbb{R}^{d_l \times d_{l-1}}$  is the weight matrix in the  $l$ -th layer,  $\sigma$  is a coordinate-wise activation function,  $c_\sigma = \mathbb{E}_{z \sim \mathcal{N}(0,1)}[\sigma^2(z)]^{-1}$ , and  $l \in [L]$ . The output of the neural network is
$$
f _ {\theta} (x) := f ^ {(L + 1)} (x) = W ^ {(L + 1)} g ^ {(L)} (x)
$$
where  $W^{(L + 1)} \in \mathbb{R}^{1 \times d_L}$ , and  $\theta = (W^{(1)}, \dots, W^{(L + 1)})$  represents all the parameters in the network. Note that here we use NTK parametrization with a width-dependent scaling factor, which is thus slightly different from the standard parametrization. Unlike the standard parameterization which only normalizes the forward dynamics of the network, the NTK parameterization also normalizes its backward dynamics.
He Initialization. The NTK theory depends on the random initialization of the network. Suppose the dimension of network parameters  $\theta$  is  $p$ . We randomly initialize all the weights to be i.i.d.  $\mathcal{N}(0,1)$  random variables. In other words, we set  $P_{\theta_b} = \mathcal{N}(0,\mathbf{I}_p)$  and let  $\theta^b$  be an instantiation drawn from  $P_{\theta_b}$ . This initialization method is essentially known as the He initialization [55].
NTK and Linearized neural network. With the above NTK parameterization and He initialization, the population NTK expression is defined recursively as follows: [64, 7]: For  $l \in [L]$
$$
\Sigma^ {(0)} (x, x ^ {\prime}) = x ^ {T} x ^ {\prime} \in \mathbb {R},
$$
$$
\Lambda^ {(l)} (x, x ^ {\prime}) = \left( \begin{array}{c c} \Sigma^ {(l - 1)} (x, x) & \Sigma^ {(l - 1)} (x, x ^ {\prime}) \\ \Sigma^ {(l - 1)} (x, x ^ {\prime}) & \Sigma^ {(l - 1)} (x ^ {\prime}, x ^ {\prime}) \end{array} \right) \in \mathbb {R} ^ {2 \times 2},
$$
$$
\Sigma^ {(l)} (x, x ^ {\prime}) = c _ {\sigma} \mathbb {E} _ {(u, v) \sim \mathcal {N} (0, \Lambda (l))} [ \sigma (u) \sigma (v) ] \in \mathbb {R}.
$$
We also define a derivative covariance:
$$
\Sigma^ {(l) \prime} (x, x ^ {\prime}) = c _ {\sigma} \mathbb {E} _ {(u, v) \sim \mathcal {N} (0, \Lambda (l))} [ \sigma^ {\prime} (u) \sigma^ {\prime} (v) ] \in \mathbb {R},
$$
The final population NTK is defined as
$$
K (x, x ^ {\prime}) = \sum_ {l = 1} ^ {L + 1} \left(\Sigma^ {(l - 1)} (x, x ^ {\prime}) \prod_ {s = l} ^ {L + 1} \Sigma^ {(s) ^ {\prime}} (x, x ^ {\prime})\right)
$$
where  $\Sigma^{(L + 1)'}(x,x') = 1$  for convenience. Let  $\langle \cdot ,\cdot \rangle$  be the standard inner product in  $\mathbb{R}^p$
Let  $J(\theta; x) \coloneqq \nabla_{\theta} f_{\theta}(x) \in \mathbb{R}^{1 \times p}$  denote the gradient of the network. The empirical (since the initialization is random) NTK matrix is defined as
$$
K _ {\theta} (x, x ^ {\prime}) = \langle J (\theta ; x); J (\theta ; x ^ {\prime}) \rangle .
$$
In practice, we may use the empirical NTK matrix to numerically compute the population NTK matrix.
One of the most important results in the NTK theory is that the empirical NTK matrix converges to the population NTK matrix as the width of the network increases [64, 114, 115, 7]:
$$
K _ {\theta} (x, x ^ {\prime}) = \langle J (\theta ; x), J (\theta ; x) \rangle \rightarrow K (x, x ^ {\prime})
$$
where  $\rightarrow$  represents "converge in probability" [7] or "almost surely" [114].
The NTK theory shows that  $J(\theta ;x) = \nabla_{\theta}f_{\theta}(x)$  is approximately a constant independent of the network parameter  $\theta$  (but still depends on  $x$ ) when the network is sufficiently wide. Therefore, when we treat  $J(\theta ;x)$  as a constant independent of the network parameter  $\theta$ , we can obtain a model obtained from the first-order Taylor expansion of the network around its initial parameters, which is called a linearized neural network [80]. Our study on epistemic uncertainty is based on the linearized neural network.
To be rigorous, we adopt the following linearized neural network assumption based on the NTK theory:
Assumption C.1. [Linearized neural network assumption] Suppose that  $J(\theta ;x) = \nabla_{\theta}f_{\theta}(x)\in \mathbb{R}^{1\times p}$  is independent of  $\theta$  during network training, which is thus denoted as  $J(x)$ . The population NTK is then given by  $K(x,x^{\prime}) = \langle J(x),J(x)\rangle = J(x)J(x)^{\top}$ .
We introduce the NTK vector/matrix as follows. For  $\pmb{x} = (x_{1},\dots,x_{j})^{\top}$ , we let
$$
\boldsymbol {K} (\boldsymbol {x}, \boldsymbol {x}) := (K (x _ {i}, x _ {j})) _ {i, j = 1, \dots , n} = J (\boldsymbol {x}) J (\boldsymbol {x}) ^ {\top} \in \mathbb {R} ^ {n \times n}
$$