STENCIL: Submodular Mutual Information Based Weak Supervision for Cold-Start Active Learning

In recent years, many natural language processing use cases have been driven by the advent of widely available pre-trained models and subsequent fine-tuning (Min et al., 2023). Indeed, this paradigm has become commonplace with the prevalence of large language models, seeing heavy use in cases such as instruction fine-tuning (Zhang et al., 2023; Ouyang et al., 2022). The success of this paradigm for downstream tasks, however, relies heavily on annotated data. To scale to today’s pre-eminent models, large corpora of supervised data are needed, resulting in prohibitively large annotation costs. One of the most common choices for reducing annotation cost is through the paradigm of active learning, which aims to select the best instances from unlabeled data for annotation that improves model performance maximally fast. Indeed, active learning has been used extensively both in classical machine learning and recent deep learning approaches (Ren et al., 2021; Settles, 2009; Beck et al., 2021).

Despite the successful use of active learning, complications in the downstream learning environment can prohibit the use of many common methods. One complication is the presence of class imbalance, which tends to reduce the efficacy of active learning methods designed for the general case. Another complication is having to deploy such strategies in cold-start scenarios. Indeed, the majority of published active learning methods assume access to an initial annotated set of data. When both complications are combined, active learning methods for dealing with a class-imbalanced cold-start environment are needed.

While the cold-start scenario has been considered for active learning, such methods tend to operate over multiple rounds of selection and retraining (Barata et al., 2021; Brangbour et al., 2022; Ni et al., 2020; Hacohen et al., 2022; Yuan et al., 2020; Jin et al., 2022; Kothawade et al., 2022a), which addresses class imbalance inefficiently. Due to the nature of cold-start scenarios, such gradual build-up is natural as the active learning strategy learns more about the data landscape before capitalizing on selecting rare-class instances. However, such capitalization can be conducted earlier in cases where prior knowledge can be incorporated, which is often the case in many fine-tuning scenarios and can be conducted via weak supervision (Ratner et al., 2017; Maheshwari et al., 2020, 2022; Rauch et al., 2022). For example, using active learning for spam detection with unlabeled spam-imbalanced data can be accelerated if the sketch of a spam instance is known. Such a sketch is also easier to provide than in other tasks (such as drawing exemplar images for vision tasks). Hence, can class-imbalanced cold-start active learning more quickly improve rare-class performance if prior knowledge of the data environment can be provided?

In this work, we present STENCIL– Submodular mutual informaTion based wEak supervisioN for Cold start actIve Learning – that effectively utilizes prior knowledge of the task via a small exemplar set of rare-class instances. STENCIL uses this set to guide the active learning selection by maximizing the Submodular Mutual Information (see Section 2) (Iyer et al., 2021a, b) between the set of selected instances and the exemplar set. As we show in Section 4, STENCIL provides an immediate improvement of $10\%-18\%$ overall accuracy and $17\%-40\%$ rare-class F-1 score with a single round of selection compared to other methods. STENCIL achieves this performance gain using as few as 15-25 exemplar instances, which are easy to provide in a variety of NLP fine-tuning tasks. Hence, STENCIL offers a sound and complimentary strategy for the cold-start round of class-imbalanced active learning, especially for NLP fine-tuning tasks where STENCIL’s weak supervision can be used.

First, we briefly introduce submodular functions and Submodular Mutual Information (SMI) (Iyer et al., 2021a, b), which is used as STENCIL’s mechanism for guided active learning selection. A set function $F:2^{\mathcal{V}}\rightarrow\mathbb{R}$ over a ground set of instances $\mathcal{V}$ assigns a real-valued score to each possible subset of $\mathcal{V}$. Unfortunately, finding the maximizing subset of a given size of $F$ is NP-hard. Instead, approximation algorithms for finding this subset are used. Notably, a $(1-\frac{1}{e})$-approximate greedy algorithm exists for finding the maximizing cardinality-constrained subset if $F$ is monotone submodular (Nemhauser et al., 1978). Specifically, $F$ is submodular if, for any $A\subseteq B\subseteq\mathcal{V}$ and $a\notin B$, $F(A\cup\{a\})-F(A)\geq F(B\cup\{a\})-F(B)$; additionally, $F$ is monotone if $F(A\cup\{a\})-F(A)\geq 0$ for any $A$ and $a\notin A$. Such a paradigm provides a useful mechanism for selecting sets under an annotation budget within active learning by using desirable submodular functions.

SMI is defined between two sets $A$ and $Q$ for a base submodular function $F$ as $I_{F}(A;Q)=F(A)+F(Q)-F(A\cup Q)$. Indeed, (Iyer et al., 2021a, b) note that SMI serves as a generalization of Shannon-entropic mutual information as one could recover it via SMI’s definition by using Shannon entropy for $F$. Notably, if $I_{F}$ also satisfies certain conditions, then $I_{F}(A;Q)$ is also monotone submodular for a fixed $Q$ (Iyer et al., 2021a). Accordingly, the same maximization framework discussed earlier extends to SMI. (Kothawade et al., 2022c) extend this framework for SMI by defining versions with restricted submodularity (Table 1); that is, $F$ is submodular for only certain subsets of $\mathcal{V}$. By defining $\mathcal{V}$ as $\mathcal{U}\cup Q$ for an auxiliary set of query instances $Q$, one can obtain a set of instances $A\subseteq\mathcal{U}$ with high information overlap with $Q$ by maximizing $I_{F}(A;Q)$, providing a query-relevant $A$ that also has the salient properties (diversity, representation, etc.) of the scoring afforded by $F$. This mechanism has been used across numerous targeted active learning applications in other modalities (Kothawade et al., 2021, 2022b; Beck et al., 2024), so its extension to text is an exciting avenue of study, especially for the cold-start scenario.

In this section, we present STENCIL, which utilizes the maximization framework mentioned in Section 2 with a query set of exemplar instances of the rare class. We provide an overview in Figure 1 and Algorithm 1. STENCIL starts by instantiating the underlying SMI function as chosen from Table 1 through a modular choice of featurization and similarity measure. Featurization can be conducted by taking GloVe embeddings (Pennington et al., 2014) or any similar choice. Subsequently, the similarity values required by each SMI function in Table 1 are computed using a choice of similarity measure, such as cosine similarity, RBF similarity, and so forth (coalesced as matrices of similarity values, or similarity kernels). Once the similarity kernels are computed, a greedy monotone submodular maximization algorithm (Nemhauser et al., 1978; Minoux, 2005; Mirzasoleiman et al., 2015; Iyer & Bilmes, 2019) selects a set of unlabeled instances $A\subseteq\mathcal{U}$ of size $B$. The set $A$ is then annotated as a set $\mathcal{L}$, and a model $M$ is then fine-tuned on $\mathcal{L}$. In our experiments, we use an average of GloVe (Pennington et al., 2014) embeddings and cosine similarity to instantiate each SMI function.

To apply STENCIL, an exemplar query set is needed. The query set comes from prior knowledge of the task and contains exemplars of the rare class. While such exemplars are hard to generate for different modalities of data such as images, we posit that it is easier to derive exemplars for reasonable text-based tasks. With spam detection, for example, prior knowledge dictates that messages that urge to open hyperlinks would likely be spam instances; in such cases, one can prepare a few example sentences of this variety to use as the query set $Q$. With a complimentary choice of featurization, STENCIL can then effectively choose a set of instances that semantically match the exemplars. Conducting selection in this way effectively returns a set of instances that are assumed to be weakly labeled as the rare class, which are then strongly labeled by an annotator.

In this section, we experimentally verify STENCIL’s single-selection ability to handle class-imbalanced cold-start active learning scenarios for text modalities versus common active learning strategies. We additionally conduct ablations across the choice of SMI function (Table 1) and the query set size, analyzing their effect on performance. We conduct experiments over three datasets: YouTube Spam Classification (Alberto et al., 2015), SMS Spam Classification (Almeida et al., 2011), and Twitter Sentiment (Wan & Gao, 2015) (see Appendix B). Such choices are natural since common knowledge of the structure of spam messages and positive sentiments can be incorporated via exemplars (shown in Appendix A). Additionally, both of the latter datasets feature natural class imbalances (SMS: $1:18$, Twitter: $1:6$). We induce an imbalance within the YouTube dataset ($1:10$) to congregate three class-imbalanced settings, and we evaluate on balanced test data to measure both rare-class and overall performance.

For baseline comparison with STENCIL’s SMI variants (Table 1), we evaluate against common active learning strategies. Random chooses $B$ random instances from $\mathcal{U}$. Entropy, Least Confidence, and Margin select the top $B$ instances from $\mathcal{U}$ with highest Shannon entropy, lowest predicted probability, and lowest classification margin, respectively (Settles, 2009). Badge calculates the loss gradients of each point in $\mathcal{U}$ using pseudo-labeling and performs k-means++ sampling (Arthur et al., 2007) to select $B$ instances from $\mathcal{U}$, which gives a diverse set of uncertain instances (Ash et al., 2019). RegEx uses regular expression matching between each data point in the exemplar query set and each data point $\mathcal{U}$. The top $B$ samples with the highest cumulative count of word matches with the patterns or phrases in the query set are selected. Finally, KMean applies $k$-means clustering with $k=B$ on the space of average GloVe embeddings for each instance and returns the 1-NN to each center, which serves as an adjacent baseline to (Hacohen et al., 2022; Jin et al., 2022). In all cases, $B$ is set to 50 for YouTube, 136 for SMS, and 144 for Twitter (roughly $1.5\%$ to $5\%$ of the full unlabeled set size for each dataset). Subsequently, we utilize an LSTM-based network architecture (Hochreiter & Schmidhuber, 1997) of $1.23M$ parameters to take the sequence as input and make classification decisions. We optimize the network utilizing SGD with learning rate $0.01$ for 25-50 epochs (YouTube: 50, SMS: 30, Tweet: 25).

We present the results of our experiments averaged across 10 trials in Table 2 and Table 3. Summarily, we see that the GCMI variant of STENCIL performs the best on the SMS dataset, the LOGDETMI variant performs the best on the YouTube dataset, and the LOGDETMI variant exhibits the best overall accuracy (on balanced test data) while the FLQMI variant exhibits the best rare-class F-1 score on the Twitter dataset. Hence, STENCIL improves the rare-class performance in the class-imbalanced cold-start setting using only one round of selection without sacrificing overall performance, seeing improvements of $10\%-18\%$ overall accuracy and $17\%-40\%$ rare-class F-1 score. Interestingly, we note that while RegEx uses the exemplar query set, its rule-based operation cannot capture diverse sets of semantically related instances as well as STENCIL. Indeed, LOGDETMI accounts for diversity through its determinants, which is advantageous in cold-start settings where a good coverage of the data is required. GCMI’s effectiveness is particularly notable in situations where rare instances exhibit high similarity, as observed in the SMS dataset. Lastly, we observe from Table 3 that, while more query examples provides better performance, there is a diminishing return on investment in terms of performance gains when increasing the query set size beyond a certain point, which means that STENCIL effectively leverages small query sets and does not specifically require large query sets.

In this work, we present STENCIL, which effectively utilizes SMI and prior knowledge of the data environment in the form of a small exemplar set of rare-class instances to manage class-imbalanced cold-start active learning scenarios. We demonstrate that STENCIL is able to make improvements of $10\%-18\%$ overall accuracy and $17\%-40\%$ rare-class F-1 score over common active learning strategies within a single round of selection. As an exciting avenue of future work, the application of STENCIL’s SMI functionality to other text-based tasks beyond classification – and the formation of query exemplars in these tasks – would greatly expand the contributions of this work.

To reproduce our results, we offer six Google Colab notebooks¹¹1https://github.com/nab170130/stencil that can be run repeatedly to amass trial outcomes for each configuration of STENCIL and for each baseline method. Per the setting and model size ($1.23M$ parameters) described in Section 4, it is sufficient to run these notebooks on the standard CPU configuration. Per dataset, amassing 10 trials requires roughly 2 hours of compute time. Hyper-parameters were chosen based on the code examples provided in DISTIL.

We utilize the following libraries in our experiments and list their licenses as follows:

• •

  DISTIL (Beck et al., 2021): MIT License

• •

  Submodlib (Kaushal et al., 2022): MIT License

• •

  PyTorch (Paszke et al., 2019): Modified BSD License

• •

  NLTK (Bird et al., 2009): Apache License Version 2.0

• •

  Transformers (Wolf et al., 2020): Apache License 2.0

• •

  YouTube Spam (Alberto et al., 2015): CC Attribution 4.0 International

• •

  SMS Spam (Almeida et al., 2011): CC Attribution 4.0 International

• •

  Twitter Sentiment (Wan & Gao, 2015): Creative Commons CC-BY-NC-ND