OTS eq2 S et: An Optimal Transport Enhanced Sequence-to-Set Model for Extreme Multi-label Text Classification (2022)

Abstract: Extreme multi-label text classification (XMTC) is the task of finding the most relevant subset labels from an extremely large-scale label collection. Recently, some deep learning models have achieved state-of-the-art results in XMTC tasks. These models commonly predict scores for all labels by a fully connected layer as the last layer of the model. However, such models can’t predict a relatively complete and variable-length label subset for each document, because they select positive labels relevant to the document by a fixed threshold or take top k labels in descending order of scores. A less popular type of deep learning models called sequence-to-sequence (Seq2Seq) focus on predicting variable-length positive labels in sequence style. However, the labels in XMTC tasks are essentially an unordered set rather than an ordered sequence, the default order of labels restrains Seq2Seq models in training. To address this limitation in Seq2Seq, we propose an autoregressive sequence-to-set model for XMTC tasks named OTSeq2Set. Our model generates predictions in student-forcing scheme and is trained by a loss function based on bipartite matching which enables permutation-invariance. Meanwhile, we use the optimal transport distance as a measurement to force the model to focus on the closest labels in semantic label space. Experiments show that OTSeq2Set outperforms other competitive baselines on 4 benchmark datasets. Especially, on the Wikipedia dataset with 31k labels, it outperforms the state-of-the-art Seq2Seq method by 16.34% in micro-F1 score. The code is available at https://github.com/caojie54/OTSeq2Set.

## introduction extreme multi-label text classification (xmtc) is a natural language processing (nlp) task of finding the most relevant subset labels from an extremely large-scale label set. it has a lot of usage scenarios, such as item categorization in e-commerce and as in many other nlp tasks, deep learning based models also achieve the state-of-the-art performance in xtmc. for example, attentionxml @xcite , x-transformer @xcite and lightxml @xcite have achieved remarkable improvements in evaluating metrics relative to the current state-of-the-art methods. these models are both composed of three parts @xcite ): text representing, label recalling, and label ranking. the first part converts the raw text to text representation vectors, then the label recalling part gives scores for all cluster or tree nodes including portion labels, and finally, the label ranking part predicts scores for all labels in descending order. notice that, the label recalling and label ranking part both use fully connected layers. although the fully connected layer based models have excellent performance, there exists a drawback which is these models can't generate a variable-length and relatively complete label set for each document. because the fully connected layer based models select positive labels relevant to the document by a fixed threshold or take top k labels in descending order of label scores, which depends on human's decision. another type of deep learning based models is seq2seq learning based methods which focus on predicting variable-length positive labels only, such as mlc2seq @xcite , sgm @xcite . mlc2seq and sgm enhance seq2seq model for multi-label classification (mlc) tasks by changing label permutations according to the frequency of labels. however, a pre-defined label order can't solve the problem of seq2seq based models which is the labels in xmtc tasks are essentially an unordered set rather than an ordered sequence. @xcite solves this problem on mlc tasks via reinforcement learning by designing a reward function to reduce the dependence of the model on the label order, but it needs to pretrain the model via maximum likelihood estimate (mle) method. the two-stage training is not efficient for xmtc tasks that have large-scale labels. to address the above problems, we propose an autoregressive sequence-to-set model, otseq2set, which generates a subset of labels for each document and ignores the order of ground truth in training. otseq2set is based on the seq2seq @xcite , which consists of an encoder and a decoder with the attention mechanism. the bipartite matching method has been successfully applied in named entity recognition task @xcite and keyphrase generation task @xcite to allievate the impact of order in targets. chen et al. (2019) and li et al. (2020) have successfully applied the optimal transport algorithm to enable sequence-level training for seq2seq learning. both methods can achieve optimal matching between two sequences, but the difference is the former matches two sequences one to one, and the latter gives a matrix containing regularized scores of all connections. we combine the two methods in our model. our contributions of this paper are summarized as follows: (1) we propose two schemes to use the bipartite matching in xmtc tasks, which are suitable for datasets with different label distributions. (2) we combine the bipartite matching and the optimal transport distance to compute the overall training loss, with the student-forcing scheme when generating predictions in the training stage. our model can avoid the exposure bias; besides, the optimal transport distance as a measurement forces the model to focus on the closest labels in semantic label space. (3) we add a lightweight convolution module into the seq2seq models, which achieves a stable improvement and requires only a few parameters. (4) experimental results show that our model achieves significant improvements on four benchmark datasets. for example, on the wikipedia dataset with 31k labels, it outperforms the state-of-the-art method by 16.34% in micro-f1 score, and on amazon-670k, it outperforms the state-of-the-art model by 14.86% in micro-f1 score. ## related work another type of deep learning-based method is seq2seq learning based methods. mlc2seq @xcite is based on the classic seq2seq @xcite architecture which uses a bidirectional rnn to encode raw text and an rnn with an attention mechanism to generate predictions sequentially. mlc2seq enhances the model performance by determining label permutations before training. sgm @xcite proposes a novel decoder structure to capture the correlations between labels. @xcite proposes a twostage training method, which is trained with mle then is trained with a self-critical policy gradient training algorithm. about bipartite matching, tan et al. (2021) considers named entity recognition as a seq2set task, the model generates an entity set by a nonautoregressive decoder and is trained by a loss function based on bipartite matching. one2set @xcite proposes a k-step target assignment mechanism via bipartite matching on the task of keyphrases generation, which also uses a nonautoregressive decoder. @xcite proposes a fast proximal point method to compute the optimal transport distance, which is named as ipot. @xcite adds the ot distance as a regularization term to the mle training loss. the ot distance aims to find an optimal matching of similar words/phrases between two sequences which promotes their semantic similarity. @xcite introduce a method that combines the student-forcing scheme with the ot distance in text generation tasks, this method can alleviate the exposure bias in seq2seq learning. the above two methods both use the ipot to compute the ot distance. ## conclusion in this paper, we propose an autoregressive seq2set model for xmtc, otseq2set, which combines the bipartite matching and the optimal transport distance to compute overall training loss and uses the student-forcing scheme in the training state. ot-seq2set not only eliminates the influence of order in labels, but also avoids the exposure bias. besides, we design two schemes for the bipartite matching which are suitable for datasets with different label distributions. the semantic optimal transport distance can enhance the performance by the semantic similarity of labels. to take full advantage of the raw text information, we add a lightweight convolution module which achieves a stable improvement and requires only a few parameters. experiments show that our method gains significant performance improvements against strong baselines on xmtc.