How to Fine-Tune BERT for Text Classification?

Chi Sun, Xipeng Qiu, Yige Xu, Xuanjing Huang

Introduction

Text classification is a classic problem in Natural Language Processing (NLP). The task is to assign predefined categories to a given text sequence. An important intermediate step is the text representation. Previous work uses various neural models to learn text representation, including convolution models Kalchbrenner et al. (2014); Zhang et al. (2015); Conneau et al. (2016); Johnson and Zhang (2017); Zhang et al. (2017); Shen et al. (2018), recurrent models Liu et al. (2016); Yogatama et al. (2017); Seo et al. (2017), and attention mechanisms Yang et al. (2016); Lin et al. (2017).

Alternatively, substantial work has shown that pre-trained models on large corpus are beneficial for text classification and other NLP tasks, which can avoid training a new model from scratch. One kind of pre-trained models is the word embeddings, such as word2vec Mikolov et al. (2013) and GloVe Pennington et al. (2014), or the contextualized word embeddings, such as CoVe McCann et al. (2017) and ELMo Peters et al. (2018). These word embeddings are often used as additional features for the main task. Another kind of pre-training models is sentence-level. Howard and Ruder (2018) propose ULMFiT, a fine-tuning method for pre-trained language model that achieves state-of-the-art results on six widely studied text classification datasets. More recently, pre-trained language models have shown to be useful in learning common language representations by utilizing a large amount of unlabeled data: e.g., OpenAI GPT Radford et al. (2018) and BERT Devlin et al. (2018). BERT is based on a multi-layer bidirectional Transformer Vaswani et al. (2017) and is trained on plain text for masked word prediction and next sentence prediction tasks.

Although BERT has achieved amazing results in many natural language understanding (NLU) tasks, its potential has yet to be fully explored. There is little research to enhance BERT to improve the performance on target tasks further.

In this paper, we investigate how to maximize the utilization of BERT for the text classification task. We explore several ways of fine-tuning BERT to enhance its performance on text classification task. We design exhaustive experiments to make a detailed analysis of BERT.

The contributions of our paper are as follows:

We propose a general solution to fine-tune the pre-trained BERT model, which includes three steps: (1) further pre-train BERT on within-task training data or in-domain data; (2) optional fine-tuning BERT with multi-task learning if several related tasks are available; (3) fine-tune BERT for the target task.

We also investigate the fine-tuning methods for BERT on target task, including pre-process of long text, layer selection, layer-wise learning rate, catastrophic forgetting, and low-shot learning problems.

We achieve the new state-of-the-art results on seven widely-studied English text classification datasets and one Chinese news classification dataset.

Related Work

Borrowing the learned knowledge from the other tasks has a rising interest in the field of NLP. We briefly review two related approaches: language model pre-training and multi-task Learning.

Pre-trained word embeddings Mikolov et al. (2013); Pennington et al. (2014), as an important component of modern NLP systems can offer significant improvements over embeddings learned from scratch. The generalization of word embeddings, such as sentence embeddings Kiros et al. (2015); Logeswaran and Lee (2018) or paragraph embeddings Le and Mikolov (2014), are also used as features in downstream models.

Peters et al. (2018) concatenate embeddings derived from language model as additional features for the main task and advance the state-of-the-art for several major NLP benchmarks. In addition to pre-training with unsupervised data, transfer learning with a large amount of supervised data can also achieve good performance, such as natural language inference Conneau et al. (2017) and machine translation McCann et al. (2017).

More recently, the method of pre-training language models on a large network with a large amount of unlabeled data and fine-tuning in downstream tasks has made a breakthrough in several natural language understanding tasks, such as OpenAI GPT Radford et al. (2018) and BERT Devlin et al. (2018). Dai and Le (2015) use language model fine-tuning but overfit with 10k labeled examples while Howard and Ruder (2018) propose ULMFiT and achieve state-of-the-art results in the text classification task. BERT is pre-trained on Masked Language Model Task and Next Sentence Prediction Task via a large cross-domain corpus. Unlike previous bidirectional language models (biLM) limited to a combination of two unidirectional language models (i.e., left-to-right and right-to-left), BERT uses a Masked Language Model to predict words which are randomly masked or replaced. BERT is the first fine-tuning based representation model that achieves state-of-the-art results for a range of NLP tasks, demonstrating the enormous potential of the fine-tuning method. In this paper, we have further explored the BERT fine-tuning method for text classification.

2 Multi-task learning

Multi-task learning Caruana (1993); Collobert and Weston (2008) is another relevant direction. Rei (2017) and Liu et al. (2018) use this method to train the language model and the main task model jointly. Liu et al. (2019) extend the MT-DNN model originally proposed in Liu et al. (2015) by incorporating BERT as its shared text encoding layers. MTL requires training tasks from scratch every time, which makes it inefficient and it usually requires careful weighing of task-specific objective functions Chen et al. (2017). However, we can use multi-task BERT fine-tuning to avoid this problem by making full use of the shared pre-trained model.

BERT for Text Classification

BERT-base model contains an encoder with 12 Transformer blocks, 12 self-attention heads, and the hidden size of 768. BERT takes an input of a sequence of no more than 512 tokens and outputs the representation of the sequence. The sequence has one or two segments that the first token of the sequence is always [CLS] which contains the special classification embedding and another special token [SEP] is used for separating segments.

For text classification tasks, BERT takes the final hidden state h\mathbf{h} of the first token [CLS] as the representation of the whole sequence. A simple softmax classifier is added to the top of BERT to predict the probability of label cc:

where WW is the task-specific parameter matrix. We fine-tune all the parameters from BERT as well as WW jointly by maximizing the log-probability of the correct label.

Methodology

When we adapt BERT to NLP tasks in a target domain, a proper fine-tuning strategy is desired. In this paper, we look for the proper fine-tuning methods in the following three ways.

1) Fine-Tuning Strategies: When we fine-tune BERT for a target task, there are many ways to utilize BERT. For example, the different layers of BERT capture different levels of semantic and syntactic information, which layer is better for a target task? How we choose a better optimization algorithm and learning rate?

2) Further Pre-training: BERT is trained in the general domain, which has a different data distribution from the target domain. A natural idea is to further pre-train BERT with target domain data.

3) Multi-Task Fine-Tuning: Without pre-trained LM models, multi-task learning has shown its effectiveness of exploiting the shared knowledge among the multiple tasks. When there are several available tasks in a target domain, an interesting question is whether it still bring benefits to fine-tune BERT on all the tasks simultaneously.

Our general methodology of fine-tuning BERT is shown in Figure 1.

Different layers of a neural network can capture different levels of syntactic and semantic information Yosinski et al. (2014); Howard and Ruder (2018).

To adapt BERT to a target task, we need to consider several factors: 1) The first factor is the preprocessing of long text since the maximum sequence length of BERT is 512. 2) The second factor is layer selection. The official BERT-base model consists of an embedding layer, a 12-layer encoder, and a pooling layer. We need to select the most effective layer for the text classification task. 3) The third factor is the overfitting problem. A better optimizer with an appropriate learning rate is desired.

Intuitively, the lower layer of the BERT model may contain more general information. We can fine-tune them with different learning rates.

Following Howard and Ruder (2018), we split the parameters θ\theta into {θ1,⋯ ,θL}\{\theta^{1},\cdots,\theta^{L}\} where θl\theta^{l} contains the parameters of the ll-th layer of BERT. Then the parameters are updated as follows:

where ηl\eta^{l} represents the learning rate of the ll-th layer.

We set the base learning rate to ηL\eta^{L} and use ηk−1=ξ⋅ηk\eta^{k-1}=\xi\cdot\eta^{k}, where ξ\xi is a decay factor and less than or equal to 1. When ξ<1\xi<1, the lower layer has a lower learning rate than the higher layer. When ξ=1\xi=1, all layers have the same learning rate, which is equivalent to the regular stochastic gradient descent (SGD). We will investigate these factors in Sec. 5.3.

2 Further Pre-training

The BERT model is pre-trained in the general-domain corpus. For a text classification task in a specific domain, such as movie reviews, its data distribution may be different from BERT. Therefore, we can further pre-train BERT with masked language model and next sentence prediction tasks on the domain-specific data. Three further pre-training approaches are performed:

1) Within-task pre-training, in which BERT is further pre-trained on the training data of a target task.

2) In-domain pre-training, in which the pre-training data is obtained from the same domain of a target task. For example, there are several different sentiment classification tasks, which have a similar data distribution. We can further pre-train BERT on the combined training data from these tasks.

3) Cross-domain pre-training, in which the pre-training data is obtained from both the same and other different domains to a target task.

We will investigate these different approaches to further pre-training in Sec. 5.4.

3 Multi-Task Fine-Tuning

Multi-task Learning is also an effective approach to share the knowledge obtained from several related supervised tasks. Similar to Liu et al. (2019), we also use fine-tune BERT in multi-task learning framework for text classification.

All the tasks share the BERT layers and the embedding layer. The only layer that does not share is the final classification layer, which means that each task has a private classifier layer. The experimental analysis is in Sec. 5.5.

Experiments

We investigate the different fine-tuning methods for seven English and one Chinese text classification tasks. We use the base BERT models: the uncased BERT-base modelhttps://storage.googleapis.com/bert_models/2018_10_18/ uncased_L-12_H-768_A-12.zip and the Chinese BERT-base modelhttps://storage.googleapis.com/bert_models/2018_11_03/ chinese_L-12_H-768_A-12.zip respectively.

We evaluate our approach on eight widely-studied datasets. These datasets have varying numbers of documents and varying document lengths, covering three common text classification tasks: sentiment analysis, question classification, and topic classification. We show the statistics for each dataset in Table 1.

For sentiment analysis, we use the binary film review IMDb dataset Maas et al. (2011) and the binary and five-class version of the Yelp review dataset built by Zhang et al. (2015).

Question classification For question classification, we evaluate our method on the six-class version of the TREC dataset Voorhees and Tice (1999) and Yahoo! Answers dataset created by Zhang et al. (2015). TREC dataset is dataset for question classification consisting of open-domain, fact-based questions divided into broad semantic categories. Compared to other document-level datasets, TREC dataset is sentence-level, and there are fewer training examples for it. Yahoo! Answers dataset is a big dataset with 1,400k train samples.

Topic classification For topic classification, we use large-scale AG’s News and DBPedia created by Zhang et al. (2015). To test the effectiveness of BERT for Chinese text, we create the Chinese training and test datasets for Sogou news corpus. Unlike Zhang et al. (2015), we use the Chinese character directly rather than Pinyin. The dataset is a combination of the SogouCA and SogouCS news corpora Wang et al. (2008). We determine the category of the news based on the URL, such as “sports” corresponding to “http://sports.sohu.com”. We choose 6 categories – “sports”, “house”, “business”, “entertainment”, “women” and “technology”. The number of training samples selected for each class is 9,000 and testing 1,000.

Data preprocessing Following Devlin et al. (2018), we use WordPiece embeddings Wu et al. (2016) with a 30,000 token vocabulary and denote split word pieces with ##. So the statistics of the length of the documents in the datasets are based on the word pieces. For further pre-training with BERT, we use spaCyhttps://spacy.io/ to perform sentence segmentation in English datasets and we use “。”,“?” and “!” as separators when dealing with the Chinese Sogou News dataset.

2 Hyperparameters

We use the BERT-base model Devlin et al. (2018) with a hidden size of 768, 12 Transformer blocks Vaswani et al. (2017) and 12 self-attention heads. We further pre-train with BERT on 1 TITAN Xp GPU, with a batch size of 32, max squence length of 128, learning rate of 5e-5, train steps of 100,000 and warm-up steps of 10,000.

We fine-tune the BERT model on 4 TITAN Xp GPUs and set the batch size to 24 to ensure that the GPU memory is fully utilized. The dropout probability is always kept at 0.1. We use Adam with β1=0.9\beta_{1}=0.9 and β2=0.999\beta_{2}=0.999. We use slanted triangular learning rates Howard and Ruder (2018), the base learning rate is 2e-5, and the warm-up proportion is 0.1. We empirically set the max number of the epoch to 4 and save the best model on the validation set for testing.

3 Exp-I: Investigating Different Fine-Tuning Strategies

In this subsection, we use the IMDb dataset to investigate the different fine-tuning strategies. The official pre-trained model is set as the initial encoderhttps://github.com/google-research/bert.

The maximum sequence length of BERT is 512. The first problem of applying BERT to text classification is how processing the text with a length larger than 512. We try the following ways for dealing with long articles.

Usually, the key information of an article is at the beginning and end. We use three different methods of truncate text to perform BERT fine-tuning.

head-only: keep the first 510 tokens512 to subtract the [CLS] and [SEP] tokens.;

head+tail: empirically select the first 128 and the last 382 tokens.

The input text is firstly divided into k=L/510k=L/510 fractions, which is fed into BERT to obtain the representation of the kk text fractions. The representation of each fraction is the hidden state of the [CLS] tokens of the last layer. Then we use mean pooling, max pooling and self-attention to combine the representations of all the fractions.

Table 2 shows the effectiveness of the above methods. The truncation method of head+tail achieves the best performance on IMDb and Sogou datasets. Therefore, we use this method to deal with the long text in the following experiments.

3.2 Features from Different layers

Each layer of BERT captures the different features of the input text. We investigate the effectiveness of features from different layers. We then fine-tune the model and record the performance on test error rates.

Table 3 shows the performance of fine-tuning BERT with different layers. The feature from the last layer of BERT gives the best performance. Therefore, we use this setting for the following experiments.

3.3 Catastrophic Forgetting

Catastrophic forgetting McCloskey and Cohen (1989) is usually a common problem in transfer learning, which means the pre-trained knowledge is erased during learning of new knowledge. Therefore, we also investigate whether BERT suffers from the catastrophic forgetting problem.

We fine-tune BERT with different learning rates, and the learning curves of error rates on IMDb are shown in Figure 2.

We find that a lower learning rate, such as 2e-5, is necessary to make BERT overcome the catastrophic forgetting problem. With an aggressive learn rate of 4e-4, the training set fails to converge.

3.4 Layer-wise Decreasing Layer Rate

Table 4 show the performance of different base learning rate and decay factors (see Eq. (2)) on IMDb dataset. We find that assign a lower learning rate to the lower layer is effective to fine-tuning BERT, and an appropriate setting is ξ\xi=0.95 and lr=2.0e-5.

4 Exp-II: Investigating the Further Pretraining

Besides, fine-tune BERT with supervised learning, we can further pre-train BERT on the training data by unsupervised masked language model and next sentence prediction tasks. In this section, we investigate the effectiveness of further pre-training. In the following experiments, we use the best strategies in Exp-I during the fine-tuning phase.

Therefore, we first investigate the effectiveness of within-task further pre-training. We take further pre-trained models with different steps and then fine-tune them with text classification task.

As shown in Figure 3, the further pre-training is useful to improve the performance of BERT for a target task, which achieves the best performance after 100K training steps.

4.2 In-Domain and Cross-Domain Further Pre-Training

Besides the training data of a target task, we can further pre-train BERT on the data from the same domain. In this subsection, we investigate whether further pre-training BERT with in-domain and cross-domain data can continue to improve the performance of BERT.

We partition the seven English datasets into three domains: topic, sentiment, and question. The partition way is not strictly correct. Therefore we also conduct extensive experiments for cross-task pre-training, in which each task is regarded as a different domain.

The results is shown in Table 5. We find that almost all further pre-training models perform better on all seven datasets than the original BERT-base model (row ‘w/o pretrain’ in Table 5). Generally, in-domain pretraining can bring better performance than within-task pretraining. On the small sentence-level TREC dataset, within-task pre-training do harm to the performance while in-domain pre-training which utilizes Yah. A. corpus can achieve better results on TREC.

Cross-domain pre-training (row ‘all’ in Table 5) does not bring an obvious benefit in general. It is reasonable since BERT is already trained on a general domain.

We also find that IMDb and Yelp do not help each other in sentiment domain. The reason may be that IMDb and Yelp are two sentiment tasks of movie and food. The data distributions have a significant difference.

4.3 Comparisons to Previous Models

We compare our model with the following a variety of different methods: CNN-based methods such as Char-level CNN Zhang et al. (2015), VDCNN Conneau et al. (2016) and DPCNN Johnson and Zhang (2017); RNN-based models such as D-LSTM Yogatama et al. (2017), Skim-LSTM Seo et al. (2017) and hierarchical attention networks Yang et al. (2016); feature-based transfer learning methods such as rigion embedding Qiao et al. (2018) and CoVe McCann et al. (2017); and the language model fine-tuning method (ULMFiT) Howard and Ruder (2018), which is the current state-of-the-art for text classification.

We implement BERT-Feat through using the feature from BERT model as the input embedding of the biLSTM with self-attention Lin et al. (2017). The result of BERT-IDPT-FiT corresponds to the row of ‘all sentiment’, ‘all question’, and ‘all topic’ in Table 5, and the result of BERT-CDPT-FiT corresponds to the row of ‘all’ in it.

As is shown in Table 6, BERT-Feat performs better than all other baselines except for ULMFiT. In addition to being slightly worse than BERT-Feat on DBpedia dataset, BERT-FiT outperforms BERT-Feat on the other seven datasets. Moreover, all of the three further pre-training models are better than BERT-FiT model. Using BERT-Feat as a reference, we calculate the average percentage increase of other BERT-FiT models on each dataset. BERT-IDPT-FiT performs best, with an average error rate reduce by 18.57%.

5 Exp-III: Multi-task Fine-Tuning

When there are several datasets for the text classification task, to take full advantage of these available data, we further consider a fine-tuning step with multi-task learning. We use four English text classification datasets (IMDb, Yelp P., AG, and DBP). The dataset Yelp F. is excluded since there is overlap between the test set of Yelp F. and the training set of Yelp P., and two datasets of question domain are also excluded.

We experiment with the official uncased BERT-base weights and the weights further pre-trained on all seven English classification datasets respectively. In order to achieve better classification results for each subtask, after fine-tuning together, we fine-tune the extra steps on the respective datasets with a lower learning rate.

Table 7 shows that for multi-task fine-tuning based on BERT, the effect is improved. However, multi-task fine-tuning does not seem to be helpful to BERT-CDPT in Yelp P. and AG. Multi-task fine-tuning and cross-domain pre-training may be alternative methods since the BERT-CDPT model already contains rich domain-specific information, and multi-task learning may not be necessary to improve generalization on related text classification sub-tasks.

6 Exp-IV: Few-Shot Learning

One of the benefits of the pre-trained model is being able to train a model for downstream tasks within small training data. We evaluate BERT-FiT and BERT-ITPT-FiT on different numbers of training examples. We select a subset of IMDb training data and feed them into BERT-FiT and BERT-ITPT-FiT. We show the result in Figure 4.

This experiment result demonstrates that BERT brings a significant improvement to small size data. Further pre-trained BERT can further boost its performance, which improves the performance from 17.26% to 9.23% in error rates with only 0.4% training data.

7 Exp-V: Further Pre-Training on BERT Large

In this subsection, we investigate whether the BERTLARGE model has similar findings to BERTBASE. We further pre-train Google’s pre-trained BERTLARGE modelhttps://storage.googleapis.com/bert_models/2018_10_18/ uncased_L-24_H-1024_A-16.zip on 1 Tesla-V100-PCIE 32G GPU with a batch size of 24, the max sequence length of 128 and 120K training steps. For target task classifier BERT fine-tuning, we set the batch size to 24 and fine-tune BERTLARGE on 4 Tesla-V100-PCIE 32G GPUs with the max sequence length of 512.

As shown in Table 8, ULMFiT performs better on almost all of the tasks compared to BERTBASE but not BERTLARGE. This changes however with the task-specific further pre-training where even BERTBASE outperforms ULMFiT on all tasks. BERTLARGE fine-tuning with task-specific further pre-training achieves state-of-the-art results.

Conclusion

In this paper, we conduct extensive experiments to investigate the different approaches to fine-tuning BERT for the text classification task. There are some experimental findings: 1) The top layer of BERT is more useful for text classification; 2) With an appropriate layer-wise decreasing learning rate, BERT can overcome the catastrophic forgetting problem; 3) Within-task and in-domain further pre-training can significantly boost its performance; 4) A preceding multi-task fine-tuning is also helpful to the single-task fine-tuning, but its benefit is smaller than further pre-training; 5) BERT can improve the task with small-size data.

With the above findings, we achieve state-of-the-art performances on eight widely studied text classification datasets. In the future, we will probe more insight of BERT on how it works.

References