SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zi Wang*
Alexander  $\mathbf{K}\mathbf{u}^{*}$
Jason Baldridge
Been Kim
Understanding which concepts models can and cannot represent has been fundamental to many tasks: from effective and responsible use of models to detecting out of distribution data. We introduce Gaussian process probes (GPP), a unified and simple framework for probing and measuring uncertainty about concepts represented by models. As a Bayesian extension of linear probing methods, GPP asks what kind of distribution over classifiers (of concepts) is induced by the model. This distribution can be used to measure both what the model represents and how confident the probe is about what the model represents. GPP can be applied to any pre-trained model with vector representations of inputs (e.g., activations). It does not require access to training data, gradients, or the architecture. We validate GPP on datasets containing both synthetic and real images. Our experiments show it can (1) probe a model's representations of concepts even with a very small number of examples, (2) accurately measure both epistemic uncertainty (how confident the probe is) and aleatory uncertainty (how fuzzy the concepts are to the model), and (3) detect out of distribution data using those uncertainty measures as well as classic methods do. By using Gaussian processes to expand what probing can offer, GPP provides a data-efficient, versatile and uncertainty-aware tool for understanding and evaluating the capabilities of machine learning models.
Deep learning models have become pervasive in a wide range of fields, yet they remain largely uninterpretable. As a result, gaining insight into which concepts deep learning models do or do not represent is of growing importance. This has led to the development of a family of methods known as "probes" [Alain and Bengio, 2016], which aim to understand the representations learned by these models. One common approach to probe the representations of a model is to train independent classifiers for each concept and measure their accuracy [Alain and Bengio, 2016, Kim et al., 2018, Belinkov, 2022]. For example, if the classifier returns 0.5 probability (for the binary case) across a range of examples, we would reasonably conclude that the model does not represent the concept.
However, simply examining the output of a classifier may not adequately reflect the intricate nature of representing a concept. For example, we can think of two different scenarios that may result in the same classifier output. If you ask a person if an olive is a fruit, only about  $60\%$  will say yes (likewise  $72\%$  for avocado,  $53\%$  for pumpkin, and  $40\%$  for acorn) [McCloskey and Glucksberg, 1978]. This does not indicate that people have a limited representation of the fruit concept, but that
*Equal contribution. {wangzi, alexku, jasonbaldridge, beenkim} @google.com, tomg@princeton.edu. Our code can be found at https://github.com/google-research/gpax.
there is intrinsic fuzziness in the label. By contrast, if you show somebody an olive in a martini glass and tell them it is a new concept called "blicket," they may be  $60\%$  sure in judging whether an olive on a tree is a blicket (as opposed to "blicket" having a specialized meaning that applies to cocktails). In this second case, the uncertainty is because they do not have enough observations to be certain about what the new concept is – a common problem in word learning [Xu and Tenenbaum, 2007].
Fuzzy concept: fruit
1
New concept: bicket
Labels:
Probability:
Confidence level:
60% (Olive is fruit)
100% (Olive is 60% fruit)
100% (Image is olive)
$60\%$  (Image is bicket)
We introduce a probabilistic approach based on Gaussian processes (GPs) [Rasmussen and Williams, 2006] that makes it possible to delineate different kinds of uncertainty in a model's concepts. Specifically, our goal is to probe a pre-trained model by constructing a probability distribution over possible binary classifiers of concepts induced by the model. Extracting this distribution allows the probe to produce two kinds of uncertainty measures - aleatory and epistemic [Fox and Ulkumen, 2011]. This allows us to distinguish cases where 1) there is fundamental fuzziness of the label (aleatory uncertainty) or 2) it does not have enough observations (epistemic uncertainty), whereas a binary classifier would predict equal probability for both and provide no further explanatory power.
Our framework also provides a natural building block for answering many downstream questions, such as detecting out-of-distribution (OOD) examples. Using the resulting uncertainty measures, we show that our method can correctly probe what a model represents, and that it achieves competitive performance in OOD detection. We validate this approach using both synthetic and real images.
Distinguishing aleatory and epistemic uncertainty has a long history in fields including cognitive science [Howell and Burnett, 1978, Fox and Ülkümden, 2011], probability theory [Jaynes, 2003], philosophy [Hacking, 2006] and more. Our work builds upon Fox and Ülkümden [2011], which introduced a framework for distinguishing the two dimensions of uncertainty from a human perspective.
In machine learning, Gaussian processes (GPs) and Bayesian neural nets (BNN) [ ] have enabled principled ways of predicting epistemic uncertainty, which led to better information-based decision making strategies [Houlsby et al., 2011, Hennig and Schuler, 2012, Wang and Jegelka, 2017, Gal et al., 2017, Wang et al., 2023]. Motivated by regression problems in reinforcement learning, Depeweg et al. [2018] decomposed uncertainty in BNN to a conditional entropy term as aleatory uncertainty and the mutual information between model weights and predictions as epistemic uncertainty. Our decomposition simplifies these quantities<sup>2</sup> and specializes them for classification. To our knowledge, aleatory uncertainty (fuzziness of concepts) has not been studied in the GP classification literature, and no previous work in the GP literature has been done to establish the correspondence between human perceptions of uncertainty and GP classification predictions.
Probing has been a popular approach to understand which concepts a model represents. There are mixed views of what probing should mean; either it is about estimating mutual information between representations and labels [Pimentel et al., 2020] or using minimum description length (MDL) [Voita and Titov, 2020] or representing the distance as squared Euclidean distance under a linear transformation [Hewitt and Manning, 2019], or using analysis with or without supervision [Saphra and Lopez, 2019, Wu et al., 2020]. Others considered different types of probing: conditional probing that compares two concepts [Hewitt et al., 2021], probing that asks where the task-relevant information "emerges" [Kunz and Kuhlmann, 2022] or which concepts are "used" [Lasri et al., 2022], and finally, probing/localizing that enables editing [Meng et al., 2022]. To the best of our knowledge, none of these probing methods offer measures of uncertainty.
2We define episteme as negative entropy, which has the same form as the mutual information  $I(y;g) = H[y] - H[y\mid g]$  used by Depeweg et al. [2018], since equivalently,  $I(y;g) = H[g] - H[g\mid y]$ .
3 GPP: Gaussian process probes
We introduce Gaussian process probes (GPP), a simple probabilistic method for probing that naturally distinguishes aleatory and epistemic uncertainty. We formally define the problem of probabilistic probing in §3.1, followed by a review of the Beta Gaussian processes (GPs) that GPP builds upon in §3.2. In §3.3 and §3.4, we detail our method and how it measures uncertainty.
Our goal is to probe a pre-trained model in a probabilistic manner by constructing a probability distribution over binary classifiers based on the model's vector representation of inputs (e.g., activations). Here, binary classifier refers to a function mapping a stimulus to the probability of a positive label.
Assumption 1. The model,  $\mathcal{M}$ , contains basis functions that map from stimuli (e.g., images, texts) to a vector representation:  $\phi : \mathcal{X} \mapsto \mathcal{A} \subset \mathbb{R}^d$ , where  $\mathcal{X}$  is the space of stimuli and  $\mathcal{A}$  the space of  $d$ -dimensional real-valued vector representations.
For example, basis functions  $\phi$  can map from images to last-layer activations of a convolutional neural network (CNN). The probabilistic probe treats basis functions  $\phi$  as a feed-forward black box, and does not require accessing the training data, gradients, or architecture used during model pre-training.
In order to define a stochastic process to describe a distribution over classifiers, the data generating process is assumed as follows.
Assumption 2. There exists a stochastic process  $\mathcal{G}(\theta)$  parameterized by  $\theta$ , and a classifier  $h: \mathcal{A} \mapsto [0,1]$  distributed according to  $\mathcal{G}(\theta)$ , such that the label  $y \in \{0,1\}$  for any vector representation  $a \in \mathcal{A}$  is independently distributed according to a Bernoulli distribution with parameter  $h(a)$ ; i.e.,  $p(y \mid a,h) = h(a)^y (1 - h(a))^{1 - y}$ .
Now the posterior inference of the stochastic process (in our case, Beta GP) requires conditioning on observations, and the stochastic process measures uncertainty for a set of queries. We denote a set of observations as  $D$ , which consists of representation-label pairs, i.e.,  $D = \{(\phi(x_i), y_i)\}_{i=1}^{|D|}$ ; here, label  $y_i$  is either 0 or 1. The queries,  $\mathbf{a}_q = [\phi(x_j)]_{j=1}^{|\mathbf{a}_q|} \in \mathbb{R}^{d \times |\mathbf{a}_q|}$ , are vector representations of a set of stimuli whose labels are to be predicted together with uncertainty measures.
Note that there are no additional assumptions about this probing dataset, e.g., the distributions of stimuli in observations or queries. We use GPP for OOD detection, as shown later. Unlike existing probes (e.g., [Hewitt et al., 2021, Xu et al., 2020]), we do not consider a stimulus and its vector representation as random variables with some underlying distributions. Instead, viewing them as deterministic allows us to probe with any stimuli (with potentially any underlying distributions).
3.2 Background: Beta Gaussian processes
GPP uses a Beta GP to define a prior distribution over classifiers. The Beta GP is a special case of Dirichlet-based GPs [Milios et al., 2018], which have been shown to either outperform or achieve similar performance than classic GP classification approximations [Rasmussen and Williams, 2006]. For each representation  $a \in \mathcal{A}$ , Beta GPs assume the binary label  $y$  is generated by sampling  $y \sim \text{Bernoulli}(g(a))$  where the Bernoulli parameter  $g(a) \sim \text{Beta}(\alpha(a), \beta(a))$ . The Beta variable  $g(a)$  can be written as two independent Gamma variables:
$$
g (a) = \frac {t _ {\alpha} (a)}{t _ {\alpha} (a) + t _ {\beta} (a)}, t _ {\alpha} (a) \sim \operatorname {G a m m a} (\alpha (a), 1) \text {a n d} t _ {\beta} (a) \sim \operatorname {G a m m a} (\beta (a), 1),
$$
where  $t_{\alpha}(a)$  and  $t_{\beta}(a)$  are independent. The Gamma distributions are then approximated with Log-normal distributions via moment matching. That means,
$$
f _ {\alpha} (a) = \log \left(t _ {\alpha} (a)\right) \sim \mathcal {N} (\log (\alpha (a)) - \frac {v _ {\alpha}}{2}, v _ {\alpha}), \text {w h e r e} v _ {\alpha} = \log \left(\frac {1}{\alpha (a)} + 1\right), \tag {1}
$$
and the same for  $f_{\beta}$ . Beta GPs use two GPs to model the two latent functions  $f_{\alpha}$  and  $f_{\beta}$ ,
$$
g = \frac {1}{1 + e ^ {- f}}, \text {w h e r e} f = f _ {\alpha} - f _ {\beta}, f _ {\alpha} \sim \mathcal {G P} (\mu , k), f _ {\beta} \sim \mathcal {G P} (\mu , k), f _ {\alpha} \perp f _ {\beta}. \tag {2}
$$
Figure 1: Illustrative example of probabilistic probing: The model of interest  $\mathcal{M}$  transforms a stimulus  $x$  to its vector representation  $a = \phi(x) \in \mathbb{R}^2$ . The probe,  $\mathcal{G}(\theta \mid D)$ , a distribution parameterized by  $\theta$  over possible classifiers (mapping from representations to predicted probabilities) conditional on the observation set  $D$ . Suppose  $D$  contains two datapoints. Sampled classifiers  $g$  from  $\mathcal{G}(\theta \mid D)$  are depicted in the middle box. The aggregation of predictions from the sampled classifiers produces the distribution for  $g(a)$ ;  $g(a)$  is the predicted probability of any representation  $a$ . Notice that the two data points in  $D$  are located close to each other, yet has opposite labels, indicating fuzziness in labels around this region. Imagine a new query data  $a$ . In the top right-most, we see that  $a$  close to  $D$  produces a distribution that is centered around  $g(a) = 0.5$ . If the representation  $a$  is far away from  $D$  (bottom right-most), the probe can have higher epistemic uncertainty due to lack of information, and the predicted probability follows the prior distribution (in this case, the prior has two modes, one at  $g(a) = 0$  and one at  $g(a) = 1$ ). Note that both distributions have an expectation at  $g(a) = 0.5$ , for drastically different reasons: fuzzy label v.s. lack of knowledge.
Note that the latent function  $f$  in Eq. 2 is still a GP given that it is the sum of two independent GPs. We denote the Beta GP as  $\mathcal{G}(\theta)$ , where  $\theta = (\mu, k)$ . While we focus on probing as binary classification in this work, it is straightforward to extend GPP to multi-class classification by switching the Beta prior to a Dirichlet prior [Milios et al., 2018] (i.e., using more latent functions than just  $f_{\alpha}$  and  $f_{\beta}$ ).
3.3 Adapting Beta GPs for GPP
There are two important details in Beta GPs that require special attention: how to set the prior and how to approximate the posterior of classifier  $g$ . This requires setting mean and kernel functions  $\mu, k$  in the prior  $\mathcal{GP}(\mu, k)$  of the latent functions  $f_{\alpha}$  and  $f_{\beta}$ , and computing their posterior  $f_{\alpha}, f_{\beta} \mid D$ .
We adapt these aspects of Beta GPs for probing settings, that also enables users to incorporate the characteristics of their end-tasks by setting intuitive hyperparameters. See more explanations and illustrations in the Appendix.
3.3.1 Setting the prior
First, we need to make sure the prior distribution approximated by GPs in our GPP matches with the Beta prior. Without loss of generality, we assume that for any  $a \in \mathcal{A}$ , the prior distribution for  $g(a)$  is  $\mathrm{Beta}(\epsilon, \epsilon)$  for some  $\epsilon > 0$ . This hyperparameter  $\epsilon$  reflects the characteristics of the task at hand; the probability of a positive label might be bi-modal with density centered at both 0 and 1 (e.g., label is not fuzzy) ( $\epsilon < 1$ ), or uniformly distribution in  $[0, 1]$  ( $\epsilon = 1$ ), or centered at 0.5 ( $\epsilon > 1$ ).
To match the Beta prior, the prior for both  $f_{\alpha}(a)$  and  $f_{\beta}(a)$  has to be  $\mathcal{N}(\log (\epsilon) - \frac{v}{2}, v)$ , where  $v = \log \left( \frac{1}{\epsilon} + 1 \right)$ . To match this normal distribution, we set the mean function to a constant  $\mu(a) = \log (\epsilon) - \frac{v}{2}$ , and constrain the kernel function to satisfy  $k(a, a) = v$ .
3.3.2 Posterior inference
This formulation offers a closed-form solution for posterior inference of the GP. Namely, any  $y_{i} = 1$  in  $D = \{(a_i,y_i)\}_{i = 1}^{|D|}$  results in updating Beta parameters of  $g(a_{i})$  to be  $\mathrm{Beta}(\epsilon +s,\epsilon)$  and any  $y_{i} = 0$  results in  $\mathrm{Beta}(\epsilon ,\epsilon +s)$ , for some  $s\geq 1$ . This hyperparameter  $s$  describes how much "strength" (i.e., weight of the observation) we add to the Beta posterior. To match the GP posterior to the Beta posterior, we synthetically add observations with heteroscedastic noise to functions  $f_{\alpha}$  and  $f_{\beta}$ . If  $y_{i} = 1$ , the observation for  $f_{\alpha}$  is  $(a_{i},\log (\epsilon +s) - \frac{v^{\prime}}{2})$  with noise  $\mathcal{N}(0,v^{\prime})$  where  $v^{\prime} = \log (\frac{1}{\epsilon + s} +1)$ , and the observation for  $f_{\beta}$  is  $(a_{i},\log (\epsilon) - \frac{v^{\prime\prime}}{2})$  with noise  $\mathcal{N}(0,v^{\prime \prime})$  where  $v^{\prime \prime} = \log (\frac{1}{\epsilon} +1)$ . And vice versa for  $y_{i} = 0$ . The posterior of a GP remains a GP with closed-form updates to its mean and kernel function. Once we have the posterior of latent functions  $f_{\alpha}\mid D$  and  $f_{\beta}\mid D$ , obtaining the posterior of classifier  $g$  is trivial by updating the GP priors in Eq. 2 to the GP posteriors.
3.3.3 The cosine kernel
The choice of kernel function is important for GPP as it directly determines the characteristics of the distribution of latent functions, which translates to the distribution of classifiers that GPP uses. We define a cosine kernel as:
$$
k \left(a, a ^ {\prime}\right) = v \frac {a ^ {\top} a ^ {\prime} + 1}{\left(\| a \| + 1\right) ^ {\frac {1}{2}} \left(\| a ^ {\prime} \| + 1\right) ^ {\frac {1}{2}}} \text {, w h e r e} v = \log \left(\frac {1}{\epsilon} + 1\right). \tag {3}
$$
Note that multiplying the constant  $v$  is the constraint we have from setting the prior for the Beta GP. Using the cosine kernel in the GP is equivalent to defining a distribution over linear latent functions in an augmented space of the original representation space  $\mathcal{A}$ , i.e.,
$$
f _ {\alpha} (a) = W _ {\alpha} ^ {\top} \psi (a) + \mu (a), \text {w h e r e} \psi (a) = \frac {\sqrt {v}}{\| a \| + 1} \left[ \begin{array}{l} a \\ 1 \end{array} \right] \in \mathbb {R} ^ {d + 1} \text {a n d} W _ {\alpha} \sim \mathcal {N} (0, I). \tag {4}
$$
And similarly for  $f_{\beta}$  with weights  $W_{\beta}$ . Here  $\mu$  is the mean function of the Beta GP. Note that we have the closed-form posterior distribution for the weights, which remains a multivariate Gaussian distribution. See Appendix for more details.
Why is the kernel defined in this way? During the pre-training stage of deep learning models, a weighted sum of the activations and a bias term are used to construct neurons of the next layer. Following Wang et al. [2023], we can view each next layer neuron as independent samples of classifiers from the same distribution, described by a Beta GP in this work. This means classifiers from the Beta GP should all consider a bias term. Hence we augment the activations by an additional dimension with value fixed at 1 (Eq. 4). Additionally, we normalize the augmented activations,  $\psi(a)$ , such that  $k(a, a) = v$  (the constraint from setting the prior in §3.2) holds for all  $a \in \mathcal{A}$  (Eq. 3).
We can interpret the cosine kernel as a scaled cosine of the angle between two vectors (standard cosine similarity) in an augmented representation space described by activations and biases.
3.4 Uncertainty measures
The probabilistic probing formulation makes it possible to compute a posterior distribution over labels for queries. Building on terms from Fox and Ulkümen [2011], we define uncertainty measures that identify the epistemic and aleatory uncertainty of the posterior prediction  $g(a)$  where  $g = \frac{1}{1 + e^{-f}} \sim \mathcal{G}(\theta \mid D), \forall a \in \mathcal{A}$ . Let the posterior of the latent function be  $f(a) \mid D \sim \mathcal{N}(\mu_D(a), k_D(a))$ .
- Episteme :=  $-\mathbb{H}[g(a)]$ , the negative entropy of the distribution of  $g(a)$ , which describes the amount of knowledge the probe has about the label of stimulus  $a$ . In GPP, we have  $\mathbb{H}[g(a)] = \mathbb{H}[f(a)] - \mu_D(a) - 2\mathbb{E}[\log(1 + e^{-f(a)})]$  and  $\mathbb{H}[f(a)] = \frac{1}{2}\log(2\pi ek_D(a))$ .
- Alea :=  $\mathbb{H}[y \mid g(a)] = \mathbb{E}[-g(a)\log(g(a)) - (1 - g(a))\log(1 - g(a))]$ , the expected entropy of the conditional distribution  $p(y \mid g(a))$ , where  $y$  is the label of stimulus  $a$ . Higher alea corresponds to more fuzziness in the label of  $a$ .
- Judged probability :=  $\mathbb{E}[g(a)] = \mathbb{E}\left[\frac{1}{1 - e^{-f(a)}}\right]$ , which is the expected probability (judged by the probe) that the label of stimulus  $a$  is positive.
Figure 2 shows the relationship of episteme, alea and judged probability for a rational agent (white region, based on [Fox and Ülkümen, 2011]).
Query
Episteme (prior): 0.042  
Judged probability: 0.5  
Alea: 0.44
(a)
Epi:0.88 Prob:0.96 Alea: 0.12
(b)
Episteme: 1.3  
Judged probability: 0.6  
Alea: 0.56
(c)
Episteme: 1.5  
Judged probability: 0.79  
Alea: 0.44
Observations
(d)
Episteme: 1.5  
Judged probability: 0.37  
Alea: 0.56
Figure 3: How episteme, judged probability and area of GPP changes as more observations are given to the CoCa model [Yu et al., 2022]. As observations are added, episteme increases while judged probability fluctuates. (a) With 1 positive observation, the judged probability for the query is 0.96, showing that the query and an observation are very close in representation space. (b) With 1 positive (rabbit) and 1 negative (dog) example, the judged probability lowers to 0.6, showing that the CoCa model would also put rabbits and dogs close in representations, but the queried rabbit is still closer to the observed rabbit. (c) With 1 more positive example of a rabbit, judged probability increases together with episteme. (d) If we set the 3rd observation in (c) to be negative, GPP gets more evidence that the query is probably negative and the judged probability decreases to 0.37.
Low episteme means we are "not sure" and high episteme that we are "highly confident" about the underlying probability. Alea, on the other hand, increases with the intrinsic randomness of the true label. The judged probability that the stimulus's label is positive reflects both episteme and alea. With low episteme, the judged probability should be closer to 0.5, meaning the probe does not have enough information to draw a plausible conclusion. It would be irrational (grey areas in Figure 2) for a person to believe something will definitely happen even if there is no evidence. Analogously, it is irrational for a probe to produce probabilities close to 0 or 1 while having low episteme. By accumulating evidence (increasing episteme), the judged prob
ability can move closer and closer to the ground truth (moving to the right in the figure). An accurate probe should produce a judged probability that matches the ground truth fuzziness of a label only if episteme is high.
Figure 2: Interaction of alea and episteme (recreation of Figure 1 from Fox and Ülkümen [2011]).
Figure 3 is a real-world example (using a CoCa model [Yu et al., 2022] and natural images, shown in the figure, as query and observations) of how our measures of uncertainty change as we increase a small number of positive and negative observations.
Table 1 presents the results of applying GPP to the concepts involving olives introduced in §1. The judged probabilities from GPP match human judgments [McCloskey and Glucksberg, 1978]: olives have a  $57\%$  chance of being fruit with a high level of confidence, and GPP is very uncertain about whether olives on a tree are bickets.
Table 1: GPP predictions for olives in §1.
Concept	Judged prob.	Episteme	Alea
Fruit	57%	2.3	0.63
Blicket	91%	0.49	0.22
4 Validating GPP
In addition to probing concepts, GPP predicts aleatory and epistemic uncertainty, which allows us to determine the fuzziness of a concept or detect out-of-domain (OOD) stimuli.
We use 3D Shapes [Burgess and Kim, 2018] and construct 3 datasets based on its concept ontology to validate these uses of uncertainty measures. In these 3 settings, we can control the ground truth level of fuzziness by artificially injecting noise to labels (§4.2, §4.3). In §4.5, we evaluate the OOD performance of GPP and competitive baselines on 3D Shapes as well as 10 datasets defined by coarse-grained labels of the ImageNet dataset [Russakovsky et al., 2015].
4.1 Experiment setups for probing