SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Neel Guha*
Mayee F. Chen*
Kush Bhatia*
Azalia Mirhoseini
Frederic Sala
Recent work has shown that language models' (LMs) prompt-based learning capabilities make them well suited for automating data labeling in domains where manual annotation is expensive. The challenge is that while writing an initial prompt is cheap, improving a prompt is costly—practitioners often require significant labeled data in order to evaluate the impact of prompt modifications. Our work asks whether it is possible to improve prompt-based learning without additional labeled data. We approach this problem by attempting to modify the predictions of a prompt, rather than the prompt itself. Our intuition is that accurate predictions should also be consistent: samples which are similar under some feature representation should receive the same prompt prediction. We propose EMBROID, a method which computes multiple representations of a dataset under different embedding functions, and uses the consistency between the LM predictions for neighboring samples to identify mispredictions. EMBROID then uses these neighborhoods to create additional predictions for each sample, and combines these predictions with a simple latent variable graphical model in order to generate a final corrected prediction. In addition to providing a theoretical analysis of EMBROID, we conduct a rigorous empirical evaluation across six different LMs and up to 95 different tasks. We find that (1) EMBROID substantially improves performance over original prompts (e.g., by an average of 7.3 points on GPT-JT), (2) also realizes improvements for more sophisticated prompting strategies (e.g., chain-of-thought), and (3) can be specialized to domains like law through the embedding functions.
Acquiring labeled data for domains like medicine and law is essential to training machine learning models or performing basic data analysis (e.g., "how many contracts contain a choice-of-forum clause" or "how many patient medical histories discuss an adverse reaction to a drug?") [15, 18, 19]. However, building large labeled datasets is difficult, and efforts like [25] show that manual labeling with domain experts is cost-prohibitive. Recent works have begun exploring if language models (LMs) could learn annotation tasks in-context [6] and replace manual labeling at scale [13, 15, 23, 31]. The promise of this approach is that LMs' in-context capabilities enable them to learn tasks from descriptions of the task (i.e., prompts). However, the challenge is that producing high performance
Embroid: Unsupervised Prediction Smoothing Can Improve Few-Shot Classification
Figure 1: The EMBROID method for prompt-patching.
prompts is still expensive, as practitioners require labeled data in order to measure the impact of modifications to a prompt [51]. Existing work has thus focused on how domain experts can optimally construct prompts for a task (prompt-engineering), while minimizing reliance on labeled data [41, 58, 64]. Yet, because language models are sensitive to even small changes in prompt language, these techniques are imperfect and still produce erroneous predictions [2, 8, 38, 51, 70].
Our work approaches the challenge of improving prompt performance without labels from an orthogonal perspective: given the predictions of any prompted LM, can we identify and correct mis-predictions using unlabeled data? We describe this as the problem of prompt-patching. In the context of data annotation tasks, prompt-patching methods should meet three goals. First, they should be theoretically explainable, so that practitioners can understand when and how to apply them. Second, they should be fast, so that practitioners can efficiently integrate them into existing workflows. Finally, they should rarely be wrong, so that they don't worsen the performance of predictions.
Our work presents EMBROID: a method for automatically identifying and correcting LM predictions with unlabeled data and no expert supervision. Recent work has shown that for many tasks, samples close-by in embedding spaces (produced by models like BERT) have the same label [7]. EMBROID applies this intuition to the prompt-patching regime. Specifically, after LM predictions for all samples have been generated, EMBROID retrieves the  $k$  most similar samples for each input, under  $N$  different embedding functions. For each embedding function, EMBROID computes a scaled-modified majority vote over the LM's predictions for the  $k$  retrieved samples. EMBROID then combines these  $N$  votes with the original LM prediction for the test sample using a simple latent variable graphical model that is learned with a fast method-of-moments estimator [14]. The intuition behind EMBROID is that good prompts are smooth with respect to their predictions over a dataset—samples which are proximate under an embedding function should receive consistent predictions. Thus, modifying the predictions of a prompt to increase neighborhood agreement can improve the accuracy of those predictions. Lastly, because a single embedding space may imperfectly capture similarities between samples, retrieving neighbors from multiple embedding spaces improves robustness [28, 40].
Because EMBROID relies on weak-supervision—the subject of recent rigorous study [7]—it is possible to theoretically analyze and explain why and when EMBROID will improve performance. In particular, we find that performance is a function of the quality of the embeddings and the performance of the initial prompt. We also empirically study EMBROID, conducting experiments over six LMs, on up to 95 tasks, with several different prompt strategies. We find that EMBROID rarely worsens performance, and often improves F1 by a substantial margin. For instance, EMBROID improves GPT-3.5 by an average of 4.9 points F1 per task, and GPT-JT by an average of 7.3 points per task. The magnitude of EMBROID's gains are such that it enables a 1.3B parameter model to outperform an instruction-tuned 6.7B parameter model. EMBROID is also complementary to advanced prompt engineering strategies, and achieves performance improvements when applied to prompts designed
using chain-of-thought [64], AMA [2], and selective annotation [58]. Finally, EMBROID can be extended to specialized domains like law, through the use of already-available domain specific embeddings.
Succinctly, our contributions in this paper are: (1) EMBROID, a simple prompt-patching framework for improving LM predictions over text classification tasks; (2) a theoretical analysis of EMBROID which explains performance improvements in terms of embedding smoothness and base accuracy; and (3) an empirical evaluation of EMBROID covering up to 95 tasks and six different LMs.
Improving LM performance Improving the in-context generalization abilities of LMs has been intensely studied. The first family of approaches focuses on adapting LMs in order to make them more amenable to prompting. This includes task-specific finetuning [26, 27, 47], training on instruction data [9, 59], RLHF [50], and weight-surgery methods which attempt to "correct" incorrect information stored in model weights [10, 22, 43, 44]. A second family of approaches explores strategies for optimizing prompts to models, either through the specific textual features of the prompt [29, 45, 64], the use of task decompositions or LM recursion [2], implicit prompt representations [36, 37], or external databases [46]. Prompt-patching, in contrast, focuses on identifying mistakes in the predictions generated from a particular prompt. The most related approaches are aggregation methods, in which the outputs of multiple prompts are combined with an ensembling method [2, 41]. We find that EMBROID outperforms many such baselines, and can be applied to enhance their outputs.
Weak supervision EMBROID leverages statistical techniques developed in the weak supervision literature. The objective in weak supervision is to generate probabilistic labels for unlabeled data by combining the predictions of multiple noisy heuristics [14, 53, 54, 57, 62, 66]. EMBROID's novelty is that it uses embeddings to construct additional synthetic predictions, which are combined with the original predictions. In contrast, recent weak supervision approaches which incorporate embeddings use them to produce more fine-grained accuracy parameters [7], detect and discard training points [34], and as the basis for label propagation with final weak supervision predictions [52].
3 Problem setup and background
Problem setup Our problem setup comprises three elements: an unlabeled dataset, predictions from a LM for each sample in this dataset, and embedding representations of our dataset. Our goal is to improve the accuracy of LM predictions, by using the embedding representations to identify and correct predictions likely to be incorrect. Because recent work has explored how predictions from multiple prompts can be combined for a task [2], we present a generalized version of EMBROID in which we have access to multiple LM predictions. In our empirical evaluation however, we show that EMBROID performs well regardless of the number of predictions per sample available.
More formally, we focus on a binary classification task where  $x \in \mathcal{X}$  denotes a sentence or paragraph and  $y \in \mathcal{Y} = \{-1, 1\}$  is the binary label. We assume we are given an unlabeled dataset  $\mathcal{D} = \{x_i\}_{i=1}^{n_u}$  of  $n_u$  points. Each point  $x$  is sampled i.i.d. from a distribution  $\mathbb{P}_x$ , and there exists a true underlying distribution  $\mathbb{P}$  on the joint  $(x, y)$ . Following the true few-shot regime [51], we assume the only labels available are those used in the prompt. We denote a language model (e.g., GPT-3) as  $\lambda_{\text{LLM}}$ , and a task-specific prompt as  $\phi$ , which prepends task instructions to input  $x$  (e.g., "Does the clause contain an audit provision? Yes or No.") We consider the "prompt" to include both the task description and the in-context samples, consistent with [70]. The prediction this prompt induces for  $\lambda_{\text{LLM}}$  over  $x$  is  $\lambda_{\text{LLM}}(\phi(x)) \in \mathcal{Y}$ . Varying  $\phi$  by changing the task description, in-context demonstrations, or punctuation will alter the prediction generated for  $x$ . For a set of  $m$  prompts  $[\phi_1, \ldots, \phi_m]$ , we denote their respective predictions on  $x$  as a vector of weak sources  $\boldsymbol{\lambda}(x) = [\lambda_{\text{LLM}}(\phi_1(x)), \ldots, \lambda_{\text{LLM}}(\phi_m(x))]$ . For convenience, we denote  $\lambda_{\text{LLM}}(\phi_i(x))$  as  $\lambda_i(x)$  or  $\lambda_i$  when the  $x$  is obvious, and similarly use  $\boldsymbol{\lambda}$  instead of  $\boldsymbol{\lambda}(x)$ . We distinguish between two regimes: in the single-prompt regime with  $m = 1$ , we have access to a LM prediction for each point  $x$ , while in the multi-prompt regime with  $m > 1$ , we have access to multiple predictions.
<sup>1</sup>We assume a task-specific mapping function which allows a practitioner to associate a text generation from an LM to a particular class prediction in  $\mathcal{V}$ .
Algorithm 1 EMBROID: Correcting LLMs with embeddings
Input: Unlabeled data  $\mathcal{D}$  LLM predictions  $\lambda (x)$  for each  $x\in \mathcal{D}$  embedding models  $\mathcal{E} =$ $\{E_1,\dots ,E_N\}$  shrinkage parameter  $\tau$  , nearest neighbors parameter  $k$    
for all unlabeled  $x\in \mathcal{D}$  do for all embedding models  $E_{j}\in \mathcal{E}$  do Compute k-nearest neighbors  $\mathrm{NN}_{j,k}(x)$  Compute smoothed neighborhood prediction  $\lambda_{\mathrm{sm},j}(x)$  using  $\lambda$ $\mathrm{NN}_{j,k}(x)$  , and  $\tau$  using eq. (2) end for   
end for   
Solve graphical model  $\operatorname *{Pr}(y,\pmb {\lambda}(x),\pmb{\lambda}_{\mathrm{sm}}(x))$  in eq. (3) with triplet method over  $\mathcal{D}$  (Algorithm 2).   
for all unlabeled  $x\in \mathcal{D}$  do Sample  $\hat{y}_x\sim \hat{\mathrm{Pr}} (y|\pmb {\lambda}(x),\pmb{\lambda}_{\mathrm{sm}}(x))$    
end for   
Output: Label set  $\hat{Y} = \{\hat{y}_x\mid x\in \mathcal{D}\}$
We assume access to  $N$  embedding models  $\mathcal{E} = [E_1, \ldots, E_N]$ , each represented as a fixed mapping  $E_i: \mathcal{X} \mapsto \mathcal{Z}_i$  from an input  $x$  to an embedding vector  $z$ . These auxiliary embedding models provide representations of  $x$  which encode different types of similarity information. Through model repositories like HuggingFace [65], it is possible to download a number of models which generate representations for text sentences (e.g., BERT or RoBERTa [12, 42]). These models have the property that semantically similar sentences are close-by in embedding space [7, 28, 40].
Weak supervision background Weak supervision uses a graphical model to combine votes from multiple noisy sources into a single prediction, by estimating the accuracy of each source. It models  $\operatorname*{Pr}(y,\lambda (x))$  as a latent variable graphical model and uses  $\hat{y} = \mathrm{argmax}_y\hat{\mathrm{Pr}} (y|\lambda (x))$  to produce label estimates, where  $\hat{\mathrm{Pr}}$  represents the learned model. The graphical model is based on a graph  $G = (V,E)$ , where  $V = y\cup \lambda$  and  $E$  consists of edges from  $y$  to each  $\lambda_{j}$ . We assume no dependencies between sources, although simple extensions can incorporate them [61]. The formal graphical model is:
$$
\Pr (y, \boldsymbol {\lambda} (x)) = \frac {1}{Z} \exp \left(\underbrace {\theta_ {y} y} _ {(I)} + \underbrace {\theta^ {\top} \boldsymbol {\lambda} (x) y}\right) \tag {1}
$$
where  $Z$  is the partition function used for normalization,  $(I)$  represents a label balance term with parameter  $\theta_{y}$  controlling the prior of  $\operatorname{Pr}(y = 1)$ , and  $(II)$  represents the source accuracy term where each  $\theta_{i}$  is an accuracy parameter for the  $i$ th source. Note that from this model, sources are conditionally independent:  $\lambda_{i} \perp \lambda_{j}|y$  for any  $i,j \in [m]$ . Our use of this model has two steps. First, we must learn the accuracy parameters of  $\operatorname{Pr}(y,\boldsymbol{\lambda}(x))$  without access to  $y$ . We use the triplet method introduced in [14], which is an efficient method-of-moments estimator for the parameters. Then, at inference we compute  $\hat{\operatorname*{Pr}}(y|\boldsymbol{\lambda}(x))$ . Appendix B contains more details.
4 EMBROID
First, EMBROID uses the embedding models  $\mathcal{E}$  to compute additional votes for each  $x$ . Let  $\mathrm{NN}_{j,k}(x)\subset \mathcal{D}$  be the  $k$ -nearest neighbors of sample  $x$  under the embedding function  $E_{j}$ . We define the smoothed neighborhood prediction vector  $\lambda_{\mathrm{sm},j}(x)\in \{-1,0,1\}^{m}$  as follows, with  $\lambda_{\mathrm{sm},j}[i](x)$  being the  $i$ th element:
$$
\begin{array}{l} \tilde {\lambda} _ {j} [ i ] (x) = \frac {1}{k} \sum_ {\tilde {x} \in \mathrm {N N} _ {j, k} (x)} \lambda_ {i} (\tilde {x}) \\ \lambda_ {\mathrm {s m}, j} [ i ] (x) = \left\{ \begin{array}{l l} 1 & \tilde {\lambda} _ {j} [ i ] (x) > \tau_ {i} ^ {+} \\ - 1 & \tilde {\lambda} _ {j} [ i ] (x) <   \tau_ {i} ^ {-} \\ 0 & \text {o . w .} \end{array} , \right. \tag {2} \\ \end{array}
$$
where  $\tau_{i}^{+}\in [-1,1]$  and  $\tau_{i}^{-}\in [-1,1]$  act as shrinkage parameters for  $\lambda_{i}$  which control the level of agreement amongst the neighbors of  $x$  necessary to generate a particular vote. The scalar  $\lambda_{\mathrm{sm},j}[i](x)$  is the average vote of  $\lambda_{i}$  amongst the neighbors of  $x$  in  $E_{j}$ . When  $\lambda_{\mathrm{sm},j}[i](x)$  is sufficiently positive, i.e.,  $\lambda_{\mathrm{sm},j}[i](x) > \tau_{i}^{+}$ , EMBROID sets  $\lambda_{\mathrm{sm},j}[i](x)$  to be a positive vote. When  $\lambda_{\mathrm{sm},j}[i](x)$  is sufficiently negative, i.e.,  $\lambda_{\mathrm{sm},j}[i](x) < \tau_{i}^{-}$ , EMBROID sets  $\lambda_{\mathrm{sm},j}[i](x)$  to be a negative vote. Otherwise,  $\lambda_{\mathrm{sm},j}[i](x)$  is set to be an abstain. The intuition is that  $\lambda_{\mathrm{sm},j}[i](x)$  will be an accurate vote over  $x$  whenever two conditions are met: (1) the LM is generally accurate, i.e.,  $\lambda_{j}$  is usually correct, and (2)  $E_{j}$  is smooth, i.e., nearest-neighbors share the same task label.
Next, we augment our base model in equation (1) to incorporate these auxiliary neighborhood predictions  $\lambda_{\mathsf{sm}} = [\lambda_{\mathsf{sm},1},\ldots ,\lambda_{\mathsf{sm},N}]\in \{-1,0,1\}^{Nm}$  computed using the embeddings:
$$
\Pr (y, \boldsymbol {\lambda}, \boldsymbol {\lambda} _ {\mathrm {s m}}) = \frac {1}{Z} \exp \left(\theta_ {y} y + \theta^ {\top} \boldsymbol {\lambda} y + \sum_ {j = 1} ^ {N} \alpha_ {j} ^ {\top} \lambda_ {\mathrm {s m}, j} y\right), \tag {3}
$$
where the vector  $\alpha_{j}\in \mathbb{R}^{m}$  represents the quality parameters for the  $j^{th}$  embedding model when used with the  $m$  different prompts. To solve this model and produce label estimates, we note that it has the same format as (1) if we concatenate  $\lambda$  and  $\lambda_{\mathrm{sm}}$  into one set of weak sources. Therefore, we can use the triplet method from [14] to learn parameters and output estimates  $\hat{\operatorname*{Pr}} (y|\boldsymbol {\lambda}(x),\boldsymbol{\lambda}_{\mathrm{sm}}(x))$  for each  $x\in \mathcal{D}$  at inference time (see Appendix B for details).
Parameters  $\theta$  and  $\alpha_{j}$  in (3) allow us to trade-off two different sources of information—one presented by directly prompting an LM to obtain a label and the other by incorporating similarity information from the embedding models—and to further account for varying error modes among the embedding models. Our use of the neighborhood predictions in (3) yields a more expressive model than the standard weak supervision framework solely on LLM predictions in (1), which we can recover when  $k = 0$ , and can thus help make corrections to the LLM predictions. In practice, we find that setting  $\tau_{i}^{+} = \tau_{i}^{-} = \mathbb{E}[\lambda_{i}]$  (i.e., the average source vote) yields good performance (Appendix G).
We analyze EMBROID, discussing the advantages of using  $\lambda_{\mathrm{sm}}$  in addition to  $\lambda$ , and show that embedding smoothness and base prediction accuracy play a critical role in information gain. Appendix F provides synthetics demonstrating these tradeoffs and comparing to weak-supervision baselines.
First, we provide a result on the generalization error of our model  $\hat{\mathrm{Pr}}(y|\lambda,\lambda_{\mathrm{sm}})$ . Define the generalization error as the expected cross-entropy loss,  $L(\lambda,\lambda_{\mathrm{sm}},\mathcal{D}) = \mathbb{E}_{y,\lambda(x),\lambda_{\mathrm{sm}}(x),\mathcal{D}}[-\log \hat{\mathrm{Pr}}(y|\lambda(x),\lambda_{\mathrm{sm}}(x))]$ . We use  $[\lambda_1,\dots,\lambda_{(N+1)m}]$  to represent  $[\lambda,\lambda_{\mathrm{sm}}]$  and denote by  $a_{\max} = \max_i \mathbb{E}[\lambda_i(x)y]$  the largest accuracy (scaled to  $[-1,1]$ ) of any source, and by  $b_{\min} = \min_{i,j} \{\mathbb{E}[\lambda_i\lambda_j], \hat{\mathbb{E}}[\lambda_i\lambda_j]\}$  the minimum expected pairwise product between any two sources. Assume that all sources are better than random, e.g.,  $\operatorname*{Pr}(\lambda_i = y) > 0.5$ . These terms and assumptions are from using the triplet method.
Proposition 5.1. Suppose that the data  $x, y, \lambda, \lambda_{\mathrm{sm}}$  follows the model in (3). The generalization error of  $\hat{\operatorname*{Pr}}(y|\lambda, \lambda_{\mathrm{sm}})$  can be decomposed into
$$
L (\boldsymbol {\lambda}, \boldsymbol {\lambda} _ {\mathrm {s m}}, \mathcal {D}) \leq \underbrace {H (y | \boldsymbol {\lambda} , \boldsymbol {\lambda} _ {\mathrm {s m}})} _ {\text {I r r e d u c i b l e E r r o r}} + \underbrace {\frac {C (N + 1) m}{n _ {u}}} _ {\text {V a r i a n c e}} + o (1 / n _ {u}),
$$
where  $C = \frac{3(1 - b_{\mathrm{min}}^2)}{8b_{\mathrm{min}}^2(1 - a_{\mathrm{max}}^2)}\left(\frac{1}{b_{\mathrm{min}}^4} +\frac{2}{b_{\mathrm{min}}^2}\right)$
In the bound above, the variance term comes from estimation error when learning the parameters via the triplet method. The irreducible error depends on quality of  $\lambda$  and  $\lambda_{\mathrm{sm}}$ . If knowledge of the LLM prediction and neighborhood prediction significantly reduces uncertainty in  $y$ , the conditional entropy term  $H(y|\lambda(x), \lambda_{\mathrm{sm}}(x))$  is low.
Information gain from using both  $\lambda$ ,  $\lambda_{\mathrm{sm}}$  We compare upper bounds on generalization error when both  $\lambda$ ,  $\lambda_{\mathrm{sm}}$  are modeled versus when only  $\lambda$  is modeled, as in (1) corresponding to classical weak
supervision. Based on the bound in Proposition 5.1, modeling both  $\lambda$  and  $\lambda_{\mathrm{sm}}$  increases the variance term by a constant multiplicative factor.
Here, we examine how the irreducible error is affected, that is, the difference  $H(y|\lambda) - H(y|\lambda,\lambda_{\mathrm{sm}})$ . Since this quantity is always nonnegative, we focus on bounding the pointwise difference in conditional entropy—which we call the information gain—for a given  $x_0$  on which the LLM is incorrect. For simplicity, suppose we have one embedding  $E$ . An embedding  $E$  is  $M$ -smooth with respect to the label if
$$
\Pr (\tilde {y} = c | y = c, \| E (x) - E (\tilde {x}) \| \leq \varepsilon) \geq M _ {E} (\varepsilon), \tag {4}
$$
where  $c \in \mathcal{V}, \varepsilon > 0$  and  $M_E(\cdot) \in [0,1]$  is decreasing in its input. Note that this definition requires knowledge of ground-truth labels, and is thus impossible to use as a metric for selecting which embeddings to use.
Define  $\beta_{i} = \operatorname{Pr}(\lambda_{i} = y)$  as the accuracy of  $\lambda_{i}$  and  $p_{\lambda} = \operatorname{Pr}(y = 1|\lambda(x_0))$  as the prediction on  $x_0$  given only access to  $\lambda$ . Let  $\varepsilon_k = \max_{\tilde{x} \in \mathrm{NN}_k(x)} \|E(x) - E(\tilde{x})\|$  be the maximum distance between  $x_0$  and its  $k$  neighbors. Without loss of generality, assume the label on  $x_0$  is  $y = 1$ .
Theorem 5.2. Assume that  $E$  is  $M$ -smooth. The pointwise information gain on  $x_0$  is
$$
H (y | \boldsymbol {\lambda} (x _ {0})) - H (y | \boldsymbol {\lambda} (x _ {0}), \boldsymbol {\lambda} _ {\mathrm {s m}} (x _ {0})) \geq
$$
$$
2 \left(1 - p _ {\lambda}\right) \left[ \prod_ {i = 1} ^ {m} \left[ 1 - \exp \left[ - 2 k \left(\beta_ {N N _ {k}, i} - 0. 5\right) ^ {2} \right] \right] - 0. 5 \right]
$$
where  $\beta_{NN_k,i} = \operatorname*{Pr}_{\tilde{x}\sim NN_k}(\lambda_i(\tilde{x}) = y)\geq \beta_iM_E(\varepsilon_k)$  is the neighborhood accuracy.
A few observations on the bound are in order.
- Improvement over WS: If the neighborhood accuracy is bounded sufficiently far from  $\frac{1}{2}$  and  $k$  is large, using EMBROID has better irreducible error than just using  $\lambda$ . For example, setting  $m = 1$ ,  $k = 10$ ,  $\beta_{\mathrm{NN}_k,i} = 0.7$ , and  $p_{\lambda} = 0.25$  gives us an improvement of 0.076 nats.
- Smoothness: If  $E$  is highly smooth, then  $M_E(\varepsilon_k)$  will be large and irreducible error will be small.
- Base prediction accuracy: If the original sources  $\lambda$  have high accuracy  $(\beta_{i})$ , irreducible error will be small.
Additionally, we observe that if  $p_{\lambda}$  is a high-quality prediction close to the true label 1, the information gain is small. Choice of the  $k$  parameter presents a performance trade-off: increasing  $k$  will increase  $\varepsilon_{k}$  and incorporate farther-away, less reliable predictions, but it will also reduce the noise of the majority vote. We also comment on the information gain when using both  $\lambda$  and  $\lambda_{\mathrm{sm}}$  over just  $\lambda_{\mathrm{sm}}$  in Appendix C.
Our empirical evaluation focuses on three questions: (1) How robust is EMBROID's performance across LMs? (2) How does EMBROID, as a prompt-patching method, compare to high performance prompt-engineering methods? (3) How sensitive is EMBROID to the embeddings and dataset size?