Paper: Dual-Encoders for Extreme Multi-label Classification (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
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).

Reviewer comment (Novelty):
The paper lacks novelty. A minor revision of the training loss may not be sufficient for a top conference paper.

Author response:
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.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.
| Method | Hard neg per query | P@1     | P@5     | R@10    | R@100   |
|--------|-------------------|---------|---------|---------|---------|
| **LF-Wikipedia-500K** | | | | | |
| Ours   | 0                 | 77.71   | 43.32   | 69.24   | 88.12   |
| Ours   | 1                 | 82.85   | 48.84   | 74.47   | 90.08   |
| Ours   | 2                 | 83.34   | 49.32   | 74.73   | 89.74   |
| Ours   | 5                 | 83.86   | 49.57   | 74.60   | 89.12   |
| Ours   | 10                | 84.77   | 50.31   | 75.52   | 90.29   |
| NGAME (with classifiers) | - | 84.01 | 49.97 | - | - |
| DEXA (with classifiers) | - | 84.92 | 50.52 | - | - |
| **LF-AmazonTitles-1.3M** | | | | | |
| Ours   | 0                 | 42.15   | 32.97   | 29.28   | 57.48   |
| Ours   | 1                 | 49.16   | 39.07   | 32.76   | 60.09   |
| Ours   | 2                 | 50.74   | 40.14   | 33.31   | 60.45   |
| Ours   | 5                 | 52.04   | 40.74   | 33.48   | 60.13   |
| Ours   | 10                | 54.01   | 42.08   | 34.19   | 61.04   |
| NGAME (with classifiers) | - | 56.75 | 44.09 | - | - |
| DEXA (with classifiers) | - | 56.63 | 43.90 | - | - |

Reviewer comment (Evaluation):
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?

Author response:
It is a standard practice in recent XMC papers to report precision for lower K (i.e. K=1-5) and report recall for higher K (i.e. K=10-100)[1,2]. For a particular K, precision and recall numbers are proportional to each other as their numerator is same (number of true positives), the only difference is that they are scaled differently, for precision we’ll divide by K but for recall we’ll divide it by total number of positives (which is a fixed quantity), we believe this is why reporting precision at small K and recall at higher K gives better resolution to compare results at different prediction budget size.

Reviewer comment (Evaluation):
Does the loss optimize precision at k or recall at k?

Author response:
For a fixed prediction budget-size K, empirically improving precision@k improves recall@k as well and vice versa since both of these share the same numerator but are only scaled differently in the denominator. The individual per-query loss terms can be scaled accordingly if we strictly want to optimize for precision@k or recall@k.

Reviewer comment (Presentation):
The manuscript's abstract and introduction do not convey the significance of the work for DE-based RAG. While the work mainly targets the XMC literature, the authors should dedicate some of their high-level discussion (i.e., abstract and/or introduction) to highlighting this significance, as otherwise this insight could be lost to readers.

Author response:
**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

Reviewer comment (Presentation):
The manuscript's main weakness is in organizing the content to provide a clear narrative. For instance, Equation (2) introduces OvA-BCE, but the manuscript delays explaining what is wrong with it until Section 5.5. This forces the reader to go back and forth in the document to understand the intent, the issues with the default approach, and the proposed solution.

Author response:
**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.

Reviewer comment (Writing):
The introduction (or abstract) should clearly state (a) the assumptions, goals, and constraints of the project, (b) the default/status quo approach and its shortcomings, and (c) the essence of the proposed solution. The reviewer deduced these points as follows and suggests complementing, clarifying disagreements, and incorporating them into the manuscript: (1a) There are a large number of classes, many of which can apply to the same sample. The models are DE, where encoder representations are converted to scores via inner products. (1b) OvA-BCE fails to train; InfoNCE discourages confident predictions. (1c) Decoupled softmax both trains and allows for confident predictions.

Author response:
**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

Reviewer comment (Experiments):
The manuscript presumes optimization of the "full" loss function over all positive and negative classes. The reviewer asks whether this is standard practice in the XMC community. While some experiments with hard negative sampling are presented in Section 5.4, there is little discussion beyond "more is better." The reviewer requests that the authors add insights and incorporate them into the manuscript.

Author response:
**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.