Uncertainty-aware Self-training for Text Classification with Few Labels

Introduction

Motivation. Deep neural networks are the state-of-the-art for various natural language processing applications. However, one of the biggest challenges facing them is the lack of labeled data to train these complex networks. Not only is acquiring large amounts of labeled data for every task expensive and time consuming, but also it is not feasible to perform large-scale human labeling, in many cases, due to data access and privacy constraints. Recent advances in pre-training help close this gap. In this, deep and large neural networks like BERT (Devlin et al., 2019), GPT-2 (Radford et al., 2019) and RoBERTa (Liu et al., 2019) are trained on millions of documents in a self-supervised fashion to obtain general purpose language representations. A significant challenge now is to fine-tune these models on downstream tasks that still rely on thousands of labeled instances for their superior performance. Semi-supervised learning (SSL) (Chapelle et al., 2010) is one of the promising paradigms to address this shortcoming by making effective use of large amounts of unlabeled data in addition to some labeled data for task-specific fine-tuning. A recent work (Xie et al., 2019) leveraging SSL for consistency learning has shown state-of-the-art performance for text classification with limited labels leveraging auxiliary resources like backtranslation and forms a strong baseline for our work.

Self-training (ST, (III, 1965)) as one of the earliest SSL approaches has recently been shown to obtain state-of-the-art performance for tasks like neural machine translation (He et al., 2019) performing at par with supervised systems without using any auxiliary resources. For self-training, a base model (teacher) is trained on some amount of labeled data and used to pseudo-annotate (task-specific) unlabeled data. The original labeled data is augmented with the pseudo-labeled data and used to train a student model. The student-teacher training is repeated until convergence. Traditionally, self-training mechanisms do not consider the teacher uncertainty or perform any sample selection during the pseudo-labeling process. This may result in gradual drifts from self-training on noisy pseudo-labeled instances (Zhang et al., 2017). Sample selection leveraging teacher confidence has been studied in curriculum learning (Bengio et al., 2009) and self-paced learning (Kumar et al., 2010) frameworks. These works leverage easiness of the samples to inform a learning schedule like training on easy concepts first followed by complex ones. Since it is hard to assess the “easiness” of a sample, especially in deep neural network based architectures, these works rely only on the teacher model loss while ignoring its uncertainties that can be leveraged for sample selection.

Intuitively, if the teacher model already predicts some samples with high confidence, then there is little to gain with self-training if we focus only on these samples. On the other hand, hard examples for which the teacher model has less confidence are hard to rely on for self-training as these could be noisy or too difficult to learn from. In this scenario, the model could benefit from judiciously selecting examples for which the teacher model is uncertain about. However, it is non-trivial to generate uncertainty estimates for non-probabilistic models like deep neural networks. To this end, we leverage recent advances in Bayesian deep learning (Gal and Ghahramani, 2016) to obtain uncertainty estimates of the teacher for pseudo-labeling and improving the self-training process.

Our task and framework overview. We focus on leveraging pre-trained language models for classification with few labeled samples (e.g., $K=\{20,30\}$) per class for training and validation, and large amounts of task-specific unlabeled data. Figure 1(a) shows an overview of the self-training framework for NLU tasks, where augmented data is obtained from hard pseudo-labels from the teacher (e.g., BERT (Devlin et al., 2019)) without accounting for its uncertainty. We extend this framework with three core components in Figure 1(b), namely: (i) Masked model dropout for uncertainty estimation: We adopt MC dropout (Gal and Ghahramani, 2016) as a technique to obtain uncertainty estimates from the pre-trained language model. In this, we apply stochastic dropouts after different hidden layers in the neural network model and approximate the model output as a random sample from the posterior distribution. This allows us to compute the model uncertainty in terms of the stochastic mean and variance of the samples with a few stochastic forward passes through the network. (ii) Sample selection. Given the above uncertainty estimates for a sample, we employ entropy-based measures to select samples that the teacher is most or least confused about to infuse for self-training corresponding to easy- and hard-entropy-aware example mining. (iii) Confident learning. In this, we train the student model to explicitly account for the teacher confidence by emphasizing on the low variance examples. Finally, all of the above are jointly used for end-to-end learning. We adopt BERT as our encoder and show that its performance can be significantly improved by $12$% for low-resource settings without using additional resources. Furthermore, we marginally outperform recent models (Xie et al., 2019) that make use of auxiliary resources like back-translation. In summary, our work makes the following contributions. (i) Develops an uncertainty-aware self-training framework for text classification with few labels. (ii) Compares the effectiveness of various sample selection schemes leveraging teacher uncertainty for self-training. (iii) Demonstrates its effectiveness for text classification with few labeled samples on five benchmark datasets.

Background

Consider $D_{l}=\{x_{i},y_{i}\}$ to be a set of $n$ labeled instances with $y_{i}$ being the class label for $x_{i}$. Each $x_{i}$ is a sequence of $m$ tokens: $x_{i}=\{x_{i1},x_{i2}\cdots x_{im}\}$. Also, consider $D_{u}=\{x_{j}\}$ to be a set of $N$ unlabeled instances, where $n\ll N$. For most tasks, we have access to a small amount of labeled data along with a large amount of unlabeled ones.

Self-training starts with a base teacher model trained on the labeled set $D_{l}$. The teacher model is applied to a subset $S_{u}\subset D_{u}$ of the unlabeled data $D_{u}$ to obtain pseudo-labeled instances. The augmented data $D_{l}\cup S_{u}$ is used to train a student model. The teacher-student training schedules are repeated till some convergence criterion is satisfied. The unlabeled subset $S$ is usually selected based on confidence scores of the teacher model. In Section 3.1, we study different techniques to generate this subset leveraging uncertainty of the teacher model. Self-training process can be formulated as:

where $p(y|x;W)$ is the conditional distribution under model parameters $W$. $W^{*}$ is given by the model parameters from the last iteration and fixed in the current iteration. The above optimization function has been used recently in variants of self-training for neural sequence generation (He et al., 2019) and data augmentation Xie et al. (2019).

Bayesian neural network (BNN) (Gal and Ghahramani, 2015) assumes a prior distribution over its weights, thereby, replacing a deterministic model’s weight parameters by a distribution over these parameters. For inference, instead of directly optimizing for the weights, BNN averages over all the possible weights, also referred to as marginalization.

Bayesian inference aims to find the posterior distribution over the model parameters $p(W|X,Y)$. Given an instance $x$, the probability distribution over the classes is given by marginalization over the posterior distribution as: $p(y=c|x)=\int_{W}p(y=c|f^{W}(x))p(W|X,Y)dW$.

This requires averaging over all possible model weights, which is intractable in practise. Therefore, several approximation methods have been developed based on variational inference methods and stochastic regularization techniques using dropouts. Here, the objective is to find a surrogate distribution $q_{\theta}(w)$ in a tractable family of distributions that can replace the true model posterior that is hard to compute. The ideal surrogate is identified by minimizing the Kullback-Leibler (KL) divergence between the candidate and the true posterior.

Consider $q_{\theta}(W)$ to be the Dropout distribution (Srivastava et al., 2014) which allows us to sample $T$ masked model weights $\{\widetilde{W_{t}}\}_{t=1}^{T}\sim q_{\theta}(W)$. For classification tasks, the approximate posterior can be now obtained by Monte-Carlo integration as:

Uncertainty-aware Self-training

Given a pre-trained language model as the teacher, we first fine-tune it on the small amount of labeled data. To this end, we use a small batch size to gradually expose the teacher model to the few available labels. Given our low-resource setting, we do not compute uncertainty estimates over the small labeled set. Instead, given the teacher model, we compute uncertainty estimates over each instance from the large unlabeled set as follows. Considering dropouts enabled before every hidden layer in the teacher model, we perform several stochastic forward passes through the network for every unlabeled sample. For computational efficiency, we perform these stochastic passes and hence the self-training over sampled mini-batches.

For each unlabeled instance $x_{u}$, given $T$ stochastic forward passes through the network with dropout, each pass $t\in T$ with corresponding model parameters $\widetilde{W_{t}}\sim q_{\theta}(W)$, generates a pseudo-label given by $p({y_{t}}^{*})=softmax(f^{\widetilde{W_{t}}}(x_{u}))$.

There are several choices to integrate this pseudo-label for self-training, including considering $E(y)=\frac{1}{T}\ \sum_{t=1}^{T}softmax(f^{\widetilde{W_{t}}}(x))$ for the soft pseudo-labels as well as discretizing them for hard labels and aggregating predictions from the $T$ passes as:

In order to incorporate the above uncertainty measures in the self-training framework, we modify the loss component over unlabeled data in the original self-training learning process (Equation 1) as:

where $W^{*}$ denotes the model parameters from the previous iteration of the self-training process.

where, $\widehat{p}^{t}_{c}=p(y_{u}=c|f^{\widetilde{W_{t}}}(x_{u})=softmax(f^{\widetilde{W_{t}}}(x_{u}))$.

Class-dependent selection. We can further modify this measure to take into account the expected class label of the instance. This helps in sampling equivalent number of instances per class, and avoids the setting where a particular class is typically hard, and the model mostly samples instances from that class. Given the pseudo-labeled set $S_{u}$, we can construct the set $\{x_{u}\in S_{u,c}:y_{u}=c\}$ for every class $c$. Now, we use the BALD measure to select instances from each class-specific set instead of a global selection.

2 Confident Learning

The above sampling strategies select informative samples for self-training conditioned on the posterior entropy in the label space. However, they use only the predictive mean, while ignoring the uncertainty of the model in terms of the predictive variance. Note that many of these strategies implicitly minimize the model variance (e.g., by focusing more on difficult examples for hard example mining). The prediction uncertainty of the teacher model is given by the variance of the marginal distribution, where the overall variance can be computed as:

where, ${y_{t}}^{*}(x)=softmax(f^{\widetilde{W_{t}}}(x))$ and the predictive mean computed as: $E(y)=\frac{1}{T}\ \sum_{t=1}^{T}{y_{t}}^{*}(x)$.

We observe the total variance can be decomposed as a linear combination of the model uncertainty from parameters $W$ and the second component results from noise in the data generation process.

In this phase, we want to train the student model to explicitly account for the teacher uncertainty for the pseudo-labels in terms of their predictive variance. This allows the student model to selectively focus more on the pseudo-labeled samples that the teacher is more confident on (corresponding to low variance samples) compared to the less certain ones (corresponding to high variance ones). Accordingly, we update the loss function over the unlabeled data in the self-training mechanism given by Equation 4 to update the student model parameters as:

In the above equation, the per-sample loss for an instance $x_{u}$ is a combination of the log loss $-log\ p(y)$ and (inverse of) its predictive variance given by $log\ \frac{1}{Var(y)}$ with $log$ transformation for scaling. This penalizes the student model more on mis-classifying instances that the teacher is more certain on (i.e. low variance samples), and vice-versa.

Experiments

Encoder. Pre-trained language models like BERT (Devlin et al., 2019), GPT-2 (Radford et al., 2019) and RoBERTa (Liu et al., 2019) have shown state-of-the-art performance for various natural language processing tasks. In this work we adopt one of these namely, BERT as our base encoder or teacher model to start with. We initialize the teacher model with the publicly available pre-trained checkpoint from Wikipedia. To adapt the teacher language model for every downstream task, we further continue pre-training on task-specific unlabeled data $D_{u}$ using the original language modeling objective. The teacher is finally fine-tuned on task-specific labeled data $D_{l}$ to give us the base model for self-training.

Datasets. We perform large-scale experiments with data from five domains for different tasks as summarized in Table 1. SST-2 (Socher et al., 2013), IMDB (Maas et al., 2011) and Elec (McAuley and Leskovec, 2013) are used for sentiment classification for movie reviews and Amazon electronics product reviews respectively. The other two datasets Dbpedia (Zhang et al., 2015) and Ag News (Zhang et al., 2015) are used for topic classification of Wikipedia and news articles respectively. For every dataset, we sample $K$ labeled instances from Train data, and add remaining to the Unlabeled data in Table 1.

Evaluation setting. For self-training, we fine-tune the base model (teacher) on $K$ labeled instances for each task to start with. Specifically, we consider $K=30$ instances for each class for training and similar for validation, that are randomly sampled from the corresponding Train data in Table 1. We also show results of the final model on varying $K\in\{20,30,50,100,500,1000\}$. We repeat each experiment five times with different random seeds and data splits, use the validation split to select the best model, and report the mean accuracy on the blind Test data. We implement our framework in Tensorflow and use four Tesla V100 gpus for experimentation. We use Adam (Kingma and Ba, 2015) as the optimizer with early stopping and use the best model found so far from the validation loss for all the models. Hyper-parameter configurations with detailed model settings presented in Appendix. We report results from our UST framework with easy sample selection strategy employing Equation 6, unless otherwise mentioned.

Baselines. Our first baseline is BERT-Base with $110$ MM parameters fine-tuned on $K$ labeled samples $D_{l}$ for downstream tasks with a small batch-size of $4$ samples, and remaining hyper-parameters retained from its original implementation. Our second baseline, is a recent work UDA (Xie et al., 2019) leveraging backtranslationA sentence is translated to a foreign language followed by backtranslation to the source language. Due to noise injected by Neural Machine Translation systems, backtranslation is often a paraphrase of the original. for data augmentation for text classification. UDA follows similar principles as Virtual Adversarial Training (VAT) (Miyato et al., 2017) and consistency training (Laine and Aila, 2017; Sajjadi et al., 2016) such that the model prediction for the original instance is similar to that for the augmented instance with a small perturbation. In contrast to prior works for image augmentation (e.g., flipping and cropping) UDA leverages backtranslation for text augmentation. In contrast to other baselines, this requires auxiliary resources in terms of a trained NMT system to generate the backtranslation. Our third baseline is the standard self-training mechanism without any uncertainty. In this, we train the teacher model on $D_{l}$ to generate pseudo-labels on $D_{u}$, train the student model on pseduo-labeled and augmented data, and repeat the teacher-student training till convergence. Finally, we also compare against prior SSL works – employing semi-supervised sequence learning (Dai and Le, 2015), adversarial training (Goodfellow et al., 2015; Miyato et al., 2017), variational pre-training (Gururangan et al., 2019), reinforcement learning (Li and Ye, 2018), temporal ensembling and mean teacher models (Laine and Aila, 2017; Tarvainen and Valpola, 2017; Sajjadi et al., 2016), layer partitioning (Li and Sethy, 2019) and delta training (Jo and Cinarel, 2019) – on these benchmark datasets on the same Test data and report numbers from corresponding works.

Overall comparison. Table 2 shows a comparison between the different methods. We observe that the base teacher model trained with only $30$ labeled samples for each class for each task has a reasonable good performance with an aggregate accuracy of $80.85\%$. This largely stems from using BERT as the encoder starting from a pre-trained checkpoint instead of a randomly initialized encoder, thereby, demonstrating the effectiveness of pre-trained language models as natural few-shot learners. We observe the classic self-training approach leveraging unlabeled data to improve over the base model by $8\%$. The UDA model leverages auxiliary resources in the form of backtranslation from an NMT system for augmentation to improve by over $10\%$. Finally, our uncertainty-aware self-training mechanism obtains the best performance by improving more than $12\%$ over the base model without any additional resources. Our method reduces the overall model variance in terms of both implicit reduction by sample selection and explicit reduction by accounting for the sample variance for confident learning. This is demonstrated in a consistent performance of the model across different runs with an aggregated (least) standard deviation of $0.57$ across different runs of the model for different tasks with different random seeds. UDA with its consistency learning closely follows suit with an aggregated standard deviation of $1.62$ across different runs for different tasks. Classic self-training without any such mechanism shows high variance in performance across runs with different seeds. In Table 4, we show the results from other works on these datasets as reported in (Li and Ye, 2018; Jo and Cinarel, 2019; Li and Sethy, 2019; Gururangan et al., 2019)Note that these models use different encoders and pre-training mechanisms.. Our UST framework outperforms them while using much less training labels per class (shown by $K$).

Ablation analysis. We compare the impact of different components of our model for self-training with $30$ labeled examples per class for each task for training and for validation with results in Table 3. Sampling strategies. The backbone of the sample selection method in our self-training framework is given by the BALD measure (Houlsby et al., 2011) that has been shown to outperform other active sampling strategies leveraging measures like entropy and variation ratios in Gal et al. (2017) for image classification. We use this measure in our framework to sample examples based on whether the model is confused about the example or not by leveraging sampling strategies in Equations 7 or 6 and optimized by self-training with Equation 11 – denoted by UST (Hard) and UST (Easy) respectively in Table 3. In contrast to works in active learning that find hard examples to be more informative than easy ones for manual labeling, in the self-training framework we observe the reverse with hard examples often contributing noisy pseudo-labels. We compare this with uniform sampling in the classic ST framework, and observe that sample selection bias (easy or hard) benefits self-training. Class-dependent selection with exploration. In this, we remove the class-dependent selection and exploration with global selection of samples based on their easiness or hardness for the corresponding UST sampling strategy. Class-dependent selection ameliorates model bias towards picking samples from a specific class that might be too easy or hard to learn from with balanced selection of samples across all the classes, and improves our model on aggregate. Confident learning. In this, we remove confident learning from the UST framework. Therefore, we optimize the unlabeled data loss for self-training using Equation 4 instead of Equation 11 that is used in all other UST strategies. This component helps the student to focus more on examples the teacher is confident about corresponding to low-variance ones, and improves the model on aggregate. Overall, we observe that each of the above uncertainty-based sample selection and learning strategies outperform the classic self-training mechanism selecting samples uniform at random.

Impact of $K$ labeled examples. From Figure 2(a), we observe the self-training accuracy to gradually improve with increase in the number of labeled examples per class to train the base teacher model leading to better initialization of the self-training process. With only $20$ labeled examples for each task for training and for validation, we observe the aggregate performance across five tasks to be $89.27$% with further improvements with more labeled data coming from IMDB and AG news datasets.

Impact self-training iterations. Figure 2(b) shows increase in self-training accuracy of UST over iterations for a single run. In general, we observe the self-training performance to improve rapidly initially, and gradually converge in $15$-$20$ iterations. We also observe some models to drift a bit while continuing the self-training process and similar for consistency learning in UDA beyond a certain point. This necessitates the use of the validation set for early termination based on validation loss.

Related Work

Semi-supervised learning has been widely used in many different flavors including consistency training (Bachman et al., 2014; Rasmus et al., 2015; Laine and Aila, 2017; Tarvainen and Valpola, 2017), latent variable models (Kingma et al., 2014) for sentence compression (Miao and Blunsom, 2016) and code generation (Yin et al., 2018). More recently, consistency-based model like UDA Xie et al. (2019) has shown promising results for few-shot learning for classification leveraging auxiliary resources like paraphrasing and back-translation (BT) (Sennrich et al., 2016).

Sample selection. One of the earlier works in neural networks leveraging easiness of the samples for learning is given by curriculum learning (Bengio et al., 2009). This is based on the idea of learning easier aspects of the task first followed by the more complex ones. However, the main challenge is the identification of easy and hard samples in absence of external knowledge. Prior work leveraging self-paced learning (Kumar et al., 2010) and more recently self-paced co-training (Ma et al., 2017) leverage teacher confidence (or lower model loss) to select easy samples during training. In a similar flavor, some recent works have also focused on sample selection for self-training leveraging meta-learning (Li et al., 2019) and active learning (Panagiota Mastoropoulou, 2019; Chang et al., 2017) based on teacher confidence. However, all of these techniques rely on only the teacher confidence while ignoring the uncertainty associated with its predictions. There are also works on anti-curriculum learning (or hard example mining) (Shrivastava et al., 2016) that leverage hardness of the samples.

Uncertainty in neural networks. A principled mechanism to generate uncertainty estimates is provided by Bayesian frameworks. A Bayesian neural network Gal and Ghahramani (2016) replaces a deterministic model’s weight parameters with distributions over model parameters. Parameter optimization is replaced by marginalisation over all possible weights. It is difficult to perform inference over BNN’s as the marginal distribution cannot be computed analytically, and we have to resort to approximations such as variational inference to optimize for variational lower bound (Graves, 2011; Blundell et al., 2015; Hernández-Lobato et al., 2016; Gal and Ghahramani, 2015).

Conclusions

In this work we developed an uncertainty-aware framework to improve self-training mechanism by exploiting uncertainty estimates of the underlying neural network. We particularly focused on better sample selection from the unlabeled pool based on posterior entropy and confident learning to emphasize on low variance samples for self-training. As application, we focused on task-specific fine-tuning of pre-trained language models with few labels for text classification on five benchmark datasets. With only $20$-$30$ labeled examples and large amounts of unlabeled data, our models perform close to fully supervised ones fine-tuned on thousands of labeled examples. While pre-trained language models are natural few-shot learners, we show their performance can be improved by upto $8\%$ by classic self-training and by upto $12\%$ on incorporating uncertainty estimates in the framework.

Appendix

In our experiments, we employ a single model for self-training. Essentially, we copy teacher model parameters to use as the student model and continue self-training. Although, some works initialize the student model from scratch.

Sample size.

Ideally, we need to perform $T$ stochastic forward passes for each sample in the large unlabeled pool. However, this is too slow. For computational efficiency, at each self-training iteration, we select $S_{u}$ samples randomly from the unlabeled set, and then select $R\in S_{u}$ samples from therein based on uncertainty using several stochastic forward passes.

2 Hyper-parameters

We do not perform any hyper-parameter tuning for different datasets and use the same set of hyper-parameters as shown in Table 5.

Also, we retain parameters from original BERT implementation from https://github.com/google-research/bert.

Similar to all other models, we add validation data to UDA to select the best model parameters based on validation loss. We retain all UDA hyper-parameters from https://github.com/google-research/uda. We use the same sequence length for every task as in our models and select the batch size as in Table 6.

Note that our UDA results are worse than that reported in the original implementation due to different sequence length and batch sizes for hardware constraints. We select the maximum batch size permissible by the V100 gpu memory constraints given the sequence length.

References