ELIAS: End-to-End Learning to Index and Search in Large Output Spaces
Authors: Nilesh Gupta
Patrick CHen
Hsiang-Fu Yu
Cho-Jui Hsieh
Inderjit S Dhillon
Venue: NeurIPS 2022 Accept (2022)
Keywords: extreme multi-label classification
large output spaces
learnable search index

Abstract: Extreme multi-label classification (XMC) is a popular framework for solving many real-world problems that require accurate prediction from a very large number of potential output choices. A popular approach for dealing with the large label space is to arrange the labels into a shallow tree-based index and then learn an ML model to efficiently search this index via beam search. Existing methods initialize the tree index by clustering the label space into a few mutually exclusive clusters based on pre-defined features and keep it fixed throughout the training procedure. This approach results in a sub-optimal indexing structure over the label space and limits the search performance to the quality of choices made during the initialization of the index. In this paper, we propose a novel method ELIAS which relaxes the tree-based index to a specialized weighted graph-based index which is learned end-to-end with the final task objective. More specifically, ELIAS models the discrete cluster-to-label assignments in the existing tree-based index as soft learnable parameters that are learned jointly with the rest of the ML model. ELIAS achieves state-of-the-art performance on several large-scale extreme classification benchmarks with millions of labels. In particular, ELIAS can be up to 2.5% better at precision@$1$ and up to 4% better at recall@$100$ than existing XMC methods. A PyTorch implementation of ELIAS along with other resources is available at https://github.com/nilesh2797/ELIAS.

ELIAS: End-to-End Learning to Index and Search in Large Output Spaces
Nilesh Gupta
Patrick H. Chen UCLA patrickchen@g eella.ddu
Hsiang-Fu, Yu Amazon rofu. yu@gaillmm
UT Austin nileshōcs. wtw.uutxxss.eo.. edu
Cho-Jui, Hsieh
Inderjit S. Dhillon
UCLA chohsieh@cs ucla. edu
UT Austin & Google
inderjit@cs. utexas. edu
Abstract
Extreme multi-label classification (XMC) a popular framework for solving many real-world problems that require accurate prediction from very large number of potential output choices. A popular approach for dealing with the large label space arrange the labels into shallow tree-based index and then learn an ML model to efficiently search this index via beam search. Existing methods initialize the tree index by clustering the label space into few mutually exclusive clusters based on pre-defined features and keep it fixed throughout the training the label
procedure. This approach results sub-optimal indexing structure space and limits the search performance the quality of choices made during the initialization of the index. In this paper, we propose novel method ELIAS which relaxes the tree-based index to specialized weighted graph-based index which is learned end-to-end with the final task objective. More specifically, ELIAS models the discrete cluster-to-labe assignments in the existing tree-based index soft learnable parameters that are learned jointly with the rest of the ML model. ELIAS achieves late-of-the-art performance several large-scale extreme classification benchmarks with millions labels. In particular, ELIAS can be up to 2.5% better precision @ and up to 4% better at recall @ 100 than existing XMC methods. A Py Torch implementation of ELIAS along with other resources is available at httss//gghhubb.oomm/nnn/77/////
Introduction
Many real-world problems require making accurate predictions from large number of potential output choices. For example, search advertising aims to find the most relevant ads to a given search query from large corpus of ads [26] [14], open-domain question answering requires finding the
right answers to given question from a large collection of text documents [8] 29], and product
recommendation requires recommending similar related products from large product catalog, based on past searches and interactions by users. eXtreme Multi-label Classification (XMC) is popular framework for solving such problems [4), which formulates these problems as a multi-label classification task with very large number of labels; here each output choice is treated as separate label. A label é often parameterized by its one-versus-all classifier vector w, and the relevance between label l and input x is formulated as w/ (xx), where is an encoding function which maps an input x to its vector representation.
"Tiis work does not relate to Hsiang-Fut Fu's position at Amazon
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Parittio Based Index
ELIAS Index
jeom
Cluster
Feod Libel
O C
Labelc
Figure 1: Traditional partition- based index vs ELIAS index; here an arrow from cluster to a label denotes the assignment of the label to the cluster, arrow width indicates the weight of the assignment. (left) Existing partition based XMC methods use shallow balanced tree as the index structure with label uniquely assigned exactly one cluster, moreover, they initialize the clusters over pre-defined features and keep them fixed throughout the training procedure. (right) ELIAS generalizes the tree based index to a sparsely connected graph index and learns the cluster-to-labe assignments end-to-ent with the task objective during training.
Evaluating w (xx) for every label l in an XMC task can get computationally expensive since the number of labels could easily be upwards of millions. To reduce the complexity, most existing methods employ a search index that efficiently shortlists a small number of labels for an input query and the relevance scores are only evaluated on these shortlisted labels. The quality of the search index plays a pivotal role the accuracy of these methods since a label outside the shortlist will be directly discarded, even if it can be correctly captured by its classifier vector wy. Moreover, the label classifier We a function of the quality of the index as during training, the label classifiers are learned with negative sampling based on the search index. Therefore, how to improve the quality of the search index becomes a key challenge in the XMC problem.
There are two main formulations of the search index: 1) partition-based approach and 2) approximate nearest neighbor search (ANNS) based approach [16]
25
10
partition-
based approach, labels are first arranged into a tree-based index by partitioning the label space into mutually exclusive clusters and then a ML model is learned route a given instance to few relevant clusters. In an ANNS- bbseed approach, fixed, black-bor -box ANNS index is learned on pre-defined label embeddings. Given an input embedding. this index is then used to efficiently query a small set of nearest labels based on some distance/similarin between the input and label embeddings. Both of these approaches suffer from a critical limitation that the index structure is fixed after it's initializedd.
This decoupling of the search index from the rest of the ML model training prevents the search index from adapting with the rest of the model during training, which leads sub-optimal performance.
To overcome this challenge, we propose a novel method called ELIAS: End-to-end Learning to Index and Search, which jointly learns the search index along with the rest of the ML model for multi-label classification in large output spaces. In particular, as illustrated in Fig. [ ELIAS generalizes the widely used partition resed index sparsely connected weighted graph-based index. ELIAS models the discrete cluster-to-labes assignments in the existing partition based approaches as soft learnable parameters that are learned end-to-end with the encoder and classification module to optimize the final task objective. Moreover, because ELIAS uses a graph-based arrangement of labels instead of tree-based arrangement, a label can potentially be assigned to multiple relevant clusters. This helps to better serve labels with a multi-modal input distribution
P2]
Through extensive experiments demonstrate that ELIAS achieves state-of oftte-of-the-ar results on multiple large-scale XMC benchmarks. Notably, ELIAS can be up to 2.5% better at precision @ 1 and up to 4% better at recall@ 100 than existing XMC methods. ELIAS's search index can be efficiently implemented on modern GPUs to offer fast inference times on million scale datasets. In particular, ELIAS offers sub-millisecond prediction latency dataset with 3 million labels single GPU.
Related Work
One-vs-all (OVA) methods: OVA methods consider classification for each label as an independent
binary classification problem. In particular, an OVA method learns L (number of classes) independent
label classifiers (we]k=]. where the job of each classifier Wy is to distinguish training points of label é from the rest of the training points. At prediction time, each label classifier is evaluated and the labels are ranked according to classifier scores. Traditional OVA methods like DiSMEC (2), ProXMI. R.)
and PPDSparse [30] represent each input instance by their sparse oforo features and learn

---

sparse linear classifiers by massively parallelizing over multiple machines. OVA methods achieve promising results on XMC benchmarks but suffer from huge computational complexity because of their linear scaling with number of labels. Subsequent XMC methods borrow the same building blocks of an OVA approach but overcome the computational overhead by employing some form of search index to efficiently shortlist only a few labels during training and prediction.
Partition based methods: Many XMC methods such as Parabel [25], Bonsai [9] XR-Linear (32),
AttentionXML B1. X-Transformer 288.. ER-Transformer [33, LightXML 18 follow this approach where the label space is partitioned into small number of mutually exclusive clusters, and then an ML model is learned to route given instance to a few relevant clusters. A popular way to construct clusters is to perform balanced k-means clustering recursively using some pre-defined input features. Traditional methods like Parabel, Bonsai, and XR-Linear represent their input by sparse bag-of- -for features and learn sparse linear classifiers with negative sampling performed based on the search index. With the advancement of deep learning in NLP,, recent deep learning based XMC methods replace sparse bag-of-word input features with dense embeddings obtained from a deep encoder. In particular, AttentionXML uses BiLSTM encoder while X- X-Transformer, XR-Transformen, and LightXML use deep transformer models such as BERT [ to encode the raw input text. In addition to dense embeddings, the state-of-the-ar XR- Transformer method uses concatenation of dense embedding and sparse bag-of- word features to get a more elaborate representation of the input, thus mitigating the information loss in text truncation in transformers.
ANNS based methods: Methods like SLICE [6), DeepXML (9. and GLaS utilize approximate nearest neighbor search (ANNS) structure over pre-defined label representations to efficiently shortlist labels. In particular, SLICE represents each input instance by its FastText 22 embedding and uses the mean of a label's 's training points as surrogate embedding for that label. It further constructs an HNSW [23] graph (a popular ANNS method) over these surrogate label embeddings. For a given input, the HNSW graph is queried to efficiently retrieve nearest indexed labels based on the cosine similarity between the input and label embedding. DeepXML extends SLICE by learning an MLP text encoder surrogate classification task instead of using a fixed FastText model to obtain input embeddings. GLaS takes a different approach and learns label classifiers with random negative sampling. After the model is trained, it constructs ANNS index to perform fast maximum inner product search (MIPS) directly on the learned label classifiers
Learning search index: There has been prior works 21 that model different types of standard data structures with neural networks. A recent paper 27 models the search index in information retrieval systems
sequence to sequence model where all the parameters of the search index is big transformer model. In a more similar spirit to our work, another recent paper [22] attempts to learn overlapping cluster partitions for XMC tasks by assigning each label to multiple clusters. Even though it serves as a generic plug-in method to improve over any existing partition based XMC method, it still suffers from the following shortcomings: 1) label assignments are not learned end-to-end with the task objective; instead, it alternates between finding the right model given the fixed label assignments and then finding the right label assignments given the fixed model, 2) all labels are assigned to pre-defined number of clusters with equal probability and get duplicated in each assigned cluster, which results in increased computational complexity of the method.
encoded in the parameters
3 ELIAS: End-to-end Learning to Index and Search
The multi-label classification problem can be formulated as following: given an input x E X. predict y {0, 1 where y is sparse L dimensional vector with y = 1 if and only if label is relevant to input x. Here, L denotes the number of distinct labels note that y can have multiple non-zero entries resulting in multiple label assignments to input The training dataset is given in the form of ((x'," : = 1,..., N). XMC methods address the case where the label space (L) is extremely large (in the order of few hundred thousands to millions). All deep learning based XMC methods have the following three key components:
Deep encoder &= : X R which maps the input to a D-dimensional dense embedding through differentiable function. For text input, a popular choice of Φ is the BERT 11] encoder where each input x is represented sequence of tokens.
3

---

BERT encoder
rejected path shortlisted path
8c
Clusters
Labels (>)
Figure 2: Illustration of ELIAS's search procedure: an input is first embedded by the text encoder e to get its embedding d(x). Only few (beam-size) clusters shortlisted based cluster relevance
Se w. ((x). All potential edges of shortlisted clusters are explored and assigned a score based on the product &, * Sc,e (Be,e is normalized form of learnable edge weight parameter de./ between cluster and label shortlisted based their assigned scores and the final label relevance computed as
(). Top-K paths (wi o(x)) * 8,,e** šc, here is sigmoid function. If label can be reached from multiple paths then path with maximal score is kept and rest are discarded.
Search Index I X + R shortlists K' labels along with score assigned to each shortlisted label for a given input x. More specifically, y = I(x) is a sparse real valued vector with only K (< L) non-zero entries and de # 0 implics that label L is shortlisted for input with shortlist relevance score jee As illustrated in Figure | many partition based methods 18 33 formulate their index label tree derived by hierarchically partitioning the label space into C clusters and then learn classifier vectors Wc = [wo]e= (w, E RR)) for each cluster which used to select only a few clusters for a given input. More specifically, given the input x, the relevance of cluster to input x is quantified by cluster relevance scores & w øxx). The top-b clusters based on these scores are selected and all labels inside the shortlisted clusters are returned as the shortlisted labels, where b(< C') is hyperparameter denoting the beam-size.
Label classifiers Wz = [[elli where We E TR. represents the classifier vector for label E and w/ d(x) represents the label relevance score of label é for input x. As explained above, w. ((x) is only computed for few shortlisted labels obtained from the search index I.
3.1 ELIAS Index
ELIAS formulates its label index as specialized weighted graph between root node G, C cluster nodes C = {c) and L label nodes y {L) As illustrated in Figure [ all cluster nodes are connected to the root node and all label nodes are sparsely connected few cluster nodes. ELIAS [a,, parameterizes the cluster-to-label edge assignments by learnable adjacency matrix A where the scalar parameter a,,z denotes the edge importance between cluster c and label
Note that A can be very large for XMC datasets and using dense A will incur O (CL) cost in each forward pass which can be computationally prohibitive for large-scale datasets. To mitigate this we restrict A to be a row-wise sparse matrix i.c. ||a, lo S к where ||-Ilo represents the 4o norm, a represents the ith row of A and is a hyper-parameter which controls the sparsity of A. During training, only the non-zero entries of A is learned and the zero entries do not participate in any calculation. We defer the details of how we initialize the sparsity structure of A to Section 3.4
Existing partition based XMC methods can be thought of special case of this formulation by adding additional restrictions that 1) each label connected to exactly one cluster node, and 2) all cluster-to-label connections have equal importance. Moreover, existing methods initialize the cluster-to-labe adjacency matrix A beforehand based on clustering over pre-defined features and keep it fixed throughout the training procedure. ELIAS overcomes these shortcomings by enabling the model to learn the cluster-to-label edge importance.
3.2 Forward Pass
ELIAS trains the entire model, including the deep encoder 6, the search index parameters W A and the label classifiers W. an end-to-end manner. We now describe the details of the forward pass of ELIAS.

---

Text representation: An input x is embedded by the encoder á into a dense vector representation 6(x). In particular, we use BERT-base as the encoder and represent ((x) by the final layer's CLS token vector.
Query search index: Recall that the goal of the search index I is efficiently compute shortlist of labels y E _t, where y is sparse real valued vector with K' (< L) non-zero entries and Òr * 0 implies that label é is shortlisted for input x with shortlist score ør. Similar to existing methods, ELIAS achieves this by first shortlisting small subset of clusters C C C based on cluster relevance scores defined by w @(x). But unlike existing methods which simply return the union of the fixed label set assigned to each shortlisted cluster, ELIAS shortlists the top-K labels based on the soft cluster-to-label assignments and backpropagates the loss feedback to each of the shortlisted paths. More specifically, ELIAS defines the cluster relevance scores & E RCC as:
Sc [Scle (1, a softmax W (x(x))).
Here hyperparameter a is multiplied by the softmax scores to allow multiple clusters to get high relevance scores. Intuitively, controls how many effective clusters can simultancously activate for given imput (in practice, we keep a xe 10).
Given cluster relevance scores $c, we define set Czoph as the top clusters with the highest cluster relevance scores, where b (< C)) is the beam size hyperparameter. In the training phase, we further define a parent set Cparent to guarantee that the correct labels of are present in the shortlist. More specifically, for each positive label of x, we include the cluster with the strongest edge connection to I Cporent. The shortlisted set defined as the union of these two sets and the selection process can be summarized follows:
Croph arg top-b($c). where b( C') is the beam size,
(2)
(} Cparent Ue farg (ac,e)
during prediction, during training
(3)
Ć = Czaph U Cparent.
After shortlisting a small subset of clusters č, all potential edges of shortlisted clusters are explored and a set P of explored paths is constructed, where P = {0 + - E E and a., > 0}. Furthermore, each path 0 - É P is assigned path score $p.c.e. where the path score $0,e. is expressed as the product of cluster relevance score 9, (defined by Eqn. [ and edge score S which quantifies the probability of label é getting assigned cluster and defined in terms of the learnable edge weight parameter dc,e as follows:
exp(ac.t) and S0.c, Se * Sc,e-
Se, min(1, 3 # a norm) where a
norm
(5)
Ee=1 exp(ac.))
Defining edge scores & in such manner allows mod- elling the desired probability distribution of label assign- ment to cluster, where few relevant labels are assigned to particular cluster with probability 1, and all other la- bels have probability O. Hyperparameter controls how many effective labels can get assigned to cluster, we choose 3 L/C. Figure 3 empirically confirms that the trained model indeed learns the desired edge score distribution with most of the probability concentrated on few labels and the rest of the labels getting assigned low probability. Moreover, this formulation also prevents labels with high softmax scores from overpowering edge assignments because as perS a relevant label é for cluster gets positive feedback for dc,t only if morm < 1/3, otherwise a.,e does not participate in the calculation of 8..e- This allows clusters to learn balanced label assignments. Note that, because of the assumption that A is row-wise sparse matrix, Eqn. can be computed efficiently in O (x) instead of O (L) time.
1.0
score
0.8
0.6
0.4
0.2
0.0
1250 Sorted label edge index
500 750
1000
Figure 3: a (edge weight) distribu- tion averaged over all clusters trained ELIAS model Amazon-670K dataset
5

---

Since there can be multiple paths in P which reach a particular label e, ELIAS defines shortlist score dE for label by the maximum scoring path in P that reaches E, i.e.
Gr max ( e : 0 - c' + +
(6)
Finally, only the top-K K entries in y are retained and the resulting vector is returned as the shortlist for input x.
Evaluating label classifiers: label classifiers [weli=] are evaluated for the K non-zero labels in y and the final relevance score between label & and input x is returned as pr = ow? o(x)) * De. here o is the sigmoid function.
3.3 Loss
ELIAS is trained on a combination of classification and shortlist loss where the shortlist loss encour- ages correct labels to have high shortlist scores (Uz) and classification loss encourages positive labels in the shortlist to have high final score (pt) and negative labels in the shortlist to have low final score. More specifically, the final loss L is defined as L = L. + XCs, where À is hyperparameter and classification loss L. is defined as binary cross entropy loss over shortlisted labels
Le =
Σ E:
(ye log(pe) + (1 ) - log(p)))),
shortlist loss Lg is defined as negative log likelihood loss over the positive labels
C,
Σ
log(ūe).
(8)
Eigle
3.4
Staged Training
Previous sub-sections described the ELIAS framework for learning the index graph along with the ML model in an end-to-end manner. Although, in principle one can optimize the network with the given loss function from a random initialization but we highlight few key challenges in doing so: 1) Optimization challenge: because of the flexibility in the network to assign a label node to various clusters, it becomes hard for a label to get confidently assigned to only few relevant clusters. As a result, the model is always chasing a moving target and for given input it is not able to be sure about any single path; 2) Computational challenge: the full cluster-label adjacency matrix A can be very large for large datasets and will incur O (CT) cost each forward pass if implemented in dense form. To address these challenges we train the ELIAS model in two stages. In the first stage, we only train the encoder o, cluster classifiers Wc, and label classifiers W,, keeping fixed and assigned based on traditional balanced partitions. We then utilize the stage- 1 trained model to initialize the sparse adjacency matrix A. In the second stage, we take the initialized A and rest of the stage 1 model, and jointly train the full model é, Wc, WL.A. A.
Stage 1:
In stage I training, similar to existing partition-based XMC methods, we partition the label space into C' mutually exclusive clusters by performing balanced k-means clustering over pre-defined label features. The adjacency matrix induced by these clusters is then used as fixed assignment for A. Keeping A fixed, we train the rest of the model (i.c. o. Wc. W&) on the loss described in Section 3.3 More details on clustering are provided in Section
the Appendix.
Initializing A: As highlighted before, to overcome the O (CL) cost associated with a full adjacency matrix A, we want to restrict A to be row-wise sparse matrix. In other words, want restrict each cluster to choose from candidate subset of x labels instead of the whole label set. Intuitively, in order for the model to learn anything meaningful, the candidate subset for each cluster should contain approximately similar labels. To achieve this, we utilize the stage model to first generate an approximate adjacency matrix A' and then select the top-x entries in each row of A' as non-zero entries for A. More specifically, we first identify top-b matched clusters for each training point x' by computing the cluster matching matrix M
[mx.elNxc as:
( {0
if c E Giopb. otherwise
(9)

---

where 8 represents the cluster relevance score and Coops represents the set of top-b clusters for 7th training point x'. After computing M, we define the approximate adjacency matrix A' [o = Y where Y = y"..""...... y.I" The element d'..e essentially denotes the weighted count of how many times the cluster got placed in top-b positions for positive training points of label l. Finally, the top x elements in each row of A' are selected as the non-zero parameters of A, i.e.
random (0, 1) i E arg top-oo-)(()) OcEÉ
(10)
otherwise
We choose a large enough R to provide the model enough degree of freedom to learn cluster-to- label assignments. In particular, ~ 10 x L/C works well across datasets without adding any computational burden. For efficient implementation on GPUs, we store matrix A in the form two tensors, one storing the non-zero indices and the other storing the values corresponding to those non-zero indices.
Stage 2: In stage 2 training, we initialize A as described above, and Φ, We from stage 1 model. We then train the full ELIAS model (i.e. o, We, WL, A) end-to-end to optimize the loss defined in Section 3.3
3.5 Sparse Ranker
State-of-the-art XMC methods like XR-Transformer (3 and X- X-Transformer 7 utilize high capacity
sparse classifiers learned on the concatenated sparse bag-of-word features and dense embedding obtained from the deep encoder for ranking their top predictions. Because of the high capacity, sparse classifiers are able represent head labels more elaborately than dense classifiers. Moreover, bag-of- words representation is able to capture the full input document instead of the truncated document that the deep encoder receives.
To compare fairly with such methods, we explore an enhanced variant of ELIAS represented by ELIAS ++, which additionally learns a sparse ranker that re-ranks the top 100 predictions of ELIAS. In particular, the sparse ranker takes the concatenated sparse bag-of-word and dense embedding input features and learns sparse linear classifiers on the top 100 label predictions made from the trained ELIAS model for each training point. Because these sparse classifiers are only trained on 100 labels per training point, they can be quickly trained by parallel lin- ear solvers like LIBLINEAR [[2]. We use the open- source PECOS 32 library to train and make predictions with the sparse ranker.
0.6
Sparse renoo
0.5
0.4
0.2
0.1
0.0
Deciles (decreasing frequency)
Figure True label's score distribu- tion of sparse ranker and ELIAS (4) over different label deciles on Amazon- 670K dataset. 1s decile represents la- bels with most training points while 1hth decile represents labels with least training points
During prediction, the top 100 predictions are first made by ELIAS and then the learned sparse ranker is evaluated these top 100 predictions. We empirically observe that the scores returned by ELIAS and sparse ranker are not well calibrated across different label regimes. As shown in Figure 4 the sparse ranker underestimates scores on tail labels while ELIAS scores are more balanced across all label regimes. To correct this score mis- mis-calibration. we learn simple score calibration module which consists of standard decision tree classifief"] that takes both of these scores and the training frequency of the label as input and predicts single score denoting the label relevance. The score calibration module is learned on a small validation set of 5,000 points. More details on the sparse ranker are in Appendix Section
3.6 Time Complexity Analysis
The time complexity for processing batch of n data-points is On Cd bk + Kd)) where Thet represents the time complexity of the bert encoder, C' represents the number of clusters in index,
tree.
becisiooTreeellesiffeer..
7

---

Table I: Performance comparison extreme classification benchmark datasets. Bold numbers represent overall best numbers that dataset while underlined numbers represent best numbers for dense embedding based methods. Methods which only use sparse bag-of-word features are distinguished by superseript, dense embedding based methods are distinguished by (4) superscript and methods that use both sparse + dense features are distinguished by (++d) superscript
Method
P81
P@3
P@5
PSP@1
PSPØ3
PSP85
P@1
P@3
P@5
PSP@1
PSP@3
PSP@5
Amazon-670K
LP-AAmazo/Titlss1311K
DiSMEÇ
39.70
36.10
30.60
34,20
35.14
23.88
17.24
Parabel(*)
44.89
39.80
36.00
25.43
29.43
32.85
32.60
21.80
15.61
23.27
28.21
32.14
XR-Lineur
45.36
40.35
36.71
Bousajt
45.58
40.39
36.60
27.08
30.79
34.11
34.11
23.06
16.63
24.75
30.35
34,86
Slicele)
33.15
29.76
26.93
20,20
22.69
24.70
30.43
20,50
14.84
23.08
27.74
31.89
47.77
42.79
39.10
33.13
37.12
25.20
18.24
29.22
34,64
39.49
(LaS(4)
46.38
42.09
38.56
38.94
39.72
41.24
AllenbioXML..
47.58
42.61
38.92
30.29
33.85
37.13
32.55
21.70
15.64
23.97
28.60
32.57
LighIXML,00 (9)
49.10
43.83
39.85
38.49
26.02
18.77
28.09
34.65
39.82
Transformereet
50.11
44.56
40.64
36.
38.39
40.99
38.42
25.66
18.34
34,98
Overlap-MMC-
50.70
45.40
41.55
36.39
39.15
41.96
ELIAS 60
50.63
45.49
41.60
32.59
36.44
39.97
39.14
26.40
19.08
30.01
36,09
41.07
ELIAS +40114
53.02
47.18
42.97
34.32
38.12
41.93
40.13
27.11
19.54
31.05
37,57
42.88
Wikipedia-500K
Amazon-33 3M
DiSMEC
70.21
50.57
39.68
31.20
33.40
37,00
47.34
44.96
80
Parabel(s)
68.70
49.57
38.64
26.
31.96
35.26
47.48
44.65
42.53
12.82
15.61
17.73
XR (R-Lineu(())
68.12
49.07
38.39
47.90
45.09
42.96
Bonsai
69.20
49.80
48.45
flicef4)
62.62
41.79
31.57
24.48
27.01
29.07
(s)
73.02
52.02
40.53
30.69
36.48
40.38
GLaS(4)
69.91
49.08
38.35
AttentionXML.C
76.95
58.42
46.14
30.85
39.23
44.34
50.86
48.04
15.52
18.45
LightXMI. (4)
77.78
58.85
45.57
XR- (e-Trasformee/<ed)
79.40
59.02
46.25
35.76
42.22
46.36
54.20
50.81
48.26
20.52
23.64
25.79
Overla--XMC/
52.70
49.92
47.71
18.79
21.90
24.10
ELIAS (4)
79.00
60.37
46.87
33.86
42.99
47.22
51.72
48.99
46. 89
16.05
19.39
21.81
ELIAS (+(++)
62.51
48.82
35.02
45.94
51.13
54.28
51.40
19.07
the embedding dimension, the beam size, k. is the row-wise sparsity of cluster-to-label adjacency matrix A. and K is the number of labels shortlisted for classifier evaluation. Assuming O(VL) and K O(VL), the final time complexity comes out to be O(n(Then + VI(2d + b))). Empirical prediction and training times benchmark datasets reported in Table of f the Appendix.
C' O(VL),, OL/CC)
4 Experimental Results
Experimental Setup We conduct experiments on three standard full-text extreme classification datasets: Wikipedia-500k Amazon-670K, Amazon-3M and one short-text dataset: AmazznnTiiless 131K which only contains titles of Amazon products as input text. For Wikipedia-500Ks Amazon- 670K, and Amazon-3M. we use the same experimental setup (i.e. raw input text, sparse features and train-test split) as existing deep XMC methods [31
33
18
7 For LF-AmazonTTitls- 1311
expertinte we use the experimental setup provided in the extreme classification repository (5). Comparison to
existing XMC methods is done by standard evaluation metrics of precision @ K (P@ K 1, and its propensity weighted variant (PSP@ and basclines with ELIAS at recall @ K (R @ K 10, 20, 100) evaluation to illustrate the superior shortlisting performance ELIAS's search index. More details on the experimental setup and dataset
= 3) [15). We also compare competing methods
statistics are presented in Appendix Section B
Implementation details Similar to existing XMC methods, we take an ensemble of 3 models with different initial clustering of label space to report final numbers. For efficient implementation on GPU, the raw input sequence is concatenated to 128 tokens for full-text datasets and 32 tokens for short-text dataset. Number of clusters C' for each dataset is chosen to be the same as LightXMI. which selects C. C 1/100.. We keep the shortlist size hyperparameter K fixed to 2000 which approximately same as the number of labels existing partition based methods shortlist assuming beam-size b = 20

---

Table 2: Precision and recall comparison of single model dense embedding- methods. ELIAS matches even outperforms the brute-force Bert-OvA baseline while existing partition based methods fail to compare well. especially at recall values.
Method
P@1
P@3
P@s
R@ 10
R@
R@ 100
P@
P@3
P@5
R@10
R@
R 100
Amazon-670K
LF-AmazonTille
BERT-OVA-1 (99)
48.50
43.41
39.67
49.53
56.60
67.90
38.17
25.66
18.44
50.29
54.71
62.80
AttentionXML-1(05
45.84
40.92
37.24
45.59
51.25
60.77
30.26
20.03
14.31
38.16
41.47
47.73
LighiXML-1100 119
47.29
42.24
38.48
47.34
53.26
62.03
37,01
24.88
17.90
48.07
52.10
59.42
XR- 19-Trrnsformer-14)
45.25
40.3
36.45
45.19
51.61
61.11
34.58
23.31
16.79
45.72
49.65
56,00
ELIAS-1(4)
48.68
43.78
40.04
50.33
57.67
68.95
37.90
25.61
18.45
50.12
54.62
62.88
Table 3: Performance analysis of different components of ELIAS, Allowing the model to learn cluster-to-lad assignments significantly improves both precision and recall performance (see row vs row 1). Sparse ranker further improves performance on top predictions (see row 4 row 2).
Method
P@1
P
P@S
R@10
R@20
RØ 100
P@1
P@3
P65
R@10
R@20
R 100
Amazon-670K
Tr-AmmzonTiite--I11
Stage 1
46.63
41.65
37.58
46.08
52.29
61.72
36.96
24.67
17.69
47.69
51.74
58,81
Stuge
$7.67
68.95
25.61
50.12
Sparse manker wio calibration
50.72
45.25
41.27
51.51
58.43
68.95
39.25
26.47
19.02
5144
55. 39
62.88
Score calibration
51.41
45.69
41.62
51.97
58.81
68.95
39.26
26.47
19.02
51.4
55.35
62.88
+ 3x ensemble
53,02
47.18
42.97
53.99
61.33
72.07
40.13
27.11
19.54
53.31
57.78
65.15
and the number of labels per cluster = 100. AdamW [20] optimizer is used to train the whole model with weight decay applied only to non-gain and non-bias parameters. Optimization update for label classifiers WL is performed with high accumulation steps (i.e. optimization update is performed at every & training steps, where k 10) since updating Wr. every step is a computational bottleneck and only few parameters inside Wt gets updated in each optimization step anyway. More details and hyperparameters for each dataset are presented in Appendix Section
B
Comparison on XMC benchmarks Table T compares our method with leading XMC methods such as DiSMEC [2], Parabel (25), XR-Linear (32.. Bonsai [1)) Slice (16), Astec (9], GlaS [33,
AttentionXML. (31], LightXML (8), AR-Transformer [33]. and Overlap-XMC (22). Most baseline
results are obtained from their respective papers when available and otherwise taken from results reported in 33 33 and extreme classification repository [5]. To allow fair comparison among methods that use the same form of input representation, we distinguish methods that use only sparse bag-of-word input features by (+) superscript, methods that use only dense embedding based input features by superscript, and methods that use both sparse + dense features by (s+d) superscript. ELIAS ++ which uses sparse + dense features achieves state-of-the-arr performance on all datasets at precision values while being either the best second best method at propensity scored precision on most datasets. The dense embedding based ELIAS (d) consistently outperforms existing dense embedding based XMC methods by significant margin and on many occasions achieves gains over previous state-of-ther the-art methods which use both sparse + dense features.
Comparison with brute-force OvA baseline To establish the classification performance that could have been achieved if there was no sampling performed by the shortlisting procedure, we implement the brute-force one-versus all baseline BERT-OvA which consists BERT encoder followed by a fully connected linear classification layer and is trained and inferred in one-versus-all fashion without any sampling. We follow the same training procedures as ELIAS for this baseline. Table Table@compare the OvA baseline with ELIAS and leading deep XMC methods such as Attention LiLhtXML. and a dense version of IR-Transformer which uses only dense embeddings, under single model (i.c. no ensemble) setup for direct comparison. Existing deep XMC methods do not compare well against the OvA model especially at recall I 100 but ELIAS matches and sometimes even marginally outperforms, the brute-force OvA baseline while enjoying faster training and inference speed due to the search index.
Component wise ablation of ELIAS Table presents a build-up ablation of performance gains made by different components of ELIAS. The stage model which fixes its adjacency matrix by clustering labels into mutually exclusive clusters performs similarly to existing single model XMC methods. Allowing the model to learn the adjacency matrix A in stage improves recall by up to 7% and precision by up to 2.5% over the stage I model. Adding the sparse ranker and ccore-calibration

---

module further improves model performance on top predictions but the gains diminish as we increase prediction set size. Finally, the ensemble of models improves performance at all evaluation metries which is well observed behaviour with all XMC methods.
Contritution R00000
60
Recall
55
B 11
Xi-Tronffomer-11
40-
Wr-Tranfommen11
25-
20
40
60
80
100
Deciles (decreasing frequency)
Figure 5: (left) Comparison of recall at different prediction set size Amazon-670K (right) Decile- wise analysis of recall@ 100 on Amazon-670K, 1st decile represents labels with most training points i.c. head labels while 10th decile represents labels with least training points i.c. tail labels
Recall comparison Next, we compare the recall capabilities of existing methods with ELIAS. The left plot in Figure plots the recall at different prediction set size for all competing methods and ELIAS. ELIAS strictly outperforms existing methods at all prediction set sizes and in particular, can be up to 4% better at recall @ 100 than the next best method. To further investigate which label regimes benefit most from ELIAS's ' search index we plot the decile wise contribution to recall @ 100 for each method. As can see, ELIAS improves recall performance decile but the most improvement come from the top deciles representing the most popular labels. We hypothesize that because the popular labels are likely to have multi-modal input distribution, existing partition based methods which assign a label to only one cluster fail to perform well on these multi-modal labels. Section C.4 contains additional discussion and results support this claim.
existing methods in each label
5 Conclusion and Discussions
In this paper, we propose ELIAS, which extends the widely used partition tree based search index to learnable graph based search index for extreme multi-label classification task. Instead of using a fixed search index, ELIAS relaxes the discrete cluster-to-labe assignments in the existing partition based approaches as soft learnable parameters. This enables the model to learn flexible index structure, and it allows the search index to be learned end-to-end with the encoder and classification module to optimize the final task objective. Empirically, ELIAS achieves state-of-the-ar performance on several large-scale extreme classification benchmarks with millions of labels. ELIAS can be up to 2.5% better at precision@)) and up to 4% better at recall @ 100 than existing XMC methods.
This work primarily explores many-shov and few-shot scenarios where some training supervision available for each label (output). It would be interesting see how we can adapt the proposed solution to zero-shot scenarios where there is no training supervision available for the labels. One potential approach could to parameterize the cluster-to-label adjacency matrix as a function of cluster and label features instead of free learnable scalars. Furthermore, one limitation of the proposed solution is that it learns a shallow graph structure over label space; this may not be ideal for scaling the method to billion-scale datasets. It would be exciting explore how one can extend ELIAS to learn deep graph structures.
10

---

References
[1] H. Abu-Libdeh, D. Altinbüken, A. Beutel, E. Chi, L. Doshi, T. Kraska, X. Li, A. Ly, and C. Olston. Learned indexes for google-scale disk-based database. CoRR, abs/2012. 125001, 12501
2020.
[2] R. Babbar and B. Schölkopf. DiSMEC: Distributed Sparse Machines for Extreme Multi-label Classification. In WSDM, 2017.
[3] R. Babbar and B. Schölkopf. Data scarcity, robustness and extreme multi-label classification. ML, 2019.
[4] S. Bengio, K. Dembezynski, T. Joachims, M. Kloft, and M. Varma. Extreme Classification (Dagstuhl Seminar 18291). Dagstuhl Reports, 8(7):62 80, 2019. ISSN 2192-5283. doi: 10. 4230/DagRep...7.666 URL http::///ddrpp dagstuhl http://p//ollleeeeee////////0//0
[5] K. Bhatia, K. Dahiya, H. Jain, P. Kar, Mittal, Prabhu, and M. Varma. The extreme classification repository: Multi-label datasets and code, 2016. URL attp://manikvarmo
org/dowmloads/ Ct////C////// XC/XMLeeoository .ttll
[6] C. w. Chang. H. F. Yu, K. Zhong, Y. Yang, and I. S. Dhillon. A Modular Deep Learning Approach for Extreme Multi-label Text Classification. CoRR, 2019.
[7] W.-C. Chang, Y. H.-F. K. Zhong, Y. Yang, and I.-S. Dhillon. Taming Pretrained Transformers for Extreme Multi-label Text Classification. In KDD, 2020.
[8] W.-C. Chang, F-X. Yu, Y.W. Chang, Y. Yang, and S. Kumar. Pre-training Tasks for Embedding- based Large-scale Retrieval. In ICLR, 2020.
[9] K. Dahiya, D. Saini, A. Mittal, A. Shaw, K. Dave, A. Soni, H. Jain, S. Agarwal, and M. Varma. DeepXML: A Deep Extreme Multi-Label Learning Framework Applied to Short Text Docu- ments. In WSDM, 2021.
[10] K. Dahiya, N. Gupta, D. Saini, A. Soni, Y. Wang, K. Dave, J. Jiao, G. K. P. Dey, A. Singh, D. Hada, V. Jain, B. Paliwal, A. Mittal, S. Mchta, R. Ramjee, Agarwal, P. Kar, and M. Varma. Ngame: Negative mining- aware mini-batching for extreme classification. arXiv, 2022. doi: 10.48550/ARXIV.22220
04452. URL https::///aat /wtpxvaxvvvvvv////////aa//aaaava org www.obsg/////eeoo/oe//..oo...ee/ /bs/ 22207004452222
[11] J. Devlin, M. W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. NAACL, 2019.
[12) R.-E. Fan, K.-W. Chang, C.-J. Hsieh, X.-R. Wang, and C.-J. Lin, Lib- linear: a library for large linear classification. In Journal of machine learning research, 2008
[13] C. Guo, A. Mousavi, X. Wu, D.-N. Holtmann Rice, S. Kale. S. Reddi, and S. Kumar. Breaking the Glass Ceiling for Embedding- -Bsedd Classifiers for Large Output Spaces. In NewrIPS, 2019.
[14] N. Gupta, S. Bohra, Y. Prabhu,
S. Purohit, and M. Varma. Generalized zero-shot extreme multi-label learning. In Proceedings of the ACM SIGKDD Conference on Knowledge Discovery and Data Mining, August 2021.
[15) H. Jain, Y. Prabhu, and M. Varma. Extreme Multi-label Loss Functions for Recommendation, Tagging, Ranking and Other Missing Label Applications. In KDD, August 2016.
[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] V. Jain, N. Modhe, and P. Rai. Scalable Generative Models for Multi-label Learning with Missing Labels. ICML, 2017.
[18] T. Jiang, D. Wang, L. Sun, H. Yang, 7. Zhao, and F. Zhuang. LightXML: Transformer with Dynamic Negative Sampling for High-Performance Extreme Multi-label Text Classification. In AAAI, 2021.

---

[19] S. Khandagale, H. Xiao, and R. Babbar. Bonsai: Diverse and shallow trees for extreme multi-
label classification. Mach. Learn., 109(11):2099-23) 2119, nov 2020. ISSN 0885-6125. doi: 10. 00.1077/510994-00-0088888
URI. htt.hhtss://do.. 0ttp://00000000/000000000003..
[20] P. D. Kingma and J. Ba. Adam: A Method for Stochastic Optimization, 2014.
[21] T. Kraska, A. Beutel, E. Chi, Dean, and N. Polyzotis. The case for learned index structures. SIGMOD, 2018.
[22] X. Liu, W.-C. Chang, H.-F. Yu, C.-J. Hsieh, and I. S. Dhillon. Label Disentanglement in Partition-based Extreme Multilabel Classification. In NeurIPS, 2021.
[23] A. Y. Malkov and D. A. Yashunin. Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs. TPAMI, 2016.
[24] T. Mikolov, I. Sutskever, K. Chen, G. Corrado, and J. Dean. Distributed Representations of Words and Phrases and Their Compositionality. NIPS, 2013.
[25] 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.
[26] Y. Prabhu, A. Kusupati, N. Gupta, and M. Varma. Extreme Regression for Dynamic Search Advertising. In WSDM, 2020.
[27] Y. Tay, V. Q. Tran, M. Dehghani, J. Ni, D. Bahri, H. Mehta, Z. Qin, K. Hui, Z. Zhao, J. Gupta, T. Schuster, W. W. Cohen, and D. Metzler. Transformer memory differentiable search index. arXiv, 2022. hRL.ttps:
wwwpxxii..oor/aaaee abs// 220..0222922
[28] T. Wei, W. W. Tu, and Y. F. Li. Learning for Tail Label Data: A Label-Specific Feature Approach. IJCAI, 2019.
[29] L. Xiong, C. Xiong, Li, K.-F Tang, J. Liu, P. Bennett, J. Ahmed, and A. Overwijk. Approx- imate nearest neighbor negative contrastive learning for dense text retrieval. arXiv preprint arXiv:2007.00806 2020.
[30] E. I. Yen, X. Huang, W. Dai, P. Ravikumar, Dhillon, and E. Xing. PPDSparse: A Parallel Primal-Dual Sparse Method for Extreme Classification. In KDD, 2017.
[31] R. You, Dai, Zhang, Mamitsuka, and Zhu. AttentionXMI.: Extreme Multi-Label Text Classification with Multi-Label Attention Based Recurrent Neural Networks. In NeurlPS, 2019.
[32] H.-F. Yu, K. Zhong, Zhang, W.-C. Chang, and I. S. Dhillon. Pecos: Prediction for enormous and correlated output spaces. Journal of Machine Learning Research, 2022.
[33] Zhang, W.-. Chang, H.-F. Yu, and I. S. Dhillon. Fast multi-resolution transformer fine-tuning for extreme multi-label text classification. In NeurlPS, 2021.
12


Reviews:
[{"review_id": "TmFn_GKaGA_", "reviewer": "Reviewer_ubcH", "summary": "The paper addresses a relevant problem in partition based approach to extreme multi-label classification (XMC), where existing methods have 2 stages – the first stage uses a shallow tree-based index to hard partition the label space and the second stage learns classifiers into labels inside that partition. The challenge is that if all the labels are not inside the partition in the first stage, then they cannot be in the output of the second stage at all. This paper introduces a variant, ELIAS, where the first stage is replaced by a weighted graph based index with soft learnable parameters that are learned together with the final task classification objective. Experimental results on popular XMC datasets show reasonable improvements in precision/recall metrics over existing XMC methods. \n\n", "questions": "The authors claim earlier in section 1 that ELIAS is better suited to handled multi-modal label distributions. However, this aspect is not really demonstrated in the paper. They hypothesize in the recall comparison section in section 4 that the increased recall might be due to the  multi-modal distribution of popular labels. Do they have any quantitative evidence to back this claim?", "limitations": "This paper addresses a technical improvement in XMC. As such, I don’t see any significant potential negative societal impact of this specific work.  ", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality: \nThe general idea of training the retrieval stage using an objective (supervised, unsupervised or self-supervised) followed by a classification stage is becoming more commonplace now, with works like REALM for language modeling, and retrieval augmented convolutional networks in computer vision. However, this paper appears to be the first work applying jointly training the retrieval and classification stages in XMC for multi-label classification. The authors have taken a natural next step in evolving this general idea.\nQuality:\nThe paper takes a reasonable approach – motivating the problem, challenges with existing approaches to XMC, proposing a solution, describing the details and testing it on datasets that show the value of this approach. \nClarity:\nThe paper is reasonably well written and easy to understand, with sufficient details. While I have not looked at the code, I am glad that they have shared the PyTorch implementation for transparency and enable reproducibility. \nSignificance:\nXMC is an important problem in a number of domains where the output space of labels is large – many information retrieval problems including search, ads, commerce etc. are in this category. As such the work addresses a significant problem in XMC.\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": ["I don’t know"]}, "extra_text": {}}, {"review_id": "LUi3e6XOQo", "reviewer": "Reviewer_BfTa", "summary": "This paper focuses on the extreme multi-label classification (XMC) problem and proposes a method called ELIAS. \n(1) ELIAS adopts a two-layer index for representing extreme-scale labels, and it learns overlapping cluster partitions by assigning each label to multiple clusters.\n(2) ELIAS adopts a two-staged training strategy, where a XMC model is trained with fixed cluster partition generated by k-balanced clustering in the first stage, the cluster partition is generated according to the weighted count of labels assigning to corresponding clusters according to the XMC model trained in first stage, and the XMC model with value of the cluster partition are trained in the second stage.\n(3) ELIAS++ adopts a sparse ranker and a calibration module for achieving further improvement compared to ELIAS.\nExperiments are conducted on benchmark dataset, where ELIAS achieves better performance compared to previous SOTA methods.", "questions": "How does the loss (i.e., classification and shortlist loss) affect the performance of ELIAS? Is it possible to add an ablation study of $\\lambda$?\n", "limitations": "None", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n+ This paper is well motivated and the proposed method achieves significant performance compared to previous SOTA methods.\n+ The writing is good and easy to follow.\n\nWeakness:\n- It is not rigorous to say that ELIAS is an end-to-end method, though Eq. 7 and Eq. 8 can be optimized in an end-to-end manner in theory. In practice, the model is trained in a two-stage manner, and the label-to-cluster assignment is generated according to Eq. 10, which is non-differentiable.  \n- Experiment results of ELIAS^{(d)} in Table 1 have not be analyzed in Table 3. Is ELIAS^{(d)} corresponds to Stage 1+Stage 2+3 x ensemble in Table 3?\n- Lack of time complexity analysis. It will be better to add a discussion about empirical wall-clock time of the two-stage training process of ELIAS. ", "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": ["I don’t know"]}, "extra_text": {}}, {"review_id": "ArBzhqlbSrd", "reviewer": "Reviewer_i6NX", "summary": "Tree-based methods are amongst the popular approaches for extreme multi-label learning problems. Existing approaches either use a decoupled strategy of training the label classifiers on top of a fixed tree index over labels, or alternate between updating the index and training the classifiers. \n\nThis paper proposes a relaxation of the tree-based index to a directed acyclic graph where each label can belong to multiple clusters with non-zero probability. The probability of label belonging to a cluster is learned together with label classifiers in an end-to-end fashion, and this leads to empirical improvement over using a fixed index over labels.", "questions": "Questions/Suggestions\n\nQ1) Brute-force-OvA: Does this model consist of a BERT encoder followed by a multi-layer perceptron with L outputs? How much time did it take to train this model? Is it possible that BERT-OvA model also suffers from similar optimization challenges as training the proposed model from scratch? If so, then it is not necessarily a strong upper-bound on performance of these models.\n\nQ2) What is the overall training time for the proposed model?\n\nQ3) How were hyper-parameters such as $\\kappa$ chosen? Was k-fold cross-validation used? \n\nQ4) Why have previous XMC papers such as SiameseXML and DeepXML been not compared with?\n\nQ5) The main set of results in the paper compare precision/recall of the proposed model with other baselines. But it might be interesting to have some results/insights/observations about the kind the structure of the index learnt using the proposed approach and how much it deviates from the fixed index from stage 1 training. Some ideas are:\n\n 5a) In the final index structure, what fraction of labels get assigned to multiple clusters? Is it possible to prune assignment of a label to multiple clusters after training without affecting the accuracy of the model?\n\n5b) In “Recall Comparison” para in Sec 4, authors hypothesize that improved performance for top 2 deciles is due to the tendency of popular labels to have multi-modal distribution. Is it the case that in the final index, popular labels tend to be assigned to multiple clusters and rare labels tend to pick a single cluster? \n\n5c) How does row-sparsity affect the performance of the model?\n\nQ6) Which dataset is used for Figure 5?\n\nQ7) Instead of using a separate sparse ranker, why is the proposed model not trained with a combination of dense and sparse features for input as done in baseline methods such X-Transformer, Overlap-XMC?", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n1) Proposed approach allows for learning label-to-cluster assignments in an end-to-end fashion which provides empirical improvements over using fixed label-to-cluster assignments.\n\n2) The paper is well-written and easy to follow. The experiments support the main claim of the paper that jointly learning classifiers and index structure in an end-to-end fashion can yield improved performance over using a fixed index structure. \n\nWeakness\n\n1) Limited Analysis: It would be interesting to see analysis of the proposed approach beyond final accuracy numbers, for example, how much does the final index structure deviate from the fixed index structure used for stage 1 training. \n\n2) Scalability: Since the proposed approach learns a two-level index, it is not immediately clear if the proposed approach can be applied to relax tree structures with three or more levels into a directed acyclic graph. This might limit the scalability of the proposed approach to larger industry scale datasets. However, this is a minor weakness as the approach is shown to scale well to and also outperform baselines on one of the largest publicly available extreme multi-label classification datasets (i.e. Amazon-3M).", "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