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: The authors address the problem of fine-tuning transformers models for extreme multi-label text classification (XMLTC). Even with pre-trained transformer models, fine-tuning it to the task of predicting for large label space is computationally expensive and requires significant GPU resources. The authors propose a new routine to train the X-Transformers-like model that significantly reduces that computational cost. The routing resembles AttentionXML learning routine, where one model is trained by tree level to predict meta-labels at that level. Training starts from the top level of the tree, and the model for the next level is initialized with weights of the final model from the previous level. Here underlying transformer model is first tuned on the task of predicting the same meta-labels. It starts with tuning using meta-labels from the top level and then continues tuning layer by layer. In addition, the authors apply a few additional techniques like cost-sensitive training (positive meta-labels are weighted by a number of relevant labels they consist) and bootstrapping of output weights for meta-labels/labels based on current transformers embeddings. After tuning, new label-partitioning is constructed, and the model is retrained with fixed transformer weights. The authors evaluate the attractiveness of their method on six popular benchmark datasets. New routing leads to a significant reduction in training time at the same often achieving better predictive performance.