Fast Multi-Resolution Transformer Fine-tuning for Extreme Multi-label Text Classification

1 Introduction . Many real-world applications such as open-domain question answering [ 1 , 2 ] , e-commerce dynamic search advertising [ 3 , 4 ] , and semantic matching [ 5 ] , can be formulated as an extreme multi-label text classification ( XMC ) problem : given a text input , predict relevant labels from an enormous label collection of size L. In these applications , L ranges from tens of thousands to millions , which makes it very challenging to design XMC models that are both accurate and efficient to train . Recent works such as Parabel [ 3 ] , Bonsai [ 6 ] , XR-Linear [ 7 ] and AttentionXML [ 8 ] , exploit the correlations among the labels to generate label partitions or hierarchical label trees ( HLTs ) which can be used to shortlist candidate labels to be considered during training and inference . While these methods are scalable in terms of the size of the label collection , most of them rely only on statistical representations ( such as bag-of-words ) or pooling from pre-generated token embeddings ( such as word2vec ) to vectorize text inputs . In light of the recent success of deep pretrained transformers models such as BERT [ 9 ] , XLNet [ 10 ] and RoBerta [ 11 ] in various NLP applications , X-Transformer [ 12 ] and LightXML [ 13 ] propose to fine-tune pre-trained transformer models on XMC tasks to obtain new state-of-the-art results over the aforementioned approaches . Although transformers are able to better capture semantic meaning of textual inputs than statistical representations , text truncation is often needed in practice to reduce GPU memory footprint and maintain model efficiency . For example , X-Transformer truncates input texts 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to contain the first 128 tokens before feeding it into transformer models . Efficiency of transformer fine-tuning poses another challenge for XMC applications . Directly fine-tuning transformer models on the original XMC task with a very large label collection is infeasible as both the training time and the memory consumption are linear in L. In order to alleviate this , both X-Transformer and LightXML adopt a similar approach to group L labels into K clusters of roughly equal size denoted by B and fine-tune transformers on the task to identify relevant label clusters ( instead of labels themselves ) . If B ≈ √ L and K ≈ √ L , then both the training time and the memory requirement of the fine-tuning can be reduced to O ( √ L ) from O ( L ) . However , as pointed out in [ 8 ] , the model performance would deteriorate due to the information loss from label aggregation . Thus , both XTransformer and LightXML still choose a small constant B ( ≤ 100 ) as the size of the label clusters . As a result , transformers are still fine-tuned on a task with around L/100 clusters , which leads to a much longer training time compared with non-transformer based models . For example , it takes X-Transformer 23 and 25 days respectively to train on Amazon-3M and Wiki-500K even with 8 Nvidia V100 GPUs . To address these issues , we propose XR-Transformer , an XMC architecture that leverages pre-trained transformer models and has much smaller training cost compared to other transformer-based XMC methods . Motivated by the multi-resolution learning in image generation [ 14–16 ] and curriculum learning [ 17 ] , we formulate the XMC problem as a series of sub-problems with multi-resolution label signals and recursively fine-tune the pre-trained transformer on the coarse-to-fine objectives . In this paper , our contributions are as follows : • We propose XR-Transformer , a transformer based framework for extreme multi-label text classification where the pre-trained transformer is recursively fine-tuned on a series of easyto-hard training objectives defined by a hierarchical label tree . This allows the transformers to be quickly fine-tuned for a XMC problem with a very large number label collection progressively . • To get better text representation and mitigate the information loss in text truncation for transformers , we take into account statistical text features in addition to the transformer text embeddings in our model . Also , a cost sensitive learning scheme by label aggregation is proposed to introduce richer information on the coarsified labels . • We conduct experiments on 6 public XMC benchmarking datasets and our model takes significantly lower training time compared to other transformer-based XMC models to yield better state-of-the-art results . For example , we improve the state-of-the-art Prec @ 1 result on Amazon-3M established by X-Transformer from 51.20 % to 54.04 % while reducing the required training time from 23 days to 29 hours using the same hardware . 2 Related Works . Sparse Linear Models with Partitioning Techniques . Conventional XMC methods consider fixed input representations such as sparse TF-IDF features and study different partitioning techniques or surrogate loss functions on the large output spaces to reduce complexity . For example , sparse linear one-versus-all ( OVA ) methods such as DiSMEC [ 18 ] , PPD-Sparse [ 19 , 20 ] , ProXML [ 21 ] explore parallelism to solve OVA losses and reduce the model size by weight truncations . The inference time complexity of OVA models is linear in the output space , which can be greatly improved by partitioning methods or approximate nearest neighbor ( ANN ) indexing on the label spaces . Initial works on tree-based methods [ 22 , 23 ] reduce the OVA problem to one-versus-some ( OVS ) with logarithmic depth trees . Down that path , recent works on sparse linear models including Parabel [ 3 ] , eXtremeText [ 24 ] , Bonsai [ 6 ] , XReg [ 25 ] , NAPKINXC [ 26 , 27 ] and XR-Linear [ 7 ] partition labels with B-array hierarchical label trees ( HLT ) , leading to inference time complexity that is logarithmic in the output space . On the other hand , low-dimensional embedding-based models often leverage ANN methods to speed up the inference procedure . For example , AnnexML [ 28 ] and SLICE [ 29 ] consider graph-based methods such as HNSW [ 30 ] while GLaS [ 31 ] considers product quantization variants such as ScaNN [ 32 ] . Shallow Embedding-based Methods . Neural-based XMC models employ various network architectures to learn semantic embeddings of the input text . XML-CNN [ 33 ] applies one-dimensional CNN on the input sequence and use the BCE loss without sampling , which is not scalable to XMC problems . AttentionXML [ 8 ] employs BiLSTMs and label-aware attention as scoring functions . For better scalability to large output spaces , only a small number of positive and hard negative labels are used in model GPU training . Shallow embedding-based methods [ 34–38 ] use word embedding lookup followed by shallow MLP layers to obtain input embeddings . For instance , MACH [ 34 ] learns MLP layers on several smaller XMC sub-problems induced by hashing tricks on the large label space . Similarly , DeepXML [ 35 ] and its variant ( i.e. , DECAF [ 36 ] , GalaXC [ 37 ] , ECLARE [ 38 ] ) pre-train MLP encoders on XMC sub-problems induced by label clusters . They freeze the pre-trained word embedding and learn another MLP layer followed by a linear ranker with sampled hard negative labels from HNSW [ 30 ] . Importantly , shallow embedding-based methods only show competitive performance on short-text XMC problems where the number of input tokens is small [ 34 , 35 ] . Deep Transformer Models . Recently , pre-trained Transformer models [ 9–11 ] have been applied to XMC problems with promising results [ 12 , 13 , 39 ] . X-Transformer [ 12 ] considers a two-stage approach where the first stage transformer-based encoders are learned on XMC sub-problems induced by balanced label clusters , and the second stage sparse TF-IDF is combined with the learned neural embeddings as the input to linear OVA models . APLC-XLNet [ 39 ] fine-tunes XLNet encoder on adaptive imbalanced label clusters based on label frequency similar to Adaptive Softmax [ 40 ] . LightXML [ 13 ] fine-tunes Transformer encoders with the OVA loss function end-to-end via dynamic negative sampling from the matching network trained on label cluster signals . Nonetheless , Transformer-based XMC models have larger model size and require longer training time , which hinders its practical usage on different downstream XMC problems . 3 Background Material . We assume we are given a training set { xi , yi } Ni=1 where xi ∈ D is the ith input document and yi ∈ { 0 , 1 } L is the one hot label vector with yi , ` = 1 indicating that label ` is relevant to instance i . The goal of eXtreme Multi-label Text Classification ( XMC ) is to learn a function f : D × [ L ] 7→ R , such that f ( x , ` ) denotes the relevance between the input x and the label ` . In practice , labels with the largest k values are retrieved as the predicted relevant labels for a given input x . The most straightforward model is one-versus-all ( OVA ) model : f ( x , ` ) = w > ` Φ ( x ) ; ` ∈ [ L ] , ( 1 ) where W = [ w1 , . . . , wL ] ∈ Rd×L are the weight vectors and Φ : D 7→ Rd is the text vectorizer that maps x to d-dimensional feature vector . Φ ( · ) could be a deterministic text vectorizer , such as the bag-of-words ( BOW ) model or Term Frequency-Inverse Document Frequency ( TFIDF ) model , or a vectorizer with learnable parameters . With the recent development in deep learning , using pre-trained transformer as the text vectorizer has shown promising results in many XMC applications [ 12 , 13 , 39 ] . WhenL is large , however , training and inference of OVA model without sampling would be prohibitive due to the O ( L ) time complexity . To handle the extremely large output space , recent approaches partition the label space to shortlist the labels considered during training and inference . In particular , [ 7 , 12 , 13 , 34 , 35 , 39 ] follow a three stage framework : partitioning , shortlisting , and ranking . First , label features are constructed to group labels into K clusters C ∈ { 0 , 1 } L×K where C ` , k = 1 denotes that label ` is included in the k-th cluster . Then a shortlisting model is learned to match input x to relevant clusters in an OVA setting : g ( x , k ) = ŵ > k Φg ( x ) ; k ∈ [ K ] . ( 2 ) Finally , a classification model with output size L is trained on the shortlisted labels : f ( x , ` ) = w > ` Φ ( x ) ; ` ∈ Sg ( x ) , ( 3 ) where Sg ( x ) ⊂ [ L ] is the label set shortlisted by g ( x , · ) . In the extreme case where only one label cluster is determined to be relevant to a input x , the training and inference cost on x would be O ( K + LK ) , which in the best case scenario is O ( √ L ) when K = √ L. For transformer based methods , the dominant time is the evaluation of Φ ( x ) . But K being too big or too small could still be problematic . Empirical results show that the model performance deteriorates when clusters are too big [ 8 ] . This is because that the signals coming from B labels within the same cluster will be aggregated and not distinguishable , whereB is the cluster size . Therefore , B can not be too big to ensure a reasonable label resolution for fine-tuning . Also , as pointed out in [ 12 ] , fine-tuning transformer models on large output spaces can be prohibitive . As a result , the label clusters need to be constructed in a way to balance the model performance and fine-tuning efficiency . In practice , both the transformer-based XMC models , such as X-Transformer and LightXML , adopt a small fixed constant as the cluster size B ( ≤ 100 ) , which means that training the shortlisting model g ( x , k ) is still very time consuming as the number of clusters K ≈ L/B .

Summary: This paper proposes a coarse-to-fine learning framework to solve the problem of extreme multi-label text classification.  Experiments show the model not only accelerates existing models 20x times but also achieves a slightly better performance on the dataset with an extremely large label set. The techniques are sound and the empirical results are promising.