Dual-Encoders for Extreme Multi-label Classification

Thank you for your thoughtful review and constructive feedback. We appreciate your positive comments on the importance of the research problem, our approach, and the thorough experimental validation. We have carefully considered your suggestions and propose the following revisions:

Clarifying Contributions: Please see our general comment for a more clear discussion on contributions and significance of this work. We acknowledge the need for a clearer articulation of the paper's unique contributions, especially in the introduction. We have revised the introduction to explicitly reflect this (please see the revised pdf)

Improving Organizational Flow: We will restructure the paper to enhance the flow and coherence of our arguments and findings. More specifically, taking into account reviewer Smm7’s suggestions, we will introduce a preliminary discussion on inability of BCE loss to train dual encoders and clearly state how we are proposing to overcome the challenges. Please let us know if there are any further suggestions to improve the organization of the paper.

We hope these revisions will improve the manuscript, addressing your concerns and making our contributions clearer to the readers. We are happy to discuss any further questions about the work and would appreciate an appropriate increase in the score if your concerns were adequately addressed.

We thank the reviewer for the review and feedback. Please find below our response to the concerns:

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...

We disagree that “there is a lack of evidence that this has been a general issue for current XMC methods and most discussions and experiments are designed to compare the Decoupled Loss and regular loss using the authors' own training framework”. We perform our evaluation on 4 XMC benchmarks which span the largest available XMC datasets. The main results of our paper (Table 2 and 3) compare our approach against current SOTA XMC approaches [1,2] and show superior performance. Existing XMC papers repeatedly show that dual encoders don’t perform well by themselves and hence need to be augmented with per-label classifiers or auxiliary parameters to achieve the best performance, see table 8 in [1] and table 2 in [2] vs results in [3] which show upto 10% absolute difference in P@1 between vanilla dual-encoder approach and dual-encoder with classifier approach. Quoting from [2], “this (dual-encoder) approach is expected to suffer if the textual descriptions are not descriptive enough which makes it challenging for an encoder to bring related data points and labels close to each other in the embedding space....To overcome this, XC methods introduce a high-capacity classifier such as 1-vs-all models where each label is endowed with a linear classifier that acts on the embedding of a data point”. In contrast, to the best of our knowledge, our work is the first one to show that vanilla dual-encoders can match or surpass existing SOTA methods when trained appropriately. This is a desirable result because a) dual-encoders can be much more parameter efficient than existing XMC methods that use per-label classifiers; b) this brings promise that dual-encoders can solve both standard retrieval and XMC problems leading to a universally applicable solution and simplifying the landscape of future solutions for these problems. Please see our general comment for a more detailed discussion on contributions and significance of this work. We also acknowledge that although this work shows that dual-encoders are capable of performing well on XMC tasks, classifier based solutions can still be attractive because of their ease of computation (which is just an embedding lookup) as compared to encoders (which requires a relatively expensive encoder forward pass).

The paper lacks novelty. A minor revision of the training loss may not be sufficient for a top conference paper.

We would like to emphasize that the technical contributions of our paper goes beyond just introducing a minor revision to existing loss. More specifically, our technical contributions include:

• Identification of Limitations in Traditional Dual-Encoder Training Losses: we analyze existing multi-label and contrastive training losses like One-versus-All Binary Cross-Entropy (OvA-BCE) 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.
• Decoupled Softmax and SoftTop-k loss: To overcome these limitations, we proposes a simple modification to the InfoNCE/Softmax loss as the Decoupled Softmax Loss and extend the loss design to soft-topk loss which is tailored to optimize prediction accuracy within a specific prediction budget size ('k') To the best of our knowledge, these limitations of standard multi-label and contrastive loss functions and the proposed changes have not been discussed in literature before. Moreover, scaling our experiments to largest XMC benchmarks require non-trivial implementation efforts which is often not highlighted but is an essential component for modern machine learning solutions.

[1] NGAME: Negative Mining-aware Mini-batching for Extreme Classification, WSDM 2023

[2] Deep Encoders with Auxiliary Parameters for Extreme Classification, KDD 2023

[3] The Extreme Classification Repository: Multi-label Datasets & Code

We hope our response addressed your concerns and are happy to discuss further questions. Please consider improving your score if you find this work adequate for acceptance after the clarifications and revisions.

Thank you for your valuable feedback. We appreciate your insights and have addressed your concerns as follows:

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...

You’re right that the application of the differentiable top-k operator as a loss function for optimizing multi-label predictions at fixed prediction budget-size is the novel contribution. The differentiable top-k operator itself has been proposed in the literature, a common approach is to use the optimal transport-based formulation (which you mentioned) but it is very hard to implement in a distributed training setup which is necessary for our experiments. This is why we instead extend the approach presented in the stackexchange answer by Thomas Ahle because of its simplicity and easily distributable operations with low communication overheads. To the best of our knowledge this formulation is not used in any existing work. We do cite the original stackexchange answer but ICLR citation style seems to only display user profile in the main text (the citation in the reference links to the original answer), we will fix this in the revised version based on Reviewer Smm7’s suggestions. We’d also like to highlight that although soft-topk loss function is one of the technical contributions we make, the major contributions of this work is in highlighting that existing popular loss functions can be inadequate for training dual-encoder models for XMC, to overcome this we revise InfoNCE and extend its loss design to SoftTopK leading to much more parameter efficient SOTA XMC solution. Please see our general comment for a more detailed discussion on contributions and significance of this work.

ablation involving the negative mining approach should include a comparison to SOTA methods

Due to lack of space we did not repeat the rows from the main results table (Table 2) which reports numbers for SOTA methods in the negative mining ablation table (Table 4). We do agree that perhaps having the existing SOTA results in Table 4 might be easier for a reader to quickly compare the approach with negative mining, we will make this change in the final version. In the table below we report the results with negative mining against SOTA methods, note that DE trained with proposed loss functions can still match or come very close to SOTA methods.

How $t_x$ in the soft top-k formulation is computed via binary search...

$t_x$ is not a hyperparameter, rather $t_x$ is a variable which is a function of the $x$ vector (i.e. the vector of all the scores) such that $\sum_{i=1}^{N}{s(x_i + t_x) = k}$. Essentially, t_x is a scalar quantity which when added to all the scores $x_i$ makes the sigmoid of new scores ($x_i + t_x$) sum upto $k$. Since sigmoid is a monotonic function, increasing $t_x$ strictly increases the sum of sigmoids and decreasing $t_x$ strictly decreases the sum of sigmoids, hence for a given $x$, binary search can be used to find the $t_x$ which makes the sum of new sigmoids equal to $k$. We hope this clarifies the doubt, we will also revise section 4.3 in the final version to make this more clear.

Thank you for your thorough and insightful review. We appreciate your recognition of the originality and significance of our work in applying Dual Encoders in Extreme Multi-Label Classification and its implications for Retrieval-Augmented Generation. Please find below our response to the concerns and suggestions raised:

Implications for RAG: Retrieval augmented generation (RAG) is an important scenario which we didn’t fully consider at the point of writing this paper, many thanks for bringing it out. In scenarios where multiple documents could contain the desired answer, the retrieval task in RAG aligns closely with extreme multi-label classification (XMC). Our work's exploration of dual-encoder methods in XMC directly informs the appropriate choice for loss formulation in such scenarios. To this end, we perform a preliminary experiment described in “RAG experiment” below, to see the promise of this work in the context of RAG.

Significance and Contributions: Please see our general comment for a more clear description of contributions and significance of this work. We acknowledge your concern regarding the lack of clarity in contributions and significance of this work in the introduction of the manuscript. To address this, we have proposed a revision of the introduction (please see the revised pdf). We plan to further refine the introduction in the final version to ensure that the broader implications of our research are clear to a wider audience. Below we slightly complement your answer to the question 1

• 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.
• 1(b) DE models empirically don’t perform well on such task but are desirable due to parameter efficiency and generalization to unseen labels; Standard loss function are inadequate - OvA-BCE won't train; InfoNCE disincentivizes confident predictions.
• 1(c)Decoupled softmax both trains and allows for confident predictions leading to SOTA DE models for XMC tasks

Narrative Clarity: We acknowledge that the manuscript would benefit from a clearer narrative structure. We will reorganize the content in the final version to more effectively present our objectives, the limitations of existing approaches, and the essence of our proposed solution. More specifically, in order to address the concern related to the lack of discussion on OvA BCE in section 4, we will introduce a preliminary discussion on the inability of BCE loss to train dual encoders and clearly state how we are proposing to overcome the challenges. We will also take into account all the formatting suggestions to improve the final version.

Discussion on negatives: optimization of the "full" loss function is not a standard practice in the XMC community but is used sometimes to establish the upper bound performance. In XMC benchmarks, we notice that an extensive use of negatives is required to train properly and is typical for classifier based methods to use ~100 negative labels per training query. In dual encoder training since the most expensive step is the forward and backward pass through the encoder, considering ~100 negative labels per training query for a batch size of 8192 would be computationally almost equivalent to computing the full loss for a dataset with upto 1M labels. Apart from showing that there is a significant dependance of the performance with the number of negatives, in section C.3 of the appendix we note that the quality of negatives can also play a modest role in determining the accuracy (specially P@1) performance of the model. In section C.2.3 we also confirm that increasing negatives for InfoNCE/Softmax loss function on XMC benchmarks doesn’t benefit the performance that much when compared to DecoupledSoftmax loss, this is because as we increase total number of negatives to consider in the loss computation we also increase the undesired competition among the positives in a multi-label setting.

RAG experiment

Inspired by the open QA benchmarks (NQ/MSMARCO) which can be treated as few-shot variants of RAG task (given a question, retrieve passage using a 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-Wikipedia-500K. More specifically, we create 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 we retrieve and we use the wikipedia page title as the query and ask a fine-tuned LLM (Llama2-7B) model to generate the corresponding wikipedia article. contd. in further comment...