CascadeXML: Rethinking Transformers for End-to-end Multi-resolution Training in Extreme Multi-label Classification
Authors: Siddhant Kharbanda
Atmadeep Banerjee
Erik Schultheis
Rohit Babbar
Venue: NeurIPS 2022 Accept (2022)
Keywords: Large Output Spaces
Extreme Classification
Multi-label text classification
label-tree based negative mining

Abstract: Extreme Multi-label Text Classification (XMC) involves learning a classifier that can assign an input with a subset of most relevant labels from millions of label choices. Recent approaches, such as XR-Transformer and LightXML, leverage a transformer instance to achieve state-of-the-art performance. However, in this process, these approaches need to make various trade-offs between performance and computational requirements. A major shortcoming, as compared to the Bi-LSTM based AttentionXML, is that they fail to keep separate feature representations for each resolution in a label tree. We thus propose CascadeXML, an end-to-end multi-resolution learning pipeline, which can harness the multi-layered architecture of a transformer model for attending to different label resolutions with separate feature representations. CascadeXML significantly outperforms all existing approaches with non-trivial gains obtained on benchmark datasets consisting of up to three million labels. Code for CascadeXML will be made publicly available at https://github.com/xmc-aalto/cascadexml.

CaseadeXML: Rethinking Transformers for End-to-end Multi-resolution Training in Extreme Multi-label Classification
Siddhant Kharbanda Atmadeep Banerjee Erik Schultheis Rohit Babbar Department of Computer Science, Aalto University, Finland firstmameelastnange @ @aalto.fi
Abstract
Extreme Multi-label Text Classification (XMC) involves learning a classifier that can assign input with a subset of most relevant labels from millions of label choices. Recent approaches, such as UR-Transformer and LightXML, leverage transformer instance to achieve state-of-the-art performance. However, in this process, these approaches need to make various trade-offs between performance and computational requirements. A major shorteoming, compared to the Bi-LSTM based AttentionXML,, is that they fail to keep separate feature representations for each resolution in label tree. We thus propose CascadeXMI. end-to-end multi- resolution learning pipeline, which can harness the multi-layered architecture of a transformer model for attending to different label resolutions with separate feature representations. CascadeXMII. significantly outperforms all existing approaches with non-trivial gains obtained on benchmark datasets consisting of up to three million labels. Code for CaL will be made publicly available at hhttp:: //github. com xmc- walto/aeeeeee cascadexml
Introduction
Extreme multi-label text classification (XMC) deals with the problem of predicting the most relevant subset of labels from an enormously large label space spanning up to millions of labels. Over the years, extreme classification has found many applications in e-commerce, like product recommendation [7] and dynamic search advertisement [16), document tagging I and open-domain question answering (6] (26] and thus, practical solutions to this problem can have significant and far-reaching impact.
The output space in extreme classification is not only large but also, the distribution of instances among labels follows Zipf's law such that large fraction of them are tail labels 17]. For example, in dataset of Wikipedia corpus containing 500K labels, only 2. % labels annotate more than 100 data points and 60% of the labels annotate less than 10 data points! This inherently makes learning meaningful representations at an extreme scale a challenging problem. Further, as the memory and computational requirements grow linearly with the size of the label space, naive treatment via classical methods and off- the-shel solvers fails to deal with such large output spaces g 54].
In extreme scenarios with millions of labels, surrogate step of shortlisting candidate labels becomes crucial to perform extreme classification and been adapted by many popular approaches 110 For each data point, these approaches first solve simpler task with coarse label-clusters or labels as the label space to create shortlist of relevant labels for the extreme classification task. This effectiivly reduces the training time and computational complexity of the extreme task to that of the shortlisting step which is OVVL), where L is the number of labels.
23
52
57
Existing works leverage label representation as centroid of their annotated instances to create Hierarchical Label Tree (HLT) 36) (55.. While earlier free-based methods like Parabel (36] and
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

AttentionXML
LightXML
XR- Trannsfforrerr
CascadeXML
Base Model
Bi-LSTM
Transformer shared
Transformer shared
Transformer
Attention Maps Single Encoder Multi-level HLT End-to-End
adaptive
adaptive
x
Full Resolution
X
Table 1: Strengths and limitations of current XMC methods: LightXML. makes end-to-end training possible, but only admits single level of the HLT. XR-Transformer allows for multiple levels, but its iterative feature learning does not train the transformer encoder at the highest resolution at all [57] Table 6]. However, by changing from multiple Bi-LSTMs to a single transformer, both LightXML and XR- ER-Transformer lost AttentionXML's ability to adapt attention maps to each resolution.
Bonsai 22 used the entire HLT for extreme classification, many recent works B
10 19
23
52
use label clusters only at certain level of the HLT as meta-labels which are in turn used to shortlist candidate labels for the extreme task. In contrast, XR-Transformer [57] and XttentionXML (55) leverage multiple levels of the HLT, such that each level corresponds to a certain label resolution in the tree structure (57
1.1 Motivation: Strengths and Weaknesses of Current Approaches
The earliest work to successfully combine label-tree based shortlisting and attention-based deep encoders AtLentionXML (55), which employs multiple Bi-LSTMs models that are trained sequen- tially for each tree resolution. More recent approaches [8] 19
52 57 replaced the model architecture
with a more powerful transformer model [46] and fine-tune a pre-trained instance such as Bert (IT). Without careful designs, such models computationally very expensive and earlier works [8] 52 could not effectively leverage transformers for both computation and performance on XMC tasks. Next, we discuss contemporary works LightXML [19 and TR-Transformer (57).
LightXML introduces the concept of dynamic negative sampling, which replaces pre-computed label shortlists with dynamically calculated shortlist that changes as the model's weights get updated. This enables end-to-end training with a single model by using the final feature representation of the transformer encoder for both the meta- and the clasffe-caasifiicatin task. Unfortunately, these two tasks appear to interfere with one another (33]. We hypothesize that this is because the meta task needs the attention maps to focus different tokens than the extreme task. Furthermore, it only uses a single-level tree, which prevents scaling to the largest datasets.
XR- AR-Transformer takes inspiration from multi-resolution approaches in computer vision like super resolution 25 and progressive growing of GANs 20 21, and enables multiple resolutions through iterative training. However, unlike progressively grown GANs, which predict only at the highest resolution, KR-Transformer needs predictions across all resolutions for progressive shortlisting pipeline, but uses representations trained at single resolution. In practice, this leads to XR- Transformer having a complex multi-stage pipeline where the transformer model iteratively trained up to a certain resolution and then frozen. This is followed by re-clustering and re-training of multiple classifiers, working at different resolutions, with the same fixed transformer features [57 Alg. 2]. For datasets with over 500K labels, the transformer is only trained up to 2 clusters [57 Table 6], and the resulting features are used for the extreme task.
Unlike AttentionXMLL, using multiple instances of transformer models becomes undesirable due to their computational overhead. This enforces LightXML and TR-Transformer to make different trade-offs when leveraging single transformer model for XMC tasks compared to AttentionXMLL see Table 1 In this paper, we present a method, CavcadeXML, that combines the strengths of these approaches, thus creating end to-end trainable multi-resolution learning pipeline which trains a single transformer model across multiple resolutions in a way that allows the creation of label resolution specific attention maps and feature embeddings.

---

Input Embeddings
forward pass backward pass
Meta-lahels:
e [1,
BERT
( (11,2) (3, 4)
5( (4)
Layers 1-5
(
4 (8) (x)
Layer 6
W1
a11
0900
Layer
((3)
(2))
Layer 8
W2
02
(43
(9(2)
Layer 9
(33)
(8)
Layer
Ws
29
Layer 11
Layer
W4
SCOrES
Figure 1: Overview over CascadeXML. The meta-classifiers W W3 use intermediate BERT
features ( to discard all meta-labels except the highest scoring at that level (marked in red), so
their descendants (gray) need not be considered the next level. More details in 3 and a algorithm 1
1.2 CascadeXML in a Nutshell
The key insights that enable calculating resolution-specific representations with single forward pass through the transformer are a) that the shortlisting tasks get easier for lower label resolution and b) that the intermediate transformer representations are already quite powerful discriminators [49] Consequently, CascadeXMI. extracts features for classification at the coarser meta-label resolution from earlier transformer layers and uses the last layer only for classification at the extreme resolution. The advantage of this is twofold: First, we postulate that the gain in classification accuracy due to enabling tree resolution-specife attention maps outweighs the loss in accuracy due to the slightly less powerful representations at the earlier transformer layers. Second, this ensures that the representation capacity of the later layers is exclusively available to the more difficult tasks.
The main idea of CascadeXML. is illustrated in Figure I] An input text is propagated through the transformer model until the first shorllistng level is reached. There, the embedding ( correspond- ing to the [CLS] token is extracted, and used to determine scores for each level-1 1 meta-label. The highest-scoring meta-labels are selected to create a shortlist of candidate labels ((1). The input is propagated further to perform another shortlisting step at the next resolution, at which new [CLS] embedding o (2) is extracted. Crucially, (2) is the result of separate self-attention operation, and its attention map a differently weighted combination of the token embeddings as compared (1)) The new embedding 0 (2) is used to refine the shortlist by selecting the top-scoring level-2 meta-labels that are children of the shortlisted level-1 I meta-labels. This process is repeated until at the final layer of the transformer, where classification performed at the full label resolution.
Contributions In this paper, our contributions are the following:
We introduce novel paradigm which can harness the multi-layered architecture of trans- formers for learning representations at multiple resolutions in an HLT. This is in contrast to the existing approaches which treat transformer merely a black box encoder,
Implemented via Ca,, it is an end-to-end trainable multi-resolution learning pipeline which is quick to train and simple to implement. It eliminates the need of bells and whistles like multi-stage procedures, reclustering, and bootstrap training to achieve state-of of--oe--n-t-- art performance, and
3

---

CascadeXML not only scales to dataset with millions of labels but also improves the current saae-off-the-rr in terms of prediction performance by upto ~ 6% on benchmark datasets. In addition to this, it is computationally most efficient to train and clocks multi-fold improvement in inference time as compared to existing transformer-based approaches.
2 Notation & Preliminaries
extreme classification, the input text instance (document or query) x E X' is mapped to a subset of labels y = [r.] out of the I. available labels identified through the integers [L] = (I....,L) L}, such that L ~ 10 Usually, we identify the labels through a binary vector y {0,1]4, where y/ = + j y. The instances and labels are jointly distributed according to unknown distribution (x, y) ~ P of which we have a i.i.d. sample D (x',y" i E [n]] available as a training set. Even though the number of labels is large, the label vectors are sparse.
Problem setup:
As part of the model pipeline, a text vectorizer 4 X' + RS projects the documents to a d-dimensional embedding space. These vectorizers can be simple bag-of-words (BOW) TF-IDF transformations Puraa, or learnable model danm that encodes the documents to the embedding space using weights learnt over the dataset. The final classification of an instance x is achieved by computing per label scores (wi, ((x)) using label-wise weight vectors Wi E Rd which are either (i) jointly learnt with + in deep architecture
10 19
32
52 [53] or (ii) solety learnt in shallow classifiers 2 [6].
Transformer Encoder: More recent approaches in extreme classification have achieved state-
of-the-art performance using pre-trained transformers as text vectorizers (99 52 57]. Popular
transformer encoders for NLP [II 28 51] consist of an embedding layer, followed by series of
stacked self-attention layers and task-specific heads that can be fine-tuned along with the encoder. For sequence of m input words, each layer in these models produces a sequence of embeddings @00 (x)())))..... (x)....
(1) (x)) where (b()) (x) E R is representation for the i'th token in the input
sequence the a' th layer, and ( (x) == (x((()) an embedding corresponding
special [CLS] token that has been pre-trained to capture a representation of the entire input [1] 28
Label Shortlisting: Computing scores (w,, ((x)) for each label I E [L] in an extremely large output space is computationally very demanding, especially when the feature space Id ((x) is high dimensional to allow for expressive features. To alleviate this problem, the labels L are grouped together under K meta-labels {C,..... Cx) R, with K se O(VL). Each meta-label is
represented as the set of its child labels, C. {1 E L]: is child of k} and each target set C [4] corresponds to a meta-target y = H(y)= = {Cx Bi Ex).)
J. A meta-label C is relevant to a
data point if it contains at least one relevant label.
Ideally, the labels grouped under one meta-label should be semantically similar to ensure the common
representations learnt for meta- and extreme- tasks are relevant to both tasks [0]. This has the added advantage that the negative labels selected by the shortlisting procedure will be the most confusing
ones, leading to hard negative sampling and improved learning of the extreme classifier [39]
The label vectors are sparse the extent that the much coarser meta-labels still contain mostly zeros, | < K. This means that it is enough to explicitly calculate the scores only for the descendants of relevant meta-labels. In practice, an algorithm will first predict scores for each meta-label and select the top-scoring ones as shortlist for further classification. The meta- and cetreme--lassifiieee tasks can either be learnt sequentially [O] or jointly through dynamic sabel-shorliiting
I9 23)
For extremely large label spaces, a single label resolution for shortlisting is not optimal either too large that the meta-task itself becomes very expensive, or its individual clusters are so large that the shortlist contains too many candidate labels (57). Thus, methods like Attention MI. and KR-Transformer use multi-level shortlisting procedure based on a Hierarchical Label Tree (HLT).
177 is
Label Trees: An HLT 5) = (X( ..... "I....,.T))) is hierarchical clustering of labels into T levels of
The clusters cl of each level form disjoint partitioning of the clusters of the next, etah c (((+)), with the convention that the T + 1 level is the full label space (T++1) = [T.]. As such, each cluster c(T)) (or meta-label) at level T represented by set of labels, each cluster at level T - 1 is a set of level-T meta-labels, and thus set of sets
successively refined clusters (1))) (c ...., Ck. ....,,,k)]
(1) )

---

of labels, and so on. In practice, label representations an aggregate of their positive instance features are leveraged create an HLT. Typically, 2 constructed using Piriar as
Vr
where Vi Σ Furar(x(1)).
21
v
kyy
The HLT often constructed by recursively partitioning the label-space using balanced k-means clustering top-down fashion 19 55 [57]. For example, HLT constructed using balanced
2-means clustering has [log, LJ levels, out of which we select T levels corresponding to different meta-label resolutions.
Method: CascadeXML
At its core, the CascadeXML model consists of three components: A pre-trained language model which allows to extract a document representation (fa) (x) E IR at different layers E [A], an HLT
5) =
H))) that divides the label space into K1 <... KT < = KT+1 successively (w()) [w] ..... w k, R
refined clusters, and set of linear classifiers
[T }
For each meta-classifier, we select a layer in order to extract resolution-specifie feature embeddings ( == (ar) (x) with 1 < 41 <... < aT < A #T+1-
Refining and Coarsening Label Vectors: In order to transfer the representation of shortlist of a label vector between different levels of the HLT, we need to define two operations. The first is refining the resolution of shortlist:
Given shortlist S (c%)) ..... c( C RQ()) of clusters of the resolution at level i, we want to calculate the representation of this shortlist in the next level, denoted At (S). Since each cluster is defined the set of its descendants, this is achieved by taking the union over the clusters in the list. The second operation is to find the cover of meta-labels that envelope given set of labels, (S)) This is achieved by identifying each of the clusters in the level above for which there exists least one descendant in the shortlist. This leads to
(8) U C (((+))
(T) = {
(t-11) E (((+1)) |sn (1-1) + #0).
(1)
CES
We define (y) ( . R.
(y) to be the level-t cover of the label set y C [L].
Multi-resolution Dynamic Label Shortlisting: We adapt dynamic label shortlisting 19 for multi-
resolution training approaches. Formally, given a shortlist S()) of clusters {c ...., k
(0)
s()) C (2( at level 4, and the corresponding classification features (()), a new shortlist ((1+1) c
(2(+1) is generated as follows: The k highest-scoring meta-labels are selected into Stop and the next shortlist ((++1) is then given the refinement:
Stop = ( : Tops (ww!) (VE)000)))) )
s((+1)= -9("" (S!00)
At the first level, however, the shortlist initialized to consider all meta labels S(1) (1), It is of essence that all true positive (meta-) labels appear at every level during training time R 19 55 However, this is difficult ensure implicitly due imperfect recall rate of classifiers. Hence, during training all meta-labels corresponding to ground-truth labels are added to the candidate set, leading to
S
(t++1) train RT ((((s)) U (()))(())
Consequently, single training step in CascadeXML looks as given in algorithm
Training Objective: At each level the tree, the goal is to correctly identify the most likely meta-labels, which can be achieved through minimization of one-vs-all loss, such BCE
L (t) (x.
Σ (()) ((w!)) ( (x)).
(x)). (y))),
(4)
|
(ESO)
5

---

Algorithm 1: Training step in CASCADEXML
Input: instance x, labels y, clusters E, features @ classifiers W, weights ( + (11) (x) E [T +1
** Forward pass through Transformer */
loss +-0 0
8(1) e(1)
Initialize shortlist
for t in 1, .......? Te....aa. T do
y(4)) - (y)
/ Ground-truti meta-labels */
P: (wh. ) Vi E S(t)
/* Sparse prediction */
loss + loss + [8 Elesto C()) vp).i"")
S Cky :j Topk (ip liE ((() 1)
/* Top-K selection */
S(++1) Ai (5(())
L (x)
** Add ground-truth positives */
Pj f- (w),
(T-1) O((+))) E(T-)-
/* Final prediction */
k;
loss +- loss [S Eies ((T+1) (p),2))
adjust + and W to reduce loss
The overall objective is a weighted sum of the individual layer losses
T+1
[(x, y)
Σ (1) (1) (x,,)),
(5)
where level T + 1 corresponds to the extreme classification task. To balance out different loss scales result of normalization by varied shortlist sizes, we re-scale losses from multiple resolutions using a simple factor (1) |/ minee/r+1| (1)(
The algorithm for inference remains the same as training, except do not teacher-force true labels R8(0()) during inference time (Eqn. 3). The inference consists of one forward pass through the transformer backbone, followed by successive shortlisting at all resolutions, and final extreme classification. For l-way branching tree with log, T. levels, each level has be meta-labels, of which [ x V are selected for the shortlist. Therefore, the time complexity is
Inference:
logs L
O (IS()) | +|S(T+))) = (
Vb()d) O dvL)
where Tam denotes the time taken to compute transformer embeddings for input.
4
Main results
Benchmark datasets Evaluation Metrics We benehmark the performance of Cas on standard datasets and evaluation metrics, as shown in Table 2] Table 3 In all, we use public XMC datasets 4 31) the statistics of which are specified in appendix A.I. For our experiments we use the same raw text input, train-test split and sparse feature representations as used in [55 [9 57 use Precision @k (P@k) and Propensity-scored P@k (PSP@k) [77 which focuses more on model's performance on tail labels, as the evaluation metrics.
Baseline Methods We compare Cas with strong baselines: DiSMEC Parabel [36] eetreme'Text [47], Bonsai [22). XR-Linear [56). AttentionXML, X- X-Trrnsformer g LightXML and Tr-nsssfformr. The baseline results for these methods were obtained from [57 Table: 2]. may be noted that, in order to incorporate the global statistical information of the dataset, XR-Transformer concatenates sparse if-idf features to those learnt via deep networks (denoted & as:
4 PPann(X
Pariar (x)
car (x)
J [x]]]
T (x)]l
For fair evaluation, we present separate comparison with methods which use (i) features from a DNN encoder only, and (ii) those obtained by combining DNN and tf-idf features.

---

Model Ensemble Settings For CascadeXML, we follow the ensemble settings of LightXML and XR- Transformer. Specifically, we ensemble a model each using BERT [TT] RoBerta. 228 and XLNet [51] for Wiki10-31K and AmazonCat-13K, For the larger datasets, we make an ensemble of three BERT models with different seeds for random initialization. On the other hand, AttentionXMI. uses ensemble of 3 models and X-Transformer uses 9 model ensemble with BERT, RoBerta, XLNNt large models with three difference indexers.
Metrics
P@1 1
P@3
P@5
P@1
P@3
P@5
P@1
P@3
P@5
Datasets
Wiki-500K
Amazon-670K
Amazon-3M
DiSMEC
70.21
50.57
39.68
44.78
39.72
36.17
47.34
44.96
42.80
Parabel"
68.70
49.57
38.64
44.91
39.77
35.98
47.42
44.66
42.55
Text
65.17
46.32
36.15
42.54
37.93
34.63
42.20
39.28
37.24
Bonsai
69.26
49.80
38.83
45.58
40.39
36.60
48.45
45.65
43.49
XR-Linear
65.59
46.72
36.46
43.38
38.40
34.77
47.40
44.15
41.87
DNN Features
AttentionXML.
76.74
58.18
45.95
47.68
42.70
38.99
50.86
48.00
45.82
LightXML 77.89
58.98
45.71
49.32
44.17
40.25
Cas
78.39
59.86
46.49
50.22
45.20
41.45
52.75
50.83
48.90
[DNN + tf-idf] Features
X- "ransformer** 77.09
57.51
45.28
48.07
42.96
39.12
51.20
47.81
45.07
XR- Transformer*
79.40
59.02
46.25
50.11
44.56
40.64
54.20
50.81
48.26
CascadeXML 81.13
62.43
49.12
52.15
46.54
42.44
53.91
51.24
49.52
Table 2: Comparison of CascadeXML to state-of-the-art methods on large scale benchmark datasets. ** and *# imply ensemble results of models and models respectively, and "-' implies that the model does not scale for that dataset due to GPU memory constraints.
CascadeXML, which leverages the multi- layered transformer architecture to learn label resolution specific feature representations, is empirically superior to all existing XMC approaches whether using: (i) tf-idf features, or (ii) DNN-based dense features, or (iii) both. This is despite CascadeXML being a simple end-to-end trainable algorithm without using any bells or whistles like multi-stage procedures [55 57). reclustering [57] or bootstrapped training [55] (7 For example, on Amazon-3M dataset, CascadeXMML significantly outperforms all DNN- approaches while taking only 24 hours train single Nvidia A100 GPU. In contrast, X- X-Transformer needs to ensemble 9 models and 23 days to train on 8 GPUs to reach the empirical performance (c.f. page 2, (57) and LightXML remains unscalable.
Empirical Performance: As demonstrated in Table 2]
We report results on Wiki10-31K and 1mazonCat-13K in [Tble 3 As a standard practice in the domain 66 19 57), use 256 tokens as an input sequence length to the model after truncation. We find this sequence length to be sufficient for competent empirical performance and hence, we do not witness any empirical benefits of leverage sparse ff-idf features for these datasets. Notably, Ca achieves best results on AmazonCat-13R dataset and performs par with XR- UR-Transformer without leveraging sparse tf-idf features. Thus, Cas- CadeXML's unique multi-resolution approach is a favorable choice not only for datasets with extremely large output spaces, but also for datasets with tens of thousand labels.
Metrics
P@1
P@3
P@5
P@
P@3
P@5
Datasets
Wiki10-31K
Amazon 13K
DISMEC
84.13
74.72
65.94
93.81
79.08
64.06
Parabel
84.19
72.46
63.37
93.02
79.14
64.51
extremeText
83.66
73.28
64.51
92.50
78.12
63.51
Bonsai
84.52
73.76
64.69
92.98
79.13
64.46
XR-Linear
85.75
75.79
66.69
94.64
79.98
64.79
DNN Features
AttentionXMI..
87.34
78.18
69.07
95.84
82.39
67.32
LightXML CascadeXML:
89.67
79.06
69.87
96.77
83,98
68.63
89.74
80.13
70.75
96.90
84.13
68.78
[DNN + tf-idf] Features
"-Transformer"
88.26
78.51
69.68
96.48
83.41
68.1
Transformer"
88.69
80.17
70.91
96.79
83.66
04
CascadeXMI.*
89.18
79.71
71.19
96.71
84.07
68.69
Table 3: Comparison of CascadeXML: to state-of-the-art meth- ods on legacy datasets.
7

---

On the concatenated feature space, CascadeXML. also significantly outperforms XR- IR-Transformer on 8 out of 9 dataset-metric combination (ref: Table 2 with notable 6% improvement on Wiki-500K dataset for P@5 metric. As XR- KR-Tranformer learns transformer features for label resolutions only up to 21 clusters [57] Table 6], incorporating tf-idf features via computation-intensi
XR-L.nnarr (56) for final extreme classification task becomes an integral part of their pipeline. On the other hand, CascadeXML does not have any such inherent limitation, and even without tf-idf features, outperforms XR- TR-Transformer on 7 out of 9 datasct-metric combinations (ref: Table 2
5
Computational cost & Ablation results
Single Model Comparison & Training Time Empirical results and training time of a single CaS model has been compared to single instance performance of state-of-the-ee DNN based XMC approaches in [Tble 4] Notably, CascadeXMI, without leveraging any sparse tf-idf features, performs at par with XR-Transformer which benefits from using the extra statistical information about the entire (not truncated) data point. When comparing training time, CascadeXML. is up to an order of magnitude faster than LightXML on a single GPU. In multi-GPU setting, CascadeXML trains in 7.2 hours on Wiki-500K dataset using only 4 NVidia V100 GPUs compared to AR-Transformet and AttentionXML which require 12.5 hours on 8 GPUs. CascadeXML not only reduces the training time across datasets but also requires half the number of GPUs to do This proves that MML is significantly more compute efficient as compared to previous DNN based XMC approaches.
Dataset
Method
P@1
P@3
P@5
T train
T train
AttentionXML
87.1
77.8
68.8
0.5
Wiki10-31K
LightXML
87.8
77.3
68.0
6.7
XR- KR-Transformer (+ tf-idf)
88.0
78.7
69.1
1.3
0.5
CascadeXMI.
88.4
78.3
68.9
0.3
0.1
AttentionXML.
75.1
56.5
44.4
12.5
Wiki-500K
LightXML
76.3
57.3
44.2
89.6
XR- TR-Transformer (+ tf-idf)
78.1
57.6
45.0
29.2
12.5
CLL
77.0
58.3
45.1
22.0
7.2
AttentionXML
45.7
40.7
36,9
8.1
Amazon-67U8 -670K
LightXML
47.3
42.2
38.5
53.0
XR- XR-Transformer (+ tf-idf)
49.1
43.8
40.0
8.1
3.4
CascadeXMI.
48.5
43.7
40.0
7.5
3.0
Table 4: Single model comparison of DNN based XMC approaches. Trrci denotes the training time
as reported on a single GPU (Nvidia V100). Trpai denotes multi-GPU training time on GPUs for AttentionXML and XR- TR-Trrasformer, and on only 4 GPUs for CascadeXML.
Inference Time Ca clocks the fastest inference speed compared to all previous DNN-based XMC models, as shown in Table 5) Notably, CascadeXML is ~ 1.5x and ~ 2x faster at inference time compared to LightXML and XR- TR-Trassfrrmer respectively.
Dataset
AttentionXML
X Transformer
XR- TR-nfrasfrrm
LightXML
CascadeXML
Wikil0-31K
20.0
48.1
39.1
27.1
12.8
AmazonCat-133k
14.4
47.6
26.1
24.1
13.1
Wiki-500K
80.1*
48.1
33.9
27.3
16.0
Amazon-670K
76.0*
48.0
30.9
23.3
16.6
Amazon-3M
130.5*
50.2
35.2
16.9
Table 5: Comparison of CascadeXML w.r.t. inference time (in milliseconds per sample) with SOTA XMC methods. Inference times were recorded on a single Nvidia V100 GPU and a single CPU with batch size of 1. Superscript implies that model parallel was used with 8 GPUs for inference.
Performance of CascadeXML on tail labels has been compared to baseline XMC methods in Table We note that for smaller datasets Wiki10-31K and AmazonCat- 13K PfI 177 significantly outperforms all other methods on PSP metrics. This is to be expected, as this method specifically optimizes for tail labels. However, for larger datasets Amazon-670K and Wiki-500K CascadeXML outperforms the strong baselines like XR-Transeeee Transformer
Performance on Tail Labels
8

---

and AttentionXML by 5-8% on PSP@3 and PSP@5 metrics while also significantly outperforming PfastreXMM over these datasets. These results conclude that our unique end-to-end nulti-resoluuionn training approach is not only empirically superior to previous approaches in P@ K metrics but also outperforms strong XMC baselines in performance over tail labels.
Methods
PSP@1
PSP
PSP@5
PSP@1
PSP
PSPe5
PSP@1
PSP
PSP@5
PSP@
PSP@3
PSP@5
WikH1031IK
AmazonCat-138
Amazon-670K
Wiki-500K
DISMEC
10,60
12.37
13.61
51.41
61.02
86
26.26
30.14
33.89
27.42
32,95
36,95
ProXML
32.31
MM
19.02
18.34
18.43
69.52
73.22
75.48
29. 30
30.80
32.43
32.02
29.75
30.19
Parabel
11.69
12.47
13:14
50.92
64,00
72.10
36.36
29.95
33.17
26.88
31,96
35.26
11.85
13:44
27.08
30.79
34.11
DNN Features (Singlc Model)
XML-CNN
9.39
10.00
10:20
52.42
62.83
67.10
17.43
21.66
24.42
XMtentionMML
16.20
17.05
17.93
53.52
68.73
76.26
32.36
35.
30.05
37.31
41.74
12.16
14.86
16.40
50.51
64.92
74.63
29.21
33.49
37.65
32.10
39.41
43.75
CascadeMLL
13.22
14.70
30.23
34.93
38.77
DNN Features (Unserbble Model)
AtemiioNML."
15.57
16.80
17.82
53.76
68.72
76.38
30.29
33.85
37.13
30.85
39.23
44.34
CascadeXMI.
13.36
15.06
77.52
31.19
36.07
40.15
32.60
46.66
[DNN ( ir-idf) Features (Fasemblee Model)
NR-Transformer CascadeXMI.
12.25
29.77
38 28.29
45.02
13.32
15.35
17.45
51.39
66.81
77.58
30.77
35.78
40.52
32.12
43.15
49.37
Table Comparison of performance of state-of of the--heerrt art methods on tail labels on benchmark datasets.
Impact of label clusters size Ear-
lier works have argued that using fine- grained clusters leads to better model per-
Dataset
Model Sctting
P@1
P@3
P@5
formance 19 23 32]. With increased
Default
48.5
43.7
40.0
number of label clusters at the penulti- mate level of HLT, the multi-resolution tasks tend get
Amazon-670K
1 Clusters
47.6
42.9
39.3
w/o Rescaling
47.8
42.8
39.2
in-sync with each other [23] and hence enables trans- fer learning across different stages of learning pipelines (10 32). Our mutli- resolution training enables Cas to use 26 label clusters for penulti- mate label resolution, double that of XR- Transformer. As TTale 7shows, increas- ing the resolution at the last shortlist- ing step results in non-trivial improve- ments, highlighting the need for archi- tectures that can efficiently handle high- resolution meta-labets
Default
76.9
58.4
45.2
Wiki-500K
2 Clusters
76.6
58.1
44.9
w/o Rescaling
76.5
57.9
44.7
Default
51.3
49.0
46.9
Amazon-3M
25 Clusters
50.8
48.7
46.5
w/o Rescaling
50.9
48.6
46.6
Table 7: Ablation experiments showing performance of a single CascadeXML model with different model settings. Default implies 2 clusters in penultimate label resolution and per-resolution loss re-scaling.
Impact of re-scaling per-resolution loss For extreme datasets, we witness more efficient training by having larger shortlist size for finer label resolutions (see appendix A.2.2). This results in varied loss scales across label resolutions leading to one resolution dominating the training. Comparing the "Default" and "w/o Rescaling" rows for respective datasets in Table 7 indicates that rescaling as introduced in section 3 helps CascadeXML. train more effectively.
Discussion
As shown above, CascadeXML substantially outperforms current state-of-the-art approaches both in terms of computational efficiency, and prediction performance. It achieves this while being simpler than IRR-ransformer and more scalable than LightXML. Revisiting Table I] the extraction of in- termediate [CLS] representations enables us to perform end-to-end meloi-resolution training with resolution-specific attention maps from a single transformer model. From Table 2, we show im- provements over both LightXMI. and KR-Transformer,, which are transformer- -hatom models sharing feature representations across multiple label resolutions. We attribute this increased performance to CascadeXML's end-to-end multi-resolution training pipeline and property of adapted feature embed-

---

dings. The presented architecture, therefore, combines the strength of different previous approaches,
without inheriting their limitations.
We postulated that the mheeta-classiffcattin tasks at different resolutions need to attend to different tokens in the input sequence, a property lacking in recent approaches [ T9 57]. This intuition is
corroborated by the fact that the attention maps learned by CascadeXML differ significantly between different resolutions. We provide graphs and supporting data in Appenddx B
Obtaining label resolution specific feature embeddings in pre-trained transformer model requires to make the compromise of extracting these at earlier transformer layers, where the token representations are not yet as refined as in the final layer [49]. In contrast to the extreme classification task, which requires placing true labels exactly the top-5 positions, the shortlisting task only requires all true meta-labels be recalled within the shortlist, which is much longer than five elements. We observe that the meta-classifiers are able achieve high-recall (as shown shown Appendix B) at the shortlist length even with the less refined features. CaS is, therefore, successful in efficiently leveraging intermediate transformer representations for the task of label shortlisting while keeping feature representations separate for each label resolution.
The structure of the model means that earlier layers in the transformer receive gradients from multiple resolutions, whereas the final layers' back propagation is limited to gradients from the extreme resolution (cf. FFigure I}. This setup ensures that portion of the representational capabilities is exclusively available for the extreme task and hence, the final transformer layers are increasingly more suitable for classification at the extreme resolution. At the same time, since the gradients from the finer tasks are also relevant for the coarser meta-label shortlisting task, so additional downstream task does not hurt the shortlisting performance at the intermediate transformer layers.
7 Other Related Works
large majority of the initial works in XMC have been focussed on learning the classifier (with fixed
features), with one of the following class of methods: (i) one-vs-rest [54] 2 [53] B
41].
(ii) label
trees 35
47 18]
(iii) decision trees (37 44 30). and (iv) label embedding based 5 45 [4]
classifiers. Beyond the above algorithmic categorization of approaches, computational considerations
on scaling the training process via negative sampling [16]. and smart initialization [12 41 has also
been studied. Furthermore, the statistical consequences of negative sampling (39] and missing labels
17
38
43
42
48
have led to novel insights in designing unbiased loss functions.
With advances in deep learning, joint learning of features and classifier has been the focus of most recent approaches. Building on the first steps for developing convolutional neural networks for text classification (24), deep extreme classification was introduced in XML-CNI (27). This work preceded the recent developments on employing BiLSTM and transformer encoders discussed in detail the Introduction section. There has been a recent surge in works to tackle XMC for instances with short
text inputs and those in which labels are endowed with textual descriptions [10 [ 32 33 40 Lately,
the XMC setting has also been extended to predict unseen labels in zero-shot learning
13 58
Conclusion
In this paper, we introduced novel paradigm for fine-tuning transformer models in XMC settings. In contrast to the existing methods, which consider transformer as a blackbox encoder, we leverage its multi-layered architecture to learn data representation corresponding to different resolutions the HLT as well as fine-grained labels at the level of the extreme classifier. The proposed instantiation of our framework in the form of CascadeXML not only yields state-of-the-art prediction performance, but is also end-to-end trainable (without intermediate reclustering steps), simpler to implement, and fast on training and inference. Beyond further research in extreme classification towards fully harnessing the representation capabilities of transformer encoders, we believe our approach can inspire future multi-resolution architectures other domains as well which leverage label hierarchy.
9 Acknowledgments
The authors would like to thank Devaansh Gupta and Mohammadreza Qaraei for useful discussions. They also acknowledge the support of CSC-IT Center for Science, Finland, as well as the Aalto Science-IT project, for providing the required computational resources. This research is supported in part by Academy of Finland grants Decision No. 348215 and 347707.
10

---

References
[1] R. Agrawal, A. Gupta, Y. Prabhu, and M. Varma. Multi-label learning with millions of labels: recommending advertiser bid phrases for web pages. Proceedings of the 22nd international conference on World Wide Web, 2013.
[2] R. Babbar and B. Schölkopf. Dismec: Distributed sparse machines for extreme multi-label classification, Proceedings of the Tenth ACM International Conference on Web Search and Data Mining, 2017.
[3] R. Babbar and Schölkopf. Data scarcity, robustness and extreme multi-label classification. Machine Learning, pages 1-23, 2019.
[4] K. Bhatia, K. Dahiya, H. Jain, P. Kar, A. Mittal, Y. Prabhu, and M. Varma. The extreme classification repository: Multi-label datasets and code, 2016.
[5] K. Bhatia, H. Jain, P. Kar, M. Varma, and P. Jain. Sparse local embeddings for extreme multi- label classification. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28, 2015.
[6] W. Chang, F. X. Yu, Y. Chang, Y. Yang, and Kumar. Pre-training tasks for embedding- large-scale retrieval. In 8th International Conference on Learning Representations, ICLR 2020.
[7] W.-C. Chang, D. Jiang, H.-F. Yu, C. H. Teo, Zhang, K. Zhong, K. Kolluri, Q. Hu, N. Shandilya, V. levgrafov, et al. Extreme multi-label learning for semantic matching in product search. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pages 2643-2651 2021.
[8] W.-C. Chang. H.-F Yu, K. Zhong, Y. Yang, and Dhillon. Taming pretrained transformers for extreme multi-label text classification. Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2020.
[9] K. Dahiya, A. Agarwal, D. Saini, G. K, J. Jiao, A. Singh, S. Agarwal, P. Kar, and M. Varma. Siamesexml: Siamese networks meet extreme classifiers with 100m labels. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2330 2330-2340. PMLR, 2021.
[10] K. Dahiya, D. Saini, A. Mittal, A. Shaw, K. Dave, A. Soni, H. Jain, S. Agarwal, and M. Varma. DeepXML: Deep Extreme Multi-Label Learning Framework Applied to Short Text Docu- ments. In WSDM, Mar. 2021.
[11] J. Devlin, M. W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In NAACL, 2019.
[12] H. Fang, M. Cheng, C.-J. Hsich, and M. Friedlander. Fast training for large-scale one-versus- all linear classifiers using tree-structured initialization. In Proceedings of the 2019 SIAM International Conference Data Mining, pages 280-288. SIAM, 2019.
[13] L. Galli and C.-J. Lin. A study on truncated newton methods for linear classification. IEEE Transactions on Neural Networks and Learning Systems, 2021.
[14] C. Guo, A. Mousavi, X. Wu, D. N. Holtmann-Rice S. Kale, S. Reddi, and S. Kumar. Breaking the Glass Ceiling for Embedding- Based Classifiers for Large Output Spaces. In NewrIPS, 2019.
[15] N. Gupta, S. Bohra, Y. Prabhu, S. Purohit, and M. Varma. Generalized zero-shot extreme multi-lahel learning. In Proceedings of the ACM SIGKDD Conference on Knowledge Discovery and Data Mining, August 2021.
[16] H. Jain, V. Balasubramanian. B. Chunduri. and M. Varma. Slice: Scalable Linear Extreme Classifiers trained on 100 Million Labels for Related Searches. In WSDM, 2019.
[17] H. Jain, Y. Prabhu, and M. Varma. Extreme multi-label loss functions for recommendation, tagging, ranking & other missing label applications. In KDD, pages 935-944, 2016.

---

[18] K. Jasinska, K. Dembezynski, R. Busa-Fekete, K. Pfannschmidt, T. Klerx, and E. Hullermeier. Extreme F-measure Maximization using Sparse Probability Estimates. ICML, June 2016.
[19] T. Jiang, D. Wang. L. Sun, H. Yang. Z. Zhao, and F. Zhuang. Lightxml: Transformer with dynamic negative sampling for perf-peeformance extreme multi-label text classification. In Proceedings of the 35th AAAI Conference on Artificial Intelligence, pages 7987-7994, 2021.
[20] Karras, T. Aila, S. Laine, and J. Lehtinen. Progressive growing of gans for improved quality, stability, and variation. In ICLR, 2018.
[21] T. Kartas, S. Laine. and T. Aila. A style-based generator architecture for generative adversar- ial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401-4414 2019.
[22] S. Khandagale, H. Xiao, and R. Babbar. Bonsai: diverse and shallow trees for extreme multi- label classification. Machine Learning, 109(1 1):2099-2119 19, 2020.
[23] S. Kharbanda, A. Banerjee, A. Palrecha, and R. Babbar. Embedding convolutions for short text extreme classification with millions of labels. arXiv preprint arXiv:2/09.07319 2021.
[24] Y. Kim. Convolutional Neural Networks for Sentence Classification. In EMNLP, Oct. 2014.
[25] W.-S. Lai, J.-B. Huang. N. Ahuja, and M.-H. Yang. Deep laplacian pyramid networks for fast and accurate super-resolution. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 624-632, 2017.
[26] K. Lee, M.-W. Chang, and K. Toutanova. Latent retrieval for weakly supervised open domain question answering. In Proceedings the 57th Annual Meeting of the Association for Computa- tional Linguistics, pages 1086-6096 Florence, Italy, July 2019. Association for Computational Linguistics.
[27] Liu, W.-C Chang, Wu, and Y. Yang. Deep learning for extreme multi-label text clas- sification. In Proceedings of the 40th international ACM SIGIR conference on research and development in information retrieval, pages 115-124, 2017.
[28] Y. Liu, M. Out, N. Goyal, Du, M. Joshi, D. Chen, Levy, M. Lewis, L.. Zettlemoyer, and V. Stoyanov. Roberta: A robustly optimized bert pretraining approach. ArXiv, 1ttp//907.119999 2019.
[29] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. ICLR, 2019.
[30] M. Majzoubi and A. Choromanska. Ldsm: Logarithm-ddepth streaming multi-label decision trees. In International Conference on Artificial Intelligence and Statistics, pages 4247-4257. PMLR, 2020.
[31] McAuley and J. Leskovec. Hidden factors and hidden topies: understanding rating dimensions with review text. In Proceedings of the 7th ACM conference on Recommender systems, pages 165-172, 2013.
[32] A. Mittal, K. Dahiya, Agrawal, D. Saini, S. Agarwal, Kar, and M. Varma. Decaf: Deep extreme classification with label features. In WSDM, March 2021.
[33] A. Mittal, N. Sachdeva, S. Agrawal, S. Agarwal, P. Kar, and M. Varma. Eclare: Extreme classification with label graph correlations. In Proceedings of The ACM International World Wide Web Conference, April 2021.
[34] A. S. Morcos, M. Raghu, and S. Bengio. Insights on representational similarity in neural networks with canonical correlation. In NeurIPS, 2018
[35] Y. Prabhu, A. Kag, S. Gopinath, K. Dahiya, S. Harsola, R. Agrawal, and M. Varma. Extreme multi-label learning with label features for warm-start tagging, ranking and recommendation. In WSDM, 2018
[36] Y. Prabhu, A. Kag, S. Harsola, R. Agrawal, and M. Varma. Parabel: Partitioned label trees for extreme classification with application to dynamic search advertising. In WWW, 2018.
12

---

[37] Y. Prabhu and M. Varma. FastXML: A Fast, Accurate and Stable Tree-classifier for eXtreme Multi-label Learning. KDD, Aug. 2014.
[38] M. Qaraei, E. Schultheis, P. Gupta, and R. Babbar. Convex surrogates for unbiased loss functions in extreme classification with missing labels. In Proceedings of the Web Conference 2021, pages 3711-3720, 2021.
[39] Reddi, S. Kale, F. Yu, D. N. H. Rice, Chen, and S. Kumar. Stochastic Negative Mining for Learning with Large Output Spaces. CoRR, 2018.
[40] D. Saini, A. K. Jain, K. Dave, J. Jiao, A. Singh, R. Zhang, and M. Varma. Galaxe: Graph neural networks with labelwise attention for extreme classification. In Proceedings of The ACM International World Wide Web Conference, April 2021.
[41] E. Schultheis and R. Babbar. Speeding-up up one-vs-all training for extreme classification via smart initialization. arXiv preprint 10X.v:2109.133 2021.
[42] R. Schultheis and R. Babbar. Unbiased loss functions for multilahel classification with missing labels. arXiv preprint arXiv:2109. /2822, 2021.
[43] R. Schultheis, M. Wydmuch, R. Babbar, and K. Demhezynski. On missing labels, long-tails and propensities in extreme multi-label classification. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 1547-1557, 2022.
[44] W. Siblini, P. Kuntz, and F. Meyer. Craftml, an efficient clustering- based random forest for extreme multi-label learning. In International Conference on Machine Learning, pages 4664 4673. PMLR, 2018.
[45] Y. Tagami. AnneXML: Approximate Nearest Neighbor Search for Extreme Multi-label Classifi- cation. In KDD, Aug. 2017.
[46] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, u. Kaiser, and I. PPlosuukin.. Attention is all you need. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[47] M. Wydmuch, K. Jasinska, M. Kuznetsov, R. Busa-Fekete, and K. Dembczynski. A no-regret generalization of hierarchical softmax extreme multi-label classification. In NIPS, 2018.
[48] M. Wydmuch, K. Josbss,a-Kobus, R. Babbar, and K. Dembczynski. Propensity-scored proba- bilistic label trees. In SIGIR, pages 2252 2252-2256, 2021.
[49] J. Xin, R. Tang, J. Lee, Y. Yu, and J. Lin. DeeBERT: Dynamic early exiting for accelerating BERT inference. In Proceedings of the 58th Annual Meeting of the Association for Com- putational Linguistics, pages 2246-2251, Online, July 2020. Association for Computational Linguistics.
[50] Y. Xiong, W.-C. Chang, C.-J. Hsieh, H.-F. Yu, and I. S. Dhillon. Extreme zero shot learning for extreme text classification. In NAACL, 2022.
[51] Z. Yang, Z. Dai, Y. Yang, J. G. Carbonell, R. Salakhutdinov. and Q. V. Le. Xlnet: Generalized autoregressive pretraining for language understanding. In NeurIPS, 2019.
[52] H. Ye, Z. Chen, D.-H. Wang, and D. B. D. Pretrained Generalized Autoregressive Model with Adaptive Probabilistic Label Clusters for Extreme Multi-label Text Classification. In ICML, 2020.
[53] E. I. Yen, X. Huang, W. Dai, Ravikumar, Pand Dhillon, and E. Xing. PPDSparse: A Parallel Primal-Dual Sparse Method for Extreme Classification. In KDD, Aug. 2017.
[54] E. I. Yen, X. Huang, K. Zhong, P. Ravikumar, and S. Dhillon. PD-Sparse: A Primal and Dual Sparse Approach to Extreme Multiclass and Multilabel Classification. In ICML, June 2016.
13

---

[55] R. You, Z. Zhang, Z. Wang. S. Dai, H. Mamitsuka, and S. Zhu. AttentionXML: Label tree-based attention-awe -waare deep model for pigh-peeformance extreme multi-label text classification. In NeurlPS, 2019.
[56] H.-F Yu, K. Zhong, J. Zhang, W.-C. Chang, and I. S. Dhillon. Pecos: Prediction for enormous and correlated output spaces. Journal of Machine Learning Research, 2022.
[57] J. Zhang, W.-c Chang, H.-f. Yu, and Dhillon. Fast multi-resolution transformer fine-tuning for extreme multi-label text classification. Advances in Neural Information Processing Systems, 34, 2021.
[58] Y. Zhang, Z. Shen, C.-H. Wu, B. Xie, J. Hao, Y.-Y. Wang, K. Wang, and J. Han. Metadata- induced contrastive learning for zero-shot multi-label text classification. In Proceedings of the ACM Web Conference 2022, pages 3162-3173, 2022.
14


Reviews:
[{"review_id": "iVl3UMOJNCu", "reviewer": "Reviewer_EmPK", "summary": "The paper proposes CascadeXML to solve the extreme multi-label text classification problem. Compared to prior works, the proposed method utilizes features from multiple layers from Bert for different granularity of the hierarchical label tree. The coarser labels with lower scores can be terminated early for more efficient learning. The proposed method is compared with several SOTA methods, including  LightXML and XR-Transformer, on benchmark datasets - Wiki-500K, Amazon-670K, and Amazon-3M.", "questions": "Please refer to the weakness section", "limitations": "NA", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I am not an expert in text classification, and thus not familiar with most of the related works. My judgment here is mostly based on the merits and the presentation of the paper on its own.\n\nThe paper in general is very well written. It explains the main drawbacks of existing works (although I am not sure if there are other works that already addressed these issues) and motivates from the drawbacks to propose CascadeXML. The core idea of CascadeXML is intuitively sound, well explained, and should be easy to reproducible. The authors also include the code in the supplementary.\n\nI did not find too many weaknesses in the paper, but I do have some questions.\n\n1. Does all compared (or the main) SOTA methods based on the same Bert model? This is not clearly stated in the paper. \n2. For early termination of the coarse labels, there are risks of early wrong predictions as well.  Is there any measurement to prevent that? ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "eACqNSUHhV-", "reviewer": "Reviewer_2RC4", "summary": "This paper studies the extreme multi-label text classification problem and established new state-of-the-art (SOTA) results with the proposed method (namely CascadeXML), which significantly simplified the multi-stages training of the previous SOTA model, XR-Transformer. CascadeXML enables end-to-end training with two key insights: (1) leverages intermediate representations of Transformer layers for learning XMC sub-problems of the label tree (2) fine-tuning at the full resolution of label space. ", "questions": "### Major Questions\n- How to decide the mapping between Transformer intermediate layer embedding \\theta^{(t)} and the level of the label tree? For example, if the depth of a label tree is greater than 12 layers (say 15 layers), where the BERT-base model has 12 layers of Transformer encoders, then the mapping is not 1-to-1?\n- For Eq(4), It seems to be that the loss function is binary cross entropy. Any comments or justification on that? Did you compare it with other loss functions such as max margin hinge loss or ranking loss?\n- Any thoughts of why the inference time is faster than other baselines in Table 3?\n- For the CascadeXML trained by sparse-dense concat features, do you retrain the label tree or is it the original label tree constructed by PIFA of sparse TFIDF input features?\n- It seems like the \"coarsening-op\" R_{\\uparrow}^{(t)}(S) (latter one in Eq(1)) is never used in the main paper. Maybe you can consider removing it for more simpler presentation and ease of notation?\n\n### Minor Comments\n- Line 111: we often write t_k \\in \\mathcal{V} in convention.\n- Line 181-182: For Algorithm 1, the short list at t+1 layer should be R_{\\downarrow}^{(t)} (S_{top}^{(t)}) \\union R^{(t+1)}(y). In Algorithm 1, the “refining-op” is missing?\n- Line 327&332: merge duplicate bibtex entry [6] and [8], which points to the same paper.\n- Line 432-433: update citation bibtex with latest info (Xiong et al., NAACL 2022)\n- Line 446-447: update citation bibtex with latest info (Yu et al., JMLR 2022)\n\n", "limitations": "No potential negative societal impact.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strength\n- Leverage intermediate representations for multi-resolution learning of XMC is novel\n- Paper writing is clear and easy to follow\n- Empirical results are quite promising \n\n### Weakness\n- Lack of ablation study for the choice of intermediate layer of Transformer embeddings \n- Lack of ablation study on loss functions\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "VjTBP594sd", "reviewer": "Reviewer_ZN1o", "summary": "The authors propose CasccadeXML, an end-to-end deep learning pipeline for extreme multi-label text classification. CascadeXML leverages the multi-layer architecture of transformer models to handle different label resolutions with separate feature representations. Specifically, for each label resolution in the hierarchical label tree (HLT), features from the corresponding transformer layer are used to predict a shortlist of candidate labels or meta-labels. The authors empirically show that CascadeXML achieves state-of-the-art performance on three extreme multi-label datasets. The computation of training and inference are both faster than the current state-of-the-art transformer-based models.", "questions": "1. How did you decide the hyperparameters, especially the corresponding transformer layer of each label resolution?\n2. Are the results sensitive to the hyperparameters?\n3. Is there a rule of thumb for deciding these hyperparameters if given a new dataset?", "limitations": "The authors said the limitations and the potential negative societal impact are discussed in the appendix, but I didn't find them. Some discussion about the lack of theoretical guarantee or the limited search space of the hyper-parameters would be good.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength\n1. The approach to leveraging the multi-layer architecture of transformers for multi-resolution learning is interesting.\n2. The experiment results are good concerning both performance and computation.\n3. The analysis of the impact of label cluster size explains the improvement of CascadeXML well.\n4. The paper is well-organized and easy to follow.\n\nWeakness\nMajor\n1. The method requires much more hyper-parameters, such as the corresponding transformer layer and the number of meta-labels of each label resolution.\n2. Lack the discussion of the sensitivity of the newly introduced hyper-parameters.\nMinor\n3. The high/low resolution mentioned in Section 1 sounds ambiguous to me. May replace them with refined/coarse.\n4. The definition of \\mathfrak{C}^{(t)} and \\mathfrak{R}^{(t)} are provided but the definition of \\mathfrak{C}^{(t)}(y) and \\mathfrak{R}^{(t)}(y) are not.\n5. Line 234: Attention -> AttentionXML", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "JgiXBytK66", "reviewer": "Reviewer_zXmp", "summary": "This work proposes CascadeXML which can leverage the multi-layered architecture to learn data representation corresponding to different resolutions in the Hierarchical Label Tree as well as fine-grained labels at the level of the extreme classifier. Each layer of the model can perform label shortlisting on a level.\n\nThe authors compare the proposed CascadeXML to strong baselines on three XMC benchmark datasets. Experiments show that the proposed method outperforms these baselines, and can achieve faster training and inference.\n", "questions": "1. Can you explain more on Eq(1)? How can we present a shortlist with its descendants, which might be a set of another shortlists? How can we get the corresponding vector representations?\n\n2. If any prediction error happens on a higher level during inference, will this error propagate easily to lower levels?", "limitations": "The authors have discussed and addressed some of the concerns I have in the paper.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\n1. The proposed CascadeXML is end-to-end differentiable and achieves better results compared to previous baselines on three benchmark datasets.\n\n2. The proposed method is faster in both training and inference, which is an important aspect of the XMC task.\n\n3. This work is well-written and easy to follow.\n\n\nWeakness:\n\n1. The idea of dynamic label shortlisting is very interesting, but I wonder how to solve the cases where HLT contains more levels than the layer number of Transformers? Also, for the HLT containing fewer levels, do we have to change the model structure? If the model is trained on HLT-1 containing $n$ layers, can we transfer it to an HLT with $m$ layers where $m \\neq n$?\n\n2. Why do the earlier layers in the Transformer correspond to the higher levels of the HLT? More explanations over this may help a bit with the answer to Q1.\n\n3. The proposed method highly relies on the multi-layer structure of the Transformer, which makes it not easy to extend to other frameworks. As mentioned above, the scalability and transferability of this method may be much more limited compared to the baselines adopted in this work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept