SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziyi Huang, Henry Lam, Haofeng Zhang *
New York, NY, USA
zh2354, khl2114, hz2553@columbia.edu
Uncertainty quantification (UQ) is important for reliability assessment and enhancement of machine learning models. In deep learning, uncertainties arise not only from data, but also from the training procedure that often injects substantial noises and biases. These hinder the attainment of statistical guarantees and, moreover, impose computational challenges on UQ due to the need for repeated network retraining. Building upon the recent neural tangent kernel theory, we create statistically guaranteed schemes to principally characterize, and remove, the uncertainty of over-parameterized neural networks with very low computation effort. In particular, our approach, based on what we call a procedural-noise-correcting (PNC) predictor, removes the procedural uncertainty by using only one auxiliary network that is trained on a suitably labeled dataset, instead of many retrained networks employed in deep ensembles. Moreover, by combining our PNC predictor with suitable light-computation resampling methods, we build several approaches to construct asymptotically exact-coverage confidence intervals using as low as four trained networks without additional overheads.
Uncertainty quantification (UQ) concerns the dissection and estimation of various sources of errors in a prediction model. It has growing importance in machine learning, as it helps assess and enhance the trustworthiness and deployment safety across many real-world tasks ranging from computer vision [67, 21] and natural language processing [113, 103] to autonomous driving [90, 91], as well as guiding exploration in sequential learning [9, 1, 119]. In the deep learning context, UQ encounters unique challenges on both the statistical and computational fronts. On a high level, these challenges arise from the over-parametrized and large-scale nature of neural networks so that, unlike classical statistical models, the prediction outcomes incur not only noises from the data, but also importantly the training procedure itself [73, 94]. This elicits a deviation from the classical statistical theory that hinders the attainment of established guarantees. Moreover, because of the sizes of these models, conventional procedures such as resampling [37, 30, 102] demand an amount of computation that could quickly become infeasible.
Our main goal of this paper is to propose a UQ framework for over-parametrized neural networks in regression that has simultaneous statistical coverage guarantee, in the sense of classical frequentist asymptotic exactness, and low computation cost, in the sense of requiring only few (as low as four) neural network trainings, without other extra overheads. A main driver of these strengths in our framework is a new implementable concept, which we call the Procedural-Noise-Correcting (PNC) predictor. It consists of an auxiliary network that is trained on a suitably artificially labeled dataset, with behavior mimicking the variability coming from the training procedure. To reach our goal, we
*Authors are listed alphabetically.
synthesize and build on two recent lines of tools that appear largely segregated thus far. First is neural tangent kernel (NTK) theory [64, 80, 7], which provides explicit approximate formulas for well-trained infinitely wide neural networks. Importantly, NTK reveals how procedural variability enters into the network prediction outcomes through, in a sense, a functional shift in its corresponding kernel-based regression, which guides our PNC construction. Second is light-computation resampling methodology, including batching [47, 99, 100] and the so-called cheap bootstrap method [75], which allows valid confidence interval construction using as few as two model repetitions. We suitably enhance these methods to account for both data and procedural variabilities via the PNC incorporation.
We compare our framework with several major related lines of work. First, our work focuses on the quantification of epistemic uncertainty, which refers to the errors coming from the inadequacy of the model or data noises. This is different from aleatoric uncertainty, which refers to the intrinsic stochasticity of the problem [92, 96, 121, 62, 36], or predictive uncertainty which captures the sum of epistemic and aleatoric uncertainties (but not their dissection) [94, 97, 12, 3, 22]. Regarding epistemic uncertainty, a related line of study is deep ensemble that aggregates predictions from multiple independent training replications [81, 73, 41, 8, 94]. This approach, as we will make clear later, can reduce and potentially quantify procedural variability, but a naive use would require demanding retraining effort and does not address data variability. Another line is the Bayesian UQ approach on neural networks [44, 2]. This regards network weights as parameters subject to common priors such as Gaussian. Because of the computation difficulties in exact inference, an array of studies investigate efficient approximate inference approaches to estimate the posteriors [43, 49, 16, 33, 32, 95, 79, 56]. While powerful, these approaches nonetheless possess inference error that could be hard to quantify, and ultimately finding rigorous guarantees on the performance of these approximate posteriors remains open to our best knowledge.
2 Statistical Framework of Uncertainty
We first describe our learning setting and define uncertainties in our framework. Suppose the input-output pair  $(X,Y)$  is a random vector following an unknown probability distribution  $\pi$  on  $\mathcal{X} \times \mathcal{Y}$ , where  $X \in \mathcal{X} \subset \mathbb{R}^d$  is an input and  $Y \in \mathcal{Y} \subset \mathbb{R}$  is the response. Let the marginal distribution of  $X$  be  $\pi_X(x)$  and the conditional distribution of  $Y$  given  $X$  be  $\pi_{Y|X}(y|x)$ . Given a set of training data  $\mathcal{D} = \{(x_1, y_1), (x_2, y_2), \ldots, (x_n, y_n)\}$  drawn independent and identically distributed (i.i.d.) from  $\pi$  (we write  $\pmb{x} = (x_1, \dots, x_n)^T$  and  $\pmb{y} = (y_1, \dots, y_n)^T$  for short), we build a prediction model  $h: \mathcal{X} \to \mathcal{Y}$  that best approximates  $Y$  given  $X$ . Let  $\pi_{\mathcal{D}}$  or  $\pi_n$  denote the empirical distribution associated with the training data  $\mathcal{D}$ , where  $\pi_n$  is used to emphasize the sample size dependence in asymptotic results.
To this end, provided a loss function  $\mathcal{L}:\mathcal{Y}\times \mathbb{R}\to \mathbb{R}$ , we denote the population risk  $R_{\pi}(h)\coloneqq \mathbb{E}_{(X,Y)\sim \pi}[\mathcal{L}(h(X),Y)]$ . If  $h$  is allowed to be any possible function, the best prediction model is the Bayes predictor [61, 93]:  $h_B^* (X)\in \mathrm{argmin}_{y\in \mathcal{Y}}\mathbb{E}_{Y\sim \pi_{Y|X}}[\mathcal{L}(y,Y)|X]$ . With finite training data of size  $n$ , classical statistical learning suggests finding  $\hat{h}_n\in \mathcal{H}$ , where  $\mathcal{H}$  denotes a hypothesis class that minimizes the empirical risk, i.e.,  $\hat{h}_n\in \mathrm{argmin}_{h\in \mathcal{H}}R_{\pi_{\mathcal{D}}}(\hat{h})$ . This framework fits methods such as linear or kernel ridge regression (Appendix B). However, it is not feasible for deep learning due to non-convexity and non-identifiability. Instead, in practice, gradient-based optimization methods are used, giving rise to  $\hat{h}_{n,\gamma}$  as a variant of  $\hat{h}_n$ , where the additional variable  $\gamma$  represents the randomness in the training procedure. It is worth mentioning that this randomness generally depends on the empirical data  $\mathcal{D}$ , and thus we use  $P_{\gamma |\mathcal{D}}$  to represent the distribution of  $\gamma$  conditional on  $\mathcal{D}$ .
Furthermore, we consider  $\hat{h}_n^* = \text{aggregate}(\{\hat{h}_{n,\gamma} : \gamma \sim P_{\gamma|\mathcal{D}}\})$  where "aggregate" means an idealized aggregation approach to remove the training randomness in  $\hat{h}_{n,\gamma}$  (known as ensemble learning [73, 18, 17, 87, 45]). A prime example in deep learning is deep ensemble [73] that will be detailed in the sequel. Finally, we denote  $h^* = \lim_{n \to \infty} \hat{h}_n^*$  as the grand "limiting" predictor with infinite samples. The exact meaning of "lim" will be clear momentarily.
Under this framework, epistemic uncertainty, that is, errors coming from the inadequacy of the model or data noises, can be dissected into the following three sources, as illustrated in Figure 1. Additional discussions on other types of uncertainty are in the Appendix A for completeness.
Model approximation error.  $\mathrm{UQ}_{AE} = h^{*} - h_{B}^{*}$ . This discrepancy between  $h_B^*$  and  $h^*$  arises from the inadequacy of the hypothesis class  $\mathcal{H}$ . For an over-parameterized sufficiently wide neural network
$\mathcal{H}$ , this error is usually negligible thanks to the universal approximation power of neural networks for any continuous functions [28, 58, 53] or Lebesgue-integrable functions [88].
Data variability.  $\mathrm{UQ}_{DV} = \hat{h}_n^* - h^*$ . This measures the representativeness of the training dataset, which is the most standard epistemic uncertainty in classical statistics [110].
Procedural variability.  $\mathrm{UQ}_{PV} = \hat{h}_{n,\gamma} - \hat{h}_n^*$ . This arises from the randomness in the training process for a single network  $\hat{h}_{n,\gamma}$ , which is present even with deterministic or infinite data. The randomness comes from the initialization of the network parameters, and also data ordering and possibly training time when running stochastic gradient descent with finite training epochs.
3 Quantifying Epistemic Uncertainty
We use a frequentist framework and, for a given  $x$ , we aim to construct a confidence interval for the "best" predictor  $h^*(x)$ . As discussed in Section 1, the over-parametrized and nonconvex nature of neural networks defies conventional statistical techniques and moreover introduces procedural variability that makes inference difficult.
Figure 1: Three sources of epistemic uncertainty.
We focus on over-parameterized neural networks, that is, neural networks whose width is sufficiently large, while the depth can
be arbitrary such as two [120, 107]. Over-parameterized neural networks give the following two theoretical advantages. First, this makes model approximation error negligible and a confidence interval for  $h^{*}(x)$  also approximately applies to  $h_B^* (x)$ . Second, the NTK theory [64] implies a phenomenon that the network evolves essentially as a "linear model" under gradient descent, and thus the resulting predictor behaves like a shifted kernel ridge regression whose kernel is the NTK [7, 80, 54, 59, 118] (detailed in Appendix C). However, to our knowledge, there is no off-the-shelf result that exactly matches our need for the epistemic uncertainty task, so we describe it below. Consider the following regression problem:
$$
\hat {R} _ {n} \left(f _ {\theta}, \theta^ {b}\right) = \frac {1}{n} \sum_ {i = 1} ^ {n} \left(f _ {\theta} \left(x _ {i}\right) - y _ {i}\right) ^ {2} + \lambda_ {n} \| \theta - \theta^ {b} \| _ {2} ^ {2}. \tag {1}
$$
where  $\theta$  is the network parameters to be trained,  $\theta^b$  is the initial network parameters, and  $\lambda_{n} > 0$  is the regularization hyper-parameter which may depend on the data size. We add regularization  $\lambda_{n}$  to this problem, which is slightly different from previous work [80] without the regularization  $\lambda_{n}$ . This can guarantee stable computation of the inversion of the NTK Gram matrix and can be naturally linked to kernel ridge regression, which will be introduced shortly in Proposition 3.1. We assume the network adopts the NTK parametrization with parameters randomly initialized using He initialization [55], and it is sufficiently wide to ensure the linearized neural network assumption holds; See Appendix C for details. Moreover, we assume that the network is trained using the loss function in (1) via continuous-time gradient flow by feeding the entire training data and using sufficient training time  $(t\to \infty)$ . In this sense, the uncertainty of data ordering and training time vanishes, making the random initialization the only uncertainty in procedural variability. The above specifications are formally summarized in Assumption C.2. With the above specifications, we have:
Proposition 3.1 (Proposition C.3). Suppose that Assumption C.2 holds. Then the final trained network, conditional on the initial network  $s_{\theta^b}(x)$ , is given by
$$
\hat {h} _ {n, \theta^ {b}} (x) = s _ {\theta^ {b}} (x) + \boldsymbol {K} (x, \boldsymbol {x}) ^ {T} (\boldsymbol {K} (\boldsymbol {x}, \boldsymbol {x}) + \lambda_ {n} n \boldsymbol {I}) ^ {- 1} (\boldsymbol {y} - s _ {\theta^ {b}} (\boldsymbol {x})), \tag {2}
$$
where in the subscript of  $\hat{h}_{n,\theta^b}$ ,  $n$  represents  $n$  training data,  $\theta^b$  represents an instance of the initial network parameters drawn from the standard Gaussian distribution  $P_{\theta^b} = \mathcal{N}(0, I_p)$  where the dimension of  $\theta^b$  is  $p$  (He initialization);  $\mathbf{K}(\mathbf{x}, \mathbf{x}) \coloneqq (K(x_i, x_j))_{i,j=1,\dots,n}$  and  $\mathbf{K}(x, x) \coloneqq (K(x, x_1), K(x, x_2), \ldots, K(x, x_n))^T$  where  $K$  is the (population) NTK. This implies that  $\hat{h}_{n,\theta^b}$  is the solution to the following kernel ridge regression:
$$
s _ {\theta^ {b}} + \underset {g \in \mathcal {H}} {\arg \min } \frac {1}{n} \sum_ {i = 1} ^ {n} (y _ {i} - s _ {\theta^ {b}} (x _ {i}) - g (x _ {i})) ^ {2} + \lambda_ {n} \| g \| _ {\mathcal {H}} ^ {2}
$$
where  $\mathcal{H}$  is the reproducing kernel Hilbert space (RKHS) constructed from the NTK  $K(x,x')$ .
Proposition 3.1 shows that 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}$ . It also reveals how procedural variability enters into the neural network prediction outcomes, highlighting the need to regard random initialization as an inevitable uncertainty component in neural networks. We provide details and deviation of Proposition 3.1 in Appendix C.
3.1 Challenges from the Interplay of Procedural and Data Variabilities
First, we will discuss existing challenges in quantifying and reducing uncertainty to motivate our main approach based on the PNC predictor under the NTK framework. To this end, deep ensemble [81, 73, 41, 8, 94] is arguably the most common ensemble approach in deep learning to reduce procedural variability. [8] shows that deep ensemble achieves the best performance compared with a wide range of other ensemble methods, and employing more networks in deep ensemble can lead to better performance. Specifically, the deep ensemble predictor [73]  $\hat{h}_n^m (x)$  is defined as:  $\hat{h}_n^m (x)\coloneqq \frac{1}{m}\sum_{i = 1}^m\hat{h}_{n,\theta_i^b}(x)$  where  $m$  is the number of networks in the ensemble,  $\hat{h}_{n,\theta_i^b}(x)$  is the independently trained network with initialization  $\theta_i^b$  (with the same training data  $\mathcal{D}$ ), and  $\theta_1^b,\dots,\theta_m^b$  are i.i.d. samples drawn from  $P_{\theta^b}$ . We also introduce  $\hat{h}_n^* (x)\coloneqq \mathbb{E}_{P_{\theta^b}}[\hat{h}_{n,\theta^b}(x)]$  as the expectation of  $\hat{h}_{n,\theta^b}(x)$  with respect to  $\theta^b\sim P_{\theta^b}$ . Taking  $m\to \infty$  and using the law of large numbers, we have  $\lim_{m\to \infty}\hat{h}_n^m (x) = \hat{h}_n^* (x)$  a.s.. So  $\hat{h}_n^* (x)$  behaves as an idealized deep ensemble predictor with infinitely many independent training procedures. Using Proposition 3.1 and the linearity of kernel ridge regressor with respect to data (Appendix B), we have
$$
\hat {h} _ {n} ^ {m} (x) = \frac {1}{m} \sum_ {i = 1} ^ {m} s _ {\theta_ {i} ^ {b}} (x) + \boldsymbol {K} (\boldsymbol {x}, \boldsymbol {x}) ^ {T} (\boldsymbol {K} (\boldsymbol {x}, \boldsymbol {x}) + \lambda_ {n} n \boldsymbol {I}) ^ {- 1} \left(\boldsymbol {y} - \frac {1}{m} \sum_ {i = 1} ^ {m} s _ {\theta_ {i} ^ {b}} (\boldsymbol {x})\right)
$$
$$
\hat {h} _ {n} ^ {*} (x) = \mathbb {E} _ {P _ {\theta^ {b}}} [ \hat {h} _ {n, \theta^ {b}} (x) ] = \bar {s} (x) + \boldsymbol {K} (x, \boldsymbol {x}) ^ {T} (\boldsymbol {K} (\boldsymbol {x}, \boldsymbol {x}) + \lambda_ {n} n \boldsymbol {I}) ^ {- 1} (\boldsymbol {y} - \bar {s} (\boldsymbol {x})) \tag {3}
$$
where  $\bar{s} (x) = \mathbb{E}_{P_{\theta^b}}[s_{\theta^b}(x)]$  is the expectation of the initial network output  $s_{\theta^b}(x)$  with respect to the distribution of the initialization parameters  $P_{\theta^b} = \mathcal{N}(0,\pmb {I}_p)$ . It is easy to see that  $\mathbb{E}_{P_{\theta^b}}[\hat{h}_n^m (x)] = \hat{h}_n^* (x)$  and  $\mathrm{Var}_{P_{\theta^b}}(\hat{h}_n^m (x)) = \frac{1}{m}\mathrm{Var}_{P_{\theta^b}}(\hat{h}_{n,\theta^b}(x))$  where  $\mathrm{Var}_{P_{\theta^b}}$  is the variance with respect to the random initialization. As for the total variance:
Proposition 3.2. We have
$$
\operatorname {V a r} \left(\hat {h} _ {n} ^ {m} (x)\right) = \operatorname {V a r} \left(\hat {h} _ {n} ^ {*} (x)\right) + \frac {1}{m} \mathbb {E} \left[ \operatorname {V a r} \left(\hat {h} _ {n, \theta^ {b}} (x) | \mathcal {D}\right) \right] \leq \operatorname {V a r} \left(\hat {h} _ {n} ^ {*} (x)\right) + \mathbb {E} \left[ \operatorname {V a r} \left(\hat {h} _ {n, \theta^ {b}} (x) | \mathcal {D}\right) \right] = \operatorname {V a r} \left(\hat {h} _ {n, \theta^ {b}} (x)\right).
$$
where the variance is taken with respect to both the data  $\mathcal{D}$  and the random initialization  $P_{\theta^b}$ .
Proposition 3.2 shows that deep ensemble improves the statistical profile of a single model by reducing its procedural variability by a factor  $\frac{1}{m}$  (but not the data variability), and achieving this reduction requires  $m$  training times. To quantify the epistemic uncertainty that contains two variabilities from data and procedure, we may employ resampling approaches such as "bootstrap on a deep ensemble". This would involve two layers of sampling, the outer being the resampling of data, and the inner being the retraining of base networks with different initializations. In other words, this nested sampling amounts to a multiplicative amount of training effort in a large number of outer bootstrap resamples and the  $m$  inner retaining per resample, leading to a huge computational burden. Moreover, the data variability and procedural variability in neural networks are dependent, making the above approach delicate. More discussion about this issue is presented in Section D.
In the following, we introduce our PNC framework that can bypass the above issues in that:
Uncertainty reduction. We train one single network and one additional auxiliary network to completely remove the procedural variability. This is in contrast to the deep ensemble approach that trains  $m$  networks and only reduces the procedural variability by an  $m$ -factor.
Uncertainty quantification. We resolve the computational challenges in "bootstrap on a deep ensemble", by combining PNC predictors with low-budget inference tools that require only as low as four network trainings.
3.2 PNC Predictor and Procedural Variability Removal
We first develop a computationally efficient approach to obtain  $\hat{h}_n^* (x)$ , the idealized deep ensemble predictor that is free of procedural variability. We term our approach the procedural-noise-correcting (PNC) predictor, whose pseudo-code is given in Algorithm 1. This predictor consists of a difference between essentially two neural network outcomes,  $\hat{h}_{n,\theta^b}(x)$  which is the original network trained using one initialization, and  $\hat{\phi}_{n,\theta^b}^\prime (x)$  that is trained on a suitably artificially labeled dataset. More precisely, this dataset applies label  $\bar{s} (x_i)$  to  $x_{i}$ , where  $\bar{s}$  is the expected output of an untrained network with random initialization. Also note that labeling this artificial dataset does not involve any training process and, compared with standard network training, the only additional running time in the PNC predictor is to train this single artificial-label-trained network.
Algorithm 1 Procedural-Noise-Correcting (PNC) Predictor
Input: Training data  $\mathcal{D} = \{(x_1,y_1),(x_2,y_2),\dots,(x_n,y_n)\}$
Procedure: 1. Draw  $\theta^b\sim P_{\theta^b} = \mathcal{N}(0,I_p)$  under NTK parameterization. Train a standard base network with data  $\mathcal{D}$  and the initialization parameters  $\theta^b$  , which outputs  $\hat{h}_{n,\theta^b}(\cdot)$  in (3).
2. Let  $\bar{s}(x) = \mathbb{E}_{P_{\theta^b}}[s_{\theta^b}(x)]$ . For each  $x_i$  in  $\mathcal{D}$ , generate its "artificial" label  $\bar{s}(x_i) = \mathbb{E}_{P_{\theta^b}}[s_{\theta^b}(x_i)]$ . Train an auxiliary neural network with data  $\{(x_1, \bar{s}(x_1)), (x_2, \bar{s}(x_2)), \dots, (x_n, \bar{s}(x_n))\}$  and the initialization parameters  $\theta^b$  (the same one as in Step 1.), which outputs  $\hat{\phi}_{n,\theta^b}^\prime(\cdot)$ . Subtracting  $\bar{s}(\cdot)$ , we obtain  $\hat{\phi}_{n,\theta^b}(\cdot) = \hat{\phi}_{n,\theta^b}^\prime(\cdot) - \bar{s}(\cdot)$ .
Output: At point  $x$ , output  $\hat{h}_{n,\theta^b}(x) - \hat{\phi}_{n,\theta^b}(x)$ .
The development of our PNC predictor is motivated by the challenge of computing  $\hat{h}_n^* (x)$  (detailed below). To address this challenge, we characterize the procedural noise instead, which is given by
$$
\hat {\phi} _ {n, \theta^ {b}} (\cdot) := \hat {h} _ {n, \theta^ {b}} (x) - \hat {h} _ {n} ^ {*} (x) = s _ {\theta^ {b}} (x) - \bar {s} (x) + \boldsymbol {K} (x, \boldsymbol {x}) ^ {T} \left(\boldsymbol {K} (\boldsymbol {x}, \boldsymbol {x}) + \lambda_ {n} n \boldsymbol {I}\right) ^ {- 1} (\bar {s} (\boldsymbol {x}) - s _ {\theta^ {b}} (\boldsymbol {x})). \tag {4}
$$
By Proposition 3.1, the closed-form expression of  $\hat{\phi}_{n,\theta^b}$  in (4) corresponds exactly to the artificial-label-trained network in Step 2 of Algorithm 1. Our artificial-label-trained network is thereby used to quantify the procedural variability directly. This observation subsequently leads to:
Theorem 3.3 (PNC). Suppose that Assumption C.2 holds. Then the output of the PNC predictor (Algorithm 1) is exactly  $\hat{h}_n^* (x)$  given in (3).
We discuss two approaches to compute  $\bar{s}(x)$  in Step 2 of Algorithm 1 and their implications: 1) Under He initialization,  $s_{\theta^b}(x)$  is a zero-mean Gaussian process in the infinite width limit [80], and thus we may set  $\bar{s} \equiv 0$  for simplicity. Note that even if we set  $\bar{s} \equiv 0$ , it does not imply that the artificial-label-trained neural network in Step 2 of Algorithm 1 will output a zero-constant network whose parameters are all zeros. In fact, neural networks are excessively non-convex and have many nearly global minima. Starting from an initialization parameter  $\theta^b$ , gradient descent on this artificial-label-trained neural network will find a global minimum that is close to the  $\theta^b$  but not "zero" even if "zero" is indeed one of its nearly global minima ("nearly" in the sense of ignoring the negligible regularization term) [35, 118, 23]. This phenomenon can also be observed in Proposition 3.1 by plugging in  $\pmb{y} = \mathbf{0}$ . Hence, the output depends on random initialization in addition to training data, and our artificial-label-trained network is designed to capture this procedural variability. 2) An alternative approach that does not require specific initialization is to use Monte Carlo integration:  $\bar{s}(x) = \lim_{N \to \infty} \frac{1}{N} \sum_{i=1}^{N} s_{\theta_i^b}(x)$  where  $\theta_i^b$  are i.i.d. from  $P_{\theta^b}$ . When  $N$  is finite, it introduces procedural variance that vanishes at the order  $N^{-1}$ . Since this computation does not involve any training process and is conducted in a rapid manner practically,  $N \gg n$  can be applied to guarantee that the procedural variance in computing  $\bar{s}(x)$  is negligible compared with the data variance at the order  $n^{-1}$ . We practically observe that both approaches work similarly well.