Paper: EHI: End-to-end learning of Hierarchical Index for Efficient Dense Retrieval (ICLR-2024)

Reviewer comment (Novelty):
There is a very closely related work based on product quantization that is not cited: Zhan et al. (ACM WSDM 2022) “Learning Discrete Representations via Constrained Clustering for Effective and Efficient Dense Retrieval” (https://doi.org/10.1145/3488560.3498443).

Author response:
We believe that this work [1] is complementary to the work of EHI for the following reason. Our work is end-to-end learning of both the embedding space as well as the indexer (encoder as well as the ANNS). The indexer itself is the ANNS structure, and we do not need another disjoint clustering approach. RepCONC as introduced in the paper quantizes document embeddings, and uses a disjoint k-means step to generate n clusters as explicitly stated in their paper Section 3.5: 
“Besides PQ, RepCONC employs the inverted file system (IVF) to accelerate vector search. After quantizing document embeddings, RepCONC uses k-means to generate 𝑛 clusters. Each document embedding belongs to the nearest cluster and is stored in the corresponding inverted list.”. However, we would also like to point out that RepCONC achieves an MRR@10 value on passage and TREC-DL task of 34.0% and 66.8% which is significantly lower than EHI. We have also added these numbers in our comparisons in the Appendix. 

[1]  Zhan et al. (ACM WSDM 2022) Learning Discrete Representations via Constrained Clustering for Effective and Efficient Dense Retrieval.

Reviewer comment (Novelty):
The paper omits references to more efficient sparse neural IR models, which are also end-to-end, such as:  
- Lassance et al. (SIGIR 2021) “Composite Code Sparse Autoencoders for First Stage Retrieval” (https://doi.org/10.1145/3404835.3463066)  
- Formal et al. (SIGIR 2021) “From Distillation to Hard Negative Sampling: Making Sparse Neural IR Models More Effective” (https://doi.org/10.1145/3477495.3531857).

Author response:
We believe there has been some misunderstanding about the term “end-to-end” as used in this paper (please refer to General comment). For instance, [1] which introduces CCSA explicitly mentions in the very first line in their paper that : “We propose a Composite Code Sparse Autoencoder (CCSA) approach for Approximate Nearest Neighbor (ANN) search of document representations based on Siamese-BERT models.” As CCSA directly uses the document representations based on the Siamese-BERT model, it is not end-to-end. 
It is also unclear why the reviewer mentions “more efficient” here, when the performance (MRR@10) of CCSA as reported in Table 2 of their paper of MSMARCO dev and TREC-DL benchmarks are 28.9% (MRR@10) and 58.3% (nDCG@10), which is significantly worse in comparison to EHI, and other benchmarks presented in the Appendix.
 
- Secondly, [2] aims to learn sparse representations of the representations to facilitate fewer FLOPS. We believe this to also be orthogonal, as each model of SPLADE++ could lead to one specific flop usage, and is not as malleable as EHI where the practitioner has the option to improve performance, by increasing the budget. Again, it is not clear how the reviewer compares efficiency here between the two models, where one - EHI learns an encoder and ANNS in an end-to-end fashion, whereas the latter - SPLADE++ removes this ANNS structure and replaces it with sparse representations for learning. Furthermore, it is not clear how the KMR curve of this SPLADE++ looks in comparison to EHI. For instance, they use the metric of FLOPS, but what would be the FLOPS if one uses the dense representation for retrieval? 

[1] Lassance et al. (SIGIR 2021). Composite Code Sparse Autoencoders for First Stage Retrieval. 
[2] Formal et al. (SIGIR 2021). From Distillation to Hard Negative Sampling: Making Sparse Neural IR Models More Effective.

Reviewer comment (Experiments):
The authors compare to ColBERT, but do not consider ColBERTv2, which performs much better.

Author response:
Ideally, we would like to compare against models which are representation-based similarity models that output single vectors. These are the models where ANNS can be readily used, and EHI helps alleviate the limitations arising from this disjoint training paradigm.  ColBERT being a late-interaction similarity model does not output a  single vector representation, and was only used a standalone baseline. Note that it is not readily possible to add an ANNS structure on top of the ColBERT architecture. ColBERT-v2 follows a similar paradigm, and has a 6x storage cost as it stores one vector per token instead of one vector per passage [1]. We would like to stress that EHI (distilbert) can achieve performances similar to ColBERT-v2 with 2% drop in MRR with a model 1/2x smaller and achieves this with more than 80% reduction in latency, and no additional storage costs over the ANNS.

[1] Wang, Liang, Nan Yang, Xiaolong Huang, Binxing Jiao, Linjun Yang, Daxin Jiang, Rangan Majumder, and Furu Wei. "Simlm: Pre-training with representation bottleneck for dense passage retrieval." arXiv preprint arXiv:2207.02578 (2022).

Reviewer comment (Experiments):
Overall, the paper misses strongly related approaches and does not compare to state-of-the-art results on first-stage rankers, either dense or sparse.

Author response:
We have added many SOTA baselines including ColBERT, SGPT, cpt-text, ANCE, DyNNIBAL and compare against these approaches in Appendix D. We have compared against other SOTA results such as DSI, NCI in this regime, which are our closest related works and have compared against their baselines as well. Some of the works you mention here which focus on sparse representations and remove the use of ANNS is tangential to our work as further elaborated in our general comment. If the reviewer believes we have missed any other strongly related approaches, please do let us know.

Reviewer comment (Presentation):
The document is not well-structured—many figures and tables cited in the main text are only in the appendix. The first 8 pages should be self-contained, which is currently not the case.

Author response:
Thanks for updating the tables – it would be better to have the discussion and actual results of the SOA systems in the main text, rather than in the Appendix; as such, the paper reads as if it were hiding results under the carpet. Given the positioning of the paper, I think it would be better to argue why such an approach is worth pursuing – and I think this is true since the efficiency-effectiveness tradeoff for dense models is understudied.

- We thank the reviewer for their insightful comments. We have added a few lines to the related works sections citing the works suggested by the reviewer but have not been able to incorporate much content in the related sections expanding further on the results of these models and why efficiency-effectiveness tradeoffs are worth pursuing due to page limitations as stated in our [previous comment](https://openreview.net/forum?id=ESq3U7z6FD&noteId=WhsaglVdYG).
As suggested, we have also incorporated some of the comparisons against SPLADE in the main text. We will expand on this further in the final version as it requires significant re-writing to accommodate. As suggested by the reviewer PLmq, we are planning to move the Tables to the Appendix anyway, but we will highlight these related works for completeness. We want to stress that we have no intentions of sweeping these results under the carpet and are happy to include them.

Reviewer comment (Theory):
There is no discussion of an important aspect: the generated identifier sequence is inherently a probabilistic process, but this is not taken into account in the loss formulation.

Author response:
We respectfully disagree with this. Both indexing similarity and intra-leaf similarity losses focus entirely on this identifier sequence. We define them formally in mathematical formulations, defining path embedding, expressing how each loss component affects training in the ablations, and further expanding into this “path embedding”/ “indexer sequence” in Appendix F.

Reviewer comment (Experiments):
The authors propose leveraging their indexing tree to sample better negatives—one of the advantages of their model—but this is only mentioned in passing, and the effect of the sampling is not studied. This is problematic because it might be the main explanation for the better performance. The authors should try to index using HNSW to measure this effect. This aspect also represents an opportunity for the paper to focus on this better learning strategy separately from the indexing/end-to-end problem.

Author response:
The hard -ve sampling aspect is only mentioned passing-by, and the effect of the sampling is not studied.

- We do discuss this in ablations (see Figure 3c). Note that even without the use of negative samples, EHI outperforms the exact search R@100 metric by ~2%, and thus helps bolster the idea that training both encoder and indexer in an end-to-end fashion is indeed helpful. We further expand on this negative sampling mechanism in Appendix F.1.

Reviewer comment (Experiments):
Please provide comparisons with state-of-the-art approaches.

Author response:
We have added many SOTA baselines including ColBERT, SGPT, cpt-text, ANCE, DyNNIBAL and compare against these approaches in Appendix D. We have compared against other SOTA results such as DSI, NCI in this regime, which are our closest related works and have compared against their baselines as well. Some of the works you mention here which focus on sparse representations and remove the use of ANNS is tangential to our work as further elaborated in our general comment. If the reviewer believes we have missed any other strongly related approaches, please do let us know.

Reviewer comment (Experiments):
The ablation study for negative sampling should be provided.

Author response:
The hard -ve sampling aspect is only mentioned passing-by, and the effect of the sampling is not studied.

- We do discuss this in ablations (see Figure 3c). Note that even without the use of negative samples, EHI outperforms the exact search R@100 metric by ~2%, and thus helps bolster the idea that training both encoder and indexer in an end-to-end fashion is indeed helpful. We further expand on this negative sampling mechanism in Appendix F.1.

Reviewer comment (Presentation):
Please make the paper self-contained; appendices should only be used to provide more details about specific aspects.

Author response:
Thanks for updating the tables – it would be better to have the discussion and actual results of the SOA systems in the main text, rather than in the Appendix; as such, the paper reads as if it were hiding results under the carpet. Given the positioning of the paper, I think it would be better to argue why such an approach is worth pursuing – and I think this is true since the efficiency-effectiveness tradeoff for dense models is understudied.

- We thank the reviewer for their insightful comments. We have added a few lines to the related works sections citing the works suggested by the reviewer but have not been able to incorporate much content in the related sections expanding further on the results of these models and why efficiency-effectiveness tradeoffs are worth pursuing due to page limitations as stated in our [previous comment](https://openreview.net/forum?id=ESq3U7z6FD&noteId=WhsaglVdYG).
As suggested, we have also incorporated some of the comparisons against SPLADE in the main text. We will expand on this further in the final version as it requires significant re-writing to accommodate. As suggested by the reviewer PLmq, we are planning to move the Tables to the Appendix anyway, but we will highlight these related works for completeness. We want to stress that we have no intentions of sweeping these results under the carpet and are happy to include them.

Reviewer comment (Experiments):
clear assumption - my first intuition after reading the paper is the story it tries to deliver lacks an good assumption/description of the initiative. The work is an incremental work built on past models, and authors claim that e2e training is better. However the baseline models it compares with are also trained e2e, while it is disjoint because the inference process, i.e. the ANNS requires doc embedding being calculate first. I would phrase the advantage to be to recapture the lost accuracy during ANNS, etc.

Author response:
- Thank you for pointing this out. We note that modern retrieval systems have a disjointed learning process that often results in representations that are sub-optimal for downstream clustering. This is primarily due to the fact that the encoder is not aware that the embeddings need to be clustered. This leads to misalignment of representations. To overcome this limitation, end-to-end training is desired - which leads to not only recapture the lost accuracy during ANNS (supported by our findings in Figure 2, Appendix E.8), but also leads to better representations learned (supported by our findings on Figure 2, Figure 3, and Appendix E). We will further highlight this in the introduction in our draft.

- Please refer above for the motivation to why e2e training is preferred. We also showcase this with findings on various benchmarks as highlighted in Figure 2, Appendix E. We would like to point out that EHI is the first work to the best of our knowledge to solve the problem of training an encoder and ANNS in an end-to-end fashion without using warm-started KNN models to facilitate learning the indexer. What we highlight in this work is an alternative training paradigm to do so and we respectfully disagree with the statement that EHI is an incremental work built on past models.

Reviewer comment (Writing):
lack of clarity in model details - it may be due to page limit but the model description is not clear enough. The softmax in the tree, the path embedding, the Wh+1, Uh+1, and p^h+1 looks not that straight forward without more clarification or a graph of training and inference process, etc.

Author response:
- For a given query/document, the root node uses the embedding information alone, and we use a linear layer ($U_1$) to transform the embeddings received by the encoder to another 768-dimensional embedding. Having these affine transformations at each height is essential as we would like to focus on different aspects of the same query/document at each height to classify which children node the item needs to go to. This is followed by another linear layer ($W_1$) to predict logits for the children at the node (2 in this case). This, followed by the softmax in the tree, leads to the probability distribution deciding which node to go towards - ($p^1 = [P_1, P_2]$). 

  This is followed by the logic at the following height, which follows roughly the same process as the root node with one significant difference. The affine transformation at this height takes both the embedding input and the one hot version of the probabilities at the previous height to make the predictions. During training, we have 1 at the index with the highest probability and 0 otherwise (assume $P_1 > P_2$, and hence one-hot version is $[1,0]$). However, once the conditional probability at the next height is computed, we multiply by this $P_1$ to have the actual probability of reaching nodes 3 and 4 (note these numberings follow the same logic as used in Figure 1). The final path embedding is the concatenation of these probability distributions leading to $[P_1, P_2, P_1 P_3, P_1, P_4]$ as depicted in Figure 14.

  During inference, the indexer follows a similar setup where, instead of only doing one forward pass with the highest probability, we do $b$ forward passes at a given height with top-$b$ highest probabilities. We have added a separate section in Appendix F describing the details of the EHI indexer further with the help of figures as requested.

Reviewer comment (Experiments):
Missing comparison with existing work that learns a kNN index jointly with the encoder model. Some examples are:  
- Gupta, Nilesh, et al. "ELIAS: End-to-End Learning to Index and Search in Large Output Spaces." NeurIPS 2022  
- Gupta, Gaurav, et al. "BLISS: A Billion scale Index using Iterative Re-partitioning." KDD 2022.

Author response:
> Missing comparison with existing work that learns a kNN index jointly with the encode model. 

- It should be noted that ELIAS and BLISS are tailored for the domain of extreme classification while we focus on the dense information literature domain (note that label classifiers, and pre-computing KNN structure for datasets the size of MSMARCO is indeed a significant overhead). ELIAS has been cited in our work, and we do point out the limitations of ELIAS [1] which requires warm starting with a knn clustered index. Regardless, we report our findings on the LF-AmazonTitles-131K  (XC dataset) below. Note that below comparison is still unfair, as ELIAS utilizes additional million auxiliary parameters in their label classifiers, while EHI does not. Furthermore, EHI uses label features for encoding, while ELIAS does not. Although EHI outperforms ELIAS by 1.67%,  it must be recognized that comparing against extreme classification based methods is unfair. 

  Furthermore, BLISS [2] as per our understanding is an efficient ANN structure which enforces load balancing. Our goal in this work has been to unify the training scheme of the encoder which learns representations, and the ANN which helps index documents for faster and efficient retrieval.
We are happy to add both of these comparisons to our paper if the reviewer think these comparisons are fair indeed.

| LF-AmazonTitles-131K                             | P@1   |
 |-------------------------------------|--------|
| ELIAS | 40.13 |
| EHI | 41.8 |

[1] Gupta, Nilesh, et al. "ELIAS: End-to-End Learning to Index and Search in Large Output Spaces." NeurIPS 2022
[2] Gupta, Gaurav, et al. "BLISS: A Billion scale Index using Iterative Re-partitioning." KDD 2022.

Reviewer comment (Experiments):
The training strategy for baseline dual-encoder models is not clear. Some recent papers propose negative mining strategies that overcome the limitations of having to maintain an up-to-date kNN index during training. Here are some examples (some are already cited by the authors but not sure why a direct comparison is omitted):  
- Monath, Nicholas, et al. "Improving Dual-Encoder Training through Dynamic Indexes for Negative Mining." AISTATS 2023.  
- Dahiya, Kunal, et al. "NGAME: Negative Mining-aware Mini-batching for Extreme Classification." WSDM 2023.  
- Hofstätter, Sebastian, et al. "Efficiently teaching an effective dense retriever with balanced topic aware sampling." SIGIR 2021.

Author response:
> The training strategy for baseline dual-encoder models is not clear. Overall, it is not clear is jointly learning kNN index and encoder model is useful to get better negatives during training or does it have advantages beyond just better hard negative mining. 

- The training strategy for the baseline dual encoder model is an in-batch hard negative sampling approach. Note that the primary purpose of our work as the reviewer have rightfully captured in their summary is the unified training paradigm of learning both the encoder and indexer without any warm starting. Note that the dynamic negative mining is something EHI achieves for almost free, as the built indexer is also used for the ANNS, and we expand on this further in Appendix F.1. 
Prior approaches including [1,2,3] do study this dynamic negative sampling approach as well, and is complementary to our approach. However, the difference between EHI and these approaches lies in the fact that EHI does not need to build a new index for clustering every few steps, and the hierarchical tree index learned by EHI is used for the downstream ANNS task. This contrasts with works such as [1,2,3], where the index is used during training and hard negative mining, but is finally discarded. Note that we have compared against DyNNIBAL in Appendix D. Furthemore, as mentioned in our new experiments, EHI with it's differentiable indexer achieves roughly 2% improvement at lower compute budgets over other off-the-hat baselines such as SCANN and Faiss (see Appendix E.8).

[1] Monath, Nicholas, et al. "Improving Dual-Encoder Training through Dynamic Indexes for Negative Mining." 
[2] Dahiya, Kunal, et al. "NGAME: Negative Mining-aware Mini-batching for Extreme Classification."
[3] Hofstätter, Sebastian, et al. "Efficiently teaching an effective dense retriever with balanced topic aware sampling."

Reviewer comment (Experiments):
The documents need to be re-indexed every few epochs to mine hard negative documents for each query. It is not clear if this step is cheaper than building an index with off-the-shelf libraries like ScANN and FAISS.

Author response:
>  It is not clear is this step is cheaper than building index with off-the-shelf libraries like ScANN and FAISS.

- It is definitely cheaper as we are simply using the same trained indexer for re-indexing. If we would like to use some off-the-hat indexer such as SCANN or FAISS, we would need to train the ANNS from scratch every once in a while. From our findings, we note that training of ANNS on MSMARCO is roughly 2x more expensive than just inference using the indexer of EHI.

Reviewer comment (Novelty):
While the index contains trainable parameters at each internal node, the basic structure of the tree index (i.e., height and branching factor) is fixed.

Author response:
> While the index contains trainable parameters at each internal node, the basic structure of the tree index i.e. height and branching factor is fixed.

- This is absolutely true. These choices are left to the practitioner similar to the network size, shape, activation layers, and other intricacies of any deep learning algorithm. However, we note that EHI is fairly robust to these choices as showcased in Table 13.

Reviewer comment (Experiments):
Joint training of encoder and nearest neighbor search index may not work very well in zero-shot settings where encoder models are trained on one domain and used on a new domain (with a new set of documents/items and without any training data).

Author response:
> Joint training of encoder and nearest neighbor search index may not work very well in zero-shot settings.

-That is correct. However, this is simply true for most current IR system, even those which use a two-stage disjoint training paradigm. It would be interesting to study how few-shot learning works with EHI, and is an interesting future direction.

Reviewer comment (Experiments):
The encoder model is initialized using a variant of `distill-bert` which is already fine-tuned using labelled data from the MS-MARCO dataset. Since MS-MARCO is one of the main large-scale evaluation datasets, it may not be fair to compare with baselines (previous papers) that use models such as `distill-bert`, `bert` or `RoBERTa` which are pre-trained only using self-supervised objectives.

Author response:
> The encoder model is initialized using a variant of distill-bert which is already fine-tuned using labelled data from MS-MARCO dataset.

- Distilbert-cos model we have used through sentencebert was also pre-trained only using self-supervised objectives [1, 2] with mined hard negatives. The comparisons against various baselines in Appendix D predominantly serve to showcase that EHI training paradigm can compete against SOTA methods of the field, while the alignment with ANNS is significantly improved as showcased in Figure 2, Appendix E.

[1] Reimers, Nils, and Iryna Gurevych. "Sentence-bert: Sentence embeddings using siamese bert-networks." arXiv preprint arXiv:1908.10084 (2019).
[2] https://github.com/UKPLab/sentence-transformers/issues/1167