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

Abstract:
Extreme multi-label text classification (XMC) seeks to find relevant labels from an extreme large label collection for a given text input. Many real-world applications can be formulated as XMC problems, such as recommendation systems, document tagging and semantic search. Recently, transformer based XMC methods, such as XTransformer and LightXML, have shown significant improvement over other XMC methods. Despite leveraging pre-trained transformer models for text representation, the fine-tuning procedure of transformer models on large label space still has lengthy computational time even with powerful GPUs. In this paper, we propose a novel recursive approach, XR-Transformer to accelerate the procedure through recursively fine-tuning transformer models on a series of multi-resolution objectives related to the original XMC objective function. Empirical results show that XRTransformer takes significantly less training time compared to other transformerbased XMC models while yielding better state-of-the-art results. In particular, on the public Amazon-3M dataset with 3 million labels, XR-Transformer is not only 20x faster than X-Transformer but also improves the Precision@1 from 51% to 54%. Our code is publicly available at https://github.com/amzn/pecos.

heading: 1 Introduction
text: 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.

heading: 2 Related Works
text: 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.

heading: 3 Background Material
text: 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 cannot 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.

heading: 4 Proposed Method: XR-Transformer
text: As noted above, the shortlisting problem (2) is itself an XMC problem with slightly smaller output size LB where B is the cluster size. In XR-Transformer, we apply the same three stage framework recursively on the shortlisting problem until a reasonably small output size is reached L
BD . We
can therefore follow the curriculum learning scheme and fine-tune the pre-trained transformers progressively on the sub-XMC problems with increasing output space { L
BD , L BD−1 , . . .}. At each fine-tuning task, the candidate label set is shortlisted by the final model at the previous task. The recursive shortlisting ensures that for any input, the number of candidate labels to include in training and inference isO(B) and therefore the total number of considered labels isO(B logB(L)). Also, we leverage the multi-step fine-tuning and use the embedding generated at the previous task to bootstrap the non pre-trained part for the current task. We now describe the model design in detail.
Hierarchical Label Tree (HLT). Recursively generating label clusters D times is equivalent to building a HLT [41] of depth D. We first construct label features Z ∈ RL×d̂. This could be done by applying text vectorizers on label text or from Positive Instance Feature Aggregation (PIFA):
Z` = v` ‖v`‖
; where v` = ∑
i:yi,`=1
Φ(xi),∀ ` ∈ [L], (4)
where Φ : D 7→ Rd it the text vectorizer. Then we follow similar procedures as [8] and [13] and use balanced k-means (k = B) to recursively partition label sets and generate the HLT in a top-down fashion. The HLT is represented with a series of indexing matrices {C(t)}Dt=1, such that C(t) ∈ {0, 1}Kt×Kt−1 where K0 = 1 and KD = L. Equivalently, once C(D) is constructed, the HLT can be built from bottom up through joining B adjacent clusters together.
Multi-resolution Output Space. Multi-resolution learning has been explored in different contexts such as computer vision [14, 42]. For instance, using an output scheme with coarse-to-fine resolutions results in better image quality for generative adversarial networks [15, 16]. As an another example in meta learning, [43] learns multiclass models via auxiliary meta classes by collapsing existing classes. Nevertheless, multi-resolution learning has not been well-explored in the XMC literature. In XR-Transformer, we leverage the label hierarchy defined by the HLT and train the transformer model on multi-resolution objectives.
The XMC task can be viewed as generating an 1-D image y ∈ {0, 1}L with binary values based on input text x. Just like a coarsified image could be obtained by a max or mean pooling of nearby pixels,
the coarse label vector can be obtained by max-pooling of labels which are nearby in the label feature space. Once the HLT is constructed using label features, the true labels at layer Y(t) ∈ {0, 1}N×Kt can be determined by the true labels of the child clusters at t+1 through a max-pooling like operation:
Y(t) = binarize(Y(t+1)C(t+1)), (5)
and Y (D)i,` = yi,` is the original label matrix. This forms a series of learning signals with coarse-to-fine resolution and can be used to generate learning tasks with easy-to-hard objectives.
Direct use of the binarized Y(t) ∈ {0, 1}N×Kt in Eq (5) results in information loss when merging several positive labels into one cluster. Ideally, a cluster containing several positive children is more relevant than a cluster with only one positive child. To add this lower level information to higher level learning objectives, we introduce the relevance matrix R(t) ∈ RN×Kt+ for layer t of the XMC sub-problem where Ri,` defines the non-negative important weight for ith instance to `th cluster. Different from cost-sensitive learning [44] for MLC (CSMLC) setting [45–47] where there is only one cost matrix explicitly derived by evaluation metrics such as F1 score, in XR-Transformer, we consider the usage of cost-sensitive learning where the relevance matrices are recursively induced by the HLT structure. Specifically, given an HLT, we recursively construct relevance matrices for t = 1, . . . , D:
R(t) = R(t+1)C(t+1), (6)
and R(D) = Y(D). Motivated by [48], we adopts the row-wise `1 normalized relevance matrix:
R̂ (t) i,j =
{ R
(t) i,j
‖R(t)i ‖1 if Y (t)i,j = 1,
α otherwise, where α is the hyper parameter to balance positive and negative weights.
Label Shortlisting. During training, XR-Transformer only focuses on discriminating the labels or clusters that have high chance of being positive. A necessary condition for a label at layer t to be positive is that its parent label at level t− 1 is positive. Therefore, an intuitive approach would be to only train on the output space shortlisted by positive clusters of the parent layer. However, in practice we found this approach sometimes leads to sub-optimal result during inference with beam search. As an effort to balance explore and exploit, we further include the top-k relevant clusters determined by the model learned on the parent layer to mimic the beam search during inference. Thus at layer t, the labels considered during training are shortlisted by the parent layer t− 1:
P(t−1) = Top(W(t−1)>Φ(X,Θ(t−1)), k), (7)
M(t) = binarize(P(t−1)C(t)>) + binarize(Y(t−1)C(t)>), (8) where the Top(·, k) operator zeros out elements in a matrix except the top-k largest values in each row. For each instance xi, only non-zero indices of Mi will be included into the training objective. We can therefore define a series of learning objectives for level t ∈ {1, 2, . . . , D} as:
min W(t),Θ N∑ i=1 ∑ `:M
(t) i,` 6=0
R̂ (t) i,`L(Y (t) i,` ,W (t)> ` Φ(xi,Θ)) + λ‖W (t)‖2, (9)
where L is a point-wise loss such as hinge loss, squared hinge loss or BCE loss, W(t), Θ are the model weights to be learned.
Text Representation. Most previous works on XMC construct text feature representation in one of two ways: statistical feature representations and semantic feature representations. Although the latter, in particular transformer models, have shown promising results on various NLP benchmarks, the self-attention mechanism makes transformers unscalable w.r.t. sequence length. To ensure efficiency, input texts are usually truncated [12, 13] which result in loss of information. On the other hand, the statistical features, such as TFIDF, are fast to construct with the whole input taken into consideration. In XR-Transformer, we use a combination of these two feature representations and each component is a complement of lost information for the other one:
Φcat(x,Θ) := [ Φtfidf (x) ‖Φtfidf (x)‖ , Φdnn(x,Θ) ‖Φdnn(x,Θ)‖ ] , (10)
Algorithm 1: Iterative_Learn(X,Y,C,θ,P) Input :X, Y, C, θ, P m,n← C.shape if n = 1 then
M← ones(N,m) if θ is not fixed then
θ∗ ← Optimize (9) with Φ = Φdnn, initialize Θ = θ else
θ∗ ← θ W∗ ← argminW ∑N i=1 ∑m `=1 R̂ (t) i,`L(Yi,`,W>` Φcat(xi,Θ)) + λ‖W‖2
else Yprev ← binarize(YC) M← binarize(PC>) + binarize(YprevC>) if θ is not fixed then
θ∗ ← Optimize (9) with Φ = Φdnn initialize Θ = θ else
θ∗ ← θ W∗ ← argminW ∑N i=1 ∑ `:Mi,` 6=0 R̂ (t) i,`L(Yi,`,W>` Φcat(xi, θ∗)) + λ‖W‖2
Return :W∗, θ∗
Algorithm 2: XR-Transformer training Input :X, Y, pre-trained transformer Φdnn(·,θ0) Ẑ` = v`/‖v`‖; where v` = ∑ i:yi,`=1
Φtfidf (xi),∀ ` ∈ [L] {Ĉ(t)}D̂t=1 ← k-means-clustering(Ẑ) Generate label hierarchy {Ŷ(t)}D̂t=1 using 5 θ∗ = θ0, P = None for t in 1, 2, 3, · · · , D̂ do
Ŵ,θ∗ ← Iterative_Learn(X, Ŷ(t), Ĉ(t),θ∗,P) P← Top(Ŵ>Φcat(X,θ∗), k)
Z` = v`/‖v`‖; where v` = ∑
i:yi,`=1 Φcat(xi),∀ ` ∈ [L]
{C(t)}Dt=1 ← k-means-clustering(Z) Generate label hierarchy {Y(t)}D̂t=1 using 5 Fix θ∗, P = None for t in 1, 2, 3, · · · , D do
W(t), _← Iterative_Learn(X,Y(t),C(t),θ∗,P) P← Top(W(t)>Φcat(X,θ∗), k)
Return :Φcat(·,θ∗), {C(t)}Dt=1, {W(t)}Dt=1
where Φdnn(·,Θ) is the transformer parametrized by Θ. Once the text representation is constructed, predictions can be made by simply applying a linear projection on top of the text representation through (1).
Training with bootstrapping. The training of XR-Transformer consists of three steps. At first, a preliminary HLT is constructed using raw statistical features. Then a pre-trained transformer model is fine-tuned recursively from low resolution output to high resolution. At each layer t, fine-tuning objective (9) is optimized with initialization Θ = θ(t−1)∗ the best transformer weights of layer t− 1. θ(0)∗ denotes the pre-trained transformer weights.
Unlike the transformer warmed-up with pre-trained weights, the projection weights W(t) is trained from scratch without good initialization. At the beginning of fine-tuning, gradient flow through these cold-start (usually randomly initialized) weights will usually worsen the pre-trained components.
We leverage the recursive learning structure to tackle this issue by model bootstrapping. Concretely, W(t) is initialized as:
W (t) init := argmin
W(t) N∑ i=1 ∑ `:M
(t) i,` 6=0
R̂ (t) i,`L(Y (t) i,` ,W (t)> ` Φdnn(xi,θ (t−1)∗)) + λ‖W(t)‖2, (11)
In practice, (11) is fast to compute since the semantic text feature for the previous layer Φcat(X,θ
(t−1)∗) is already computed and thus (11) can be solved very quickly on CPUs with a variety of parallel linear solvers, such as LIBLINEAR [49].
Once the fine-tuning is complete, the refined HLT is constructed with the text representation that combines statistical text feature and fine-tuned semantic text embeddings. Then the ranking models are trained on top of the combined text features for the final prediction. The detailed training procedure is described in Algorithm 1 and 2.
Inference. The inference cost of XR-Transformer consists mainly of two parts: cost to compute transformer embedding and to retrieve relevant labels through beam search. Therefore, the inference time complexity is O(Tdnn + kd log(L)), where k is the beam size, d is the concatenated feature dimension and Tdnn is the time to compute Φdnn(x) for a given input. Note that even the inference is done with beam search through the refined HLT, the transformer text embedding only need to be computed once per instance.
Connections with other tree based methods. Although methods such as AttentionXML [8] also train on supervisions induced by label trees, the final model is a chain of sub-models which each on is learned on single-resolution. In particular, given a hierarchical label tree with depthD, AttentionXML will train D different text encoders on each layer of the tree where as XR-Transformer trains the same transformer encoder progressively on all layers of the tree. This difference leads to a longer inference time for AttentionXML than XR-Transformer since multiple text encoders need to be queried during inference, as shown in the comparison in the inference time in Appendix A.4.2.

heading: 5 Experimental Results
text: We evaluate XR-Transformer on 6 public XMC benchmarking datasets: Eurlex-4K, Wiki10-31K, AmazonCat-13K, Wiki-500K, Amazon-670K, Amazon-3M. Data statistics are given in Table 1. For fair comparison, we use the same raw text input, sparse feature representations and same train-test split as AttentionXML [8] and other latest works [12, 13]. The evaluation metric is Precision@k (P@k), which is widely-used in XMC literature [3, 8, 12, 13, 18, 28]. The results of Propensity-score Precision@k (PSP@k) are defer to Appendix A.4.3, which focus more on tail labels’ performance.
Baseline Methods. We compare XR-Transformer with state-of-the-art (SOTA) XMC methods: AnnexML [28], DiSMEC [18], PfastreXML [41], Parabel [3], eXtremeText [24], Bonsai [50], XMLCNN [33], XR-Linear [7], AttentionXML [8], X-Transformer [12] and LightXML [13]. We obtain
most baseline results from [8, Table 3] and [7, Table 3] except for the latest deep learning based algorithms [8, 12, 13]. To have fair comparison on training time, we use the same hardware (i.e., AWS p3.16xlarge) and the same inputs (i.e., raw text, vectorized features, data split) to obtain the results of AttentionXML, X-Transformer and LightXML. The hyper-parameter of XR-Transformer and more empirical results are included in Appendix A.3.
Model Performance. The comparisons of Precision@k (P@k) and training time are shown in Table 2 and Table 8, respectively. The proposed XR-Transformer follows AttentionXML and LightXML to use an ensemble of 3 models, while X-Transformer uses an ensemble of 9 models [12]. More details about the ensemble setting can be found in Appendix A.3. The proposed XR-Transformer
framework achieves new SOTA results in 14 out of 18 evaluation columns (combination of datasets and P@k), and outperforms competitive methods on the large datasets. Next, we show the training time of XR-Transformer is significantly less than other DNN-based models.
Training Cost. Table 8 shows the training time for these DNN-based models. To have fair comparison, all the experiments are conducted with float32 precision on AWS p3.16xlarge instance with 8 Nvidia V100 GPUs except for LightXML, which was run on single V100 GPU since multi-GPU training is not implemented. XR-Transformer consumes significantly less training time compared with other transformer based models and the shallow BiLSTM model AttentionXML. On Amazon3M, XR-Transformer has 20x speedup over X-Transformer while achieving even better P@k. Finally, in table 4, we compare XR-Transformer with LightXML under the single model setup (no ensemble), where XR-Transformer still consistently outperforms LightXML in P@k and training time.
Comparison of Different Semantic Embeddings. To provide more empirical justifications, that the improvement in performance comes from better semantic embedding rather than the introducing of TF-IDF features, we further tested models using Pre-Trained Transformer and word2vec embeddings concatenated with the same TF-IDF features.
Table 5 summarizes the performance of these models on all 6 datasets. In particular, word2vec is using token embedding from word2vec-google-news-300 and for Pre-Trained we use the same setting as XR-Transformer (3-model ensemble). On large datasets such as Wiki-500K/Amazon670K/Amazon-3M, Pre-Trained +TF-IDF has marginal improvement compared to the baseline TF-IDF features. Nevertheless, our proposed XR-Transformer still enjoy significant gain compared to Pre-Trained +TF-IDF. This suggests the major improvement is from learning more powerful neural semantic embeddings, rather than the use of TF-IDF.
Effect of Cost-Sensitive Learning. In Table 6, we analyze the effect of cost sensitive learning on four XMC datasets with the largest output spaces. On most datasets, cost sensitive learning via aggregated labels yields better performance than those without. We also show that cost-sensitive learning is not only beneficial to XR-Transformer, but also useful to its linear counterpart XRLinear [7]. See Appendix A.4.1 for more results.
Effect of Label Resolution and Text Representation. Next, we compare the effect of label resolution on the quality of the finetuned transformer embeddings. We finetune transformer models in a non-recursive manner on a two layer HLT with different leaf cluster size. Then the fine-tuned transformer embeddings are used along or in combination with TF-IDF features to produce the predictions with refined HLT. From Figure 2 we can observe that a larger cluster size will result in worse semantic features. Figure 2 also shows that combining semantic features Φdnn with statistical features Φtfidf could in general improve the model performance.

heading: 6 Conclusion and Future work
text: In this paper, we have presented XR-Transformer approach, which is an XMC architecture that leverages multi-resolution objectives and cost sensitive learning to accelerate the fine-tuning of pre-trained transformer models. Experiments show that the proposed method establishes new stateof-the-art results on public XMC datasets while taking significantly less training time compared with earlier transformer based methods. Although the proposed architecture is designed for XMC, the ideas can be applied to other areas such as information retrieval or other DNN models such as CNNs/ResNets. Also, more extensive study is required to understand why the coarse-to-fine scheme would lead to not only faster training but better overall quality. A hypothesis is that the problem is being solved at multiple scales hence leading to more robust learning of deep transformer models.

Review summaries:
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.

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.

The authors propose an approach for extreme multi-label classification  that is based on a transformer network and leverages a reduction approach that breaks down the problem to smaller ones via a hierarchical approach that constructs a label tree. The authors use both the representation of the transformer as well as classical TF-IDF ones as the final one for the model. The approach is validated in 6 large datasets and compared with other state-of-the-art approaches.