Label-Wise Document Pre-Training for Multi-Label Text Classification

We design a multi-label pre-training task and model Label-Wise Pre-Training (LW-PT). When given a target document $T$ and it’s label set $\{1,...,k,...,l\}$, we sample a positive example $C^{+}$ which has the same label $k$ as target and several negative examples $C^{-}$ which doesn’t have the label $k$. This two parts are combined as candidates $\{C_{i}\}_{i=1}^{n}$ with size $n$. The training goal is to discriminate the positive example from all $\{C_{i}\}_{i=1}^{n}$. We do this for each label $k$ in the label set of targer document to make training samples.

Figure 1-a shows the architecture of pre-training task and model LW-PT. We denote the target encoder as T-Encoder and candidates encoder as C-Encoder. They have different parameters.

Therefore, as for the label $k$, we calculate the representation for the target document $T$ and candidates $\{C_{i}\}_{i=1}^{n}$, and obtain $Q^{t}_{k}\in\mathbb{R}^{2d}$ and $Q^{c_{i}}_{k}\in\mathbb{R}^{2d}$.

Further, the similarity of any two documents is directly calculated by dot product and trained using the negative log-likelihood loss function.

Obviously, we can choose different encoder for it. In MLTC task, different labels may attend on differnt part of the document. Therefore, we design a label-wise encoder. Normally, Label-Wise LSTM (LW-LSTM) can meet the requirements. However, LSTM is not effective for long documents. Inspired by Hierarchical Attention Network (HAN) [19], we also introduce Hierarchical Label-Wise LSTM(HLW-LSTM) to model long documents with several sentences. The details are described in the following.

After pre-training, we use T-Encoder and C-Encoder in the downstream model for MLTC task. As illustrated in Figure 1-b, we concatenate the output of T-Encoder and C-Encoder to get document representation for each label $k$. Then, each label-wise representation is also concatenated. And we obtain $M\in\mathbb{R}^{l\times 4d}$ as the final representation for a given document $D$.

After that, we simplely use a MLP layer to predict the probability of each label $\hat{y}$.

where $W\in\mathbb{R}^{(l\times 4d)\times l}$ is a trainable parameter.

We use the cross-entropy loss function for MLTC task.

where $\hat{y}_{k}$ is the predict probability of label $k$ for current document. And $y_{k}$ is the ground truth of whether label $k$ appeared in current document.

It should be noted that T-Encoder and C-Encoder are fine-tuned in the MLTC task at the same time.

We use two MLTC datasets RMSC[20] and AAPD[18] with different length and language of documents. The former has longer documents with several sentences in Chinese. And the latter is shorter in English. Table 1 shows some statistics.

• •

  RMSC-V2[20]: The dataset is collected from a popular Chinese music review website²²2https://music.douban.com. For each music, it includes a set of human annotated styles, and associated reviews. 22 styles are defined in it. The dataset contains over 7.1k samples, 288K reviews, and 3.6M words. However, the initial version didn’t contains a certain split. Thus, we split the dataset into trian/valid/test by the same ratio as the original paper.

• •

  AAPD[18]: The dataset is collected from a English academic website³³3https://arxiv.org in the computer science field. Each sample contains an abstract and the corresponding subjects. 54 labels are defined and 55,840 samples are included.

We use four widely used evaluation metrics, the same as [5, 18, 20].

• •

  One-Error: One-error calculates the fraction of samples whose top-1 predicted label is not in the ground truth.

• •

  Hamming Loss: Hamming loss calculates the fraction of the wrong labels to the total number of ground truth labels.

• •

  Macro-F1: Macro F1 takes the average of each label’s F1 score.

• •

  Micro-F1: Micro F1 calculates the F1 score over all sample-label pairs.

For RMSC-V2 dataset, we use jieba ⁴⁴4https://github.com/fxsjy/jieba toolkit tokenizer and train word embedding by Skip-gram model[12]. The embedding dimision and hidden size is 100. For AAPD dataset, we also train word embedding by Skip-gram model[12]. The embedding dimision and hidden size is 256.

We use 2 layer LSTM for LW-LSTM encoder. And a hierarchical LSTM layer for HLW-LSTM encoder. In addition, Adam[11] optimizer with learning rate 0.001 is used. We add layer normalization[2] and dropout with probability 0.2. In pre-training procedure, 3k iterations are runned and the number of candidates documents is 3. Batch size is 128. In fine-tuning, 20 epochs are runned.

It should be noted that the pre-training procedure is on dataset of train.

We use several MLTC baselines to compare with our method.

• •

  Binary Relevance (BR)[3] transforms the multi-label problem into several single-label classification problems and ignore the correlation between labels.

• •

  Classifier Chains (CC)[16] transforms the multi-label problem into a chain of several single-label classification problems.

• •

  Label Powerset (LP)[17] transforms the multi-label problem to a multi-class problem with one multi-class classifier trained on all unique label.

• •

  CNN[10] attempts to use a convolutional neural network with several different size of kernels, and predicts each label with logistic regression.

• •

  LSTM[9] uses a 2-layer LSTM neural network with self-attention on the last layer to get document representation, and predicts each label with logistic regression.

• •

  Hierarchical Attention Network (HAN)[19] uses a hierarchical attention network on word and sentence level to get document representation, and predicts each label with logistic regression.

• •

  HAN+LG[20] introduces a label graph matrix into HAN.

• •

  CNN-RNN[5] utilizes CNN and RNN on Seq2Seq framework to get labels one by one.

• •

  SGM+GE[18] computes a weighted global embedding based on all labels as opposed to just the top one at each timestep.

• •

  set-RNN[15] presents an adaptation of RNN sequence models to the problem, where the target is a set of labels, not a sequence.

• •

  reg-LSTM[1] proposes a simple BiLSTM architecture with regularization.

• •

  MAGNET[13] uses a graph attention network to capture the attentive dependency structure among the labels.

The complete results on RMSC-V2 and AAPD datasets are shown in Table 2. We calculate the One-Error, Hamming-Loss, Macro F1 score and Micro F1 score for several models. Compared to the previous approaches, our LW-LSTM encoder or HLW-LSTM encoder with pre-training mechanism presents a outstanding performance. Adding the fine-tuning procedure further improve the performance on a substantial margin.

It should be noted that LW-LSTM model denotes directly use the LW-LSTM encoder in the downstream task model without pre-training mechanism. And the same as HLW-LSTM model. We can see that pre-training and fine-tuning have a significant contribution to the performance. But for which label-wise encoder to use, it’s determined by the dataset. For long documents with several sentences (e.g. RMSC), we suppose to choose HLW-LSTM encoder with hierarchical attention. And for short documents with less sentences (e.g. AAPD), we suppose to use LW-LSTM encoder.

We can see that the Macro F1 score has an outstanding improvement when pre-training mechanism is used. This is because the pre-training approach improves the learning of labels with less frequency by capturing the correlation between labels. And Macro F1 score is the average of each label’s F1 score. These experimental results is consistent with our ideas.

In this section, some ablation studies to demonstrate the method proposed above are made on RMSC-V2 dataset.

To examine the ability of discovering the label correlations, we cite three correlated labels “alternative”, “metal” and “rock” in music tagging task of RMSC-V2 dataset for example. In Table 5, for target song (i.e. document) “Janes Addiction - Nothings Shocking” possesses labels of “alernative” and “rock”. We compute its top 50 similar songs via cosine similarity of document representation learnt respectively from three labels “alternative”, “metal” and “rock”. For each set of 50 similar songs, we calculate the label frequency (i.e., the proportion of songs that contain a specific label) and demonstrate the statistics in Table 5. The proportion smaller than 10% are not presented.

From Table 5 we can find that, for the target song, among its top 50 similar songs derived through “alternative”-wise document representation, 76% have label “rock”, 42% “alternative” and 26% “metal”. The similar observations can also be found for “metal”-wise and “rock”-wise document representation. This reveals that the proposed label-wise document representation is capable of capturing label correlations. Meanwhile, we can see all three columns has similar results, which means the label-wise representation on the original three labels are also similar. It is also interesting to notice that, even the target song doesn’t hold the label “metal”, it can also be accurately represented using the “metal”-wise representation due to label “metal” is correlated with “alternative” and “rock”.

As mentioned above, our pre-training mechanism captures the label correlation. Thus, the labels have lower frequency is also predicted accurately. To demonstrate it, we count the frequency of each label and the corresponding F1 score. As shown in the Figure 2, frequency are discretized and average F1 score of labels on RMSC-V2 are calculated with HLW-LSTM encoder. And also LW-LSTM encoder for AAPD dataset. Obviously, the model have pre-training and fine-tuning has the best performance at most times and are better for lower frequency labels.

We present some examples for HLW-LSTM models with or without pre-training and fine-tuning mechanism. As shown in Table 6, five typical examples are selected from RMSC-V2 dataset in music tagging task. In the first three examples, models with pre-training accurately predict the whole label set. In the fourth, further using the fine-tuning mechanism obtains a accurate prediction. In the last, our pre-training and fine-tuning models both predict a extra label “pop” which is not appeared in the ground truth. However, we visit the website ⁵⁵5https://music.douban.com/subject/1774742/ and find this song also have some labels similar to “pop” but not appear in the dataset. Note that the labels are made by human and maybe not accurate and complete. But for the label “indie” which is not predicted, this is where the model needs futher improvement.