SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nicholas Roberts, Xintong Li, Dyah Adila, Sonia Cromp, Tzu-Heng Huang, Jitian Zhao, Frederic Sala
University of Wisconsin-Madison {nick11roberts, fredsala}@cs.wisc.edu {xli2224, adila, cromp, thuang273, jzhao326}@wisc.edu
Machine learning models—including prominent zero-shot models—are often trained on datasets whose labels are only a small proportion of a larger label space. Such spaces are commonly equipped with a metric that relates the labels via distances between them. We propose a simple approach to exploit this information to adapt the trained model to reliably predict new classes—or, in the case of zero-shot prediction, to improve its performance—without any additional training. Our technique is a drop-in replacement of the standard prediction rule, swapping arg max with the Fréchet mean. We provide a comprehensive theoretical analysis for this approach, studying (i) learning-theoretic results trading off label space diameter, sample complexity, and model dimension, (ii) characterizations of the full range of scenarios in which it is possible to predict any unobserved class, and (iii) an optimal active learning-like next class selection procedure to obtain optimal training classes for when it is not possible to predict the entire range of unobserved classes. Empirically, using easily-available external metrics, our proposed approach, LOKI, gains up to  $29.7\%$  relative improvement over SimCLR on ImageNet and scales to hundreds of thousands of classes. When no such metric is available, LOKI can use self-derived metrics from class embeddings and obtains a  $10.5\%$  improvement on pretrained zero-shot models such as CLIP.
The use of pretrained models is perhaps the most significant recent shift in machine learning. Such models can be used out-of-the-box, either to predict classes observed during pretraining (as in ImageNet-trained ResNets [14]) or to perform zero-shot classification on any set of classes [32]. While this is exciting, label spaces are often so huge that models are unlikely to have seen even a single point with a particular label. Without additional modification, pretrained models will simply fail to reliably predict such classes. Instead, users turn to fine-tuning—which requires additional labeled data and training cycles and so sacrifices much of the promise of zero-shot usage.
How can we adapt pretrained models to predict new classes, without fine-tuning or retraining? At first glance, this is challenging: predictive signal comes from labeled training data. However, relational information between the classes can be exploited to enable predicting a class even when there are no examples with this label in the training set. Such relational data is commonly available, or can be constructed with the help of knowledge bases, ontologies, or powerful foundation models [39].
How to best exploit relational structure remains unclear, with a number of key challenges: We might wish to know what particular subset of classes is rich enough to enable predicting many (or all) remaining labels. This is crucial in determining whether a training set is usable or, even with the aid of structure, insufficient. It is also unclear how approaches that use relational information interact with the statistical properties of learning, such as training sample complexity. Finally, performing adaptation requires an efficient and scalable algorithm.
Figure 1: Classification regions in the probability simplex of 3-class classifiers faced with a 100-class problem. The probability simplex using arg max prediction can only output one of three classes. LOKI uses the entire probability vector to navigate the class metric space, leading to more prediction regions. (Left) regions from arg max prediction. (Centers, Right) classification regions from LOKI.
This work addresses these challenges. It proposes a simple and practical approach to learning in structured label spaces, with theoretical guarantees. First, we offer a simple way to translate the soft outputs (i.e., probability vectors) produced by any supervised learning model into a more general model that can exploit geometric information for label structure. In other words, our approach, called LOKI, $^{1}$  is a simple adaptor for pretrained models. LOKI can be applied via a fixed linear transformation of the model outputs. LOKI's simplicity makes it applicable to a broad range of settings while enabling very high-cardinality predictions subject to a potentially small model output budget—we provide a visualization of this key idea in Figure 1.
Theoretically, we provide a rich set of results for the metric-based adaptation setting. First, we introduce a learning-theoretic result in terms of the sample complexity of the pretrained model. It captures a key tradeoff between the number of classes and metric structure, the problem dimensionality, and the number of samples used to train the model prior to adaptation. Next we exhaustively study the properties of training sets, determining for a wide class of relational structures the minimum number (and structure) of subsets that enable reliable prediction. Finally we show how to exploit this result in an active learning-like approach to selecting points that will improve deficient training datasets.
Experimentally, we show that using structural information improves prediction in high-cardinality settings. We demonstrate the strength of the active learning-based approach to dataset expansion over random baselines. Finally, and most excitingly, we show that even in zero-shot models like CLIP, where it is possible to produce probability vectors over any possible class, the use of our adaptor leads to a  $19.53\%$  relative improvement.
First, we introduce the problem setting, notation, and mathematical tools that we will use. Afterward, we discuss how LOKI relates to prior work.
Problem Setting As in conventional supervised learning, we have a dataset  $(x_{1},y_{1}),\ldots ,(x_{n},y_{n})$  drawn from a distribution on  $\mathcal{X}\times \mathcal{Y}$ , where  $\mathcal{X}$  and  $\mathcal{Y}$  are the input and label spaces. In our setting,  $N\coloneqq |\mathcal{V}|$  is finite but large; often  $N\gg n$  so that many labels will simply not be found in the training dataset. We let  $\Lambda \subseteq \mathcal{V}$  with  $|\Lambda | = K$  be the set of observed labels. For convenience of notation, we also define  $\lambda \coloneqq [\lambda_i]_{i = 1}^K,\pmb {y}\coloneqq [y_i]_{j = 1}^N$  to be the vectors of elements in  $\Lambda$  and  $\mathcal{V}$ .
In addition to our dataset, we have access to a relational structure on  $\mathcal{V}$ . We assume that  $\mathcal{V}$  is a metric space with metric (distance)  $d: \mathcal{V}^2 \to \mathbb{R}$ ;  $d$  encodes the relational structure of the label space. Specifically, we model this metric space using a graph,  $G = (\mathcal{V}, \mathcal{E})$  where  $\mathcal{E} \subseteq \mathcal{V}^2 \times \mathbb{R}_+$  is a set of edges relating the labels and the standard shortest-path distance  $d: \mathcal{V}^2 \to \mathbb{R}_{\geq 0}$ . In addition to its use in prediction, the metric  $d$  can be used for evaluating a model by measuring, for example,  $\frac{1}{n} \sum_{i=1}^{n} d^2(f(x_i), y)$ —the analogue of the empirical square loss.
Frechet Mean Estimator Drawing on ideas from structured prediction [5, 35], we use a simple predictor that exploits the label metric. It relies on computing the Fréchet mean [8], given by
$$
m _ {\boldsymbol {y}} (\boldsymbol {w}) := \underset {y \in \mathcal {Y}} {\arg \min } \sum_ {i = 1} ^ {K} \boldsymbol {w} _ {i} d ^ {2} (y, \boldsymbol {y} _ {i}), \tag {1}
$$
<sup>1</sup>Refers to the 'locus' (plural: loci) of the Fréchet mean.
where  $\boldsymbol{w} \in \mathbb{R}_{\geq 0}^{K}$  is a set of weights. The Fréchet mean generalizes the barycenter to metric spaces and is often used in geometric machine learning [26].
Locus of the Fréchet mean. The locus of the Fréchet mean is the set of all Fréchet means under different weights [29]. We write it as  $\Pi (\pmb {y})\coloneqq \cup_{\pmb {w}\in \Delta^{K - 1}}m_{\pmb{y}}(\pmb {w})$
$\Pi (\pmb {y})$  can be thought of the set of all labels in  $\mathcal{V}$  that are reachable by the Fréchet mean given  $\{\pmb {y}_i\}_{i = 1}^K\subseteq \mathcal{V}$  and different choices of its parameter  $\pmb{w}$ . Intuitively, we can think of the locus for a given dataset as describing how usable it is for predicting beyond the observed labels. Trivially, if  $\{\pmb {y}_i\}_{i = 1}^K = \mathcal{V}$ , then  $\Pi (\pmb {y}) = \mathcal{V}$ . We are primarily interested in the case in which  $\{\pmb {y}_i\}_{i = 1}^K\subset \mathcal{V}$  yet we still have  $\Pi (\pmb {y}) = \mathcal{V}$ , or that  $|\Pi (\pmb {y})|$  is at least sufficiently large.
2.1 Relation to Prior work
LOKI is primarily related to two areas: zero-shot learning and structured prediction.
Zero-Shot Learning Like zero-shot learning (ZSL), LOKI is capable of predicting unobserved classes. Our framework is most closely related to generalized ZSL, which uses side information to predict both observed and unobserved classes. Many types of external knowledge are used in ZSL, including text [28, 38, 9], attributes [22, 23, 7], knowledge graphs [42, 34, 10], ontologies [25], and logical rules [33]. Our work is most closely related to ZSL approaches that rely on knowledge graph information. Often, ZSL methods that use knowledge graph information rely on the use of graph neural network architectures [42, 17, 18]. However, we note that these architectures can be heavyweight and can be challenging to scale up to extremely large graphs, whereas LOKI does not have architectural requirements and scales linearly in the size of the graph when  $K$  is small.
Structured Prediction Structured prediction (SP) operates in label spaces that are endowed with algebraic or geometric structure [1, 21]. This includes problems such as predicting permutations [19], non-Euclidean and manifold regression [31, 36], and learning on graphs [12]. LOKI is the most immediately related to the latter, however, any finite metric space can be represented as a graph, which further lends to the flexibility of our approach. Even in discrete structured prediction settings, the cardinality of the label space may be combinatorially large. As such, LOKI can be viewed as a simple method for adapting classifiers to structured prediction.
The Fréchet mean has been used in structured prediction—but in approaches requiring training. In [5],  $\hat{f}(x) = \arg \min_{y \in \mathcal{Y}} \frac{1}{n} \sum_{i=1}^{n} \alpha_i(x) d^2(y, y_i)$ , where  $\alpha(x) = (K + n\nu I)^{-1} K_x$ .  $K$  is the kernel matrix for a kernel  $k: \mathcal{X} \times \mathcal{X} \to \mathbb{R}$ , so that  $K_{i,j} = k(x_i, x_j)$ ,  $(K_x)_i = k(x, x_i)$ .  $\nu$  is a regularization parameter. In other words, the weight  $w_i$  corresponding to  $y_i$  is the average produced by solving a kernel regression problem at all points  $x_k$  in the dataset where  $y_k = y_i$ . It has also been used in weak supervision (WS) for metric-equipped label spaces [41, 37], where the goal is to produce labeled data for training structured prediction models.
3 Framework
We introduce our framework—LOKI. We show how LOKI can be used to adapt any supervised classifier over a set of  $K$  classes to a much richer set of possible class predictions. It does so by weighting the Fréchet mean by the classifier's per-class prediction probabilities or logits, allowing it to predict any class in the locus of the Fréchet mean—potentially far more classes than the initial  $K$ . Next, we show how LOKI can be expressed as a fixed linear transformation of a model's outputs. Finally, we show that LOKI relates to standard classification.
3.1 LOKI: Adapting Pretrained Models
We describe our approach to adapting pretrained classifiers-trained on a set of classes  $\Lambda$  to the metric geometry of the label space  $\mathcal{V}$ , enabling the prediction of unobserved classes.
We model unobserved classes  $\mathcal{Y} \setminus \Lambda$  using the Fréchet mean among observed classes weighted by their prediction probabilities  $P(y = \lambda_i | x$  and  $y \in \Lambda$ ). We denote the vector of model outputs as
$\mathbf{P}_{y|x} \coloneqq [\mathbf{P}_{\lambda_i|x}]_{i=1}^K = [P(y = \lambda_i|x \text{ and } y \in \Lambda)]_{i=1}^K$ . Then predictions using LOKI are given by
$$
\hat {y} \in m _ {\boldsymbol {\lambda}} (\mathbf {P} _ {y | x}) = \underset {y \in \mathcal {Y}} {\arg \min } \sum_ {i = 1} ^ {K} \mathbf {P} _ {\lambda_ {i} | x} d ^ {2} (y, \lambda_ {i}).
$$
3.2 LOKI as a linear transformation of model outputs
Most standard classifiers output a vector of prediction probabilities,  $\mathbf{P}_{y|x}$ , whose entries correspond to the confidence of predicting a specific class. Predictions are typically given by  $\hat{y} \in \arg \max_{i \in [K]} (\mathbf{P}_{y|x})_i$ . LOKI generalizes this prediction rule when viewed as a linear transformation of  $\mathbf{P}_{y|x}$ . Consider the LOKI prediction rule  $\hat{y} \in m_{\lambda}(\mathbf{P}_{y|x}) = \arg \min_{y \in \mathcal{Y}} \sum_{i=1}^{K} \mathbf{P}_{\lambda_i|x} d^2(y, \lambda_i) = \arg \max_{j \in [N]} (\mathbf{D}\mathbf{P}_{y|x})_j$ , where  $\mathbf{D}_{j,i} = [-d^2(y_j, \lambda_i)]; \mathbf{D} \in \mathbb{R}^{N \times K}$  is the matrix of negative squared distances between the observed classes and the rest of the label space. Thus LOKI can be used within standard classification pipelines when the model output  $\mathbf{P}_{y|x}$  is multiplied by the fixed matrix  $\mathbf{D}$ .
3.3 Generalizing Standard Classification
We provide a simple intuition for our approach. The fact that LOKI reduces to standard classification among the observed classes has several implications. This includes the idea that under our framework, forms of few-shot, zero-shot, hierarchical, and partial label learning all reduce to standard classification when additional metric information is introduced.
Generalizing the arg max prediction rule In the absence of this metric information—a situation that we model using the complete graph and setting  $\Lambda = \mathcal{V}$ —our framework also recovers standard classification. Indeed, both in terms of error modeling and in terms of inter-class similarity, the intuition of standard classification and of the 0-1 loss are captured well by the unweighted complete graph—simply treat all differences equally. This graph is given as  $K_{N} := (\mathcal{V}, \mathcal{V}^{2} \times \{1\})$ —i.e., every label is equidistant from every other label. Plugging this into LOKI, we obtain the following:
$$
\hat {y} \in m _ {\pmb {\lambda}} (\mathbf {P} _ {y | x}) = \underset {y \in \mathcal {Y}} {\arg \min} \sum_ {i = 1} ^ {K} \mathbf {P} _ {\lambda_ {i} | x} d ^ {2} (y, \lambda_ {i}) = \underset {y \in \mathcal {Y}} {\arg \min} \sum_ {i = 1} ^ {K} \mathbf {P} _ {\lambda_ {i} | x} \mathbf {1} \{y \neq \lambda_ {i} \} = \underset {i \in [ K ]} {\arg \max} \mathbf {P} _ {\lambda_ {i} | x},
$$
which is exactly the standard classification prediction rule.
Generalizing the 0-1 loss via the expected squared distance The expected squared distance  $\mathbb{E}[d^2 (y,\hat{y})]$  is the standard loss function for many structured prediction problems. Note that accuracy fails in such settings—since it cannot distinguish between small and large errors. This is most clearly seen in the extreme case of regression, where test accuracy will virtually always be zero no matter how good a trained model is. At the other extreme, the complete graph, this loss function becomes the standard 0-1 loss:  $\mathbb{E}[d^2 (y,\hat{y})] = \mathbb{E}[1\{y\neq \hat{y}\} ]$  . As an adaptor for structured label spaces, we use the empirical version of this loss to evaluate LOKI. Note that the expected squared distance subsumes other metrics as well. For example, when  $\mathcal{V} = \mathbb{R}$  , we can derive the standard MSE by setting  $d(y,\hat{y}) = |y - \hat{y} |$  , which is just the standard L1 distance metric. Other scores such as recall at top-k can be similarly obtained at the cost of  $d(\cdot ,\cdot)$  being a true metric. In other words,  $\mathbb{E}[d^2 (y,\hat{y})]$  is an very general metric that supports any metric space, and we use it throughout this work.
4 Theoretical Results
Challenges and Opportunities The arg max of per-class model probabilities is a ubiquitous component of classification pipelines in machine learning. In order to predict unobserved classes using metric space information, LOKI replaces this standard component. As a simple but significant change to standard pipelines, LOKI opens up a new area for fundamental questions. There are three main flavors of theoretical questions that arise in this setting:
1. How does the performance of LOKI change as a function of the number of samples?
2. What minimal sets of observed classes are required to predict any class in the metric space?
3. How can we acquire new classes that maximize the total number of classes we can predict?
Excitingly, we provide a comprehensive answer to these questions for the most common metric spaces used in practice. First, we provide a general error bound in terms of the number of samples, observed classes, problem dimension, and the diameter of the metric space, that holds for any finite metric space. Second, we characterize the sets of observed classes that are required to enable prediction of any class, and show how this set differs for various types of metric spaces of interest. Finally, we provide an active learning algorithm for selecting additional classes to observe so as to maximize the number of classes that can be predicted and we characterize the types of metric spaces for which the locus can be computed efficiently. These results provide a strong theoretical grounding for LOKI.
4.1 Sample Complexity
What is the interplay between predicting unobserved classes based on metric space information and standard learning-theoretic notions like the sample complexity needed to train a model? Our first result illustrates this tradeoff, and relates it to the squared distance loss. Suppose that we only observe  $K \leq N$  of the classes at training time, and that we fit a  $K$ -class Gaussian mixture model. We use LOKI to adapt our pretrained classifier to enable classification of any of the  $N$  classes. We have that (a formal statement and interpretation are in the Appendix),
Theorem 4.1 (Informal LOKI sample complexity). Let  $\mathcal{V}$  be a set of classes represented by  $d$  dimensional vectors under the Euclidean distance, and let  $\Lambda \subseteq \mathcal{V}$  be the set of  $K$  observed classes. Assume that  $n$  training examples are generated by an identity covariance Gaussian mixture model over classes  $\Lambda$ , and that test examples are generated over all classes  $\mathcal{V}$ . Assume that we estimate a Gaussian mixture model on the training set and obtain probability estimates  $\hat{\mathbb{P}}(y_i|x)$  for  $i \in [K]$  for a sample  $(x,y_{*})$  from the test distribution. Then with high probability, under the following model,
$$
\hat {y} _ {*} \in m _ {\Lambda} ([ \hat {\mathbb {P}} (y _ {i} | x) ] _ {i \in [ K ]}) = \underset {y \in \mathcal {Y}} {\arg \min} \sum_ {i \in [ K ]} \hat {\mathbb {P}} (y _ {i} | x) d ^ {2} (y, y _ {i})
$$
the sample complexity of estimating target  $y_{*}$  from the test distribution  $\mathcal{D}_{test}$  with prediction  $\hat{y}_{*}$  is:
$$
\mathbb {E} _ {(x, y _ {*}) \sim \mathcal {D} _ {\text {t e s t}}} [ d ^ {2} (y _ {*}, \hat {y} _ {*}) ] \leq O \left(\frac {d}{\alpha} \sqrt {\frac {\log K / \delta}{n}} \left(\frac {1}{\left(R ^ {1 - \frac {2}{d}} - \frac {\log R}{R}\right)} + \sqrt {d}\right)\right)
$$
where  $\alpha$  is related to the sensitivity of the Fréchet variance to different node choices.
4.2 Optimal Label Subspaces
Next we characterize the subset of distinct labels required to predict any label using our label model with respect to various types of metric spaces.
We first consider label spaces whose metric is a tree graph—such metrics are, for example, related to performing hierarchical classification and weak supervision (where only partial labels are available). We consider a special type of tree called a phylogenetic tree, in which only the leaves can be designated as labels—phylogenetic trees are commonly used to relate the labels of image classification datasets. Afterwards we perform a similar analysis for grid graphs, which are important for label spaces that encode spatial information. Finally, we discuss the case in which no useful metric information is available, i.e., the complete graph.
Our goal in this section is to characterize the properties and size of  $\{\lambda_i\}_{i=1}^K$  in each of these metric spaces such that we still have  $\Pi(\Lambda) = \mathcal{V}$ . We characterize 'optimal' subsets of classes in each of the spaces under a certain notions of optimality. We provide several relevant definitions pertaining to this concept, starting with a notion of being able to predict any possible class using observed classes.
Definition 4.2 (Locus cover). Given a set  $\Lambda \subseteq \mathcal{V}$  for which we construct a tuple of its elements  $\Lambda$ , if it holds that  $\Pi(\Lambda) = \mathcal{V}$ , then  $\Lambda$  is a locus cover.
Definition 4.2 captures the main idea of LOKI—using some set of observed classes for which we can train classifiers, we would like to be able to predict additional unobserved classes using the geometry that relates the observed and unobserved classes. Namely, elements of  $\Pi(\Lambda)$  are 'reachable' using LOKI. We refine this Definition to describe the trivial case that defaults to standard classification and the nontrivial case for which LOKI moves beyond standard classification.
Definition 4.3 (Trivial locus cover). If  $\Lambda = \mathcal{V}$ , then  $\Lambda$  is the trivial locus cover.
This Definition captures the notion of observing all of the classes in the label space. Here, all of the elements of  $\mathcal{V}$  are trivially reachable using LOKI.
Definition 4.4 (Nontrivial locus cover). A locus cover  $\Lambda$  is nontrivial if  $\Lambda \neq \mathcal{V}$
LOKI is more useful and interesting when faced with a nontrivial locus cover—under Definition 4.4, we can use some subset of classes  $\Lambda$  to predict any label in  $\mathcal{V}$ .
Definition 4.5 (Minimum locus cover). Given a set  $\Lambda \subseteq \mathcal{V}$ , if  $\Lambda$  is the smallest set that is still a locus cover, then it is a minimum locus cover.
In cases involving an extremely large number of classes, it is desirable to use LOKI on the smallest possible set of observed classes  $\Lambda$  such that all labels in  $\mathcal{V}$  can still be predicted. Definition 4.5 characterizes these situations—later, we obtain the minimum locus covers for all trees and grid graphs. It is worth noting that the minimum locus cover need not be unique for a fixed graph.
Definition 4.6 (Identifying locus cover). Given a set  $\Lambda \subseteq \mathcal{V}$ , if  $\Lambda$  is a locus cover where  $\forall y \in \mathcal{V}, \exists \mathbf{w} \in \Delta^{|\Lambda| - 1}$  such that  $m_{\Lambda}(\mathbf{w}) = \{y\}$ , then  $\Lambda$  is an identifying locus cover.
The Fréchet mean need not be unique—as an arg min, it returns a set of minimizers. In certain metric spaces, the minimum locus cover can yield large sets of minimizers—this is undesirable, as it makes predicting a single class challenging. Definition 4.6 appeals to the idea of finding some set of classes for which the Fréchet mean always returns a unique minimizer—this is desirable in practice, and in some cases, more so than Definition 4.5.
Definition 4.7 (Pairwise decomposable). Given  $\Lambda \subseteq \mathcal{V}$ ,  $\Pi(\Lambda)$  is called pairwise decomposable when it holds that  $\Pi(\Lambda) = \cup_{\lambda_1, \lambda_2 \in \Lambda} \Pi(\{\lambda_1, \lambda_2\})$ .
In many cases, the locus can be written in a more convenient form—the union of the locus of pairs of nodes. We refer to this definition as pairwise decomposability. Later, we shall see that pairwise decomposability is useful in computing the locus in polynomial time.