SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Collaborative Alignment of NLP models
Fereshte Khani
fkhani@microsoft.com
Marco Tulio Ribeiro
Google DeepMind*
marcotcr@gmail.com
Despite substantial advancements, Natural Language Processing (NLP) models often require post-training adjustments to enforce business rules, rectify undesired behavior, and align with user values. These adjustments involve operationalizing "concepts"—dictating desired model responses to certain inputs. However, it's difficult for a single entity to enumerate and define all possible concepts, indicating a need for a multi-user, collaborative model alignment framework. Moreover, the exhaustive delineation of a concept is challenging, and an improper approach can create shortcuts or interfere with original data or other concepts.
To address these challenges, we introduce CoAlign, a framework that enables multi-user interaction with the model, thereby mitigating individual limitations. CoAlign aids users in operationalizing their concepts using Large Language Models, and relying on the principle that NLP models exhibit simpler behaviors in local regions. Our main insight is learning a local model for each concept, and a global model to integrate the original data with all concepts. We then steer a large language model to generate instances within concept boundaries where local and global disagree. Our experiments show CoAlign is effective at helping multiple users operationalize concepts and avoid interference for a variety of scenarios, tasks, and models.
NLP models have showcased remarkable capabilities, yet they are not exempt from imperfections. Unacceptable values embedded in their training data, persistent errors, or violations of business rules highlight the need to teach certain concepts to these models. A concept relates a set of inputs to desired behaviors, e.g. in sentiment analysis, a concept may dictate that "religion does not connote sentiment" (e.g., "I'm Muslim" is Neutral). Similarly, in natural language inference (NLI), the broader concept of "downward monotonicity" specifies entailment relations when certain parts of expressions are made more specific (e.g., "All cats like tuna" entails "All small cats like tuna").
The standard way of teaching concepts to models is adding new training data that exemplifies the concept, e.g. adding neutral sentences containing religious words [1], or adding entailment pairs that illustrate downward monotonicity [2]. However, it is hard to ensure that the data provided does not lead to shortcuts, i.e. spurious correlations or heuristics that allow models to make predictions without capturing the true underlying concept, such as "all sentences with religious terms are neutral", or "going from general to specific leads to entailment". Further, the model may overfit - fail to generalize from the instances provided to the actual concept, e.g. only recognizing pairs in the form ("all X...", "all ADJECTIVE X..."), and not pairs like ("all animals...", "all cats..."). Finally, both shortcuts and overfitting can lead to interference with the original data or other concepts, e.g. leading to failures on "I love Islam" or pairs like ("Some cats like tuna", "Some small cats like tuna"). In sum, operationalizing concepts is challenging, since users typically cannot think of all concept boundaries and interactions in advance.
*Work done while at Microsoft
Figure 1: CoAlign loop for operationalizing a single concept. (a) The user starts by providing some seed data from the concept and their labels, (b) they are used to learn a local concept model. (c) GPT-3 is then prompted to generate new examples, prioritizing examples where the local model disagrees with the global model. (d) Actual disagreements are shown to the user for labeling, and (e) each label improves either the local or the global model. The loop c-d-e is repeated until convergence, i.e. until the user has operationalized the concept and the global model has learned it.
One possible solution is to ask domain experts to create data that covers the concept as comprehensively and precisely as possible, e.g. the GLUE diagnostics dataset [3] or the FraCaS test suite [4]. However, these datasets are often expensive to create, small (and thus not suitable for training), and not exhaustive, as even experts still fail to think about all aspects and nuances of a concept [5]. Another solution is to have users provide data incrementally while they receive feedback from the model, e.g. adversarial training [6] or adaptive testing [5]. These do not require users to think about everything in advance, and can expose and correct model weaknesses. However, adversarial training does not include the notion of concepts explicitly, and adaptive testing does not address the interaction between different concepts or between a concept and the original data. As a result, users may not explore concept boundaries efficiently, and thus do not know when they have sufficiently covered a concept or whether they have introduced interference with other concepts.
In this paper, we introduce Collaborative Alignment of NLP Models (CoAlign). CoAlign leverages the collective knowledge of multiple users to cover many concepts instead of relying on a single user. CoAlign allows users to collaborate with AI systems and other users in specifying concepts and teaching them to models. Our key insight is making concepts explicit by learning a local model for each concept, in addition to a global model that integrates the original data and all additional concepts. When operationalizing a new concept, we rely on a large language model (LLM) to generate new instances where the local and global model disagree, and ask users to label such examples (Figure ??). Intuitively, these examples are either cases where the local model is not yet fully specified, or where the global model still makes errors on the concept (due to overfitting or shortcut reliance). As users label these examples, both models are updated until convergence, i.e. until the concept has been learned in a way that does not conflict with prior data or prior concepts. In a sense, each local model is ever-improving cheap expert in its respective concept. The speed of prediction of local models and diversity of examples generated by the LLM enable users to explore the boundaries between concepts and existing data, an exploration that could be challenging for users to do without aid.
Our experimental results demonstrate the effectiveness of CoAlign in operationalizing concepts and handling interference. We first show CoAlign outperforms AdaTest [5], a SOTA tool for debugging NLP models that also uses GPT-3 by revealing more bugs and fixing them without interference. We then demonstrate that CoAlign operationalize concepts even when user starts with biased data, outperforming a model that relies solely on data collection. We compare the data selection mechanism of CoAlign to random selection and uncertainty sampling by running simplified version of CoAlign where instead of using GPT-3 we iteratively select examples from an unlabeled pool. We show CoAlign outperforms both baselines in teaching an NLI model about downward- and upward-monotone concepts [3], as well as teaching a sentiment analysis model about Amazon products reviews. Finally, in a pilot study, we demonstrated that CoAlign helped users clarify their concepts.
Let  $x$  be a text string, and  $y$  a categorical label, e.g. sentiment (positive, negative, neutral). We assume there is a true function  $f(x) = y$  that we aim to approximate with a model  $\hat{f}(x)$ . Assume
we have access to a "base" dataset  $D_0 = \{(x_1, y_1), \ldots, (x_n, y_n)\}$ , e.g. of movie reviews, from base distribution  $P_0$ . We refer to the model trained on  $D_0$  as the base model  $\hat{f}_0$ .
A concept  $C_i$  is associated with a distribution  $P_i$  over the input space, e.g. a concept might place probability mass exclusively on sentences containing religious words. We say  $x \in C_i$  if  $P_i(x) > 0$ . Since it's hard for users to be exhaustive, we do not assume users can generate from  $P_i$ , but that they can label any  $x$  with  $f(x)$ , and as to whether it is in  $C_i$  or not. We further assume users can provide a very small number samples in the support of  $P_i$ .
Without loss of generality, we assume that we have  $k$  users developing a model collaboratively, each with their own concept. Our goal is to train a model  $\hat{f}$  that does well on both the base distribution and all concepts, i.e., minimizing  $\frac{1}{k + 1}\sum_{i = 0}^{k}\mathbb{E}_{x\sim P_i}[\hat{f} (x)\neq f(x)]$ .
3 Collaborative Alignment of NLP models
In this section we describe how a single user operationalizes their concept by producing a dataset  $D_{i}$  in the context of an existing global model  $\hat{f}$  trained on the base dataset  $D_{0}$  and previous concepts  $D_{1:i-1}$ . If  $\hat{f}_{0}$  is already "aligned" within the concept, i.e.  $\hat{f}_{0}(x) = f(x)$  for all  $x \in C_{i}$ , we would be done and there would be no need for  $D_{i}$ . Thus, what we really want is for  $D_{i}$  to specify the boundary around "failures", cases not currently handled by  $\hat{f}_{0}$ . We abstract away the choice model and learning procedure, assuming that the model can learn the concept given  $D_{i}$  (in practice, we use neural networks that are expressive enough).
3.1 Sampling from the concept
Since by assumption we cannot sample from the concept distribution  $P_{i}$ , it is a challenge to find regions of  $P_{i}$  where  $\hat{f}_0$  fails. To address this, we use GPT-3 [7] as a generator  $\mathcal{G}$  to simulate a random walk within the concept. To do so, we construct a prompt with  $m$  in-concept examples, and use this prompt as input to  $\mathcal{G}$  to generate more samples. Then, we ask the user in the loop to accept or reject each generated sample  $x'$  ( $x'$  is accepted if  $x' \in C_i$ ), and also to label the accepted  $x'$  with  $f(x')$ . The value of  $m$  controls the tradeoff between precision and recall, with high  $m$  generating more in-concept examples and low  $m$  exploring the space more broadly.
Under some conditions it can be shown that  $\mathcal{G}$  simulates a Markov chain with stationary distribution  $P_{i}$  (Appendix A), but the weaker condition of connectivity suffices for finding the concept failures, i.e. there must be a path between any  $x^{\prime},x^{\prime \prime}\in C_i$  with nonzero transition probabilities according to  $\mathcal{G}$  and the prompt. That is, if the concept is connected, with enough time we should be able to find regions of  $P_{i}$  that are not already learned by  $\hat{f}_0$ .
While sampling from  $\mathcal{G}$  eventually leads to the yet-unlearned concept regions, it is an inefficient use of human effort, as it does not use the user labels to guide generation (i.e. the user has to label many examples that the current model already handles correctly). A better approach would be ask the user to label in a way that maximizes the expected information gain for the concept, to which we now turn.
3.2 Local Concept Models
Complex functions can be approximated by simpler functions in a local neighborhood, as evidenced by theoretical results (e.g. Taylor expansion) and empirical applications [8, 9]. Since a concept is a natural local neighborhood, we use this insight and learn a local model  $\hat{f}_i$  to approximate  $f(x)$  in  $C_i$ .
We present a toy example for intuition in Figure 2, where we show toy distributions  $P_0$  (Figure 2a) and  $P_0$  with an additional concept  $P_1$  (Figure 2b). In 2b,  $\hat{f}_0$  learned on samples from  $P_0$  (dashed line) predicts most of  $P_1$  correctly, except for a small region in the bottom left quadrant. However, we would need many random samples from  $P_1$  in order to learn that region, and reach the best model  $\hat{f}$  (solid line in Figure 2b). In contrast, we can learn a good local model  $\hat{f}_1$  for  $P_1$  with a trivial number of samples (Figure 2c). This local model can be used to produce a disagreement region between  $\hat{f}_1$  and  $\hat{f}_0$  (Figure 2d), and sampling from that region would lead to failure discovery much faster.
(a) The initial data  $(P_0)$  is elliptical Gaussian
Figure 2: (a) A model is trained on two elliptical Gaussian (b) The spherical Gaussian are showing a user concept, who wants to find bugs in the model and teach the model about the new concept. However, since the model (dashed line) have high accuracy on the user concept it is hard to find bugs and teaching the model requires data from the low probability region. (c) We fit a classifier to the Spherical Gaussian which can be done with only a few data points (i.e., local model), (d) we then focus on disagreements between these two models to find bugs in the user concept.
(b) The user concept  $(P_{i})$  is spherical Gaussian
(c) local function for the user data
(d) disagreement regions between local and global models
More generally, we define a score function as the disagreement between the local and global function. This score function is used to steer generation such as to maximize the score of generated samples  $x'$ , by adding instances to the prompt for  $\mathcal{G}$  with probability proportional to their score (similar to Ribeiro and Lundberg [5], who use a different score function). We note that models may present false agreement on some samples, i.e.  $\hat{f}_i(x') = \hat{f}_0(x') \neq f(x')$ . To account for this, we also sample from the agreement region sometimes, with a probability that decays over time as we gain more confidence in the local model.
3.3 Operationalizing a concept: from disagreement to convergence
The local and global models disagree on regions where the concept has not yet been learned (global is wrong) or the local model does not fit the user's concept correctly (local is wrong). Thus, every label from the disagreement region results in an improvement in whichever model is incorrect. As labeling progresses, we update both the concept model  $\hat{f}_i$  and the global model  $\hat{f}$  until we reach convergence (i.e.  $\hat{f}$  has learned the concept  $C_i$ ). Note that while  $\hat{f}_i$  is trained on data from  $P_i$ ,  $\hat{f}$  is trained to do well on the base distribution  $P_0$ , as well as on all concepts  $P_{1:k}$ .
(a) seed
Figure 3: This figure illustrates the main steps in CoAlign: (a) The user starts by providing a small number of data points and their labels within their domain (represented by the black circle), (b) we fit a simple model (shown as a linear model) to represent the user's concept, (c) we use the generator  $(\mathcal{G})$  to generate data points towards the region where the local model disagrees with the global model, (d) a diverse set of data points are selected for the user to label, (e) based on the user's feedback, the local and global models are updated until convergence.
(b) local fit
(c) generation
(d) selection
(e) update
We present pseudo-code for operationalizing a concept in Algorithm 1, and illustrate it in Figure 3. The local and global models are initialized in lines 1-2 (Figure 3a and b). In line 5 (Figure 3c),  $\mathcal{G}$  is prompted such as to generate in-concept samples that are likely to be in the disagreement region. Then, in line 6 (Figure 3d), we select a batch of  $b$  examples for labeling according to the disagreement score of the generated instances. This generation-labeling loop is repeated for  $L$  iterations, after which both the local and global models get updated in lines 7-8 (Figure 3e). For the sake of interactivity, in practice we do not train from scratch when updating the models, and instead finetune them from a previous checkpoint. After a few iterations,  $\hat{f}$  and  $\hat{f}_i$  converge (i.e. the user has operationalized the concept and the global model has learned it). If the number of generated samples  $b * L$  is large
enough, we can assume that no disagreement on generated sample between  $\hat{f}$  and  $\hat{f}_i$  means they have converged, and thus we stop and output the generated data  $D_{i}$  (line 9).
Algorithm 1: Operationalizing a new concept  $i$
input:Base dataset  $\overline{D_0}$  , Concept datasets  $D_{1:i - 1}$  . A small set of samples from concept  $D_{i}$    
1 Init local and global: Train  $\hat{f}_i$  on  $D_{i}$  , and train  $\hat{f}$  on  $D_{0:i - 1}$  .   
2 do   
3 for L iterations do   
4 Generation: Prompt  $\mathcal{G}$  with subset from  $D_{i}$  chosen with probability  $\propto |\hat{f} (x) - \hat{f}_i(x)|$  .   
5 Labeling: Select b samples with prob.  $\propto |\hat{f} (x^{\prime}) - \hat{f}_i(x^{\prime})|$  . Users reject  $x^{\prime}$  if out of concept, or add to  $D_{i}$  with label  $f(x^{\prime})$  .   
6 Update local and global: Train  $\hat{f}_i$  on  $D_{i}$  , and train  $\hat{f}$  on  $D_{0:i}$  .   
7 while  $D_{i}$  was updated this round; output:A dataset  $D_{i}$
3.4 Handling interference between concepts
When aligning a model to multiple users (or training a model on different distributions), two types of conflicts can happen: (1) Literal disagreements: where two users disagree in labeling the same input. (2) Interference: where over-generalization in one concept interfere with another concept (e.g., if a user adds a lot of neutral examples for the "religion does not connote sentiment" concept, then training the model on these new examples can cause the model to predict neutral for sentences like "I love Islam" which interfere with a concept that says "love X" is positive). As stated in Section 2, we assume there is a single true function  $f$  (thus no literal disagreement) and we only focus on handling interference. Handling interference is crucial since any local change in ML models can interfere with other parts of the model [10-12].
Having local functions (cheap experts) enable us to check interference efficiently. Every time that a user operationalizes a concept according to Algorithm 1, we check the resulting global model  $\hat{f}$  against the local models for all previous concepts. To do so, we re-run Algorithm 1, only asking the user for labels if there is a newfound disagreement region. In practice, this means that a user adding a new concept needs to make sure it does not break any concepts from other users, a process similar to regression testing in software engineering.
While we want to handle interference, having multiplicity of concepts can be beneficial in refining a concept. In particular, we want the global model to not overfit to a concept (i.e., only memorize the training data) and generalize well; however generalizing well is dependent on other concepts and the previous data. Take, for instance, a world where only bananas are yellow. A user might guide an ML model to recognize bananas solely based on their yellow color. Now if a new user introduces another yellow object, like corn, the model must discern other distinguishing features. Merely combining training data for two such concepts doesn't suffice (as shown in [11]); the boundaries must be distinguished. Furthermore, interference can be beneficial by exposing false agreement regions between the global model and any individual concept model. In other words, while both  $\hat{f}$  and  $\hat{f}_i$  may be relying on the same shortcut to solve some regions of  $P_i$ , it is unlikely that this shortcut does not cause interference between  $\hat{f}$  and all local models  $\hat{f}_j$  for  $j \neq i$ . In this case, the interference caused by adding  $C_j$  is actually beneficial, as it further refines concept  $C_i$ .
In practice the original dataset  $(D_0)$  is often very large and we cannot fine-tune the model on  $D_{0:k}$ . Instead of choosing the whole  $D_0$ , every time we sample data points with highest disagreement between  $\hat{f}_0$  and  $\hat{f}$  from  $D_0$ . In other words, we treat  $\hat{f}_0$  as a user with concept distribution  $P_0$ , this enable us to deal with interference with original model as well (as an example in Figure 2 we might only choose the data points from elliptical Gaussian in disagreement region).
2We like to note that while we haven't introduced a specific mechanism to resolve literal disagreements, our method can surface such disagreements. This can pave the way for resolution through discussions, voting, or even tweaking the model to reflect multiple perspectives, especially in cases where there isn't a consensus among users.
4 Linear regression analysis
We now study CoAlign in a noiseless overparametrized linear regression scenario. This particular setup is used in recent literature to gain some insights into the behaviors of deep networks [10, 11, 13, 14]. For more details and proofs see Appendix B.
Setup. Let  $x \in \mathbb{R}^d$  where only a subset the data points are valid. We assume  $\theta^{\star} \in \mathbb{R}^{d}$  such that  $y = \theta^{\star} \top x$ . Let  $S_{i}$  denote the smallest subspace containing all valid data points of  $C_i$ . Given  $k$  examples in  $C_i$ , let  $S_{i}^{obv}$  denote the subspace observed by the training data,  $S_{i}^{uno}$  denote the unobserved subspace (thus  $S_{i} = S_{i}^{obv} + S_{i}^{uno}$ ), and  $S_{i}^{inv}$  is the subspace that concept  $i$  does not have any variation in it. We consider overparametrized noiseless linear regression, where the number of features  $(d)$  exceeds the number of training data, enabling us to always interpolate all training data. We assume local and global models infer the min  $L_2$  norm interpolant.
As a running example, let  $x \in \mathbb{R}^3$  and consider a concept where data points belonging to that concept satisfy  $x_1 = x_2$ . Let's assume we observed  $x = [1, 1, 0]$  with label  $y = 2$ . In this case we have:  $S_1^{obv} = \{[1, 1, 0]\}$ ,  $S_1^{uno} = \{[0, 0, 1]\}$  and  $S_1^{inv} = \{[1, -1, 0]\}$ . The min-norm solution interpolating the concept is  $\hat{\theta}_i = [1, 1, 0]$ .
Operationalizing a concept. The min-norm solution can also be constructed by introducing constraints that projection of  $\hat{\theta}_i$  on  $S_i^{uno}$  and  $S_i^{inv}$  has to be zero. In the above example, the min-norm solution is the unique answer of the following linear equations ([0, 0, 1] $\theta = 0$ , [1, -1, 0] $\theta = 0$ , [1, 1, 0] $\theta = 2$ ). When we teach the global model about a concept, the naive combination of data points could violate these constraints, leading to disagreements between local and global models. For our running example, suppose the original data point is  $x = [0, 1, 1]$ ,  $y = 2$ . Combining the concept data point with the original data point for inferring global model results in  $\hat{\theta}_{\mathrm{global}} = \left[\frac{2}{3}, \frac{4}{3}, \frac{2}{3}\right]$ , which leads to a disagreement between local and global model for data points varying in the  $[0, 0, 1]$  direction.
In case of disagreement, two scenarios could arise: (1) the direction is non-zero and the local model needs more specification (happens a lot at the beginning), or (2) the direction is zero, but an explicit constraint is needed to prevent the global model from assuming other values. In our running example, the generator can help us to find a data point that the two model disagree  $^3$ . Let's assume the data point is  $x = [0, 0, 1]$  where local model predicts 0 but global model predicts  $\frac{2}{3}$ . We show this data point to the user and either the user confirms  $[0, 0, 1]$  as non-zero or makes the zero explicit, resulting in a global prediction of  $\hat{\theta}_{\mathrm{global}} = [0, 2, 0]$ .
Once we learn the local concept (i.e., all unobserved directions are indeed zero), how many of  $S_{i}^{uno}$  directions need to be added as explicit constraints? Intuitively, we do not need to query user for any direction in  $S_{i}^{uno}$  that has already been observed in  $S_{0}^{obv}$  or are orthogonal  $S_{0}^{obv}$  since there is no interference. The following proposition indicates the maximum number of disagreements for teaching global model about the local concept.
Proposition 1. If  $\mathrm{proj}_{S_i^{uno}}(\theta^\star) = 0$ , then the maximum number of disagreement between local and global models is  $\dim (\mathrm{proj}_{S_0^{obv}}(S_i^{uno} \cap (S_i^{uno} \cap S_0^{obv})^\perp))$ .