Dual-Encoders for Extreme Multi-label Classification
Authors: Nilesh Gupta
Fnu Devvrit
Ankit Singh Rawat
Srinadh Bhojanapalli
Prateek Jain
Inderjit S Dhillon
Venue: ICLR 2024 poster (2024)
Keywords: retrieval
dense retrieval
dual encoders
extreme multi-label classification
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Dual-encoder (DE) models are widely used in retrieval tasks, most commonly studied on open QA benchmarks that are often characterized by multi-class and limited training data. In contrast, their performance in multi-label and data-rich retrieval settings like extreme multi-label classification (XMC), remains under-explored. Current empirical evidence indicates that DE models fall significantly short on XMC benchmarks, where SOTA methods linearly scale the number of learnable parameters with the total number of classes (documents in the corpus) by employing per-class classification head. To this end, we first study and highlight that existing multi-label contrastive training losses are not appropriate for training DE models on XMC tasks. We propose decoupled softmax loss - a simple modification to the InfoNCE loss - that overcomes the limitations of existing contrastive losses. We further extend our loss design to a soft top-k operator-based loss which is tailored to optimize top-k prediction performance. When trained with our proposed loss functions, standard DE models alone can match or outperform SOTA methods by up to 2\% at Precision@1 even on the largest XMC datasets while being 20× smaller in terms of the number of trainable parameters. This leads to more parameter-efficient and universally applicable solutions for retrieval tasks. Our code and models are publicly available [here](https://github.com/nilesh2797/dexml).

Published as a conference paper at ICLR 2024
DUAL- FOR EXTREME MULTI-LABEL CLASSIFICATION
Nilesh Gupta Devvrit Khatri Ankit Singh Rawat
Srinadh Bhojanapalli'
Prateek Jain Inderjit D
The University of Texas at Austin Google
*Google Research
ABSTRACT
Dual-encoder (DE) models are widely used in retrieval tasks, most commonly stud- ied on open QA benchmarks that are often characterized by multi-class and limited training data. In contrast, their performance in multi-label and data-rich retrieval settings like extreme multi-label classification (XMC), remains under-explored. Current empirical evidence indicates that DE models fall significantly short XMC benchmarks, where SOTA methods (Dahiya et al., 2023a by linearly scale the number of learnable parameters with the total number classes (documents in the corpus) by employing per-class classification head. To this end, we first study and highlight that existing multi-label contrastive training losses not appropriate for training DE models on XMC tasks. We propose decoupled softmax loss a simple modification to the InfoNCE loss that overcomes the limitations of existing contrastive losses. We further extend our loss design to soft top-k operator-based loss which is tailored to optimize top-k prediction performance. When trained with our proposed loss functions, standard DE models alone can match or outperform SOTA methods by up to 2% at Precision @ even on the largest XMC datasets while being 20x smaller in terms of the number of trainable parameters. This leads to more parameter-effictt
and universally applicable solutions for retrieval tasks. Our code and models are publicly available here
INTRODUCTION
Dual-encoder (DE) models have been highly successful in dense retrieval tasks for open-domain question answering (openQA) systems (Lee al. 2019 Karpukhin et al.) 2020 Qu al., [202T). where they efficiently retrieve relevant passages or documents from a large corpus given a user s query, These models map both queries and documents into a shared embedding space, enabling efficient retrieval using fast similarity search methods (Guo et al., 2020: Johnson et al., 202T). Notably, openQA benchmarks are often characterized by multi-class (most queries are tagged with a single positive document) and limited (each document in the corpus has zero or very few tagged queries i.e. zeroffew-shot scenario) training data. On these benchmarks, models are required to generalize and retrieve relevant documents even though they might not have appeared in the training set.
Another important retrieval scenario that frequently arises in real-world applications such as search engines [Mitra & Craswell (0188) and recommendation systems (Covington et al.] 2016 Zhang ct al. 2019) is where we want to perform retrieval for multiple documents/items from a large collection based a significant number of seen examples per document/item, i.e., a many-shot scenario. Such tasks are often formulated as extreme multi-label classification (XMC) tasks, where each document/item also appears in other emerging applications such as retrieval augmented generation (RAG) (Lewis)
be retrieved is considered as separate label (Bhatia et al., 2016). XMC
al. 22020). In particular, when multiple documents can potentially contain the desired answer, the retrieval task in RAG closely aligns with an XMC problem. Typically, XMC algorithms need to both "memorize" and "generalize". That is, for each label they need to memorize the type of queries that are relevant; e.g., for product to product recommendation scenario, the algorithm should memorize which products can lead click on a particular product using the provided product-produt co-click data. At the same time, the algorithm should generalize on unseen queries.
"Correspondence: ni Teeeeh@ac www.ea............. edd

---

Published as a conference paper at ICLR 2024
It is a prevailing belief in the XMC community that due to the semantic gap and wewwledge-intensiig nature of XMC benchmarks, DE by themselves are not sufficient to attain good performance (Dahiyal et al. 2023b). As a measure to overcome the semantic gap and enable memorization, bfateoo-the-an art
(SOTA) XMC methods augment DE with either per-label classifiers (Dahiya et al. 2023a Mittal
et al. a02 auxiliary parameters (Dahiya et al.] 2023b). We explore this belief by performing
a simple experiment on synthetic dataset where a random query text is associated with a random document text, and the task is to memorize these random correlations during retrieval. We find that DE models are able to perform this task with perfect accuracy at least on up to 1M scale datasets (cf. Section
C.I). disputing the previously held belief in the literature.
In light of the aforemmentioee synthetic experiment, this work aims to answer the following question: "Are pure DE models sufficient for XMC tasks?" A Pure DE model is desirable be- cause: 1) unlike XMC methods, DE methods are parameter efficient i.e. the model size (consequently the number of train- able parameters) does not scales linearly with the number of labels (see Figure 1. 2) XMC methods require re-training or model updates on encountering new labels. In contrast, DE methods can generalize new labels based their features. Interestingly, our work shows that pure DE models can indeed match or even outperform SOTA XMC methods by up to 2% even on the largest public XMC benchmarks while being 20x smaller in model size. The key the improved performance is the right loss formulation for the underlying task and the use of extensive negatives to give consistent and unbiased loss feedback. We refer to DE models trained with the proposed approach as DEXML (Dual-Encoded for eXtreme Multi-Label classification).
58-
MGAVE
56-
54
Decouieniootmma SPM hard
10
4B
4B-
CANC
20 ony
1.0
bumber of training parameters
Figure 1: Number of trainable pa- rameters used by different models and their Precision performance on Fr-AmazonTitles-- 1,3M dataset (Bhatia ct al. 2016
We begin by analyzing existing multi-label and contrastive training losses like One-versus-Al Binary Cross-Entropy (OvA-BCH and InfoNCE, highlighting how they may be inadequate for training dual-encoder models in XMC. Specifically, we note that OvA-BCE does not train effectively and InfoNCE disincentivizes confident predictions. To overcome these limitations, we propose a simple modification to the InfoNCE loss the DecoupledSoftmax loss which removes the undesirable correlation between positive labels during training. We further extend our loss design to Soft Top-k loss which is tailored to optimize prediction accuracy within specific prediction budget size (k). This can be particularly relevant to XMC applications where the goal is to make a fixed number of highly accurate predictions from a vast set of possible labels. Moreover, we note that XMC problems involve a long tail of labels, which at the query level manifests as large number of highly relevant negatives. To this end, in order to establish the full potential of DE methods, we first study our proposed loss using all the negatives in the loss term. Naturally, including all negatives is challenging for million-label scale datasets. To this end, we provide a memory
distributed implementation that can use multiple GPUs and modified gradient cache (Gao et al.) 2021 approach to scale DE training to largest public XMC benchmarks even with a modest GPU setup (see Table T [ We further study the performances and implications of approximations of the proposed loss functions using standard hard negative mining approaches. In summary, the contributions of this work include:
We demonstrate that dual-encoder methods alone can achieve SOTA performance on XMC tasks leading to a more parameter-efftcieet and generalizable approach for XMC. This also unifies the approaches for XMC and retrieval problems, paving the way for simplified and universally applicable solutions.
We analyze and show existing multi-label and contrastive training losses not appropriate for training dual-encoder models in XMC settings. Specifically, we note that OvA-BCE does not train effectively and InfoNCE disincentivizes confident predictions.
We propose a simple modification to the InfoNCE loss, namely DecoupledSoftmax loss, which overcomes the limitations. We further extend our loss design to SoftTop-k loss which is tailored to optimize prediction accuracy within a specific prediction budget size.
Lastly, we explore applications of our proposed loss functions which shows promise for our proposed approach over standard encoder training in RAG.
RAG task in Section C99
Before providing a detailed account of our contributions in Section 4 and 5 begin by discussing relevant literature in Section 2 and introducing necessary background in Section 3
2

---

Published as a conference paper at ICLR 2024
2 RELATED WORK
Extreme classification methods. There is large body of literature utilizing sparse features or
dense pre-trained dense features to develop XMC methods (Babbar & Schölkopf] 2017 Prabhu et al.
2018
Prabhu& & Varma 2014]ooulin
[Juuln et al., [017] Yen et al. 2016|[2017 Guppt et al. [021]] Yu et al. Lately, task-specific dense features, e.g. based on BERT (Devlin et al. | 99. have been also
2022
utilized to design solutions for XMC problems (see, e.g.. /hang et al. 20188] You et al. 2019:] Gupta
et al.
2019 [Guo et al. [019 Medini et al., 2019) Chang et 20200 et al. 2021] Ye et al.
2020 Zhang et al. [02 Gupta et al., 2022 and references therein). However, these approaches
treat labels as featureless ids and employ classification networks that assign a unique classification
vector to each label. Given their increasingly popularity in information retrieval literature (Lee et al.,
2019; Karpukhin et al., 2020) Luan et al. [021] Qu et al.] 2021 Xiong et al. (02T). DE methods that employ two encoders to embed query and items into a common space have been increasingly explored for XMC settings well. By leveraging the label features such as label text (Mittal et al.
2021a Dahiya et al. 202I). images (Mittal et al., 2022), and label (correlation) graph (Mittal et al. 2021b Saini et al. (021). such methods have already shown improved performance for the settings involving tail labels. That said, the current SOTA methods such as NGAME (Dahiya et al., (023a). do not solely rely on a DE model. For instance, NGAME's two-stage approach involves first training query representations via a DE approach, and then utilizing a classification network in the second stage. DEXA (Dahiya al., 2023b), is a recently proposed approach which builds on NGAME to improve the encoder embeddings by augmenting them with auxiliary parameters.
Loss functions for multi-label classification problems. One of the prevalent approaches solve multi-label classification problem involves reducing it to multiple (independent) binary (Dembozyński al.] 2010) or multiclass (Joulin et al.) 2017 Jernite et al., 2017) classification problems. Subse- quently, one can rely on wide range of loss functions designed for the resulting problems, aimed at
minimizing the top-k classification error (Lapin et al., 2016 Berrada et al,, 2018:) Yang & Koyejo 2020; Petersen et al. (022). Since different queries have varying number of relevant documents, 2020:Petersen
Lapin et al. (2016); 6); Yang & Koyejo (2020) identified desirable loss functions that are simultaneously calibrated (Bartlett et al. |0006] Zhang 2004) for multiple values of k: Another popular approach to solve a multi-label classification problem relies on (contextual) learning to rank framework (Liu let al. 2009) where (given query) the main objective is to rank relevant documents above irrelevant ones. Various surrogate ranking loss functions, aimed at optimizing the specific ranking metrics such as Precision @k and NDCG@K k. have been considered in the literature (see e.g.. [sunier al. 2009
Kar et al. 2015 Li et al. 2017 Jagerman et al. 2022 Su et al.
2022 and references therein).
Negative sampling. Modern retrieval systems increasingly deal with large number of items, leading to increased training complexity. Negative sampling (Bengio & Senecal 22088] Mikolov et al. 2013a
Mnih & Teh 2012. Mnih & Kavukcuoglu 2013) helps mitigate this via employing a loss function restricted to the given relevant documents and subset of negatives or irrelevant documents. Selection of informative negatives is critical for training high-quality models. For classification networks,
Bengio & Senecal (2008) studied negative sampling for softmax cross entropy loss by relying on importance sampling. Recently, many works have extended this line of work, e.g., by leveraging quadratic kernels (Blanc & Rendle, 2018) and random Fourier features (Rawat al., 2019). As for the generic loss function and model architecture, is common to employ sparse retrieval mechanism, e.g..
BM25 (Robertson et al. 1995 2004), to obtain negative documents (Luan et al.
102 Karpukhin 2017 Karpukhin
et al. 2020). Another strategy is to leverage in- in-batch negatives (Henderson et al.
al. 2020 for given query in a mini-batch, items appearing as positive labels for other queries in the mini-batch are treated as negatives. However, in-batch negatives often constitute "easy" negatives
(with very small contribution the overall gradient) and do not aid the training much (see, e.g., Karpukhin et al.) 2020: Luan et al., 2021). Same holds for the negatives sampled according to uniform or unigram distribution (Mikolov al. 2013b). |eddi al. (2019) proposed sampling large uniform subset of items and weighting those according to their hardness for each query. Taking this approach to an extreme, for each query, ANCE (Xiong et al. (02T) selects the hardest negatives
from the entire corpus according to the query-document similarity based on a document index.
3 BACKGROUND: MULTI-LABEL CLASSIFICATION
As discussed in Section D the (many-shot) retrieval problem explored in this work is essentially multi-label classification problem. Assuming that Q and DD denote the query and document (or
3

---

Published as a conference paper at ICLR 2024
label) spaces, respectively, the underlying query-document relevance is defined by the distribution where P(d[g) captures the true relevance for the guery-document pair (g. d) E Q x D. In this paper, we assume that the document space is finite with a total [D] L documents. Let the training data consists of N examples S := {1 NeN S x {0, 1)2, where, for j E [L] Yig = iff jth document in D is relevant for query 9i- We also denote the set of positive labels for di as P. P [j Yi. }}
Dual-encoder models and classification networks. Note that learning retrieval model is equivalent to learning a scoring function (or simply scorer) 8a : Q x D - R. which is parameterized by model parameters 0. DE model consists of a query encoder fø : Q - F and document encoder hig: : D + RR that map query and document features to d-dimensional embeddings, respectively. Accordingly, the corresponding scorer for the DE model is defined as: sø d) = (a za) (os(d) (d)), where 0 = ( T/] and zq. 4d E IR Note that it is common to share encoders for query and document, i.e., Φ eb. We also focus on this shared parameter setting in our exploration.
Different from DE models, classification networks ignore document features. Such networks consist of query encoder Iø Q + E and classification matrix W E IR' with ith row the classification weight vector for the ith document. The scorer for the classification network then becomes so so d) (4), Wa), where wa denotes the classification vector associated with document d and 0 = ( W.. For case of exposition, we do not always highlight the explicit dependence trainable parameters 0, Φ, vb while discussing DE models and classífication networks.
Loss functions. Given the training set S, one learns a scorer by minimizing the following objective: (1)
() Σ iEIM y;:
1
N
where l is a surrogate loss-function for some specific metrics of interest (c.g., Precision @k k and Recall k).. Informally, ((g, y; serves as a proxy of the quality of scorer for query q, as measured by that metric with y as ground truth labels. One of the popular loss functions for multi-lahel classification problem is obtained by employing one-vs-all (OvA) reduction convert the problem into L binary classification tasks (Dembczyúski et al. (2100). Subsequently, we can invoke a binary classification loss such as sigmoid binary cross-entropy (BCE) for L tasks, which leads OvA-BCE:
s, = - -,,, jE(L]
yy log
1 1 (0q,dd)) (1 4;) log
+ cs(q,d,)
Alternatielyy, one can employ a multi-label to multi-class reduction (Menon et al. 2019) and invoke softmax cross-entropy (SoftmaxCE) loss for each positive label:
Elg.y, 8) = E JE(L] Yj log (e((9dd)) / el(q))))
(3)
Our empirical findings in Section 5.5 reveal that DE models fail to train with BCE loss. We hypothesize that this might be due to the stringent demand of the BCE loss, which requires all negative pairs to exhibit high absolute negative similarity and all positive pairs to have a high absolute positive similarity a feat that can be challenging for DE representations. In the section below we analyze SoftmaxCF loss (an extreme case of InfoNCE) in more detail and highlight its shortcomings.
4 IMPROVED TRAINING OF DUAL-ENCODER MODELS
In this section, first discuss the loss function typically used for training dual-encoders, and its shortcomings for multi-label problems. We then suggest simple fixes to address the problems, argue why these changes are necessary, and discuss practical implications of our modifications. Finally, we discuss our proposed soft top-k loss formulation and discuss its implementation.
4.1 LIMITATIONS OF STANDARD CONTRASTIVE LOSS FUNCTIONS FOR EXTREME MULTI-LABEL PROBLEMS
Existing dual-encoder methods generally rely on contrastive losses for training. For the rest of the paper, we will anchor technique, experiments, and discussion on the popular InfoNCE contrastive learning loss (Oord et al. 2018), but the observations and the approach apply to other contrastive learning loss functions considered in Xiong et al. (2021): Dahiya et al. (202) 2023a). In multi-label
4

---

Published as a conference paper at ICLR 2024
setting, given batch of queries B and their positive labels, InfoNCE loss takes the form:
e'(()dy))
((s; Σ )i)
Σ
Σ
Mij Mij-loy log ed(aidd)) + + etla.d
9.EB
9:EB JELL]
where sq; ) is the similarity score for query q and document dj. D denotes the set of (hard) negatives created for the batch B. That is, for in-batch negatives, DD- represents union of all positive labels for all the queries yi E B. In the extreme case when we consider all labels in D 44 becomes the same as the SoftmaxCE loss in (3). A quick analysis of the gradients of this loss function reveals that:
[11 1P..PPo
P otherwise
Jj
Σx ex(nnd))) where s(q, d) x
Ozd:
This implies that the loss function is optimized when all positive labels score du [PTT while all negative labels achieve σu = O. While this approach appears to be appropriate for elevating the ranking of positive labels above negatives, it raises concerns when applied to imbalanced real-world XMC datasets. In these datasets, not all positive labels possess equal ease of prediction from the available data. To illustrate this issue, consider scenario where one positive label, denoted 1, is straightforward to predict for given query qe while the remaining positive labels are characterized by ambiguity. Consequently, the model assigns high score to the unambiguous label 7. However, this leads to an increase in the softmax score for label I, resulting in on exceeding [PT Consequently, the loss function begins penalizing this particular positive label. Given the inherent label imbalance XMC datasets, the uniform scoring of all positive labels, and the subsequent penalization of confidently predicted positives, do not align with the ideal modeling strategy.
Decoupled multi-label InfoNCE formulation. Based on the above observations, we re-formulate the multi-label InfoNCE loss as:
es(a;,,))
(DecoupledSoftmax)
f(g.. y1:) Y(: = Σ Moy log cx(g.dd)) + ElE(L - yut) eslaidd)
(5)
SE(L]
where s(g. d) model assigned score to the query-document pair. Note that we exclude other positives from the softmax's 's denominator when computing the loss for a given positive, and add all the negatives in the denominator to eliminate the bias/variance the estimation of that term. Analyzing the gradients with the proposed change gives:
0e,
(1 o.))T
if P. otherwise
Ozd.
EjEP, (alaiit) Exer, d Lq.
(((()))))
here oij d0())d + EneR, (x(a)d where s(g, d) = x ad
It simple see DecoupledSoftmax is optimized when we have σa 1 i P. (for positives) and when the score for a negative s(g, di) is significantly smaller than all positive scores, i.e., s(g, d,) for j E P.. Thus, DecoupledSoftmax avoids penalizing positives that are confidently predicted by the model and allows the training to better handle the imbalanced multi-label nature of XMC datasets.
We empirically validate the efficacy of our proposed modifi- cation in addressing the aforementioned issues. To achieve this, we construct a synthetic dataset that accentuates the chal- lenges and demonstrate how modification mitigates these challenges, in contrast to its absence. Specifically, we create this synthetic dataset follows: We randomly sample 1000 training query texts and 5000 label texts. For the initial 100 training queries, we intentionally designate the first token as "+*" (a randomly selected token) and associate all these queries with the first five labels i.c. labels 0, 3, and 4. However, for the Oth label, we append the token "ts" to its label text. This augmentation aims facilitate the model's learning of the pattern that whenever "t*" is present in a query, label O should
Teat tresiiions
62
Sotrmox
360
Trainiing Esoch
Figure 2: Decoupled softmax vs stan- dard softmax on synthetic dataset.
5

---

Published as a conference paper at ICLR 2024
be confidently predicted. The test set comprises 1000 new, randomly generated query texts, wherein we again replace the first token with "t3" and exclusively tag these queries with the Oth label.
As illustrated in Figure 2 our proposed modification results in 100% P@L this dataset, while the standard softmax approach achieves approximately 20% P@L. This stark contrast highlights the softmax's 's inability capture this straightforward correlation.
We note that similar situations arise in real-world XMC datasets. For instance in EURLex-4K dataset, labels such as "afghanistan." 'international sanctions," f and 'cfsp' frequently co-occur in queries. However, "afghanistan" is easily predictable, as training queries consistently contain the "afghanistan" keyword, whereas the predictability of the other labels is not trivial. Thus, the observations in Figure 2are are not limited to artificially crafted scenarios, but extend to practical XMC datasets, emphasizing the relevance of our proposed modification. Figure visualizes the gradients for two labels from EURLex-4K dataset for all positive training queries encountered during training. This plot reveals that regular softmax loss gives very high variance gradient feedback, but decoupled softmax gives much more consistent gradient feedback and positive feedback.
Sottrar
2000
4000
Poctive rriliing Samppl
Pautive
Figure 3: Gradient analysis of two labels [left] "afghanistan"" Echerry-pickede and [righr]"data datall"ddaaaaaaaaa "rtaanmmmmision"" (randomly-chosen) EURLex-4K dataset for all positive training queries encountered during training.
4.2 MEMORY EFFICIENT TRAINING
Training a DE model with the above loss function for a very large number of labels is challenging. As the pool of labels to be considered in the loss computation grows, the computational requirements to process these examples increase significantly. - as for each query and label, the model needs to compute the embeddings and store all the intermediate activations facilitate backpropagation. In order to overcome these memory constraints, we utilize a modified version of gradient cache approach (Gao al. 2021) perform distributed DE training with large pool of labels under limited GPU memory. Due to space constraints, we describe our implementation and its runtime/memory analysis in detail in Section B.5 of the appendix.
Although, with the memory-efficient implementation we are able to train DE with the full loss even on the largest XMC benchmarks and establish the capabilities of DE models for XMC tasks with modest GPU setup (see Table 9). the training is still expensive and is not desirable for scaling to even bigger datasets. To this end, we explore (see Section (5.4) approximations of our proposed approach with standard negative mining techniques similar to Xiong et al. (2021) and evaluate the performance of the learned DE model with the increasing number negative samples. Furthermore, in Section CE following the embedding cache approach proposed in Lindgren et al. (2021). also study how the quality of the sampled negatives affects the model performance.
4.3 PROPOSED DIFFERENTIARIE TOP-K OPERATOR- -BASED LOSS
Many-shot retrieval systems are applicable in a variety of fixed-size retrieval settings like say product- product retrieval similar to LF-Amazon Titlonniiieee--I- fitles-1.3M 1.3M dataset where the goal is to retrieve top-k products for given input product. While our modification to InfoNCE loss does tend to give higher scores to positives than the negatives, it might be bit too harsh in certain cases. For example, consider there are five positives and is difficult to rank all the five positives over all the negatives. Suppose the goal is to optimize Recall @ 10. Then long the five positives are ranked above all but five negatives, the loss should be small (ideally zero). However, the proposed loss in 5 would still penalize such solutions.

---

Published as a conference paper at ICLR 2024
To explicitly optimize such top-k metrics, we propose to use differentiable soft top-k operator to define the loss function. The soft top-k operator takes a score vector as input, and essentially serves as filter, assigning values close to 1 for the top-k scores and values close to 0 for the remaining scores. is designed to be differentiable to allow gradients to flow through it during backpropagation.
Once we have the soft top-k score vector, we apply log- likelihoood loss over this vector. This loss serves to provide the model with feedback on its performance in placing all positives among the top-K predictions. When all positive labels are among the top-k predictions, the loss is zero. However, if any positive label falls outside the top-k scores, the loss provides a non-zero feedback signal that guides the model to adjust its parameters during backpropagation. That is, given query qs. let the score vector s; := (s(q, d1),... ..... s(q,, dz)) E R' and z; := Soft (op-k(s;) be the output of the soft top-k operator. With y, {0, 1 as the label vector, the loss function for qi takes the form
(SoftTop-k) ((;. - Σ SE((.] Vij log Rig.
(6)
Since zij depends on the full score vector 8;, this loss also provides feedback to all the positives and negatives in the dataset. To complete the loss function formulation, we specify the differentiable soft top-k operator below. Note that the operator is based on a proposal by Ahle 2022).
Soft differentiable top-k formulation. We first look variational definition of hard top-k operator. {s(2z + fx)] where saa) = [a >> 0]
Let x { be the input vector, then Hard the step function that outputs if 0 and 0, otherwise. tx a threshold so that exactly elements exceed it, i.e., E sax; + (x) = k;; for simplicity, let x comprise unique elements.
Based on Hard operator, we can now define Soft operator by essentially replacing the above hard step function with a smooth a function such as the standard sigmoid function. We can change the smoothness of our soft approximator by changing the function. We use scaled sigmoid function i.c. (zz) = sigmoid (axr) where a > 0 is a hyperparameter that controls the smoothness. In general, we might not have closed-form solution for ix depending our choice of However, for monotonic a, we can find ix using binary search: start with an under- and over-estimate of tx and at every iteration reduce the search space by half based comparison of o + (x) and k.
Computing gradient of SoftTop-k operator. By definition, Sof- = {a + 4))) where lx is computed using binary search. Now,
0tx 0 (x: + tx)
Σ o( +tx) =
Dz. Σ, 0'( + (x)
Hence, @SoftTop- o (xx+ [=
0tx d' d' + (x)+ [1]1]
o' d' + tx) E' +xx)
+xx)
(7)
Ox
0m
Combining definition of Soft Top-k operator along with associated loss function in 66 and by using gradient of Soft in (7), we now get an approach- - Top- DE - to train DE models in extreme multi-label kind of scenarios. Note that here again we have to compute gradient wrt all negatives; thus, requiring the memory efficient and distributed backpropagation implementation from Section 4.2
5 EXPERIMENTS
In this section we compare DE models trained with the proposed loss functions to XMC methods. We show that the proposed loss functions indeed substantially improve the performance of DE models, even improving over SOTA in some settings. We also present ablations validating choices such as loss functions, number of negatives, and further provide a more exhaustive analysis in Section of the Appendix
5.1 DATASETS AND EVALUATION
We experiment with diverse XMC datasets, EURLex-4K, AF-AAmzzoTTiles-133 131K, LF-iiiipedia- 500K, and LLF-AmazoTTiles--IA each containing label texts. These datasets cover various applica- tions like product recommendation, Wikipedia category prediction, and EU-law document tagging.
Note that, with slight abuse of notation, we use Soft Soo-Top-k() and Soft to refer to soft version of the top-k operator and the corresponding loss function, respectively.
7

---

Published as a conference paper at ICLR 2024
Table 1: Performance comparison on large XMC benchmark datasets. Methods suffixed with x3 use an ensemble of learners. Bold entries represent overall best numbers while underlined entries represent best numbers among pure DE methods. Results for most existing XMC baselines are from either their respective papers or the XMC repository (Bhatia et al. 2016), blank entries indicate source does not have those numbers.
Label Text
Label Classifier
Method
P@1
P@5
N@5
PSP@5
P@1
P@S
N@5
PSP@
LF Wikipedha-500K
LF-AmazonTileee 1.3M
DistilBERT OvA Classifier
X
V
82.00
48.54
73.44
48.87
48.72
39.09
44.79
25.91
ELJAS
X
81,96
48.90
73.71
48.92
49.26
39.29
45.44
27.37
XR- URRTransformer
V
81.62
47.85
72.43
50.14
46.59
27.79
ECLARE x3
V
V
68.04
35.74
56.37
38.29
50.14
40.00
46.68
30.56
SiameseXML x3
V
67.26
33.73
54.29
37.07
49.02
38.52
45.15
32.52
NGAME
V
84.01
49.97
75.97
57.04
56.75
44.09
52.41
35.36
NGAME DE
X
77.92
40.95
57.33
45.82
35.48
36.80
DEXA
V
84.92
50.52
76.80
58.33
56.63
43,90
52.37
34.86
DEXA DE
X 79.99
42.52
57.68
51.92
38.86
37.31
DecoupledSoftmax DEXML
X 85.78
50.53
77.11
58.97
58.40
45,46
54.30
36.58
Softoo-55 DEXML
X
82,93
51.11
76.61
59.55
54.41
43.82
51.70
32.20
We follow standard setup guidelines from the XMC repository (Bhatia al. For EURLex-4K, we adopt the XR-Transformer setup due to unavailable raw texts in the XMC repository version. Dataset statisties are in Table 6 and summary is in Table 5
(016) for LF_* datasets.
Model evaluation
employs standard XMC and retrieval metrics (P@) 1, P@5, PSP@5, nDCG@5 R@ 10, @@ 100). Details are in Section B.4
5.2 BASELINES AND SETUP
We compare our approach with SOTA XMC methods such DEXA (Dahiya et al. 2023b),
NGAME (Dahiya et al. 2023a XR-Transformer ZZang et al. 2021 ELIAS (Gupaa et al. 2022).
ECLARE (Mittal et al. 2021b), SiameseXML (Dahiya et al.) (202)) To keep the comparison fair, transformer (Sanh et al. 2019) as used in NGAME in all of our dual-encoder experiments and share the encoder parameters for both query and document. We also compare against encoder only version of NGAME and DEXA (represented by NGAME DE and DEXA DE), i.e., retrieval using only encoder representation. direct comparison between pure classification model trained with all negatives and our proposed approach, similar to the BERT-OvA approach in Gupta et al. [2022] Table 2), we also report results for the Distil-BERT OvA classifier model which stacks classification layer with L labels top of dist wiilbbrrtbaa baase transformer and fine-tunes both the classification layer and the transformer by using the full one-vs-all BCE loss in 2 See Appendix B for more details on the experimental setup.
we use the same 66M parameter
order to have
5.3 COMPARISON WITH XMC METHODS
Table prreeents the main comparison results, showcasing the performance of our approach relative to existing methods. On large-scale datasets such LF- L.F-Wikppddia-500X and I.F-AmazonTiteee 1.3M, our approach significantly outperforms all existing XMC methods. More specifically, it can be upto 2% better than the next best method on P@ and up 2.5% better on PSP@k while being 20x smaller (in terms of trainable parameters) than existing SOTA method on LF-Amazon Titoonniieee--te 1.3M (cf. Figure T]. Moreover, compared to the DE versions of NGAME and DEXA our approach can be 6% better at P@k.
smaller datasets such as EURLex-4K, TR-rransformer outperforms our approach but as we demonstrate in Section[[] its gains are mainly because it uses ensemble of 3 learners and sparse classifier-based ranker. Such add-ons can be incorporated into any method. On AF-AmzoonTtlls-- 131K, our approach is comparable to DEXA/NGAME P@ 1, which can be partially explained by a significant jump in performance at P@ achieved by using a label propensity aware score fusion module in NGAME; see Dahiya et al. (2023 Table 8).
P@5 but there is considerable gap
5.4 DecoupledSoftmax WITH HARD NEGATIVES
Since our proposed DecoupledSoftmax loss uses all negatives in its formulation, we investigate approximation of this loss by sampling only a few negatives per train query the batch; see Table
8

---

Published as a conference paper at ICLR 2024
Table 2: Performance comparison on small extreme classification datasets. Please refer Table for notations.
Label Text
Label Classifier
Method
P@I
P@5
N@5
PSP@5
P@
P@5
N@s
PSP@5
EURLex-4K
LF- FF-AAazooTiilss 131K
DistiBERT OvA Classifier
K
V
85.25
59.96
69.83
53.72
37.55
18.49
40.75
40.18
ELIAS
X
/
86,88
61.73
71.65
55.62
39.26
19.02
42.23
41.89
XR- ERR-ransformerx:
X
V
88.41
63.18
38.10
18.32
40.71
39.59
NGAME
46.01
21.47
48.67
49.43
NGAME DE
X
42.61
20.69
48.71
DEXA
46.42
21.59
49.00
49.65
DEXA DE
X
44.76
21.18
49.50
DecoupledSoftmax DEXML
X
86,78
60.19
70.37
54.78
42.52
20.64
46.33
47.
SofTop-5 DEXML
X
83,42
60.95
70.07
56.84
41.11
21.23
46.58
48.70
Following Xiong et al. 2021), for mining the hard negatives we compute the top 100 nega- tive shortlist at regular intervals and sample the negatives from this shortlist. For each query we train on the full pool of documents sampled in the batih, as a result, for a mini-batch of size B and number of hard negatives sampled per query m, we compute the loss on almost B(1 + m) documents. Clearly from Table 3 increasing the number of effective labels consid- ered in the loss significantly increases the per-
Table Ablation of number of negatives sampled per query in the DecoupledSoftmax DEXML training.
Batch size
Hard neg
Effective doc pool
P@
P@5
R@ 10
R@ 100
per query
I.F.Amazon Titles 1.3M
81992
32.97
8192
1
~ 16384
49.16
39.07
32.76
60.09
21576
50.74
40.14
8192
5
~ 49152
52.04
40.74
33.48
60.13
90112
42.08
LF-Wikipcdia-5000
formance of the approximate versions. In Sec-
2048
0
2048
77.71
43.32
69.24
88.12
2018
1
4096
82.85
48.84
74.47
90.08
tion C.2.3 of the appendix we also compare the results of InfoNCE loss with hard negatives and
2048
2
6144
83.34
49.32
74.73
89.74
2018
12288
83.86
49.57
74.60
89.12
2018
10
22528
84.77
50.31
75.52
90.29
DecoupledSoftmax loss with hard negatives.
5.5 COMPARISON ACROSS DIFFERENT LOSS FUNCTIONS
Next, we compare different loss formulations as dis- cussed in Sections 3 and the EURLex-4K and LF.
Table 4: Ablation of DEXML loss functions
AmazonTitles-I3 K. datasets (cf. Table 4 The DE model fails to train with BCE loss in We hypothe- size that this might due to the stringent demand the BCE loss, which requires all negative pairs to exhibit high negative similarity and all positive pairs to have a high positive similarity feat that be challenging for DE representations. DecoupledSoftmax loss in (5) performs the best for top-1 predictions. Moreover, the SoftTop-5 and SoftTop-100 100 losses (cf. (6)) yield optimal results for top-5 and top-100 predictions, respectively. These results show that Soft Top-k loss can outperform other loss formulations when optimizing for specific prediction set size. We further analyze the
Loss
P@1
P@5
R@
R 100
EURLex-IK
BCE
0.1
0.07
0.14
1.84
Sofimax CE
80.05
58.36
72.41
92.57
Decoupled
5
86.78
60.19
72.56
91.75
SoftTop-3 6
83.42
60.95
74.21
91.30
SoltTop-
37.41
56.97
93.72
LF-AmazonTilles
5.96
25.20
Softmax
41.77
20.87
56.91
69.49
DecoupledSpftrme
42.52
20.64
56.36
68.52
SoltTop-5
41.11
21.23
57.72
69.46
Soft 1orTTo--100 100
33.38
18:27
53.87
71.52
label score distributions of models learned with different loss functions in Section Clof the Appendix.
6 CONCLUSIONS & LIMITATIONS
Our key contribution is simple yet effective training procedure for DE methods that allows them to match or surpass SOTA extreme classification methods on standard benchmarks for many-shot retrieval. One limitation of our empirical evaluation is that we have only evaluated Ommilionn-sized label spaces. Verifying our gains on O(billion) scale is an interesting and challenging avenue for future work. Additionally, designing algorithms that improve DE performance XMC tasks without requiring all the negatives another interesting direction. We investigated well known many-shot learning problem and studied it in abstract form, so don't envision significant additional negative implications for society.
proposal up to

---

Published as a conference paper at ICLR 2024
REFERENCES
Thomas Ahle(2022).
Differentiable top-k function.
Mathematics
Stack
Exchange.
URL.
http///maat.sttaeeccccaeaaaae ctm. q/ 4506773
htl.:httss//matt.stccekeece
httpk//eccaanee.ooo///5656777770 com/g/450677 (version: 2022-08-07).
Rohit Babbar and Bernhard Schölkopf. Dismec: Distributed sparse machines for extreme multi-label classification. In Proceedings of the Tenth ACM International Conference Web Search and
Data Mining. WSDM 17. pp. 721-729, New York, NY, USA. 2017. Association for Computing Machinery. ISBN 9781450346757.
doi: 00.1145/3018661.30111 10. 11 0014441/11165 3018741
URLhttps://ddo org/
Peter L Bartlett, Michael Jordan, and Jon McAuliffe. Convexity, classification, and risk bounds. Journal of the Americun Statistical Association, 101(433): 2006.
Yoshua Bengio and Sean-SEbastien Senecal. Adaptive importance sampling accelerate training of
a neural probabilistic language model. IEEE Transactions on Neural Networks, 19(4)7713-722
2008. doi: 10. 10.1009/NN.2007..1133111
Leonard Berrada, Andrew Zisserman, and M. Pawan Kumar. Smooth loss functions for deep top-k classification. In International Conference on Learning Representations, 2018. URL
https::/
net mett/
K. Bhatia, K. Dahiya, H. Jain, P. Kar, A. Mittal, Y. Prabhu, and M. Varma. The extreme classi- fication repository: Multi-label datasets and code, 2016. URL http::/manikvarmat org/
Guy Blanc and Steffen Rendle. Adaptive sampled softmax with kernel based sampling. In Jennifer Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 590-599. PMLR, 10-15 Jul 2018.
Wei-Cheng Chang, Hsiang-Fu Yu, Kai Zhong, Yiming Yang. and Inderjit S. Dhillon. Taming pretrained transformers for extreme multi-label text classification. In Proceedings of the 26th ACM SIGKDD International Conference Knowledge Discovery and Data Mining, KDD '20, pp. 3163-3171, New York, NY, USA, 2020. Association for Computing Machinery. ISBN 9781450379984, doi: 0..1455339448663403366
URL
3394486. 3403368
Paul Covington, Jay Adams, and Emre Sargin. Deep neural networks for youtube recommendations.
In Proceedings of the 10th ACM Conference Recommender Systems, RecSys '6, pp. 191-198, New York, NY, USA. 2016. Association for Computing Machinery, ISBN 978 145044039, doi: 10.1 N0.1145/295910.222 2959999.. URL/https: //do:i///ooo////// (addii.org/lii 10. 145 01055/22551000 2959190
K. Dahiya, N. Gupta, D. Saini, A. Soni, Y. Wang. K. Dave, J. Jiao, K. Gururaj, P. Dey, A. Singh, D. Hada, Jain, B. Paliwal, A. Mittal, Mehta, R. Ramjee, S. Agarwal, Kar, and M. Varma. Ngame: Negative mining- aware mini mani-batching for extreme classification. In Proceedings of the ACM International Conference on Web Search and Data Mining, March 2023a.
K. Dahiya, S. Yadav, S. Sondhi, D. Saini, S. Mehta, J. Jiao, S. Agarwal, P. Kar, and M. Varma. Deep encoders with auxiliary parameters for extreme classification. In Proceedings of the ACM SIGKDD Conference on Knowledge Discovery and Data Mining, August 2023b.
Kunal Dahiya, Ananye Agarwal, Deepak Saini, Gururaj K. Jian Jiao, Amit Singh, Sumeet Agarwal, Purushottam Kar, and Manik Varma. Siamesexml: Siamese networks meet extreme classifiers with 100m labels. In Marina Meila and Tong Zhang (eds.), Proceedingy of the 38th International Conference on Machine Learning, volume 139 of Proceedinggs of Machine Learning Research, pp. 2330-2340. PMLR, 18-24 Jul 2021. URL hntps:: http/tttts:/ppoceeeiii
mlr. www..preps./viee
dahlya2la.htm.
Krzysztof Dembezyáski, Weiwei Cheng, and Eyke Hüllerreeer. Bayes optimal multilabel classifi- cation via probabilistic classifier chains. In Proceedings of the 27th International Conference International Conference on Machine Learning. ICML' 10, pp. 779-286. Madison, WI, USA, 2010.
Omnipress. ISBN 0781605589077.
10

---

Published as a conference paper at ICLR 2024
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-traie Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguisties: Human Language Technologies, Volume (Long and Short Papers); pp. 2019. Association for Computational Linguistics. doi: 10. 1/553/////11191423. URL https wwt.ooooooooro/119911
Minneapolis, Minnesota, June
Luyu Gao, Yunyi Zhang, Jiawei Han, and Jamie Callan, Scaling deep contrastive learning batch size under memory limited setup. In Proceedings of the 6th Workshop on Representation Learning for NLP (RepLANLP 2022), pp. 316-321, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v 1/20:///122222 http://e........po...p repl4nlp-1.311 URL https: / htt///aaaaantooooooy.ooooo org/
2021. rep441p-114 1.31
Chuan Guo, Ali Mousavi, Xiang Wu, Daniel Holtmann-Rice, Satyen Kale, Sashank Reddi, and Sanjiv Kumar. Breaking the Glass Ceiling for Bibeedding-Basee Classifiers for Large Output Spaces. Curran Associates Inc., Red Hook, NY, USA, 2019.
Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiy Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning, 2020. URLhttpss: //axxiv http://:///aaivoooo///b/////////
httrggabs//990 1908. 10396
N. Gupta, S. Bohra, Y. Prabhu, 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.
N. Gupta, P.H. Chen, H-F. Yu, C-J. Hsieh, and I. Dhillon. Elias: End-to-end learning to index and search in large output spaces. In Neural Information Processing Systems, November 2022.
Vivek Gupta, Rahul Wadbude, Nagarajan Natarajan, Harish Karnick, Prateek Jain, and Piyush Rai. Distributional semantics meets multi-label learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3747-3754, 2019.
Matthew Henderson, Rami AI-Rfou, Brian Strope, Yun-Hsuan Sung, László Lukács, Ruiqi Guo, Sanjiv Kumar, Balint Miklos, and Ray Kurzweil. Efficient natural language response suggestion 2017.
for smart reply. arXiv preprint
Rolf Jagerman, Zhen Qin, Xuanhui Wang, Michael Bendersky, and Mare Najork. On optimizing top-k metrics for neural ranking models. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR '22, pp. 2303-2301 New York, NY, USA, 2022. Association for Computing Machinery. ISBN 978 450387323, doi: 10.1 00.144/334749553555
URL :ttp::::://o.../ www.g/...000. 1145//1.1115/4774999 3531 849
Yacine Jernite, Anna Choromanska, and David Sontag. Simultaneous learning of trees and rep- resentations for extreme classification and density estimation. In Doina Precup and Yee Whye Tch (cds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1674. PMLR, 06-11 Aug 2017. URL htttp::/pproeeeding
ml htt.Fppee////////eeiittallaal html
Ting Jiang, Deqing Wang, Leilei Sun, Huayi Yang. Zhengyang Zhao, and Fuzhen Zhuang, Lightxml: Transformer with dynamic negative sampling for high-performance extreme multi-label text classification. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 7987-7994, 2021.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion- scale similarity search with gpus. IEEE Transactions on Big Data, 7(3):535-547, 2021. doi: 10. 10.110///DAATA.211...220 292
Armand Joulin, Edouard Grave, Piotr Bojanowski, and Tomas Mikolov. Bag of tricks for efficient text classification. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers, pp. 427 4311, Valencia, Spain, April 2017. Association for Computational Linguistics. URL https: https/////aaaaatoooooooooooo
E1 E17-200 207-068
11

---

Published as a conference paper at ICLR 2024
Purushottam Kar, Harikrishna Narasimhan, and Prateek Jain. Surrogate functions for maximizing precision at the top. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference Machine Learning. volume 37 of Proceedings of Machine Learning Research, pp. 199,198, Lille, France, 07-09 Jul 2015. PMLR. URL
ngs.
Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Densè passage retrieval for open-domain question answering. In Proceed- ings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), PP. 6769-6781, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v1/ 4020.emnlppmmale 550. https:
org 2000..2020... emn main. 550
Maksim Lapin, Matthias Hein, and Bernt Schiele. Loss functions for top-k error: Analysis and insights. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1468-1477, 2016.
Kenton Lee, Ming- Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. In Anna Korhonen, David R. Traum, and Lluís Mårquez (eds.), Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pp. 6086-6096. Association for Computational Linguistics, 2019.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmenter genera- tion for mnowledge-intensee nlp tasks. Advances in Neural Information Processing Systems, 33: 9459-9474, 2020.
Yuncheng Li, Yale Song, and Jiebo Luo. Improving pairwise ranking for multi- image classi- fication. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), July 2017.
Erik Lindgren, Sashank Reddi, Ruiqi Guo, and Sanjiv Kumar. Efficient training of retrieval models using negative cache. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 4134- 4146. Curran Associates, Inc., 2021. //R/itttss///rroceeddrs new.pps.cc/ppprrr files/ http:///2222/111121155cccc666660
Paper pdf
Tie-Yan Liu al. Learning to rank for information retrieval. Foundations and Trends@ in Information Retrieval, 3(3):225-311, 2009.
Yi Luan, Jacob Eisenstein, Kristina Toutanova, and Michael Collins. Sparse, dense, and attentional representations for text retrieval. Transactions of the Association for Computational Linguistics. 9:329-345, 2021.
doi: 10.1162/tacl a 00369. URL tacl-1.2C
http//aaaaatoooo.oorgg.o.. 2021
Tharun Kumar Reddy Medini, Qixuan Huang. Yiqiu Wang, Vijai Mohan, and Anshumali Shrivastava. Extreme classification in log memory using count-min sketch: A case study of amazon search with 50m products. Advances Neural Information Processing Systems, 32, 2019.
Aditya K Menon, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Multilabel reductions: what is my loss optimising? In H. Wallach, H. H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, F. Fox, and R. Garnett (eds.), Advances Neural Information Processsig Systems, volume 32. Curran Asso-
ciates, Inc., 2019. URL httt:pss//pooceeeiissse.eo www.:e/wwwe................r.... wtwpeiissscccpppeeee cc/paper fi filesspappz htop:ffiiiaaas/oooggddee222222e2
Paper.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation word representa- tions in vector space. arXiv preprint 1301.:130..3781, 2013,
Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahra- mani, and K.Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 26.
wttp::/wwwccccccecceeeeeeeeeeeee Curran Associates, Inc., 2013b. URL httpaapppp//ppooeeeeeeeee.eeeeea
fi http://pprrnnioo11eeeeaaaaaaaaa1
12

---

Published as a conference paper at ICLR 2024
Bhaskar Mitra and Nick Craswell. An introduction to neural information retrieval. Foundations and Trends@ in Information Retrieval, 13(1):1-126, 2018. ISSN 1554-0669. doi: 10.1561/1500000061. URL http:// dx doi wwtp.o........10..00000000
Anshul Mittal, Kunal Dahiya, Sheshansh Agrawal, Deepak Saini, Sumeet Agarwal, Purushottam Kar, and Manik Varma. Decaf: Deep extreme classification with label features. In Proceedings of the 14th ACM International Conference on Web Search and Data Mining, WSDM "21, pp. 49-57, New York, NY, USA, 2021a. Association for Computing Machinery, ISBN 978 450382977 doi: 10.1 10.1145/443966334441
URL //doi. www.10..//.............. 145/ / 344.433344441 3441807
Anshul Mittal, Noveen Sachdeva, Sheshansh Agrawal, Sumeet Agarwal, Purushottam Kar, and Manik Varma. Eclare: Extreme classification with label graph correlations. In Proceedings of the Web Conference 202/, www '21, pp. 3721-3732, New York, NY. USA, 2021b. Association for Computing Machinery. ISBN 9781450383127,
doi: 10.1145/3442381.3444999
URL
wtt.:////...11/334411 34981
Anshul Mittal, Kunal Dahiya, Shreya Malani, Janani Ramaswamy, Seba Kuruvilla, Jitendra Ajmera, Keng- Han Chang. Sumeet Agarwal, Purushottam Kar, and Manik Varma. Multi-modal extreme classification. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12383-12392, 2022. doi: 10.1009/ 2022.01207.
Andriy Mnih and Koray Kavukcuoglu. Learning word embeddings efficiently with noise- contrastive estimation. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Wein- berger (eds.). Advances ciates, Inc., 2013. URL http:hss//poccccssiiginoooo/ rwwrips wttppps.ccpppeeeeeeee cc/paper__..p.....eppe fi fillespaape 2013/file/db2/0 61 org.a//00000.001..11776000000000
Neural Information Processing Systems, volume 26. Curran Asso-
821
Andriy Mnih and Yee Whye Teh. A fast and simple algorithm for training neural probabilistic language models. In Proceedings of the 29th International Coference on International Conference on Machine Learning, ICML'11,, pp. 419 426, Madison, WI, USA, 2012. Omnipress. ISBN
9781450312851.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprini 1rX.vv1807.03745 2018.
Felix Petersen, Hilde Kuehne, Christian Borgelt, and Oliver Deussen. Differentiable top-k classifica- tion learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari. Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 17656-17668. PMLR, 17-23 Jul 2022. URL J http:::tp//...........tt http:::///roccceee ngs ml www.srssssve....... v1 http:///w..eeeeeeeeeeeeee2 html
Yashoteja Prabhu and Manik Varma. Fastxml: A fast, accurate and stable tree-classifier for extreme multi-label learning. In Proceedings of the 20th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD '1, pp. 263 272, New York, NY, USA, 2014. Association for Computing Machinery. ISBN 9781450329569,
doi: 10.1145/262333..222666 URL UtRp:stss://doii http:////ooo..o//..11o 114.0........... /1w.11////22332322 2623651
Yashoteja Prabhu, Anil Kag. Shrutendra Harsola, Rahul Agrawal, and Manik Varma. Parabel: Partitioned label trees for extreme classification with application to dynamic search advertising. In Proceedings of the 2018 World Wide Web Conference, WWW '18, pp. 993-1002, Republic and Canton of Geneva, CHE, 2018 International World Wide Web Conferences Steering Committee. ISBN 978 450356398. doi: 10. 10.1145/3178876.3 855998.
3178876. 3185999
Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Wayne Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. RockeiQA: An optimized training approach to dense passage retrieval for
open-domain question answering. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 5835,5847, Online, June 2021. Association for Computational Linguistics. doi: 18.13v/3/v1/ n0ac_naacl-o 466. htI.httta:.. : antho wotthelogy.oom wtw.lloo.oorg//211111 naacl ma n. 466
13

---

Published as a conference paper at ICLR 2024
Ankit Singh Rawat, Jiecao Chen, Felix Xinnan X Yu, Ananda Theertha Suresh, and Sanjiv Kumar. Sampled softmax with random fourier features. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https:///pooeeeeee neuri www.rrips.ccpppaeeee htcc///775bbaaacccc77eeeeee22222 -aap pdf
wwl.fiess/ppp wtllee/pppp//019..
Sashank J. Reddi, Satyen Kale, Felix Yu, Daniel Holtmann-Riee Jiecao Chen, and Sanjiv Kumar. Stochastic negative mining for learning with large output spaces. In Kamalika Chaudhuri and Masashi Sugiyama (eds.), Proceedings of the Twenty-Second International Conference Arti- ficial Intelligence and Staristics, volume 89 of Proceedings of Machine Learning Research, pp. mlr. www.pprssssv... v
1940 1949. PMLR, 16-18 Apr 2019. URL
redci 199...9m.m
Stephen Robertson, Hugo Zaragoza, and Michael Taylor. Simple bm25 extension to multiple weighted fields. Proceedings of the Thirteenth ACM International Conference on Information and Knowledge Management, CIKM '04, pp. 42-49, New York, NY, USA, 2004. Association doi: 10.11451103117111033 111.0//11..11111111111111111111. 103111
for Computing Machinery. ISBN 1581 138741,
URLhttps
Stephen E Robertson, Steve Walker, Susan Jones, Micheline M Bancock-Beaulinn Mike Gatford, et al. Okapi at trec-3. Nisl Special Publication Sp, 109:109, 1995.
Deepak Saini, Arnav Kumar Jain, Kushal Dave, Jian Jiao, Amit Singh, Ruofei Zhang, and Manik Varma. Galaxe: Graph neural networks with labelwise attention for extreme classification. In Proceedings of the Web Conference 2021, WWW '21, pp. 3733-3744, New York, NY, USA, 2021, Association for Computing Machinery. ISBN 9781450383127. doi: 00114453442381.3449999 3449937
URL https: http:/............oo.o.. 1etee//11..115444422220
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1990.01 0000000000000.1 108, 2019.
Jianlin Su, Mingren Zhu, Ahmed Murtadha, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Zlpr: novel loss for multi-label classification. arXiv preprint arXiv:2208.02955: 2022.
Nicolas Usunier, David Buffoni, and Patrick Gallinari. Ranking with ordered weighted pairwise classification. In Proceedings of the 26th annual international conference on machine learning, pp. 1057-1064, 2009.
Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang. Jialin Liu, Paul N. Bennett, Junaid Ahmed, and Arnold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. In International Conference on Learning Representations, 2021. URL https //openreview. net. wtwpee//ooooomodddfffrrrrra
Forest Yang and Sanmi Koyejo. On the consistency of top-k surrogate losses. Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 10727-10735. PMLR, 13-18 2020. URL htt.attss//prrreednngeg mlr www.ppeesssvviii httesss/vll//e /ttpss//119/aaan2.....
Hui Ye, Zhiyu Chen, Da-Han Wang, and Brian D. Davison. Pretrained generalized autoregressive model with adaptive probabilistic label clusters for extreme multi-label text classification. In Proceedings of the 37th International Conference on Machine Learning, ICML'20. JMLR.org, 2020.
Ian E. H. Yen, Xiangru Huang, Kai Zhong, Pradeep Ravikumar, and Inderjit Dhillon. Pd-sparse: A primal and dual sparse approach to extreme multiclass and multilabel classification. In Proceedings of the 33rd International Conference on International Conference on Machine Learning Volume 48, ICML' 16. pp. 3069-3077. JMLR.org, 2016.
Ian E.H. Yen, Xiangru Huang, Wei Dai, Pradeep Ravikumar, Inderjit Dhillon, and Eric Xing. Ppdsparse: A parallel primal-dual sparse method for extreme classification. In Proceedings of the 23rd ACM SIGKDD International Conference Knowledge Discovery and Data Mining,
14

---

Published as a conference paper at ICLR 2024
KDD '17, pp. 545-553, New York, NY, USA. 2017. Association for Computing Machinery. ISBN 9781450348874, doi: 10.11445720079983309988888 /RL//tttpp::/do..oo////00
3097983 3098083
Ronghui You, Zihan Zhang, Ziye Wang. Suyang Dai, Hiroshi Mamitsuka, and Shanfeng Zhu. Attentionxml: Label tree-based attention-aware deep model for high-performance extreme multi- label text classification. Advances Neural Information Processing Systems, 32, 2019.
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.
Jiong Zhang, Wei-Cheng Chang, Hsiang-Fu Yu, and Inderjit Dhillon. Fast multi-resolution trans-
former fine-tuning for extreme multi-label text classification. Advances Neural Information Processing Systems, 34:7267-7280, 2021.
Shuai Zhang, Lina Yao, Aixin Sun, and Yi Tay. Deep learning based recommender system: A survey and new perspectives. ACM Comput. Surv. 52(1), feb 2019. ISSN 0360-0300. doi: 10.1145/3285029,
Tong Zhang. Statistical analysis of some multi-category large margin classification methods. Journal of Machine Learning Research, (20ct):1225-125,, 2004.
Wenjie Zhang, Junchi Yan, Xiangfeng Wang, and Hongyuan Zha. Deep extreme multi-label learn- ing. In Proceedings of the 2018 ACM on International Conference on Multimedia Retrieval, ICMR '18, pp. 100-107, New York, NY, USA, 2018. Association for Computing Machinery. ISBN 9781450350464, 3206025. 3206030
doi: 10. 22144///222002522200000 URLhttps: olo..ooog/...1...
15

---

Published as a conference paper at ICLR 2024
Table 5: A snapshot the datasets used our experiments. Each row represents a different dataset with sample query and associated labels/documents separated by ":"
Datantt
Query
Labstelleuneets
BUllka-k
singliis qpa sonrnellaia
str_tranapertz aictines catalints pont ptatteeeua aidd covitatioon. benders secvicka, pt Legpeeet interwats spaint stabe wttteassdtteessppes company
13IX
sopix Hartie dy, huper hasds wioky Mistowlo aystes cillo uiee borsoley
hactng?
daicc anancce reteed pestonur talggo elo purtious inkepaiic euuie...
inlqpeddnnes polit the From port land enteny cosssiy of witro die
Rizelar Fectain Alrond innay, DE. Daze, 12-Coursy
thinatsin aectoin Nave, Dananol nudge Dlute Froty -1-bnne tars chack 100 riber
"Power Chenca 121, Chokios Prunat Hotee Fuege fierge
Mickland 5 siqnature Easca Fancy tnsalced Riadd Mita 2.5 (U81) sichitzitions Teseacch
tuper they Counch Proteis Triple Ctr
1.30r harrs
byeatize 180100 apdrobyand 1008 dayy veotain isolace ausrrtt venilla 1bss chocolate
A FUTURE WORK AND REPRODUCIBILITY
While our work shows that dual encoder model can be better or as competent as SOTA extreme classification methods without needing to linearly scale the model parameters with the number of labels, there are still challenges that need to be addressed. Below we highlight a few key future work directions:
Scaling to larger label spaces: We have established the performance of small dual encoders on up O(million) label spaces. However, it would be interesting study up to how big datasets can the limited-size dual encoders can retain their performance and if there is any fundamental capacity threshold that is essential for scaling to practical web-scale datasets.
Improving DE performance without all negatives: We have observed a clear relationship between DE performance and the availability of negatives. While it is computationally challenging to have all negatives, designing algorithms that can boost DE performance on shot retrieval problems without requiring large pool of documents in loss computa- tion is another promising area for exploration.
DE encoder with in- built fast search procedure: DE approaches rely on external approx- imate nearest neighbor search routines to perform fast inference when searching for top documents for a given query. In contrast, many hierarchical extreme classification methods come in-built with fast search procedures, it would be interesting to explore such hierarchical approaches for DE models.
Reproducibility: Our PyTorch based implementation is available the following link We plan to open-souree the code after the review period along with the trained models to facilitate further research in these directions. Below we provide additional details to clarify our approach and experimental setup.
B IMPLEMENTATION DETAILS
B. DATASET
In this subsection we provide more details of the datasets considered in this work,
The EURLex-4K dataset is small extreme classification benchmark dataset. It was sourced from the publications office of the European Union (EU). Specifically, it's a collection of documents concerning European Union law. The dataset comprises approximately 4,000 distinct labels, represented by Euro Voc descriptors. Euro is multilingual thesaurus maintained by the EU, covering fields of EU interest and offering consistent set of terms for the categorization of documents. There are over 15,000 training examples in the dataset. Each document within the EURLex-4K dataset is provided as a text string, consisting of the title and the content (or body) of the document.
EURLex-4K:
LF- AitzonTitlet-----aiiiia fiitles-11KK Originated from Amazon, this dataset encapsulates product rela- tionships, aiming to predict associated products based on a queried product. It comprises about 131. ,000 labels (products for recommendation from the catalog), and around 300,000
16

---

Published as a conference paper at ICLR 2024
training query products. The products recognized via their respective Amazon product titles, which function as the principal descriptive feature.
50K:: A conventional large-scale extreme classification benchmark, this dataset comprises document tagging of Wikipedia articles. Human annotators have tagged query documents with multiple categories, amounting to total of 500,000 labels and approximately 1.8 million training queries. It is marked by the presence of certain super- head classes.
albeit on larger scale and with a denser network. It includes a total of 2.2 million training points and 1.3 million labels. The distribution of these pa-purchasing links can be significantly skewed. Some popular products relate to a considerable number of other products, while less popular ones maintain fewer connections.
LF-Amazon' Tities--iMll- This dataset is similar to Amazonnitlss
Table 6: Dataset statistics
Num Train Points
Num Test Points
Num Labels
Avg Labels per Point
Avg Points per Label
Dataset
EURLex-4K
15,449
3,865
3801
5.30
20.79
F--AmazonTitlls-- 131K
294,805
134,835
131,073
2.29
5.15
LF- LF-Wikipedia-5006
1,779.881
769.421
501,070
4.75
16.86
LF- A-Amazoo Titles- 1.3M
2,248,619
970,237
1,305,265
22.20
38.24
B.2 TRAINING HYPERPARAMETER
In all our experiments we train the dual encoders for 100 epochs with Adamw optimizer and use the linear decay with warmup learning rate schedule. Following the standard practices, use 0 weight deacy for all non-gain parameters (such as layernorm, bias parameters) and 0.01 weight decay for all the rest of the model parameters. Rest of the hyperparameters considered in experiments are described below:
max len: maximum length of the input to the transformer encoder, similar to (Dahiya et al. 2023a) we use 128 max_len for the long-text datasets (EURLex-4K and LF. Wikipedia-500K) and 32 max_len for short-text datasets [LF-AmazonTTtles-- 131K and J--AmazooTTllle--.333
LR: learning rate of the model
batch_sie siie: batch-size the mini-batches used during training
drepout: dropout applied to the dual-encoder embeddings during training
multiplicative factor to control steepness of a function described in Section
4.3
nf: micro-batch size hyperparameter that controls how many labels get processed time when using gradient caching
т: temperature hyperparameter used to scale similarity values (i.e. s(ys dz)) during loss computation
Table 7: Hyperparameters
Dataset
len
LR
batch size
dropout
?)
EURLex-4K
128
0.0002
1024
0.1
4096
0.05
LF- AmazonTitles- 131K
32
0.0004
4096
2048
0.05
LF. Wikipedia-5008
128
0.0004
4096
0.1
2048
LF- AmazonTitles 1.3M
32
0.0008
8192
0.1
2
1024
17

---

Published as a conference paper at ICLR 2024
Droader call wiiaaoring reerred
Migeer
Coup C/Ce
Lesp
Excader wiout sking stare aciaa
Gnader
Teip X//0 thea
Figure 4: Illustration distributed implementation with gradient caching applied on label embedding computa- tion. Here solid black line indicate forward pass direction and solid red lines indicate gradient backpropagation direction. L is the number of labels considered in the loss computation, B is the batch size of queries, is micro batch-size hyperparameter which controls how many labels are processed at time.
B.3 BASELINES AND EVALUATION METRICS
For existing extreme classification baselines (apart from ELJAS), we obtain their numbers from their respective papers or the extreme classification repository (Bhatia et al. (016). Since the ELIAS paper do not report results on most of the datasets considered in this work, we ran the publicly available ELIAS implementation to report the results.
B.4 EVALUATION METRICS
We evaluate the models using standard XMC and retrieval metrics such as P@ I, P@5, P nDCG@5, R@@ 10, and @@ 100. P@k (Precision at k) measures the proportion of the top-k predicted labels that are also among the true labels for given instance. PSP@k (Propensity Scored Precision k) is a variant of precision that takes into account the propensity scores of the labels, which are indicative of the likelihood of a label being positive. nDCG@K @k (normalized Discounted Cumulative Gain at k) measures the quality of the ranked list of labels up to the position k. R@k (Recall at k) the proportion of the true labels that included in the top-k predicted labels.
B.5 MEMORY EFFICIENT DISTRIBUTED IMPLEMENTATION
In this section we describe our distributed implementation in more detail. For multi-GPU setting, where we have GPUs, batch size of B queries, and pool size of L. labels considered in the loss computation.
We begin by distributing the data and labels evenly across the GPUs. Each GPU receives B/G queries and L/G labels. However, for extreme multi-label problems with > 1M labels, even this reduced load can be too large for the L/G labels memory-wise wise - for an individual GPU during the gradient caching method (Gao et al., 202I). More
backpropagation. To handle this. concretely, on each GPU, we first divide the large L/G label input batch into manageable sub-batches of size 1). Then the encoder performs L/nIG forward passes compute label embeddings without
18

---

Published as a conference paper at ICLR 2024
constructing the computation graph (i.e. with torch. no mo_grad ()). Then, the loss is computed based on these embeddings, and the gradients for each embedding are computed and stored in a cache. Once we get the gradients with respect to the label embeddings, the encoder is run again each 17 sized sub-batch separately (but this time we retain the computation graph), and the stored embedding gradients are used to backpropagate through the encoder parameters. This allows for constant memory usage during the encoder's bawwardbackward pass at the cost of doing the forward pass twice.
To compute the loss, we perform an all-gather operation on query embeddings across all GPUs. This enables each GPU compute the dot product similarity matrix for the full queries and L L/G labels. The computation of the loss requires certain common terms, such as the denominator in the case of softmax-style losses. To compute these common terms, we perform an all-reduce operation across all GPUs. Finally, each GPU computes the gradients with respect to its batch of L./G label and B B/G query embeddings and backpropagates the respective gradients through the encoder parameters. By computing the loss and gradients in this distributed manner and by employing gradient cache, we can effectively handle large pools of labels without overwhelming the memory capacity of individual GPUs.
In the following description we assume that the loss is standard softmax loss but appropriate modifi- cations can be done to make it work for other losses discussed in Section 4 Figure Iillustrates the overall processing of a mini-batch in a 2 GPU setup. More generally, let there be G GPUs, B queries { ) and L labels { considered in the loss computation.
Each GPU receives B/G queries and L/G labels. More concretely, we define the set of B/G and the set of labels as (1)) L/G
queries received by the gt GPU as {x. ji=1
We can then define the embedding matrices for each GPU for B/G queries and L/G labels
as Qg = ( B/G) = encoder. Note that, computing Lg for large L requires gradient caching as described in Section 4.2
and Lg (((4") ) respectively, where ø is the query/document
We then gather all query embeddings at each GPU, denoted as Q (Gather(Qg)f.
We compute the similarity matrix for B queries and L/G labels as S, = QLT
Finally, we compute the softmax loss for each query-label pair (i.j) (on respective
GPUs) as Lij - (Si. 1 (ogsumexp)=1 is distributed across GPU, we will have to perform all_reduce operation to compute 1 (St,)) term in the loss.
Note, since S = [St, S2, .... Sa]
B.5.1 MEMORY AND RUNTIME ANALYSIS OF PROPOSED IMPLEMENTATION
Below we provide analysis of memory and compute requirements of the implementation with or without gradient cache based approach: let's assume we have G gpus, our batch size is B, number of labels L, gradient cache sub-batch size embedding dimension d and the memory required for processing (combined forward and backward pass of transformer encoder during training) a single query and label is M and Mf respectively. Without gradient caching, the total memory requirements for processing the whole mini-batch during training will be: O +d) w
BL, ), notice that with transformer encoders LM is the largest bottle- neck for XMC datasets with L in order of millions. With gradient caching the memory requirement Ld +BL ). Similarly for compute let's assume the compute required for the forward and backward pass of query be Teae- and T respectively, the compute required for 9
(B(M,, B(M4g"
is: O
G
the forward and backward pass of label be Tenc-s f and T, Tr Then, without gradient caching
the total compute requirement is: O B(rane 44pame +(d)+ trene ++0)+BL) With gra-
G dient caching the compute requirement will be: O B(T +T - +d) 1(2T 1 4ppne-h +d)+ BB.)
Empirically even with 16 A GPUs, the naive approach without gradient caching runs easily out of GPU memory even on AmazonTitles -itlees-lllkK dataset as it demands about ~TB combined GPU memory just for encoding all the labels. Table provides the ablation of memory used and runtimes
19

---

Published as a conference paper at ICLR 2024
of the memory-efficte approach with different global batch sizes and gradient cache sub-batch sizes on TF-AmazonTilles 131K dataset (experiments ran with 4 A100 GPU setup):
Table 8: Ablation of memory used and runtimes of the grad-cache based memory efficient implemen- tation
Global batch size
Gradient cache sub-batch size
Max Memory per GPU
Avg time per batch
Avg time per epoch
1024
64
6.03 GB
17.8s
1.42 hrs
1024
256
6.31 GB
7.1s
0.56 hrs
1024
1024
14.89 GB
5.9s
0.47 hrs
1024
4096
48.51 GB
5.6s
0.44 hrs
4096
256
15.75 GB
7.9s
0.16 hrs
4096
1024
16.16 GB
6.7s
0.14 hrs
4096
4096
51.85 GB
6.0s
0.12 hrs
16384
256
53.12 GB
12.2s
0.06 hrs
16384
1024
53.12 GB
8.2s
0.04 hrs
16384
4096
58.77 GB
6.6s
0.03 hrs
B.6 HARDWARE AND RUNTIMES
We run our experiments maximum 16 A100 GPU setup each having 40 GB GPU memory. In principle, our experiments can also be performed on a single GPU setup. but it'll take significantly longer to train. Table reports the training times and number of GPU used during training for each dataset.
Table 9: Training times and resources used for each dataset in our experiments
Parameter
EURLex-4K
LF-Amazon Titles-131k 13l
LF-Wikipedia-500k
FF-AmazonTills-I..m
Number of GPU
Training time (hrs)
0.27
B.7 DISTRIBUTED log SOFTTOP-k: IMPLEMENTATION
Below we provide the code snippet for efficient distributed implementation of log SoftTop-k operator (for loss computation it is more desirable to directly work in log domain) in PyTorch, here it is assumed that xs (the input to our operator of shape total number of labels and G is the number of GPUs used in the distributed setup.
I./G where the batch size, the
Topk code adapted from pttps://gisse github. com/ thomasable/4 2ie85ee5842401b007a00 fffe33a118
si sigmoid torch, s sigmoid
def sigmoid_graa :
sig_x sigmoid(x)
return sig x (1 sig x)
import Function class DistLegTopK (Functionn:
From
@staticmethed
10
def forward xs, k, alpha, n_iter=32):
11
logits, log_sig = _idsst_find_tsxse, k, alpha=alpha, n_iter=n_iter
retarn_log_sig-Tree
ctx.save for back_baccwadd (terch. tensor logits)
13
return log_sig
14
15
@staticmethed
def backward (ctx, grad_outputs Compute vjp, that is grad_output.
20

---

Published as a conference paper at ICLR 2024
alpha, logits ctx. saved tensors
19
sig_x eiggoi
20
v alpha+sig_x (1-sig_x)
s = v.sum (dim=11, keepdims-True)
alst.all_reducel//
uv grad_output alpha (1-sig_x)
24
uv_sum uv. sum (dim-l, keepdims-True)
25
reduue (uv_sum,
26
t1 = uv_sum v /
return tl.add (uv), None, None, None
in-place addition
29 @torch.no_grant
30
def _eist_fin_tte (xs, k, alpha-1, n_iter=64, return_logssigrFalls.
assert alpha >
b, xs.shape
33
n dest.get_woottee.... world_size(
34
if isinstance (k, int):
35
assert 0 < k <
elif isinstance (k, torch.LongTensor)>
assert
.))( and (k n) .al(()
88
Lo should be small enough that all sigmoids are in the area. # Similarly Hi is large enough that all are in their area. xs mi_ xs.min (dim=l, keepdims-True).ce values ma. (dim=1, keepdims=True) values dist.all reduce (xs min, op=dist Reducetp..cc MINN) dist.all__rdune (x__,, 10 -xs_max 10/alpha
19
40
42
43
ReduceeO MAX)
44
hi -xs_min + 10/alpha for in range (n_ i
47
mid (hi 12)
48
sigmoid_sum = sigmoid (alpha* (xs + mid) ) J. (dim=l) cist.all_reduce (sigmoid_ssss op=dist ReduccOp. SUK)
49
mask = sigmoid_sum
51
lo[mask] mid [mask] hi [-mas]] mid mmask]
52
53
ts (1c + /2
logits alpha* (xs ts)
56
if return log log__iog____o sig:
57
sig = logits - torch. (logits, torch, zeros like ( logitsl:, :11)) return logits, log_sig
59
clse:
60
return logits, sigmoid (logits)
C ADDITIONAL RESULTS AND ANALYSIS
C. MEMORIZATION CAPABILITIES OF DE
In order to verify if DE models are capable of memorizing random text co-relations even in the presence of an extreme semantic gap between query and label text, we create synthetic datasets of different sizes with completely random query and label texts, and then train and evaluate retrieval performance DE models. More specifically, for synthetic dataset of size N, we sample random query texts (by independently sampling 16 random tokens to form one query text) and similarly sample N random label texts. We assign 7th query to 7th label, train the DE model with in-batch and sampled hard negatives this dataset, and then report the retrieval performance on the same dataset. As seen in Table I0 even at a million scale DE models perform this task with perfect accuracy which shows that they are well capable of performing random text memorization.
21

---

Published as a conference paper at ICLR 2024
Table 10:
Performance of DE on random text pair synthetic dataset at different scales
Dataset size (N)
P@I 1
P@5
R@n 10
R@
P@1
P@5
R@ 10
R@ 100
DE with in in-batch negative
DE with hard negative
100K
100.00
20.00
100.00
100.00
100.00
20.00
100.00
100.00
500K
98.60
20.00
100.00
100.00
100.00
20.00
100.00
100.00
IM
81.21
18.67
95.71
99.15
99.93
20.00
99.99
99.99
C.2
DecoupledSoftmax vs STANDARD SOFTMAX
C.2.1
ANALYTICAL GRADIENT ANALYSIS
Softmax gradient analysis:
C
4 Σ jEP:
Eig where P, is the set of positive documents of label
bi log( here oij Σx es(n.da)
where s(q, d) zzÿ 2d
0lig
(1 -1)))))
ifl ifl=j otherwise
Ozd,
Ol,
(1 1
if P. otherwise
Ozd
This suggests that the multi-label extension of standard softmax optimized when:
i P. otherwise
P
Tu
0
DecoupledSoftmax gradient analysis:
C
4 Σ
Eij, where P; is the set of positive documents of label
JEP:
(lg))d)))
4j = log((; (zi))) here oij- = (s(g)dd)) Exep, e)())d)))
where s(g, d) = x zd
- oz))
if j if + and E P; otherwise
Olij
0
0zd:
Macid Deep cland Lqg.
(1 J.)
i otherwise
azd.
EjEP, mmoridy) Exer, otta 45 Tqg
C.2.2 EMPIRICAL RESULTS ON XMC BENCHMARKS
Table II compares DecoupledSoftmax with SoftmaxCE loss across all XMC datasets considered in this work.
22

---

Published as a conference paper at ICLR 2024
Table 11: Side by side comparison of DE models trained with DecoupledSoftmax vs Softmax loss on XMC datasets.
Dataset
P@I
P@5
Decoupled Softmax
Softmax
Decoupled Softmax
Softmax
EURLex-4K
86.78
80.05
60.19
58.36
Fr-AmmzonTiiles--- -iiees-11KKK
42.52
41.77
20.64
20.87
JF-Wikipedia-5000
85.788
79.85
50.53
49.78
Tr-AAmzzonTitlss-.....
58.40
52.42
45,46
43.43
C.2.3 COMPARING DecoupledSoftmax AND SOFTMAX WITH HARD NEGATIVES
Table 12 compares approximations of DecoupledSoftmax and SoftmaxCE loss with hard negatives on the largest LF-AmazonTitles 1.3M dataset.
Table 12:
Performance comparison of DecoupledSoftmax Softmax with sampled hard negatives, note that Softmax with sampled hard negatives same as InfoNCE with hard negatives
Batch size
Hard neg
Effective doc pool
P@1
P@S
R@ 10
R@100
P@1
P@5
R@10 10
R@ 100
per query
DecoupledSoftmax
Softmax
8192
16384
49.16
39.07
32.76
60.09
47.60
38.85
32.12
61.49
8192
2
24576
50.74
40.14
33.31
60.45
49.11
40.16
32.77
62.05
8192
10
90112
54.01
42.08
34.19
61.04
51.19
42.06
33.31
62.76
8192
All
L.M
58.40
45.46
36.49
64.25
52.42
43.42
34.00
65.84
C.3 QUALITY OF NEGATIVES IN DE TRAINING
As mentioned in Section 5.4 we perform negative sampling following (Xiong et al. 2021 we
mine hard negatives by computing the shortlist of nearest indexed labels for each training query regular intervals and sample negatives from this shortlist. Although, this approach is computationally efficient suffers from the problem of having deal with stale negatives since the model parameters constantly change but the negative shortlist gets updated only after certain training steps. To study the importance the quality of hard negatives we implement an approach similar (Lindgren et al.) 2021) where we maintain an active cache of all label embeddings and use this to first identify the hardest negatives and then use these negatives along with in-batch labels to perform loss computation. More specifically, we maintain a label embedding cache of L x d size where T. is the number of labels and is the embedding dimension. We initialize this matrix with the embeddings of all the labels at the start of training. For a mini-batch of B queries and L' sampled labels in the mini-batch, we first update the embedding cache of all mini-batch labels, then we compute the loss using embeddings of B training queries and the label embedding cache and identify the labels that receive the highest gradients. We take the top- K labels that are not already part of the mini-batch labels. Then we use these top-k labels the hard negative for the mini-batch and then compute the loss using the T/ mini-batch labels and K hard negative labels. This approach allows us to mine almost fresh hard negatives for each mini-batch In Table 13 we compare the results obtained using this approach and ANCE (Xiong et al. (2221) style hard negatives. Note that for a fair comparison, we compare them in settings where an equivalent total number of negatives are sampled per batch. These results suggest that there can be moderate gains made with an increase in the quality of negatives and further research into mining better negatives can be helpful in training better DE models with lower cost.
C.4 RESULTS WITH CONFIDENCE INTERVALS ON SMALL DATASETS
Since XMC datasets are fairly big benchmarks running multiple experiments with different seeds is very expensive. Below in Table [4] we provide mean and 95% confidence intervals for DecoupledSoftmax DE results for the smaller EURLex-4K and L- Titles-- 131K datasets.
23

---

Published as a conference paper at ICLR 2024
Table 13: Comparing DecoupledSoftmax trained with ANCE
Xiong et
2021 style stale negatives
fresh negatives mined using embedding cache (Lindgren et al. 2021
Dataset
P@L 1
P@5
R@ 10
R@ 100
P@I
P@5
R 10
@
100
ANCB (Xiong et al.
2021
Embedding cache (Lindgren ct al. 2021
BURLex-4K
84.61
59.15
71.08
90.21
85.87
59.93
71.54
90.34
AitzonTitles- 13IK
40.99
20.59
57.20
70.08
41.72
20.25
45.54
65.54
LF- Ff-Wikipedia-500
82.85
48.84
74.47
90.08
84.53
48.66
72.97
88.93
Table 14: Results with 95% confidence intervals for DecoupledSoftmax DE trained EURLex-4K and Ti-emazonTidee--11 131K dataset
P@S
RØ
R@100 100
P@5
R@
100
EURLex-4K
LF-Amazon -31K
86.50
0.24
00.24:0.19
72.62 0.26 91.76 0.10 42.37 0.13
20.57 +
+ 0.09 08.51 0.04
C.5
Score DISTRIBUTION ANALYSIS
In Figure plot the precision-recall curve on EURLex-4K for DE models trained using DecoupledSoftmax and Softmax loss. This plot reveals that DE model trained with DecoupledSoftmax offers better precision
Softmax Softmax
vs recall tradeoff. In Figure [6 we analyze
Precision
the distribution of scores given by DE mod- els trained with different loss function the test set of EURLex-4K dataset. More specifically, we plot the distribution of scores of positive and negative labels for Softmax loss 3 DecoupledSoftmax loss 5 and SoftTop- 5 loss 6 These plots reveal that the positives and negatives are better separated when we use DecoupledSoftmax and Top- loss, when compared to the Softmax loss.
0.4
0.2
Figure curve for Decoupled softmax vs stan- dard softmax on EURLex 4K dataset
DecoupledSoftmar (pos)
12
SaftTop-5 (pos)
DecoupledSoftmaa (neg)
SeftTop-5 ineg!
Softmax (pos)
(pos)
Softmax (neg)
10
Softmax (neg)
Density
Density
0.2
0.4
0.6
0.8
0.2
0.4
0.6
Score
Score
Figure 6: Plot of score distribution of DE models trained with different loss functions. left plot compares DecoupledSoftmax 5 vs softmax right plot compares SoftTop-5 [ vs softmax
C.6 DECILEWISE COMPARISON
We investigate the performance variation across different label regimes (many-shot versus few-shot labels) among various approaches. These include the pure classification-be based method (DistilBERT OvA Classifier), DecoupledSoftmax DE models trained with in- batch negatives (In-batch DE), hard negatives (Hard negative DE), and all negatives (DecoupledSoftmmx DE and Soft ToftTop-k DE).
24

---

Published as a conference paper at ICLR 2024
Consistent with Dahiya et al. (2023a). devise different label deciles based the count of training examples per label, wherein the first decile signifies labels with the most training examples and the anticipated, Figure 7 shows that the pure classification-baseed approach achieves subpar performance on the last deciles. However, intriguingly, DE methods trained with hard all negatives exhibit performance comparable to the classification-basse approaches even on head deciles. DE methods trained with our approach does well across all deciles and hence improve the performance over existing DE methods. On LF- LF-Amazon Titles- 1.3M, when comparing the DistilBERT OvA classifier with DE models, we notice similar trend that the pure classification based approach and DE approach are comparable at head deciles but tail deciles DE starts significantly outperforming the classifier approach. When comparing DE trained with all negatives and the ones trained with in-batch or hard-negatives, we notice substantial gap on head deciles but the gap diminishes on tail deciles.
last decile signifies those with the least. On LF-Wikipedia-500k
Decile-wise ise contribution to P385
Deci berill--bbee contribution to PgS
Cecoudie-Saaiv
InbusshiCt
Heed repelive DE Dou Bont Ovt Gasiter
hord regutive EE
2
Decii
Decile
Figure 7: Decile- wise analysis of P@5 on L.- J.F-Wikipddia-500 (left) and FF-AmazonTitless 1.3M (right), 1 decile represents labels with most training points i.e. head labels while 10th decile represents labels with least training points i.e. tail labels
C.7 BETTER GENERALIZATION OF DE MODELS ON TAIL LABELS
In this section, we study whether DE models offer improved learning when dealing with similar but distinct labels when compared to the classifier approach. We use LF LF-Wikipedia-500K dataset for this experiment. We first quantify how much a label is similar to the rest of the labels in the dataset, we do this by computing the normalized embeddings of each label (from the trained DE model) and searching for the top 10 nearest labels for each label, we use the mean similarity score of the top 10 nearest labels an indication of how similar this label is compared to rest of the labels. Based on this mean similarity score, we divide our labels into 3 equal bins, where the first bin corresponds to the *most dissimilar* labels and the last bin corresponds to the *most similar* labels. We compute the contribution of labels from each bin to final P@5 numbers for both the pure classifier-basee approach BERTOV baseline) and pure dual encoder- based approach (DecoupledSoftmax DE). We then plot the relative improvement between the classifier-based approach and the dual-encoder approach, more specifically we plot the relative improvement P@5 contribution relative to the classifier method. We also perform the same analysis but this time create label bins for only tail labels (we define a tail label as any label with < 5 training points
the dual-encoder method
From Figure g we infer that there doesn't 't seem to be any meaningful pattern when look at the results of the analysis on all labels but when we look at the analysis just on tail labels there appears to be trend that the relative improvement more similar labels and less dissimilar labels. This demonstrates that for labels with less training data, dual-encoders can perform relatively better on similar but distinct labels compared to pure classifier-based approach which treats each label as an atomic entity.
25

---

Published as a conference paper at ICLR 2024
Delative improvement in Pg5 contribution
2.00
1.0
1.75
improweiment Fg
0.6
1.00
0.75
0.4
0.50
0.25
DesoupicdSotmmax DE
DecoupledSoftmax DE DiSIIBETTVA
DistiBERTOvA
0.00
0.0
Label bins based similarity other labels
Label bins based an similarity to other labels
Figure 8: Plot to study dual encoder training performs better on similar but distinct labels on LF- LF-Wikipedia-SONM dataset. We first divide labels into equal bins, where the first bin corresponds to the *most dissimilar* labels and the last bin corresponds to the *most similar* labels (based on the mean similarity score of the top 10 nearest embedded labels). We then compute the contribution of labels from each bin to final P@5 numbers for both pure classifier-based approach (DistilBERRTOvA baseline) and pure dual encoder-based approach (DecoupledSoftmax DE). In the right subplot, plot the relative improvement between the classifier based approach and the dual encoder approach when considering all labels, more specifically we plot the relative improvement P@5 contribution of the dual encoder method relative to the classifier method. In the left perform the same analysis but this time we create label bins for only tail labels (we define a tail label as any label with S 5 training points.
C.8 COMPARISON WITH XR- TRANSSFORRR ON EURLEX-4K
As highlighted (Gupta et al. 2022 methods like XR-Trransformer classifiers learned on the concatenated sparse tf- idf features and dense embedding obtained from the BERT encoder for ranking their top predic- tions. [5] compare XR- Transformer's performance without ensembling and without using tf-idf features. We also, report results by adding the sparse ranker layer (implemented similarly to the approach described in Section 3.5 of (Gupta et al. 2022) on top of DecoupledSoftmax DE top 100 predictions. As can that XR- TRR-Trasformer'ss 's better performance EURLex- 4K dataset can be greatly attributed to the use of ensemble of 3 learners and high-capacity sparse classifiers.
Table 15: Comparison with XR-Transformer FURLex-4K dataset
high-capacity sparse
XR Transformer x3 ensemble
88.41
75.97
63.18
XR Transfomerr
87.19
73.99
61.60
XR Trannssffrrrr w/o tf-idf ranker
83.93
69.48
56.62
DecoupledSoftmax DE
86.78
73.40
60.19
DecoupledSoftmax with tf-idf ranker
86.55
75.21
62.26
C.9 APPLICATIONS OF PROPOSED LOSS FUNCTIONS FOR RAG
Inspired by the open QA benchmarks (NQ/MSMARCO) which can be treated as few-shot variants of RAG task (given question, retrieve passage using DE model. then given (question, retrieved passage) a reader model generates answer), we construct a naturally many-shot and multi-label RAG scenario from the existing XMC Wikipedia dataset LF- L.F-Wikipedia-500 the task of generating Wikipedia content from a Wikipedia page title given some tags associated with the page. The tags are essentially the evidence documents that retrieve and we use the Wikipedia page title as the query and ask fine-tuned LLM (Llama2-7B) model to generate the corresponding Wikipedia article. Following are some sample query prompts:
More specifically,
create
"Given the associated wikipedia tags [living people), (american film actresses), (actresses from chicago illinois)], Write wikipedia like article on [Aimee_ [Aimee_Garcia]: ...
"Given the associated wikipedia tags [(1928 births), (auschwitz concentration camp vivors), (american soccer league 1933 83 coaches)), Write a Wikipedia like article on
26

---

Published as a conference paper at ICLR 2024
"Given the associated wikipedia tags [[given names), (germanic given names), (norwegian feminine given names)], Write a wikipedia like article on [Solveig]:"
We fine-tune the LLM model on 10K training samples and evaluate the LLM augmented with different retrieval approaches on separate 1K test samples. We compare the generated content with the gold content using ROUGE metrics.
Retrieval
ROUGE-1
ROUGE-3
ROUGE-L
No retrieval
0.1622
0.0468
0.1236
InfoNCE DE
0.1806
0.0532
0.1371
DecoupledSoftmax DE
0.1838
0.0540
0.1387
Table 16: Performance comparison of different retrieval methods on proposed RAG scenario using ROUGE scores.
Although this a very preliminary experiment that needs to be thoroughly tested before making any concrete conclusions. it does provide promise for our proposed changes in RAG settings and opens up new applications for XMC approaches.
27


Reviews:
[{"review_id": "evoKkJZmAv", "reviewer": "Reviewer_Smm7", "summary": "The manuscript considers the use of Dual Encoders (DE) models for Extreme Multi-Label Classification (XMC). In previous work using softmax loss functions, the presence of other positive classes decreases the negative loss that can be gained from a training example. This work considers decoupled softmax, in which the negative loss gained for each positive class is independent of other positive classes. The motivation for this new loss function is that some classes may \"obviously\" apply, which should lead to confident predictions for those \"obvious\" cases. In experiment with synthetic dataset where one of the samples is very obviously marked, both softmax and decoupled softmax methods latch on this obvious annotations early on, but standard softmax approaches then entice the model to give the same score to all positive samples, obvious or not, thus decreasing the model's confidence in the obvious case. In non-synthetic settings, decoupled softmax reduces the variance in gradient feedback.\n\nDirectly training with the full decoupled softmax loss is challenging when there are many labels: appendix presents how to do this training with the full loss in a memory-efficient way. Still, for higher number of labels (e.g., documents retrieval), there is a need to approximate this full loss, and some preliminary results are given in this direction. A natural extension is presented for SoftTop-k.\n\nThe Appendices were not considered as part of this review.", "questions": "### 1\n(a) What are the assumptions, goals and constraints specifying what this project is about; (b) what would be the default/status quo approach and what's wrong with it; and (c) what is the essence of the solution you propose?\n\nI think that I managed to figure out the answer to those questions, but these things should be clearly identifiable from the introduction (or even abstract). Here's my own crude attempt at it: please complement it, clarify any disagreements, and propose edits to the manuscript.\n\n- 1(a) There are a large number of classes, and many of them can apply to the same sample. The models being considered are DE where the representations from each encoder will be converted to a score using an inner product.\n- 1(b) OvA-BCE won't train; InfoNCE disincentivizes confident predictions.\n- 1(c) Decoupled softmax both trains and allows for confident predictions.\n\n\n### 2\nThe bulk of the manuscript presumes the optimization of the the \"full\" loss function over all positive and negative classes. Is this standard practice in the XMC community? Some experiments sampling hard negatives are presented in Section 5.4, but there is no real discussions besides \"more is better\". Do you have any insights to add? Could those be added to the manuscript?\n\n\n*The remaining points are more minor*\n\n### 3\nFigure 1: why express the x axis in millions, instead of replacing the $10^2$ and $10^3$ ticks by $10^8$ and $10^9$?\n\n### 4\nPlease avoid notations such as \"O(100)\", \"O(million)\" and \"O(billion)\".\n\n### 5\nFigure 2: \"Precision@1\" here has a special meaning. There are 5 samples that are marked as positives, but one of them is \"more positive\", and \"Precision@1\" here means \"the more positive sample must be in first position\". This special setting is relatively clear in the text, but looking quickly at Figure 2 and its caption can be misleading. Please consider inventing a different term/notation.\n\n### 6\nThere are multiple missing punctuations after mathematical expressions. For example, a period should be added at the end of the first paragraph of Section 3, and another one should be added after the equation at the bottom of page 4. As a side note, I would personally add a lot of commas to the English text, but I understand that this may be more a matter of style.\n\n### 7\nThe text appears to use \\cite or \\citep everywhere, but many should actually be \\citet. For example, in the second paragraph of Section 4.2, \"similar to (Xiong et al., 2021)\" should become \"similar to Xiong et al. (2021)\", and \"in (Lindgren et al., 2021)\" should be \"in Lindgren et al. (2021)\".\n\n### 8\nThe style of the stackexchange citation in the paragraph following Equation (6) should be revised. Personally, I would have expected something like \"... on a proposal by Ahle (2022).\"\n\n### 9\nSome variables are overloaded; for example $d$ is used both for documents and dimensions. I thought that there were more cases, but I can't find them anymore, so I may be confused with another paper. Please check.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "To the best of my knowledge, the work is original and significant. Although the distinction between using a single multi-class cross-entropy vs using multiple binary cross-entropies (OvA-BCE) have long been well understood, it isn't immediately obvious that the DE setting would be so different, and that the OvA-BCE loss would fail to train.\n\nIn the context of Retrieval-Augmented Generation (RAG), when the retriever is not trained in an end-to-end manner with the answerer, one often conceptualize the retrieval problem as \"get that one good document that is know to suffice to answer the question\". However, there could be multiple such documents that contain the desired answer: this retrieval task can be understood as a particularly extreme case of XMC. The take home message that I personally get from this manuscript is that I should exercise caution along those lines if I ever attempt to train DE retrieval models with contrastive augmentations. (continued in weaknesses)...", "weaknesses": "... (continued from strengths) However, I wouldn't have come to this realization from the manuscript's abstract nor introduction, and that nugget would have been lost on me had I encountered the manuscript outside of a reviewing context. I understand that this work mainly targets the XMC literature (of which - full disclosure - I personally don't know much), but I still believe that the authors should dedicate some of their high-level discussions (i.e., abstract and/or introduction), to the significance of their work for DE-based RAG.\n\nMore generally, the manuscript's main weakness is with organizing the content to give a clear narrative. For example, Equation (2) introduces OvA-BCE, then eschew saying what's wrong with it until Section 5.5. As a reader, I had to go back-and-forth in the document to just understand what the authors are set in doing, what goes wrong with the default approach, and what's the authors' solution. More details in Question 1 below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "cv18MOsqcD", "reviewer": "Reviewer_BerB", "summary": "This paper explores the use of dual encoder models for extreme multilabel classification (XMC). Dual encoder models, which as the name suggests, use two encoder models, have been generally effective for a variety of other types of tasks involving zero-shot and few-shot learning, but have underperformed the state-of-the-art on XMC tasks. Dual encoder methods are, on the other hand, desirable for extreme classification in principle because they can be much more parameter efficient than extreme classification methods whose parameter counts have a strong dependence on the number of classes. \n\nThe authors propose a new loss function, the differentiable top-k, for dual encoder models on XMC tasks that make them competitive with the state-of-the-art methods on these tasks--often outperforming them by a significant margin.", "questions": "- Can the authors elaborate on how $t_x$ in the soft top-k formulation is computed via binary search? As this is a crucial hyperparameter of the loss function, does this require a logarithmic number of retraining runs? How many runs are required to set this, what are the compute costs involved with setting this hyperparameter, and how strongly does the final performance depend on this choice? \n- In Table 5, why is precision at 5 reported while recall at 100 is reported? Why not both precision and recall at 5 and 100? Is this a standard choice? \n- Does the loss optimize precision at k or recall at k?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The work is well-motivated and addresses the practical problem of getting dual encoder methods to work in the extreme multilabel setting. \n- The proposed contribution is simple, as it is just a loss function paired with either a negative mining approach or a memory-efficient implementation using all negatives. \n- The ablation of the different loss variants -- soft top-5 and soft top-100 -- is compelling and shows that the method can effectively optimize precision or recall at 5 or 100, respectively.", "weaknesses": "- I might have missed something, but I think it should be made clearer earlier on in the paper that the differentiable top-k operator had been proposed previously [1]. The authors also link to the author of the stackexchange answer, but it would be ideal to cite the specific answer at the link (please correct me if this appears somewhere in the main text, but I couldn't find it). Relatedly, is there other work that uses the formulation by Thomas Ahle? For example, how does this formulation compare to the one linked in the stackexchange post [2]? In my opinion, this brings the novelty of the contribution *when posed as a new loss function* into question. \n- Again, I might have missed something, but I think the ablation involving the negative mining approach should include a comparison to SOTA methods, as the negative mining approach may be required for scaling the method up even further. \n\n[1] https://math.stackexchange.com/questions/3280757/differentiable-top-k-function/4506773#4506773 \n\n[2] https://arxiv.org/pdf/2002.06504.pdf", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "WsFffu5bQg", "reviewer": "Reviewer_jMAd", "summary": "This paper explores the performance of dual-encoder models in extreme multi-label classification (XMC) tasks. The authors first reveal the shortcomings of traditional dual-encoder training loss, which may over-penalize the correct prediction of \"easy-positive\" labels. To address this, the Decoupled Multi-label Loss is proposed, aiming to minimize the undesirable correlation between positive labels during training. A memory-efficient training framework is also introduced in this paper.", "questions": "- Please address the questions mentioned in the Weaknesses section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The authors highlight a neglected problem in the XMC dual-encoder training stage. The proposed Decoupled Loss effectively solves this problem.\n- The theoretical part of the paper is well-presented, with Section 3 providing clear symbol definitions.", "weaknesses": "- The main motivation for this paper is the imperfect design of current dual-encoder training loss. However, there is a lack of evidence that this has been a general issue for current XMC methods. Most discussions and experiments are designed to compare the Decoupled Loss and regular loss using the authors' own training framework. Some experiments are implemented using a synthetic dataset (Fig. 2) or pre-selected labels (Fig. 3). After reading the entire paper, I believe the proposed loss can solve the mentioned problem, but I am not convinced that this problem is a universal issue in the XMC community.\n- The paper lacks novelty. A minor revision of the training loss may not be sufficient for a top conference paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "tdXMSVnKXZ", "reviewer": "Reviewer_1qgK", "summary": "Dual-encoder models have achieved substantial success in dense retrieval tasks for open-domain question answering, especially in zero-shot and few-shot scenarios. However, their performance in many-shot retrieval problems, where abundant training data is available, such as extreme multi-label classification (XMC), has received limited attention. Existing evidence indicates that dual-encoder methods tend to underperform compared to state-of-the-art extreme classification methods that scale the number of learnable parameters linearly with the number of classes in such tasks. Some recent extreme classification techniques combine dual-encoders with a learnable classification head for each class to excel in these scenarios. This paper explores the potential of \"pure\" dual-encoder models in XMC tasks and provides insights for XMC.", "questions": "As in Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Important Research Problem: The paper addresses a significant and relevant problem in the field of machine learning - extreme multi-label classification. This is a challenging task with practical applications, and the choice of topic adds value to the existing literature.\n\nInteresting Approach: The idea of using dual-encoders for solving the problem is intriguing and adds a novel dimension to the research. This innovative approach can potentially open up new avenues for tackling similar problems in the future.\n\nComprehensive Experimental Validation: One of the strengths of the paper is its thorough experimental validation. The fact that the proposed approach has been tested on multiple datasets indicates a comprehensive evaluation of its effectiveness. This enhances the credibility of the findings and their potential applicability to real-world scenarios.", "weaknesses": "I do not find obvious flaws. One notable weakness is the lack of clarity in articulating the paper's contributions. It's important for the reader to clearly understand what novel insights or advancements are being offered. The paper should explicitly state the unique contributions and why they matter in the context of extreme multi-label classification. This clarity is crucial for both researchers and practitioners in the field.\nTo enhance the paper, it would be beneficial to provide a more structured and explicit statement of the research's contribution and significance in introduction section. This will help the readers better grasp the key takeaways from the study. Additionally, the paper could benefit from improved organization and flow to ensure that the reader can easily follow the arguments and findings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)