EHI: End-to-end Learning of Hierarchical Index for Efficient Dense Retrieval
Authors: Ramnath Kumar
Anshul Mittal
Nilesh Gupta
Aditya Kusupati
Inderjit S Dhillon
Prateek Jain
Venue: Accepted by TMLR ()

Abstract: Dense embedding-based retrieval is widely used for semantic search and ranking. However, conventional two-stage approaches, involving contrastive embedding learning followed by approximate nearest neighbor search (ANNS), can suffer from misalignment between these stages. This mismatch degrades retrieval performance. We propose End-to-end Hierarchical Indexing (EHI), a novel method that directly addresses this issue by jointly optimizing embedding generation and ANNS structure. EHI leverages a dual encoder for embedding queries and documents while simultaneously learning an inverted file index (IVF)-style tree structure. To facilitate the effective learning of this discrete structure, EHI introduces dense path embeddings that encodes the path traversed by queries and documents within the tree. Extensive evaluations on standard benchmarks, including MS MARCO (Dev set) and TREC DL19, demonstrate EHI's superiority over traditional ANNS index. Under the same computational constraints, EHI outperforms existing state-of-the-art methods by +1.45% in MRR@10 on MS MARCO (Dev) and +8.2% in nDCG@10 on TREC DL19, highlighting the benefits of our end-to-end approach.

Published in Transactions on Machine Learning Research (08/2024)
EHI: End-to-end Learning of Hierarchical Index for Efficient Dense Retrieval
Ramnath Kumar Google Inc.
remnathk@@google.co
Anshul Mittal +1
me@@anshaummittll..om
Microsoft Google Inc.
Nilesh Gupta
mileshguutta27978@maa..oo
UT Austin Google Inc.
Aditya Kusupati
kusupatii@goggle.cmm
Google Inc.
Inderjit Dhillon Google Inc. UT Austin
isd isd@google. com
Prateek Jain
prajain crejaln@@oogle.coom com
Google Inc.
Reviewed on OpenReview: https: openreview. net/ forum? id= GeLLOGsVVV
Abstract
Dense embedding- retrioval is widely used for semantic search and ranking. However, conventional two-slage approaches, involving contrastive embedding learning followed by approximate nearest neighbor search (ANNS), can suffer from misalignment between these stages. This mismatch degrades retrieval performance. We propose End-to-end Hierarchical Indexing (EHI), novel method that directly addresses this issue by jointly optimizing embedding generation and ANNS structure. EHT leverages dual encoder for embedding queries and documents while simultaneously learning an inverted file index (IVF)-style tree structure. To facilitate the effective learning of this discrete structure, EHI introduces dense path embeddings that encodes the path traversed by queries and documents within the trec. Extensive evaluations on standard benchmarks, including MS MARCO (Dev set) and TREC DL19, demonstrate EHI's superiority over traditional ANNS index. Under the same computational constraints, EHI outperforms existing -tate-of-the-art the-art methods by +1.45% in MRR@10 on MS MARCO (Dev) and +8.2% in nDCG@10 on TREC DL19, highlighting the benefits of our end-to-end approach.
1 Introduction
Semantic search (Johnson et al., 2019) aims retrieve relevant or semanlically simvlar documents /ittettsss///s/ssssseese///// /items for given query. In the past few years, semantic search has been applied to numerous real-world applications like web search, product scarch, and news search (Nayak, 2019; Dahiya et al., 2021). The problem in the
"Equal Contribution
Work done while Google Inc.

---

Published in Transactions on Machine Learning Research (08/2024)
simplest form can be abstracted as: for given query q, retrieve the relevant document(s) d(g) from static set of documents {di, d2...., dx} J s.t. d(q) arg max SIM(q, dj). Here SIM is a similarity function that has high fidelity to the training data B = {(q d,, 91y)}. Tuple (4s, dj. Mey) indicates if document dj is relevant (Uig = 1) or irrelevant (Vi = 1) for a given query q: E Q.
Dense embeskding-based retrieval (Johnson et al., 2019; Jayaram Subramanya et al., 2019; Guo et al., 2020) is the state-of-the-art (SOTA) approach for semantic search and typically follows a two-stage process. In the first stage, it embeds the documents and the query using a deep network like BERT (Devlin et al., 2018) (Additional details about the encoder used in EHI is described in Section 3.2, and the training methodology expanded in Section 3.5). That is, it defines similarity SIM(g. d) := (E) Ed(d) as the inner product between embeddings Ea(4) and Ev(d) of the query and the document d, respectively. Ex(-) is dense embedding function learned using contrastive losses (Ni et al., 2021; Menon et al., 2022).
In the second stage, approximate ncarest neighbor scarch (ANNS) retrieves relevant documents for a given query. That is, all the documents are indexed offline and are then retrieved online for the input query. ANNS in itself has been extensively studied for decades with techniques like ScaNN (Guo et al., 2020), IVF (Sivic & Zisserman, 2003), HNSW (Malkov & Yashunin, 2020), DiskANN (Jayaram Subramanya et al., 2019) and many others being used heavily in practice.
1.1 Motivation for EHI
The prevailing two-stage approach to dense retrieval, involving separate training of the embedding encoder and the approximate nearest neighbor search (ANNS) structure, suffers from inherent limitations:
Embedding- NNS Misalignment:: The lack of joint optimization leads to a potential mismatch between the learned embedding space and the requirements of the ANNS structure. This misalignment can cause suboptimal clustering of data points, hindering the ANNS algorithm's ability to efficiently identify semantically similar examples. For example, the documents might be clustered in six clusters to optimize encoder loss. However, due to computational constraints, ANNS might allow only five branches/clustere merging clusters unnaturally and inaccurately. Figure (()) helps illustrate the issue with the disjoint training regime typically used in dense retricval literature. If the disjoint ANNS tries to cluster these documents into buckets, the resulting model will achieve suboptimal performance since the encoder was never aware of the task of ANNS. Instcad, EHI proposes the learning of both the encoder and ANNS structure in single pipeline. This leads to efficient clustering, as the embeddings being learned already share information about how many clusters the ANNS is hashing.
thus splitting or
Disregard for Query Distribution: Conventional ANNS methods focus on generic retricval efficioncy, potentially neglecting the specific characteristics of the query distribution. This can result in an indexing structure that is ill-suited to the patterns present in the actual queries encountered during deployment (Jaiswal et al., 2022). To provide an illustrative instance, let us consider scenario in which there exists a corpus of documents denoted by D, with two classes of documents categorized D2], following a certain probability distribution [di, d2). Suppose, di >> d2, then under the utilization ion of the unsupervised off-the-shelf ANNS indexer such as ScaNN algorithm, clustering into two partitions, the algorithm would distribute the two document classes in an approximately uniform manner across the resultant clusters. Concurrently, the documents of class D1 would be allocated to both of these clusters, as guided by the partitioning procedure, regardless of the query distribution. However, supposing scenario in which an overwhelming majority of the queries target D2, the goal emerges rigorously isolate the clusters associated with D and D2 since we would like to improve latency in the downstream evaluation and visit fewer documents for each query. In this context, the EHI emerges a viable solution it takes the query distribution into consideration, effectively addressing this concern by facilitating a clear delincation between the clusters corresponding to the distinct domains. Figure 1(b) depicts the query distribution of EHI model trained on the MS MARCO dataset. It. is important here to highlight that for tail queries, which are barely searched in practice, EHI is able to grasp the niche concepts and hash as few documents in these buckets as indicated by the very few average number of documents visited. Moving towards more popular queries, or the torso queries, we note rise in the average number of documents visited as these queries are more broad in general. Moving even further to
2

---

Published in Transactions on Machine Learning Research (08/2024)
58000
6000
1000
Learned Embedding siructure by
learning embeddings alignment with ANNs structure
400
Query popularity
(a) Misalignment of representations
(b) Query Distribution MSMARCO
Figure Motivation of EHI.
the head queries, we again nolice a drop in the average number of documents visited, which showcases the efficiency and improved latency the EHI model.
1.2 Overview and Evaluation
Motivated by the aforementioned issues, we propose EHT- End-to-end learning of Hierarchical Index that jointly trains both the encoder and the search data structure; see Figure 3. To the best of our knowledge, EHI is the first end-to-end learning method for dense retrieval. Similar to two-stage approaches, EHI employs tree-structured inverted file-like index, but unlike DE+ANNS (the conventional two-stage approach where standard Dual Encoder (DE) is trained independently, followed by an Approximate Nearest Neighbor Search (ANNS) using the learned embeddings), EHI integrates the training process to simultaneously learn the discrete document assignments within its index and refine the encoder parameters without any warm-starting, We can learn balanced and roughly uniform clusters with just the query, document pairs information which is generally available in the supervised training paradigm of most encoders, and does not require any proxy information through KNN other clustering algorithms for learning the INF-style tree structure. EHI's core innovation lies in the introduction of dense path embeddings. These embeddings represent the traversal paths of queries and documents within the tree index, enabling the learning of discrete document assignments. By optimizing the index such that semantically similar (query, document) pairs share similar path embeddings, EHT facilitates efficient. retrieval by elustering relevant pairs within the same leaf nodes.
We conduct an extensive empirical evaluation our method against SOTA techniques on standard benchmarks. Our experiments on the popular MS MARCO benchmark (Bajaj et al., 2016) demonstrate that EHI shows improvements of 0.6% in terms of nDCG@10 compared to dense-retrieval with ScaNN baselines when only 10% of documents are searched. We attribute these improved embeddings to the fact that EHI enables an additional signal to the encoder from the indexer as well the integrated hard negative mining as it can retrieve irrelevant or negative documents from indexed leaf nodes of a query. Here, the indexer parameters are always fresh, unlike techniques akin to ANCE (Xiong et al., 2020). Similarly, EHI provides upto 8.2% higher nDCG@10 than state-of-the-arl (SOTA) baselines on the MS MARCO TREC DL19 (Craswell et al., 2020) benchmarks for the same compute budget. EHI also achieves SOTA exact search performance on both MRR@10 10 and nDCG@10 metrics with up 80% reduction in latency, indicating the effectiveness of the joint learning objective. Similarly, we outperform SOTA architectures such NCI NQ320k by 0.6% and 1.8% on Recall@10 and Recall@100 metrics. (see Section 4.2).
1.3
Contributions
Our work introduces several key advancements to the field of dense retrieval:
Novel End-to-End Paradigm. We propose End-to-end Hierarchical Indexing (EHT), an end-to-end learn- ing framework that jointly optimizes the embedding encoder and the search index structure. EHI addresses the inherent misalignment conventional two-stage approaches, enabling more efficient. and accurate retricval (Section 3, Figure 2). To the best of our knowledge, we are the first propose this without any warm starting soft labels.
3

---

Published in Transactions on Machine Learning Research (08/2024)
Rigorous Empirical Validation. We conduct comprehensive evaluations on the industry-standaar MS MARCO benchmark, demonstrating EHI's superiority over state-of-the-art ANNS techniques such as ScaNN, Faiss. Furthermore, we also showcase that EHI is competitive with other SOTA baselines such as ColBERT, SGPT, cpt-text, ANCE, DyNNIBAL. Our experiments underscore EHT's focus on enhancing retrieval accuracy within a fixed computational budget (Section 4.2, Appendix 4.3).
Framework Flexibility. EHI is designed to be agnostic to specific encoder architectures, similarity metrics, and hard negative mining strategies. This flexibility underscores its potential for broad applicability within the dense retrieval domain.
2 Related Works
Dense retricval (Mitra et al., 2018) underlies a myriad of web-scale applications like search (Nayak, 2019),
recommendations (Eksombatchai et al., 2018; Jain et al., 2019), and powered by (a) learned representa-
tions (Devlin et al., 2018; Kolesnikov et al., 2020; Radford al., 2021), (b) ANNS (Johnson et al., 2019;
Sivic & Zisserman, 2003; Guo et al., 2020) and (c) LLMs in retrieval (Tay et al., 2022; Wang et al., 2022;
Guu et al., 2020). Figure 2 illustrates the difference between dual encoder trained with an ANNS (DE +
ANNS), Generative Retrieval approaches (such DST, NCT) and EHI.
DE ANNS
Rentritive-Retrieval
EHI
Stage-1
Stage-1
Stage-2
Encoder
BERT Model
Encoder
Encoder
Stage-2
K-Maans
ANNS
Decoder
Indexee
Lears next sequence prediction with generated semantic target
2
Semantic
Figure 2: Comparison between EHI's end-to-end training paradigm in comparison to other approaches such Dual encoder + ANNS and Generative Retrieval approaches which require a multi-stage training. We denote stages which are not optimized for the task objective by grey shade. Note that unlike DE + ANNS, and Generative Retrieval, EHI optimizes for recall in both encoder and indexer.
Representation learning. Powerful representations are typically learned through supervised and un/self- supervised ised learning paradigms that use proxy tasks like masked language modeling (Devlin et al., 2018) and autoregressive training (Radford et al., 2018). Recent advances in contrastive learning (Gutmann & Hyvärinen, 2010) helped power strong dual encoder-based dense retrievers (Ni et al., 2021; Izacard et al., 2021; Nayak, 2019). They consist of query and document encoders, often shared, which are trained with contrastive learning using limited posítively relevant query and document pairs (Menon et al., 2022; Xiong et al., 2020). While most modern-day systems use these learned representations as is for large-scale ANNS, there is no need for them to be aligned with the distance metrics or topology of the data structures. Recent works have tried to address these concerns by warm-starting the learning with a clustering structure (Gupta et al., 2022) but fall short of learning jointly optimized representations alongside the scarch structure. We explicitly state the challenges of directly adapting hots like ELIAS, originally designed for classification, to the dense retrieval task. However, we include a comparison against ELIAS on a classification benchmark with label features in Appendix C.3. Other works such as RepCONC (Zhan et al., 2022), and SPLADE (Formal et al., 2022) also work on the efficiency aspect of retrieval, where they focus on quantization of the representations using regularizers which explicitly work on reducing FLOPS.
4

---

Published in Transactions on Machine Learning Research (08/2024)
Approximate nearest neighbor search (ANNS). The goal of ANNS is to retrieve almost nearest neighbors without paying exorbitant costs of retrieving true neighbors (Clarkson, 1994; Indyk & Motwani, 1998; Weber et al., 1998). The "approximate" nature comes from rruningbbssed search data structures (Sivic & Zisserman, 2003; Malkov & Yashumin, 2020; Beygelzimer et al., 2006) as well as from the quantization based cheaper distance computation (Jegou el al., 2010; Ge et al., 2013). This paper focuses on ANNS data structures and notes that compression is often complementary. Search data structures reduce the number of data points visited during the search. This is often achieved through hashing (Datar et al., 2004; Salakhutdinov & Hinton, 2009; Kusupati et al., 2021), trees (Friedman et al., 1977; Sivic & Zisserman, 2003; Bernhardsson, 2018; Guo et al., 2020) and graphs (Malkov & Yashumin, 2020; Jayaram Subramanya et al., 2019). ANNS data structures also carefully handle the systems considerations involved in a deployment like load-balancing, disk I/O, main memory overhead, elc., and often tree-based data structures lend to prove highly performant owing to their simplicity and flexibility (Guo et al., 2020). For a more comprehensive review of ANNS structures, please refer to Cai (2021); Li et al. (2020); Wang et al. (2021). Works such as CCSA (Lassance el al., 2021) propose alternate ANN structures for efficient retrieval via constrained clustering. Other works such as TDM (Zhu et al., 2018) proposes learning a trke-like ANNS but the ANNS is learnt in a disjoint fashion where k-means algorithm is used for learning the clustering without any supervision. Works such as JTM (Zhu et al., 2019), and Deep Retrieval (Gao et al., 2020b) also attempt to optimize the learn the indexer with a neural network- - but they require pre-computation of optimal paths taken by the query and document for learning. Prior works including PQN (Yu et al., 2018) and GPQ (Jang & Cho, 2020) have explored end-to-end ANNS techniques in the visual domain which primarily rely on hashing- based approaches, and not an hierarchical tree-like like structure. EHI differentiates itself through its novel tree structure and dense path embeddings, tailored for the unique characteristies of text data and semantic search. Furthermore, EHI handles significantly larger number of buckets (e.g., 7000) compared to the 48- hashing used in approaches like GPQ, demonstrating its suitability for the scale of textual datasets with millions of documents. Similarly, the scale of the experiments are also significantly larger as PQN, GPQ showcase their efficacy on ~170K images to be indexed (where the overhead of exact search might not be very large), while we experiment on 10M documents, where exact search is a significant overhead during inference.
EHI works in end-to-cnd fashion without any pre-computation of optimal paths, and is the first end-to-end hierarchical indezer framework to the best of our knowledge.
Generative Retrieval for Semantic Search. Recently, there have been some efforts towards modeling retricval as a sequence-to-scquence problem. In particular, Differential Search Index (DSI) (Tay et al., 2022) and more recent. Neural Corpus indexer (NCT) (Wang el al., 2022) method proposed encoding the query and then find relevant document by running a learned decoder. However, both these techniques, at their core, use a seporately computed hierarchical k-means-basse clustering of document embeddings for semantically assigning the document-id. That is, they also index the documents using an ad-hoc clustering method which might not be aligned with the end objective of improving retrieval accuracy. In contrast, EHT jointly learns both representation and k-ary tree- tree-based search data structure end-to-end. This advantage is reflected on MS MARCO dataset EHI is uplo 7.12% more accurate (in lerms of nDCG@10) compared to DSI. Recently, retrieval has been used to augment LLMs also (Guu et al., 2020; Izacard & Grave, 2020b;a; Izacard et al., 2022). We would like to stress that the goal with LLMs is language modcling while retricval's goal precise document retrieval. However, retrieval techniques like EHI can be applied improve retrieval subcomponents in such LLMs.
3 End-to-ent Hieararchical Indexing (EHI)
Problem definition and Notation. Consider problem with a corpus of N documents D = {di, .... dx }. a set of Q training queries Q = (.),..,, 42); and training data (7i, da, sia), where yik E (-1,1} is the label for a given training (query, document) tuple and yik = 1 denotes that dy is relevant to q. Given these inputs, the goal is to learn a retriever that maps a given query to a set of relevant documents while minimizing the computation cost. While wall-clock time is the primary cost metric, comparing different methods against it is challenging due to very different setups (language, architecture, parallelism, etc.). Instead, we rely on recall vs.
5

---

Published in Transactions on Machine Learning Research (08/2024)
% searched curves, widely considered reasonable proxy for wall-clock time modulo other eetup/environnan changes (Guo et al., 2020).
3.1
Overview of EHI
EHI
Pie
Tsianese Van
Ishone walman cost
Pazt
:28
Walmart is selling the iPhone SE for $299.
Encoder
Vur
Tentecing (Par Pixz Pane)
nudlous
176
Pan
Linmeted Parl
Figure 3: EHI is end-to-end hierarchical indexer which comprises an encoder and hierarchical tree the indexer where the entire pipeline is learnable and differentiable. Here, variables V98, V123, and Vs76 are dense represental ions (embeddings) of the text and Pas, P123, and P376 are path embeddings of the respective samples. To efficiently train EHI without any warm starting, we use combination of objectives Lslameses Lindexing: Lintra-leaf (sce Section 3 for dotails).
At a high level, EHT has three key components: Encoder &g, Indexer Ip and Retriever. Parameters 0 of the encoder and á of the indexer are the trainable parameters of EHI. Unlike most existing techniques, which train the encoder and indexer in a two-step disjoint process, we train both the encoder and indexer parameters jointly with an appropriate loss function; see Section 3.5. Learning the indexer generally a discontimuous function a combinatorial problem that also requires multiple rounds of indexing the entire corpus. However, by modeling the indexer using a hierarchical tree and its "internal represenlation" compressed path embedding, demonstrate that the training and retrieval with encoder+ fndexxr can be executed efficiently and effectively.
In the following sections, we provide details of the encoder and indexer components. In Section 3.4, we detail how encoder-indexe can be used to retrieve specific documents for a given query, which is used for inference and negative mining (further elaborated in Appendix C.2) during training. Section 3.5 provides an overview of the training procedure. Finally, Section 3.6 summarizes how documents are ranked after retricval.
3.2 Encoder Eg: Dense embedding of query/documents
Our method agnostic to the architecture used for dual encoder. But for simplicity, we use standard dual encoder (Ni et al., 2021) to map input queries and documents a common vector space. That is, encoder Ep parameterized by 0, maps query (q E Q) and document (d E D) to common vector space: Eg(g) E _m, and Ed(d) E R²" where m is the embedding size of the model (768 here). While such encoder can also be multi-modal as well as multi-vector, for simplicity, we mainly focus on standard textual data with single embedding per query/document.. We use the standard BERT archileclure for encoder &g and initialize parameters using a pre-trained Sentence-BERR distilbert model (Reimers & Gurevych, 2019). Our base model has layers, 768 dimensions, 12 heads with 66 million parameters. We then fine-tune the final layer of the model for the target downstream dataset.
3.3 Indexer Iz: Indexing of qucry/document in the hierarchical data structure
EHT's indexer (Iz) is a tree with height H and branching factor B. Each tree node contains classifier that provides a distribution over its children. So, given a query/ muery/dccumen, we can find out the leaf nodes that the query (derymdoeument indexes into, well as the probabilistic path taken in the tree.

---

Published in Transactions on Machine Learning Research (08/2024)
The final leaf nodes reached by the query are essential for retrieval. But, also propose to use the path taken by a query/ soryyddoument in the tree as an embedding of the query/document - which can be used in training through the loss function. However, the path a query/docunee (in height H) vector space, owing to B leaf nodes, making it computationally intractable even for a small H and B. For instance, with B 10 and = 5, a relatively small tree, there would be 100,000 (10*) possible paths.
takes is an object in an exponentially large
Instead, below, we provide significantly more compressed path embedding denoted by ((: and parameterized by Ø embeds any given query document in relatively low-dimensional (B H)) vector space (which would lead to just 50 in the above example). For simplicity, we denote the query and the document path embedding as Ta(g) = T(Ex(4); o) and (d( - T(Ea(d); o), respectively.
We construct path embedding a query/document as:
T(Ea(a)) ((Eo())) = |p :p -1 :p'),
Where p" E: [0, 1 denotes the probability distribution of children nodes for a parent at height h. For a given lcaf 1, say path from root node is defined as 1= i7 .i. i(") where if (1... B] for E [H]. The probability at a given height in path is approximated using a height-specific simple feed-forward neural network parameterized by Wh+1 E (BBhtm)xB and U E (B ( (B+ (m is the embedding size). That is,
h Softmax W 44 Kyi Uk+1)) [''
(1)
where Kn. =F(((()):)0 : (4)), and one-hot-vector o(i) is the i-th canonical basis vector and F is non-linear transformation given by F(x; U,) = x + Re.
In summary, the path embedding for height 1 represents a probability distribution over the leaves. During training, we compute path embedding for higher heights for only the most probable path, ensuring that the summation of leaf node logits remains a probability distribution. Also, the indexer and path embedding function T(: 6) has the following collection of trainable parameters: & = [WH,... W1, ......y U1}, which we learn by optimizing loss function based the path embeddings; see Section 3.5. Please refer to Appendix E for additional details about the indexer in EHI.
3.4 Retriever: Indexing items for retrieving
Indexing and retrieval form backbone for any search structure. EHI efficiently encodes the index path of the query and documents in (B =)ddimensia embedding space. During retrieval for query q, EHI explores the tree structure to find the "most relevant" leaves and retrieves documents associated with those leaves. For retrieval, it requires encoder and indexer parameters (0 along with Leaf, document hashmap M.
The relevance of a leaf I for a query q is measured by the probability of a query reaching leaf at height H (P(g. 1, H)). Recall from previous section that path to a leaf I defined 1 = [t.. iff". where z/t E [1... B) for h [H]- The probability of reaching leaf for given query q E Q to an arbitrary leaf E Leaves can be computed as P(g, (, ""[ "i'' using equation 1. But, we only need to compute the most probable leaves for every query during inference, which we obtain by using the standard beam-search procedure summarized below:
1. For all parent node height h 1, compute probability of reaching their children $ Uceen11a(p) P(q, c, Yper
2. Keep top 3 children based on score S and designate them as the parents for the next height.
Repeat steps 1 and 2 until the leaf nodes are reached.
Once we select 8 leaves EHI retricves documents associated with cach Icaf, which is stored in the hashmap M. To compute this hash map, EHI indexes each document d E D (similar to query) with , = 1. Here, 3 == 1 is a design choice considering memory and space requirements and kept as a tuneable parameter. It stores document-to-leaf associations, allowing us to strategically focus the search on the most relevant subsot of documents. The hashmap M used store the information of which documents reach which leaves in our tree. This is essentially the result of the indexing step. Note that to add new ad-hoc documents after
7

---

Published in Transactions on Machine Learning Research (08/2024)
training EHI, one would need to follow the indexing step listed above to store the new documents in the appropriate buckets and update the hashmap M.
During retrieval, we find out which leaves the query goes to, and we only use those documents which have reached similar leaves. This is how we sample a few documents to search over instead of searching over the entire document space. Algorithm 2 in the appendix depicts the approach used by our Indexer for better understanding.
3.5
Training EHI
Given the definition of all three EHI components encoder, indexer, and retriever we are ready to present the training procedure. As mentioned earlier, the encoder and the indexer parameters (0; are optimized simultaneously with our proposed loss function, which is designed to have the following properties: a) Relevant documents and queries should be semantically similar, b) documents and queries should be indexed together if they are relevant, and c) documents should be indexed together iff they are similar.
Given the encoder and the indexer, we design one loss term for each of the properties mentioned above and combine them to get the final loss function. this end, we first define the triplet loss as:
L(a, b, c) [ac ab + +
(2)
where we penalize if similarity between query q and an irrelevant document d_ (v(g, 1) is within y margin of the corresponding similarity between g and relevant document d+ (vlg, d+) 1). We now define the following three loss terms:
1. Semantic Similarity: the first term is standard dual-encoder contrastive loss between relevant document d4 i.e., p(g, d+) = +1- and an irrelevant document with v(g, d-) * 1.
Liiamese L(E)))).
(3)
2. Indexing Similarity: the second term is essentially a similar contrastive loss over the query, relevant-doc, irrelevant-doc triplet, but where the query and documents are represented using the path-embedding (- given by the indexer Ig.
Lindexing L(Tx(g)) Tod_;))
(4)
3. Intra-leaf Similarity: to spread out irrelevant docs, third loss applies triplet loss over the sampled relevant and irrelevant documents for a query 4- Note that we apply the loss only if the two does are semantically dissimilar according to the latest encoder, i.e., SIM(a, b, Ea(o) (b(h) pre-specified threshold 0.9.
for
[staa]
Tintra-leaf = =IMM(4,,,,, (a(L)
(5)
The final loss function L is given as the weighted sum of the above three losses:
C X1 Lsiameso Lindexin Ag
(6)
Here is set to 0.3 for all loss components, and Å₁ A2, Ag are tuneable hyper-parameters. Our trainer (see Algorithm 1) learns 0 and & by optimizing using standard techniques; for our implementation we used AdamW (Loshchilov Hutter, 2017). Note that the loss function only uses in-batch documents' encoder embeddings and path embeddings, i.e, we are not even required to index all the documents in the tree structure, thus allowing efficient joint training of both encoder and indexer. To ensure fast convergence, we use hard negatives mined from the indexed leaves of given query q for which require documents to be indexed in the tree. But, this procedure can be done once in every step where r is hyper-parameter set to by default across our experiments. We will like to stress that existing methods like DSI, NCI, or ANCE not only have to use stale indexing of documents, but they also use stale or even fixed indexers like DSI, NCI learns a fixed semantic structure over docids using one-time hicrarchical clustering. In contrast, EHI jointly updates the indexer and the encoder in each iteration, thus can better align the embeddings with the tree/indese indexer.
8

---

Published in Transactions on Machine Learning Research (08/2024)
3.6 Re-ranking and Evaluation
This section describes the re-ranking step and the lest-time evaluation process after retrieval. In Section 3.4, we discussed how each document is indexed, and we now have learned mapping of d x L, where d is the corpus size, and I the number of leaves. Given a query at test time, we perform a forward pass similar to the indexing pipeline presented in Section 3.4 and find the top-b leaves (b here is the beam size) the given query reaches. We collate all the documents that reached these leaves (set operation to avoid any repetition of the same documents across multiple leaves) and rank them based on an appropriate similarity metric such cosine similarity, dot product, manhattan distance, etc. We use the cosine similarity metric for ranking throughout our experiments (see Section 4.2).
4 Experiments
We conduct comprehensive empirical evaluation of EHI on established dense retrieval benchmarks, Our experiments have two primary objoctives:
(a) Validate the End-to-end Paradigm: We seek to demonstrate the superiority of EHT's joint optimization of the embedding encoder and Indexer compared the conventional disjoint training approach that utilizes off-tho-shelf ANNS methods (c.g., ScaNN, Faiss-IVF) alternative generative retrieval or sparse retrieval approaches. This comparison provides strong evidence for the benefits of our proposed paradigm shift.
(b) Justify Design Decisions: Through systematic experimentation, we aim to elucidate the rationale behind the design choices within EHI. These insights will solidify the theoretical foundations of our method and guide future research in this domain.
Detailed training hyperparameters for EHI are provided in Appendix B.
0.372
0.88
0.890
0.370
0.86
MRR @
0.368
0.84
@
0.885
DE Exact Search ScaNN
0.366
+0E + ScaNN
0.82
+DE ScaNN
DE + Faise-P9 IVF
0.880
Faiss-IVF
0.364
*+EHI (Ouns)
0.80
+ EHil (Ouns
*HH (Durs)
10 % Documents visited
100
Documents visited
% Documents visited
(a) MSMARCO: MRR@10
(b) TREC DL19: nDCG@10
(c) NQ320k: R@10
Figure 4: EHI is significantly more accurate than DE + ScaNN or Faiss-IVF, especially when restricted visit small fraction of documents. See Figure 10 in Appendix for results Scifact, Fiqa.
4.1 Experimental Setup
Datasets. We evaluate EHI on four standard but diverse retrieval datasets of increasing size: SciFact (Wad- den et al., 2020), FIQA (Maia et al., 2018), NQ320k (Kwiatkowski et al., 2019), and MS MARCO (Bajaj et al., 2016). Appendix A provides additional details about these datasets.
We establish a robust set of basclines to comprchensively evaluate EHI's performance and isolate the impact of our core contributions:
State-of-the-Art (SOTA) Comparisons. We evaluate EHI against leading dense retrieval methods such as ColBERT-v2 (Santhanam et al., 2021), SPLADE (Formal et al., 2022), and many other basclines on the MS MARCO benchmark (see Table 1, Section 4.3). These comparisons, while acknowledging architect.ral differences, demonstrate EHI's strong performance and its potential within the broader research landscape. Importantly, EHI achieves this competitivenoss using a significantly small encoder (66M parameters), high- lighting the effectiveness of our end-to-end index learning approach. We note that re-ranking and other late interaction methods employed in ColBERT-v2 and SPLADE currently outperform EHT in exact search

---

Published in Transactions on Machine Learning Research (08/2024)
scenarios on the MS MARCO benchmark. This presents promising future research direction to explore how re-ranking losses and other techniques could be integrated into EHT. Note that methods like ColBERT-v2, due to their late-interaction similarity computation, are less readily adaptable to ANNS-based retrieval.
Controlled Ablation Studies. To directly assess the value of EHI's end-to-end end indexer, we construct baselines using the same DistilBERT encoder paired with: (a) Dual Encoder (DE with ANCF, hard-negatives; Menon et al. (2022)) + Exact Search (for theoretical upper bound),
(b) DE + ScaNN (representative of popular ANNS libraries; Guo et al. (2020)), and (c) DE + Faiss-IVF (for comparison with another IVF-style index; Johnson et al. (2019)).
Generative Retrieval. We include baselines like DSI (Tay et al., 2022) and NCI (Wang et al., 2022) to provide context, acknowledging their potential limitations in this setting. We report DSI numbers on MS MARCO using an implementation validated by the authors. However, we note that NCI fails to scale to large datasets like MS MARCO.
Sparse Neural Retrieval. We include sparse neural IR baselines like CCSA (Lassance el al., 2021) and SPLADE (Formal et al., 2022) for a comprehensive comparison in Table 1. However, it's important to note their potential differences in our context. Methods like CCSA, with their two-step disjoint training, face inherent challenges similar to those discussed earlier. While SPLADE offers notable improvements by replacing ANNS with sparse representations and using a FT.OP-minmmzzing regularization term, this approach necessitates training multiple models for different FLOP targets of the practitioner. In contrast, EHI achieves this fiexibility and elasticity within a single model. Furthermore, as Lassance et al. (2023) observes, SPLADE models can exhibit slower inference speeds compared to dense IR models which could be alleviated with static pruning, but the carlier limitation still holds water. However, the distillation-based improvements achieved by SPLADE are promising, and exploring the integration of distillation techniques with EHI represents an exciting direction for future research.
4.2 Results
Evaluations and findings on Scifact and Fiqa datasets is presented in Appendix C. We evaluate EHI on the MS MARCO passage retrieval task, a widely recognized benchmark for semantic search. Results are presented for both the standard dev set and the TREC DL-19 set (Craswell et al., 2020). We also report our findings on other datasets such as NQ320k (Kwiatkowski et al., 2019). We compare against the standard Sentenco-BERT model (Huggingface, 2019), trained on the MS MARCO dataset, which is further fine-tuned to the various datasets compared in this paper.
EfficinnyyAAccuracy Tradc-off: EHI demonstrates a compelling advantage in the efficiency-accuracyy trade-off. For instance, on the standard MS MARCO dev set, EHI matches or exceeds the accuracy of Exact Search while visiting 80% fewer documents (see Figure 4(a)). This significantly outperforms NcNN and DEiFaass-IVV -VFF, which require visiting nearly double the documents for comparable accuracy. Furthermore, on the TREC-DL 19 benchmark, EHI is able to match surpass the nDCG@10 of baseline Exact Search with an 78% reduction in latency (see Figure 4(b)). On the NQ320A dataset, EHT matches or surpasses the accuracy of baseline Exact Search with a 60% reduction in latency (see Figure 4(c)).
Gains in compute bottlenecks: We emphasize the significance of EHI's 3.6% nDCG@10 improvement (with only 1% documents visited) on the well-optimized MS MARCO benchmark. Such gains underscore the practical value of our approach. Similarly, when restricted to visiting 1% of the documents, EHI achieves 8.2% higher nDCG@n @10 than DE+ScaNN and Faiss-VVVV on the TREC-DL benchmark.
Comparison with Generative Retrieval Methods: We note that the DSI base model with 250M parameters is almost four times the size of the current EHI model. After multiple weeks of DSI training with doc2query + atomic id + base model, DSI achieves a MRR@10 metric of 26% on the MS MARCO dev set, which is 8.58% lower than EHI with just 1% visited documents (see Table 1). Note that despite significant efforts, we could not scale NCI code (Wang et al., 2022) on MS MARCO due to the dataset size; NCI paper does not provide metrics on MS MARCO dataset. To compare against NCI, we report our findings on the NQ320k dataset. Note that EHI is able to significantly outperform DSI and NCT (without query generation)
10

---

Published in Transactions on Machine Learning Research (08/2024)
Table 1: Performance metrics (%) evaluated on the MS MARCO dev dataset. The best value for each metric is indicated in bold font.
Method
NERC10 (Dev)
aDCG010 (Dev)
parameters
DPR Thakar ot al. (2021)
17.7
110M (BERT busee)
BM25 (official) Khattab & Zaharia (2020)
16.7
22.8
BM25 (Anserini) Khattab & Zaharia (2020)
18.7
cpt-taxt S Neelalzantan et al. (2022)
19.9
300M (ept-ttect S)
opt-text M Neelskantau et al. (2022)
20.6
1:28 (ept-texct M)
opt-text L Neelakanten et sl. (2022)
21.5
GB (ept-text
cpt-taxt XI. Noclakantan et al. (2022)
22.7
175 (ept-teex XI.)
DSI (Atounic Docid + Doc2Queryy + Buse Model] Tay et (2022)
26.0
32.28
250M (115-buse)
DSI (Naive String Dockd Doc2Query + XL Model) Tay al. (2022)
21.0
3B (TS-XL)
DST (Naive String Docid Doe2Query 4 XXT. Model) Tay et al. (2022)
16.5
11B (TS- (T5-XXI))
DSI (Semmntic String Docid + Doc2Query + XL Model) Tay et al. (2022)
20.3
27.86
3B (TS-XL)
OCSA Lasssnce et al. (2021)
28.9
110M (BERI- -base)
HNSW Malkow & Yoshunin (2018)
28.9
RoBEICTa- base + In-batch Negutives Momath et al. (2021)
24.2
123M (RoBEITT- base
HoBERTi-base Cniform Negatives Monath al. (2023)
125M (RoBERTe- base)
RoBERTa-base DyNNIBAL Monach et al. (2023)
33.4
123M (RoBERTa-lee
RoBBEITa--bss: Stochastie Negutives Monath et al. (2021)
33.1
123M (RoBERTa- base)
HoBEItIio-ase- +Nggative Cache Monath al. (21223)
33.1
(RoBERTI base)
RoBERTa-base Exhaustive Negatives Monsth es al. (2023)
31.5
123M (RoBERTa-less
SGPT-CF-2.7B Mumnighooff (2022)
27.8
2.77 (GPT-Neo)
SGIT-CE-661B Munnnighoft (2022)
6,1 (GPT----G))
SGPTT-E-55SB Muennighoff (2022)
39.9
5.8B (GPT-J)
doc2query Khattab Zaharia (2020)
21.5
DespTT Thakur
(2021)
24.3
110M (BERTT -luse)
doeTTTTqqery Klattab & Zabaria (2020)
27.7
SPARTA Thskor es al. (2021)
35.1
110M (BERT-base)
docuSqu Thakur et al. (2021)
33.8
ANCE
33.0
38.8
EII (distilbert-co:; Ours)
33.8
39.4
66M (distilBerr)
RepCONC Zhau et al. (20122
31.0
123M (RoBERTx bass))
TAS-B Thakur al. (2021)
40.8
110M (BERT-base)
GeaQ Thakur et al. (2021)
40.8
ColBERT (re-rauk) Khattab Zakaria (2020)
34.8
110M base)
CalBERT (end-to-end) Khattab & Zaharia (2020)
36.0
40.1
110M (BERT-bess
BM25 + CE Thakur et al. (2021)
41.3
SPLADE (simple training) Formul es al. (2022)
34.2
110M (BERT -uuse)
SPLADE DistiMME Formel
(2022)
35.8
66M (distiBERR))
SPLADE SelfDissil Formal et al. (2022)
36.8
66M (dit
SPLADE EaserabieDistil Formal et al. (2012)
36.9
66M
(distlBERTT)
EHI (distibert-dol; Ours)
37.2
(distIBERT)
SPLADE SoCoondenee-SelfDDia Formal et al. (2022)
37.5
60M (dittiBERT)
SPLADE + CoCandmeer-Eeeeee Ensoibbl:Dist Formal et al. (2022)
38.0
66M (distilBERT)
CoIBER1-22 Santhamam et (2012)
39.7
(BLRII-bae)
despite NCT utilizing 10x larger encoder! Furthermore, even with query generation, NCT is 0.6% and 1.8% less accurate than EHI Recall@10 and Recall@100 metrics, respectively. (sce Table 3)
Our observations about EHI are statistically significant as evidenced by p-value tests Section 4.5. Additional experiments such as the results on other benchmarks such as scifact, fiqa, negative mining, comparison against Elias (Gupta et al., 2022), and qualitative analysis on the leaves of the indexer learned by the EHI model is depicted in Appendix C.
4.3 Comparisons to SOTA
In this section, we compare EHI with SOTA approaches such ColBERT, SGPT, cpt-text, ANCE, DyN- NIBAL, which use different dot-product metrics or hard-ncgative mining mechanisms, etc. Note that EHI proposed in this paper is a paradigm shift of training and can be clubbed with any of the above architectures. So the main comparison is not with SOTA architectures, but it is against the existing paradigm of two-stage dense retricval as highlighted in Section 4.2. Nevertheless, we showcase that EHI outperforms SOTA ap- proaches on both MS MARCO dev set as well as MS MARCO TREC DL-19 setting as highlighted in Table and Table 2 respectively. Table 3 depicts the performance of EHI in comparison to other SOTA approaches.
11

---

Published in Transactions on Machine Learning Research (08/2024)
Table 2: Performance metrics (%) evaluated on the MS MARCO TREC-DL 2019 Craswell et al. (2020)
dataset. The best value for each metric is indicated in bold font.
Method
MRR010
mDCG810
parameters
BM25 Robertson al. (2009)
68.9
50.1
CCSA Lassance et al. (2021)
58.3
110M
(BRRT- -hase))
RepCONC Zhan et al. (2022)
66.8
123M (RoBERTa basee)
atass-attention BERT (12-layer) Menon et. al. (2022)
82.9
74.9
110M
BERT-base
Dual Encoder BERT (6-layer) Menon et al. (2022)
83.4
67.7
66M
(distiBEETT
DistiBERRT + MSF Menon et al. (2022)
78.1
69.3
66M
(dis))BR)))
SPLADE (simple training) Formal al. (2022)
69.9
110M
(BERT-base)
DistIBERT + Margin MSE Menon al. (2022)
86.7
71.8
66M
(distIBERT)
DistiBEETT + RaukDistil-B Menon et al. (2022)
85.2
66M
(distiBBERT)
DistiBERT + Softmma CE Menon al. (20222
66M
(distiBRRT)
DistiBBERT + MESE Menon
(2022)
85.2
66M
(distilBERT)
SPLADE DistiMSE Formal
(2022)
72.9
66M
(distBERTT)
SPLADE SelfDistil Formal et al. (2022)
72.3
66M
(distiBEETT
SPLADE EusembleDistil Formal (2022)
72.1
66M
(distiBERT)
SPLADE CoCondenser-SSfffD Formal et al. (2022)
73.0
66M
(distiBBERT)
SPLADE CoCondenser- CoCondlouee-BBseellllllll Formal et al. (2022)
73.2
66M
(disUBERT))
EiIl (distilbert-cos; Ours)
97.7
88.4
66M
(distiBERTT)
Table 3: Performance metrics (%) evaluated on the NQ320k dataset (Kwiatkowski et al., 2019). The best value for each metric indicated bold font.
Method
R@10
RØ100
# parameters
BM25 Robertson et al. (2009)
32.48
50.54
BERT ANN (Faiss) Johnson et al. (2019)
53.63
73.01
110M SBERT-base
BERT
BruteForce Wang et al. (2022)
53 42
73.16
110M (BERT-base
BM25 DocTryuury Nogueira et al. (2019)
61.83
76.92
ANCE (FirstP) Xiong al. (2020)
80.33
91.78
ANCE (MaxP) Xiong et al. (2020)
80.38
91.31
SEAL (Base) Bevilacqua et (2022)
79.97
91.39
139M
(BART- -Baee-
SEAL (Large) Bevilacqua et al. (2022)
81.24
91.93
406M (BART-large
RoBERTa-base In-batch Negatives Monath et al. (2023)
69.5
84.3
123M
(RoBEITTa-base))
RoBERITa-base- + Uniform Negatives Monath et al. (2023)
72.3
84.8
123M
(RoBENTTa-base)
RoBERTa-base + DyNNIBAL Monath al. (2023)
75.4
86.2
123M
(RoBEIITa-base)
RoBERTaabasse + Stochastic Negatives Monath al. (2023)
75.8
86.2
123M
(RoBERTa-base)
RoBERTa-hase -eegia---taseeeeeNeeie Negative Cache Monath et al. (2023)
85.6
123M
(RoBERTa-base)
RoBERTa-base + Exhaustive Negatives Monath al. (2023)
80.4
86.6
123M
(RoBERTa-bas.
DSI (Base) Tay et al. (2022)
56.6
250M (15-base)
DSI (Large) Tay al. (2022)
62.6
3B (T5-XL)
DSI (XXL) Tay et al. (2022)
70.3
11B (T5-XXL)
NCI (Base) Wang et al. (2022)
85.2
92.42
220M (T5-base)
NCI (Large) Wang et al. (2022)
85.27
92.49
770M (TT-L)
NCI w/ gg-ft (Base) Wang et al. (2022)
88.48
94.48
220M (T5-
NCI w/ qg-ft (Large) Wang et al. (2022)
88.
94.53
770M (T5-L)
EHI (distilbert-cos; Ours)
88.98
96.3
66M (disuiBERTT)
12

---

Published in Transactions on Machine Learning Research (08/2024)
Scifact
MS Marco
1.0
0.61
Fiqa
0.394
nDCG @ 10
Recall @ 100
0.9
Recall 100
0.56
0.392
0.390
0.8
0.51
EHH Indecer
0.1|2)
EHI (Ours)
1==
0.388
10
30 % Documents visited
20
Faiss-VV
SCANN
0.7
20
30
40
50
Exact Search
60
0.490
A1 0
50
A
100
Branching Factor
% Documents visited
(a) Efficacy of EHI Indexer
(b) Effect of Branching factor
(c) Effect of Loss components
Figure 5: Ablations studies depicting the various properties of EHI's training paradigm.
4.4 Robustness to initializations
In this section, we showcase the robustness of our proposed approach- - EHT to various initializations of the model. Note that analysis in this manner is uncommon in efficient retrieval works, as showcased in ScaNN, DSI, and NCI. However, such approaches worked with two disjoint processes for learning embeddings, followed by the ANNS. Since we attempt to solve the problem in end-to-end framework in differentiable fashion through neural networks, it is worth understanding the effect of different initializations seeds on our learning. Overall, run the same model for the best hyprraarameters over five different seeds and report the mean and standard deviation numbers below in Table 4.
Table 4: Performance metrics (%) evaluated the various datasets across various ranges of visited documents. The best value for each metric, corresponding to specific number of visited documents, is indicated in bold font. Here, n denotes the number of sceds run for the given model. We report the mean and standard deviation of the metrics for EHI and showcase a robustness to initialization.
Metric
Method
20%
30%
502
70%
80%
96%
100%
Search
91.77
Selbet - R31000
ScnNN
67.80
80.76
83.78
87.49
89.82
00.6
91.8
9L:27
01.27
91.77
Paies
71.81
79.42
53,90
86.50
80.6
91.1
91.77
91.77
91.77
91.77
EMIl (n 5)
80.08
87.05
90.06
91.75
$3.03
93.26
93.34
93.44
93.51
80.51
Search
53.78
Pige - 911000
ScNNN
51.0
51.0
52.33
52
52.65
53.30
52.89
52.97
59.84
53.78
Paies
52.75
57.27
53.65
53.85
58.65
53.63
51.78
53.78
59.78
58.78
EIll (n 5)
$7.23
57.74
$7.91
57.98
58.07
58.08
58.11
58.11
58.1
58.1
Search
39.30
MS Marro - aDCGR10
ScNNN
Rase-IVF
38. 86
38.00
30.1
00
39:25
39.23
39.3
39.5
39.37
30.35
59.50
39.30
30.35 96
30.38
$9:38
59.38
39.58
30.59
39.39
30.30
EMIl (n
39.17
39.34
39.39
39.41
30.41
39.42
39.42
39.42
39.42
30.42
Search
33.77
MSManco- MRRR10
ScnNN
Rase-IVF
33.
53.34
33.54
33.58 58
33.67
53.66
35.72
35.7
33.76
33.75
33.75
$3.75
33.74
39.77
$3.77
53.76
33.76
39.77
53.77
33.77
EIIl (n 5)
33.62
33.7
33.77
33.78
33.78
33.79
33.79
33.79
33.79
33.779
4.5 Statistical Significance Tests
In this section, we aim to show that the improvements achieved by EHI over other baselines as depicted in Section 4.2 are statistically significant. We confirm with hypothesis tests that all the results on datasets considered in this work are statistically significant with a p-value less than 0.05.
Table 5: Statistical Significance p-test
Dataset Scifact (R@100)
Fiqa (R8100)
MSMARCO (nDCG010)
MSNARCO (MRR010)
p-value
0.0016
7.97F-07
0.0036
0.0018
13

---

Published in Transactions on Machine Learning Research (08/2024)
MSMarco
MSMarco
1300
# Queries
700
-Doc.
# 1200
EHI
EHI
Ideal (Uniform)
Ideal (Uniform)
1100
600
0
2000
4000
6000
0
2000
4000
6000
Leaf ld
Leaf ld
(a) MS MARCO: Document load balancing
(b) MS MARCO: Query load balancing
Figure 6: Leaf distribution of documents and queries of MS Marco in our trained EHI model. Our model learns splits closer to ideal split (uniform), this suggests that all the leaves are efficiently used, and our model does not learn very lob sided clusters. This roughly uniform spread of both queries (Figure 6(b)), as well as documents (Figure 6(a)) provides a more fino-grained clustering structure which not only takes advantage of the leaves but also shows stale of the art performance compared to baselines as discussed exlensively Section 4.2.
4.6 Load balancing
Load balancing is crucial objective that we aim to accomplish in this work. By achieving nearly uniform clusters, we can create clusters with finer gramularity, ultimately reducing latency in our model. In Figure 6(a), we can observe the distribution of documents across different leaf nodes. It is noteworthy that the expected number of documents per leaf, calculated as Σ Le= Prcs, where p; represents the probability of a document in bucket index i and 6; denotes the number of documents bucket i, yields an optimal value of approximately 1263.12. Currently, our approach attains an expected number of documents per leaf of 1404.65. Despite a slight skewness in the distribution, we view this as advantageous since we can allocate some leaves store infrequently accessed tail documents from training dataset. Additionally, our analysis demonstrates that queries are roughly evenly divided, highlighting the successful load balancing achieved, as illustrated in Figure 6.
4.7 Learning to balance load over epochs
In the previous section, we showed that our learned clusters are almost uniformly distributed, and help improve latency as shown in Section 4.2. In this section, we study the progression of expected documents per leaf as defined above over training progression on the Scifact dataset. In Figure demonstrates the decline of the expected documents per leaf over the training regime.
4.8 Discussions
In the previous section, we demonstrated the effectiveness of EHI against multiple baselines on diverse benchmarks. this section, we report results from multiple ablation studies to better understand the behavior of EHT.
Fraction of visited documents as proxy for latency. Ideally, we would want to compare query throughput against recall/MRR, but obtaining head-to-head latency numbers is challenging different systems are implemented using different environments and optimizations (for instance, off-the-shelf indexers such as SCANN, FAISS work better with CPU optimization while EHI takes advantage of GPUs). Thus, following standard practice in the ANNS community, we use fraction of documents visited visited/searchee /saarched as a proxy for latency (Jayaram Subramanya et al., 2019; Guo et al., 2020).
Stand-alone Indexer comparisons. To showcase the alignment between the emboddings and the down- stream clusters learned, conduct experiment where initialize the encoders with the encoder trained
14

---

Published in Transactions on Machine Learning Research (08/2024)
Scifact
210
Expected Doc / Leaf
Uniform
190
EHI
170
150
130
110
10
30
50 Epoch
70
90
Figure 7: Load balancing of documents over the training regime of EHI. EHI learning objective encourages well balances leaf nodes well as state of the art accuracy.
via EHI. We find the indexer trained by EHI is better aligned with the representations and leads to better performance when visiting fewer fraction of documents (see Figure 8). This underscores the crucial alignment and efficiency gains achieved through EHI's joint optimization.
To isolate the EHT indexer's impact, we fix encoder embeddings (pre-trained EHT model) and compare EHl indexing against ScaNN, and Faiss-IVF. Figure 5(a) illustrates that EHI outperforms other baselines by up to 4.6% when visiting a limited number of documents, highlighting the EHI indexer's intrinsic efficiency.
Effect of branching factor. Figure 5(b) shows recall@100 of FHI on SeiFact with varying branching factors. We consider versions of EHI, one with exact-search, and another where where we restrict EHl to visit about 10% visited document. Interestingly, for EHI Exact Search, the accuracy decreases with higher branching factor, while it increases for the smaller beam-size of 0.1. We attribute this to documents in leaf node being very similar to cach other for high branching factors (fewer points per leaf). We hypothesize that EHI is sampling highly relevant documents for hard negatives leading to a lower exact search accuracy.
Ablation w.r.t loss components. Next, on FIQA dataset, we study performance of EHI when one of the loss components equation 6 is turned off: see Figure 5(c). First, we observe that EHI outperforms the other three vanilla variants, implying that each loss term contributes non-trivially to the performance of EHI. Next, we observe that removing the pecument-similarityy basedd loss term (Aa), Eq. equation 5, has the least effect the performance of EHIl the other two-loss terms already capture some of its desired consequences. However, turning off the contrastive loss on either encoder embedding (A1), Eq. equation 3, or path embedding (A2), Eq. equation 4, loss leads to significant loss in accuracy. This also indicates the importance of jointly and accurately learning encoder and indexer parameters.
MS Mance
0.330
8
2.85
830.87
lscall
MiR
11
++11 Indear
+
0.334
0.75
fF
10
20
30
30 Doocument voited
20
10
20
30
Documents visited
Doouments visited
(a) Scifact: R10100
(b) MSMARCO: MRR@10
(c) TREC DL19: nDCG&10
Figure 8: EHI indexer is significantly more accurate than EHI encoder + off-the-shelf indexers such ScaNN, Faiss-IVF especially when restricted to visit small fraction of documents. Note that all the algorithms have been initialized with the same EHT encoder for computing embeddings.
15

---

Published in Transactions on Machine Learning Research (08/2024)
Scifact
100
90
Recall @
DE + Exact Search
80
EHI (H = 2, B = 8)
EHI (H = 3, B 4)
EHI (H 1, 40)
70
10
50 % Documents visited
100
Figure 9: EHI upon extending to deeper trees.
4.9 Scaling to deeper trees
In this section, we further motivate the need to extend to deep trees and showcase how this extension is critical to be deployed in production. We show that the hierarchical part in EHI can be tailored for large-scale tasks (>1B documents or otherwise) and is absolutely necessary for two reasons:
Sub-optimal convergence and blow-up in path embedding: When working with hundreds of billions of scale datasets as ones generally used in industry, a single height indexer would not work. Having more than a billion logits returned by a single classification layer would be suboptimal and load to poor results. This is a common trend observed other domains, such as Extreme Classification, and thus warrants hierarchy in our model design. Furthermore, the path embedding would blow up if we use single height indexer, and optimization and gradient alignments would become the bottleneck in this scenario.
Sub-linear scaling and improved latency: Furthermore, increasing the height of EHI can be shown scale to very large leaves and documents. Consider indexer is single height tree that trics to index cach document of D dimension into L leaves. The complexity of forward pass in this setling would be O(P + DL). However, for very large values of L, this computation could blow up and lead to suboptimal latency and blow up in memory. Note, however, that by extending the same into multiple leaves (H), one can roduce the complexity at cach height to a sublincar scale. The complexity of forward pass the hierarchical k-ary tree would be O(P + HDL1/H). Note that since > H, the hierarchical k-ary tree can be shown to have a better complexity and latency, which helps in scaling documents of the scale of billions observed in general applications in real-world scenarios. For instance, notice significant speedup with incrcasing height; EHI trained Scifact with equal number of leaves,
where encoder has P parameters, and the
for example at (B 64, H = 1), (B 8, H 2), and (B = 4, H = 3), we notice per-query latency of 2.48ms, 2.40ms, and 1.99ms respectively at the same computation budget.
We study the accuracy EHI when extending to multiple heights of the tree structure to extend its effectiveness in accurately indexing extensive web-scale document collections. Traditional indexing methods that rely on a singlo-height. approach can be computationally impractical and sub-optimal when dcaling with billions more documents. To address this challenge, EHI treats height as hyperparameter and learns the entire tree structure end-to-end. This extension to hierarchical k-ary tree is absolutely necessary for scalability.
In the above results presented in Figure 9, we showcase that one could increase the depth of the tree to allow multiple linear layers of lower logits say 1000 clusters at each height instead of doing so at a single level. This is preciscly where the hicrarchical part of the method absolutely necessary for scalability. Furthermore, we also extend EHI model on other datasets such as Scifact to showease this phenomenon as depicted in Table 6. We do confirm with p-test with p-value of 0.05 and confirm that the results across different.
16

---

Published in Transactions on Machine Learning Research (08/2024)
permutations are nearly identical, and the improvements R@50 and R@100 over other permutations are not statistically significant (improvements < 0.5%). Furthermore, improvements of other permutations of (B,H) on the R@10 and R&20 metries over the H=1 baseline are also statistically insignificant. Our current experiments on S /MMMARCO not only portray that our method does indeed scale to deeper trees but also showcase that EHI can exlend deeper trees with little to no loss in performance and improved latency. We also maintain a roughly similar performance and uniform splitting of documents even when extended deeper trees as depicted in Figure 9.
Table 6: Performance metries (%) evaluated on the Scifact dataset across various permutations of braching factor (B) and height (H). The bost value for each metric, corresponding to specific number of visited docunents, is indicated in bold font.
Scifact
R810
R@20
R850
R0100
Two-tower model (with Exact Search)
75.82
82.93
87.9
91.77
EHI (B=10, 1=1)
84.73
87.57
93.87
95.27
EHI (B-4, 2-2)
86.23
88.7
93.43
94.77
EHl (B=6, H=2)
83.8
88.63
92.77
95.27
EHI (B-8, H
84.46
88, 47
93.27
95.27
EHl (B-4, H-3)
84.96
88.7
94.77
EHl (B=6, H=3)
85.36
89.3
93.6
94.77
EHI (B-8, H-3)
83.86
87.47
92.27
95.27
5 Conclusions, Limitations, and Future Work
We presented EHT, a novel framework that fundamentally shifts the paradigm of dense retrieval by jointly optimizing the query/document encoder and the search indexer. EHI's core components the encoder, indexer, and retricver work in tandem generate compact path emboddings. These path embeddings, representing the path taken by queries and documents within the index tree, are crucial for EHI's effective joint training. Extensive evaluations standard benchmarks demonstrate EHI's superior efficiency and accuracy compared to traditional approaches. While the success of path embeddings highlights their potential, deeper theoretical understanding and formal guarantees would be valuable. Additionally, exploring the integration EHT with hierarchical representations like Matryoshka Embeddings (Kusupati et al., 2022) or techniques like RGD (Kumar et al., 2023), and extending to billion scale (see Appendix D) could further enhance performance, particularly for complex and tail queries, and be of interest the broader rescarch community. We believe EHI lays a strong foundation for future innovations dense retrieval, promising more efficient and accurate semantic search systems.
Broader Impact Statement
Our proposed approach helps improve large scale retricval and has potential to be applicd to multiple industry grade retrieval solutions. Improving scale and recall quality of retrieval itself might have multiple societal conquences, but noe of which are specific to our method.
References
Payal Bajaj, Daniel Campos, Nick Craswell, Li Deng, Jianfeng Gao, Xiaodong Liu, Rangan Majumder, Andrew McNamara, Bhaskar Mitra, Tri Nguyen, et al. Ms marco: A human generated machine reading comprehension dataset. ariie preprint xrXi::1611.09268 2016.
Erik Bernhardsson. Annoy: Approzimate Nearest Neighbors in C+ C++/Python, 2018. URL https://pypi. prg/project/annoyy Python package version 1.13.0.
17

---

Published in Transactions on Machine Learning Research (08/2024)
Michele Bevilacqua, Giuseppe Ottaviano, Patrick Lewis, Scott Yih, Sebastian Riedel, and Fabio Petroni. Au- toregressive search engines: Generating substrings as document identifiers. Advances in Neural Information Processing Systems, 35:31668 31683, 2022.
Alina Beygelzimer, Sham Kakade, and John Langford. Cover trees for nearest neighbor. In Proceedings of the 23rd international conference on Machine learning, pp. 97 17-104, 2006.
Deng Cai. revisit of hashing algorithms for approximate nearest neighbor search. IEEE Transactions on Knowledge and Data Engineering, (5(6))2337-2343 2021. doi: 10.1109 TT///TDEE...........2..
Kenneth Clarkson. An algorithm for approximate closest-point queries. In Proceedings of the tenth annual symposium on Compulational geomelry, pp. 164, 1994.
Nick Craswell, Bhaskar Mitra, Emine Yilmaz, Daniel Campos, and Ellen M Voorhees. Overview of the trec 2019 deep learning track. arXiv preprint arXiv:2003. 07820, 2020.
Kunal Dahiya, Deepak Saini, Anshul Mittal, Ankush Shaw, Kushal Dave, Akshay Soni, Himanshu Jain, Sumeet Agarwal, and Manik Varma. Deepxml: A deep extreme multi-label learning framework applied short text documents. In Proceedings of the 14th ACM International Conference Web Search and Data Mining, pp. 31 39, 2021.
Kunal Dahiya, Nilesh Gupta, Deepak Saini, Akshay Soni, Yajun Wang, Kushal Dave, Jian Jiao, Gururaj K, Prasenjit Dey, Amit Singh, et al. Ngame: Negative mining-aware mini-batching for extreme classification. In Proceedings of the Sizteenth ACM International Conference on Web Search and Data Mining, pp. 258 266, 2023.
Mayur Datar, Nicole Immorlica, Piotr Indyk, and Vahab S Mirrokni. Locality hashing scheme based on p-stable distributions. In Proceedings of the twentieth annual symposium on Computational geometry, pp. 253 262, 2004.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04800 2018.
Chantat Eksombatchai, Pranav Jindal, Jerry Zitao Liu, Yuchen Liu, Rahul Sharma, Charles Sugnet, Mark Ulrich, and Jure Leskovec. Pixie: A system for recommending 3+ billion items to 200+ million users in real-time. In Proceedings of the 2018 world wide web conference, pp. 1775 1784, 2018.
Thibault Formal, Carlos Lassance, Benjamin Piwowarski, and Stéphane Clinchant. From distillation to hard negative sampling: Making sparse neural ir models more effective. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 2353-2359, 2022.
Jerome H Friedman, Jon Louis Bentley, and Raphael Ari Finkel. An algorithm for finding best matches in logarithmic expected time. ACM Transactions on Mathematical Software (TOMS), 3(3):209 226, 1977.
Luyu Gao, Zhuyun Dai, Zhen Fan, and Jamic Callan. Complementing lexical retricval with semantic residual embedding. corr abs 20s99 (2020). arXiv preprinl vrXiv:2004.13900 2020a.
Weihao Gao, Xiangjun Fan, Chong Wang, Jiankai Sun, Kai Jia, Wenzhi Xiao, Ruofan Ding, Xingyan Bin, Hui Yang, and Xiaobing Liu. Deep retrieval: Learning a retrievable structure for largo-scale recommendations. Xix preprint ar Xix::2077007000 2020b.
Tiezheng Ge, Kaiming He, Qifa Ke, and Jian Sun. Optimized product quantization. IEEE transactions on pattern analysis and machine intelligence, 36(4):744-755, 2013.
Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelcrating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning, pp. PMIRR, 2020.
18

---

Published in Transactions on Machine Learning Research (08/2024)
Nilesh Gupta, Patrick Chen, Hsiang-Fu Yu, Cho-Jui Hsich, and Inderjit Dhillon. Elias: End-to-end learning to index and search in large output spaces. Aduances in Neural Information Processing Systems, 35: 19798 19809, 2022.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new cstimation principle for unnormalized statistical models. In Proceedings of the thirleenth international conference on artificial intelligence and statistics, pp. 297-304- JMLR Workshop and Conference Proceedings, 2010.
Kelvin Gun, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retricval augmented language model pre-training. In International conference on machine learning, pp. 3929 3938. PMLR, 2020.
Sebastian Hofstätter, Sheng-Chiel Lin, Jheng-Hong Yang, Jimmy Lin, and Allan Hanbury. Efficiently teaching an effective dense retriever with balanced topic aware sampling. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 113 122, 2021.
Huggingface. HuggingFace Sentencee-rransformer for MSMarco, 2019. URL https://huggingface.oo co/
sentence- sentence-trans
Piotr Indyk and Rajeev Motwani. Approximate nearest neighbors: towards removing the curse of dimen- sionality. In Proceedings of the thirticth anauual ACM symposium on Theory of computing, pp. 604 613, 1998.
Gautier Izacard and Edouard Grave. Distilling knowledge from reader to retriever for question answering. arXiv preprint arXiv:2012 04584, 2020a.
Gautier Izacard and Edouard Grave. Leveraging passage retrieval with generative models for open domain question answering. arXiv preprinl arXiv:2007.0188, 2020b.
Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Towards unsupervised donse information retricval with contrastive learning. arXiv preprint arXiv:2112.09116 2021.
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi- Yu, Armand Joulin, Sobastian Riedel, and Edouard Grave. Few-shot learning with retrieval augmented language models. Xixiv preprint arXiv:2208. 03299,, 2022.
Himanshu Jain, Venkatesh Balasubramanian Bhanu Chunduri, and Manik Varma. Slice: Scalable linear extreme classifiers trained 100 million labels for related searches. Tn Proceedings of the Twelfth ACM International Conference on Web Search and Data Mining, pp. 528 536, 2019.
Shikhar Jaiswal, Ravishankar Krishnaswamy, Ankit Garg, Harsha Vardhan Simhadri, and Sheshansh Agrawal. Ood-diskann: Efficient and scalable graph anns for out-of-distribution queries. or orXiv preprinl arXiv:2211.19850, 2022.
Young Kyun Jang and Nam Ik Cho. Generalized product quantization network for semi-supervise image retrieval. In Proceedings of the IEEE/CVP Conference on Compuler Vision and Pallern Recognilion, pp- 3420 3429, 2020.
Suhas Jayaram Subramanya, Frm Devvrit, Harsha Vardhan Simhadri, Ravishankar Krishnawamy, and Rohan Kadekodi. Diskann: Fast accurate billion-point nearest neighbor search on a single node. Advances in Neural Information Processing Systems, 32, 2019.
Herve Jegou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions on pottern analysis and machine intelligence, 33(1):117 128, 2010.
Jeff Johnson, Matthijs Douze, and Hervé Jégon. Billion-scale similarity search with GPUs. IEEE Transactions Big Data, 7(3):535 547, 2019.
19

---

Published in Transactions on Machine Learning Research (08/2024)
Omar Khattab and Matei Zaharia. Colbert: Efficient and effective passage search via contextualized late interaction over bert. In Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval, pp. 39 48, 2020.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigoerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. In Computer Vision-ECCV ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 23-28, 2020, Proceedings, Part V 16, pp. 491-507. Springer, 2020.
Ramnath Kumar, Kushal Majmundar, Dheeraj Nagaraj, and Arun Sai Suggala. Stochastic re-weighted gradient descent via distributionally robust optimization. arXiv preprint AXX::2306.09222 2023.
Aditya Kusupati, Matthew Wallingford, Vivek Ramanujan, Raghav Somani, Jac Sung Park, Krishna Pillutla, Prateek Jain, Sham Kakade, and Ali Farhadi. Lle: Accurate, multi-purpose learnt low-dinensional binary codes. Advances in Neural Information Processing Systems, 34:23900-233 23901,, 2021.
Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ramanujan, William Howard-Snyder, Kaifeng Chen, Sham Kakade, Prateek Jain, et al. Matryoshka representation learning. Advances in Neural Information Processing Systems, 35:30233 00002223--04,,,, 2022.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Coliins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia. Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Lingwistics, 7:453 466, 2019.
Carlos Lassance, Thibault Formal, and Stéphane Clinchant. Composite code sparse autoencoders for first stage retrieval. In Proceedings of the 44th International ACM STGIR Conference on Research and Development in Information Retrieval, pp. 2136-2140, 2021.
Carlos Lassance, Simon Lupart, Hervé Dejean, Stéphane Clinchant, and Nicola Tonellotto. A static pruning study on sparse neural retrievers. In Proceedings of the 46th International ACM SIGIR Conference Research and Development in Information Retrieval, pp. 1771-1775, 2023.
W Li, Y Zhang, Y Sun, W Wang, W Zhang, and Lin. Approximate nearest neighbor search on high dimensional data mata-experimentts analyses, and improvement. IEEE Transactions on Knowledge and Data Engineering, 2020.
lyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXio preprinl ar Xivo: Nit:11111.01011, 2017.
Macedo Maia, Siegfried Handschuh, André Freitas, Brian Davis, Ross McDermott, Manel Zarrouk, and Alexandra Balahur. Www'18 open challenge: financial opinion mining and question answering. In Compunion proceedings of the the web conference 2018, pp. 1941 1942 2018.
Yu A Malkov and DA Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE Transactions on Pattern Analysis & Machine Intelligence, 42(04): 224-836, 2020.
Yu Malkov and Dmitry A Yashunin. Efficient and robust approximate nearest neighbor scarch using hierarchical navigable small world graphs. IEEE transactions on pallern analysis and machine intelligence, 42(():824- 836, 2018.
Aditya Menon, Sadeep Jayasumana, Ankit Singh Rawat, Seungyeon Kim, Sashank Reddi, and Sanjiv Kumar. In defense of dual-encoders for neural ranking. In International Conference on Machine Learning, pp. 15376-15400. PMLR, 2022.
Bhaskar Mitra, Nick Craswell, et al. An introduction to neural information retrieval. Foundations and Trends in Information Retrieval, 13(1):1-126, 2018.
20

---

Published in Transactions on Machine Learning Research (08/2024)
Nicholas Monath, Manzil Zaheer, Kelsey Allen, and Andrew McCallum. Improving dual-encoder training through dynamic indexes for negative mining. In International Conference on Artificial Intelligence and Statistics, pp. 9308 9330. PMLR, 2023.
Niklas Muennighoff. Sgpt: Gpt sentence embeddings for semantic scarch. arXiv preprint ar arXiu:2202.0890; 2022.
Pandu Nayak. Understanding searches better than ever before. Google AI Blog, 2019. URL https: //blog. http:///rrrttcs/ssaa/sss/earrrda -agidaanddddddddddddddnnn bert/.
Arvind Neelakantan, Tao Xu, Raul Puri, Alec Radford, Jesse Michael Han, Jerry Tworek, Qiming Yuan, Nikolas Tezak, Jong Wook Kim, Chris Hallacy, et al. Text and code emboddings by contrastive pre-training. arXiv preprinl or 2t::2201.100000 2022.
Jianmo Ni, Chen Qu, Jing Lu, Zhuyun Dai, Gustavo Hernández Abrego, Ji Ma, Vincent Zhao, Yi Luan, Keith Hall, Ming- Wei Chang, et al. Large dual encoders are generalizable retrievers. arXiv preprint arXiv:2112.0P1 2.0899, 2021.
Rodrigo Nogueira and Kyunghyun Cho. Passage re-ranking with bert. arXiv preprint arXiv:1901.04055, 2019.
Rodrigo Nogueira, Wei Yang, Jimmy Lin, and Kyunghyun Cho. Document expansion by query prediction. arXiv preprint arXie:1904.00 2019.
Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Wayne Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. Rocketqa: An optimized training approach to dense passage retrieval for open-domain question answering, arXiv preprint XrXi::2010. 08191, 2020.
Alec Radford, Karthik Narasimhan, Tim Salimans, Nya Sutskever, et al. Improving language understanding by generative pre-training. OpenAl, 2018.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp.
2021.
Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks. arXie preprint XrX1x:1908. 10084, 2019.
Stephen Robertson, Hugo Zaragoza, al. The probabilistic relevance framework: Bm25 and beyond. Foundations and Tyends in Information Retrieval, 3(4):333-338) 2009.
Ruslan Salakhutdinov and Geoffrey Hinton. Semantic hashing. International Journal of Approcimate Reasoning, 50(7):969 978, 2009.
Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. Colbertv2: Effective and elficient retrieval via lightweight late interaction. arXie preprint arXiv:2112.01 02.01888 2021.
Josef Sivic and Andrew Zisserman. Video google: A text retrieval approach object matching in videos. In Computer Vision, IEEE International Conference on, volume 3, pp. 1470-1470. IEEE Computer Society, 2003.
Yi Tay, Vinh Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupla, et al. Transformer memory as a differentiable search index. Advances in Newral Information Processing Systems, 35:21831-21643, 2022.
Nandan Thakur, Nils Reimers, Andreas Rücklé, Abhishek Srivastava, and Iryna Gurevych. Beir: A heteroge- nous benchmark for zero-shot evaluation of information retrieval models. arXiv preprint arXiv:1144006663 2021.
21

---

Published in Transactions on Machine Learning Research (08/2024)
David Wadden, Shanchuan Lin, Kyle Lo, Lucy Lu Wang, Madeleine van Zuylen, Arman Cohan, and Hannanch Hajishirzi. Fact fiction: Verifying scientific claims. arXiv preprint arXiv:2004. 14974, 2020.
Mengzhao Wang, Xiaoliang Xu, Qiang Yue, and Yuxiang Wang. A comprehensive survey and experimental comparison of graph-hased approximate nearest neighbor search. Proceedings of the VLDBB Endowment, 14 (11):1964-1978, 2021.
Yujing Wang, Yingyan Hou, Haonan Wang, Ziming Miao, Shibin Wu, Qi Chen, Yuqing Xia, Chengmin Chi, Guoshuai Zhao, Zheng Liu, et al. A neural corpus indexer for document retrieval. Advances in Neural Information Processing Systems, 35:25600 -5666, 2022.
Roger Weber, Hans-Jörg Schek, and Stephen Blou.. quantitative analysis and performance study for similarity-search methods high-dimensional spaces. In VLDB, volume 98, pp. 194 205, 1998.
Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Arnold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. arXiv preprint #xXiv:2007.00000 2020.
Tan Yu, Junsong Yuan, Chen Fang, and Hailin Jin. Product quantization network for fast image retrieval. In Proceedings of the European Conference Computer Vision (ECCV), pp. 186 201, 2018.
Jingtao Zhan, Jiaxin Mao, Yiqun Liu, Jiafeng Guo, Min Zhang, and Shaoping Ma. Learning discrete representations via constrained clustering for effective and efficient dense retrieval. In Proceedings of the Fifteenth ACM International Conference Web Search and Data Mining, pp.
2022.
Han Zhu, Xiang Li, Pengye Zhang, Guozheng Li, Jie He, Han Li, and Kun Gai. Learning tree- tree-based deep model for recommender systems. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 1079-1088, 2018.
Han Zhu, Daqing Chang, Ziru Xu, Pengye Zhang, Xiang Li, Jie He, Han Li, Jian Xu, and Kun Gai. Joint optimization of tree-based index and deep model for recommender systems. Advonces in Neural Information Processing Systems, 32, 2019.
22

---

Published in Transactions on Machine Learning Research (08/2024)
A
Datasets
In this scction, we briefly discuss the open-source datascts used in this work adapted from the Beir benchmark Thakur et al. (2021).
Scifact: Scifact (Wadden et al., 2020) is a fact-checking benchmark that verifies scientific claims using evidence from research literature containing scientific paper abstracts. The dataset has ~ 5000 documents and has standard train-test split. We use the original publicly available dev split from the task having 300 test queries and include all documents from the original dataset as our corpus. The scale of this dataset is smaller and included to demonstrate even splits and improvement over basclines even when the number of documents is in the order of 5000.
Fiqa: Fiqa (Maia et al., 2018) is an open-domain question-annwwringg task in the domain of financial data by crawling StackExchange posts under the Investment topic from 2009-2017 as our corpus. It consists of 368 documents and publicly available test split from (Thakur et al., 2021) As the test set, we use the random sample 500 queries provided by Thakur al. (2021). The scale of this dataset is 10x higher than Scifact, with almost 37,638 documents our corpus.
The MSMarco benchmark Bajaj et al. (2016) has been included since it is widely recognized the gold standard for evaluating and benchmarking large-scale information retrieval systems (Thakur et al., 2021; Ni et al., 2021). It is collection of real-world search queries and corresponding documents carefully curated from the Microsoft Bing search engine. What sets MSMarco apart from other datasets is its scale and diversity, consisting of approximately 9 million documents in its corpus and 532,761 query-passage pairs for fine-tuning the majority of the retrievers. Due to the increased complexity in scale and missing labels, the benchmark is widely known to challenging. The dataset has been extensively explored and used for fine-tuning dense retrievers in recent works (Thakur et al., 2021; Nogueira & Cho, 2019; Gao al., 2020a; Qu et al., 2020). MSMarco has gained significant. attention and popularity in the research community due to its realistic and challenging nature. Its large-scale and diverse dataset reflects the complexities and nuances of real-world search scenarios, making it an excellent testbed for evaluating the performance of information retrieval algorithms.
MS Marco:
NQ320k: The NQ320k benchmark kowiat et al., 2019) has become a standard information retrieval benchmark used to showcase the efficacy of various SOTA approaches such as DSI (Tay et al., 2022) and NCI (Wang et al., 2022). In this work, we use the same NQ320k preprocessing steps as NCI. The queries are natural language questions, and the corpus is Wikipedia articles in HTML format. During preprocessing, we filter out tags and other special characters and extract the title, abstract, and content strings. Please refer to Wang et al. (2022) for additional details about this dataset.
Table 7 presents more information about our datasets.
Table 7: Datasel details.
Dataset
Train Pairs
# Test Queries
# Test Corpus
Avg.
Test Doc per Query
Scifact
920
300
5183
1.1
Fiqa-2018
14,166
648
57,638
2.6
MS Marco
532,761
6,980
8,841,823
1.1
B Model and Hyperparameter details
We followed standard approach called grid search to determine the optimal hyperparameters for training the EHI model. This involved selecting a subset of the training dataset and systematically exploring different. combinations of hyperparameter values that perform best on this held-out set.. The key hyperparameters
23

---

Published in Transactions on Machine Learning Research (08/2024)
focused on were the Encoder (Es ) and Indexer (Iş) learning rates. Other hyperparameters to tune involved the weights assigned to each component. of the loss function, denoted as A A in Equation 6.
The specific hyperparameter values used in our experiments are detailed in Table 8. This table provides a comprehensive overview of the exact settings for each hyperparameter. The transformer block used in this work the sentence transformers distilBERT model, which has been trained on the MSMarco dataset (https: //uuggingfcce. co/ centence-tranc mameer//mmarcoo wewtarooodiiitti wwttssseettttsoo-eeeeeeeeeeeeeee cos-v5). We fine-tune the encoder for our task and report our performance metrics.
Table 8: Hyperparameters used for training EHI various datasets. Note that number of epoch and refresh rate (r) was set to 100 and 5, respectively. EHI initializes with distilBERT with encoder embedding representation as 768
Description
Scifact
Fiqa
MSMarco
NQ320k
argument name
Generul Settings
Batch size Number of leaves
64
64
4096
1024
batch_size
40
5000
7000
1000
leaves
Optimizer Settings
Encoder Learning rate
-e
2e-1
2.5e
1e-4
enc 1r
Classifier Learning rate
0.016
9e
Se-
cl_ir
Loss factor 1
0.2
0.03
0.11
0.2
A1
Loss factor 2
0.8
0.46
0.84
0.4
A2
Loss factor 3
0.2
0.54
0.60
0.11
A
C
Additional Experiments
This section presents additional experiments conducted on our proposed model, referred to as EHI. The precise performance metrics of our model, as discussed in Section 4.2, are provided in Appendix C.1. This analysis demonstrates that our approach achieves state-of-the-art performance even with different initializations, highlighting its robustness (Appendix 4.4).
Furthermore, delve into the secument-tt-leaavee ratio concept, showcasing its adaptability to degrees greater than one and the advantages of doing so, provided that our compulational requirements are met (Appendix C.4). This flexibility allows for the exploration of more nuanced clustering possibilities. We also examine load balancing in EHI and emphasize the utilization of all Icaves in a well-distributed manner. This indicates that our model learns elficient representations and enables the formation of finer-grained clusters (Appendix 4.6).
To shed light on the learning process, we present the expected documents per leaf metric over the training regime in Appendix 4.7. This analysis demonstrates how EHI learns to create more evenly distributed clusters as training progresses, further supporting its effectiveness.
Finally, we provide additional insights into the semantic analysis of the Indexer in Appendix C.4, highlighting the comprehensive examination performed to understand the inner workings of our model better.
Through these additional experiments and analyscs, we reinforce the efficacy, robustness, and interpretability of our proposed EHI model, demonstrating its potential to advance the field of information retrieval.
SciFact. We first start with the small-scale SciFact dataset. Figure 10(a) and Table 9 compares EHI to three DE basclines. Clcarly, EHI's recall-compute curve dominates that of Dc+ScaNN and DE+Faiss-IVE For example, when allowed to visit risit/search about 10% of documents, EHI obtains up to 115.64% higher Recall@ 100. Furthermore, EHI can outperform DE+Exact Search with a 60% reduction latency. Finally, representations from EHI's encoder with exact scarch can be much 4% more accurate (in terms of Recall@ 100) than bascline d Exact Search, indicating effectiveness of EHI's integrated hard negative mining.
24

---

Published in Transactions on Machine Learning Research (08/2024)
Table Performance metrics (%) evaluated on the Scifact dataset across various ranges of visited documents. The best value for each metric, corresponding specific number of visited documents, indicated bold font.
Metric
Method
10%
20%
30%
40%
50%
60%
70%
80%
90%
100%
Exact Search
82.73
ScaNN
66.2
74.11
78.28
79.94
81.28
81.61
81.94
82.28
82.28
82.28
R@10
Faiss-IVF
70.74
78.27
80.07
82.23
82.9
82.73
82.73
82.73
82.73
82.73
EHI
76.97
81.93
83.33
83.67
84.73
85.07
84.73
84.73
84.73
84.73
Exact Search
86.9
ScaNN
68.32
77.97
0.8247
84.63
85.97
86.3
86.63
86.97
86.97
86.97
R@20
Faiss-IVIF
73.26
81.43
84.07
86.73
87.4
87.23
86.9
86.9
86.9
86.9
EHI
81.03
85.27
86.5
86.17
87.23
87.57
87.57
87.57
87.57
87.57
Exact Search
92.53
ScaNN
72.39
82.7
87.2
90.2
91.53
91.87
92.2
92.53
92.53
92.53
R@50
Faiss-IVF
76.66
85.4
87.37
90.03
91.03
90.87
90.87
90.87
90.87
90.87
FHI
82.87
88.67
91.73
92.07
93.2
93.53
93.53
93.53
93.87
93.87
Exact Search
94.1
ScaNN
72.62
83.03
88.53
91.53
92.87
93.2
93.53
93.87
93.87
93.87
R@100
Faiss-IVF
77.16
87.3
89.93
92.6
93.27
93.43
93.43
93.77
94.1
EIII
83.53
89.33
92.9
93.23
94.43
94.77
94.77
94.93
95.27
95.27
FIQA. Here also we observe a similar trend SciFact; Figure 10(b) and Table 10. That is, when restricted to visit only 15% documents (on an average), EHI outperforms ScaNN and Faiss-IVF in Recall@100 metric by 5.46% and 4.36% respectively. Furthermore, EHT outperforms the exact search in FIQA with a 84% reduction in latency or documents visited. Finally, when allowed to visit about 50% of the documents, EHI is about 5% more accurate than Exact Search, which visits all the documents. Thus indicating better quality of learned embeddings.
C.1 Results on various benchmarks
In this section, we depict the precise numbers we achieve with various baselines as well as EHI as shown in Table 9, Table 10, and Table 11 for Seifact, Fiqa, and MS Marco respectively.
0.95
0.60
Recall 100
100
0.90
0.55
@
0088
0.50
DE Exact Search
DE Exact Search
+EDE ScaMN
ScaNN
0.45
+cDE Faiss-IVF
* DE Faiss-IVF
0.75
++H1 (Oure)
+ EHi (Ours)
100
% Documents visited
Documents visited
(a) SciFact: 100
(b) FIQA: Recal1@100
Figure 10: (a), (b): Recall@100 100 EHI and basclines on Scifact and Fiqa dataset. EHI is significantly more accurate than dual encoder + ScaNN or Faiss-IVE baselines, especially when computationally restricted to visit small fraction of documents. For example, on Scifact datasot, at 10% document visit rate, EHI has 11.7% higher Recall@100 than baselines.
To understand how the performance of EHI model on smaller number of in Recall@k, we report Recall@10, Recall@25, and Rocall@50 metrics on the Scifact datasct. Figure depicts the performance of the EHI model in this setting. On the Recal1@10, Recall@20 and Recall@50 metric, EHI outperforms exact search metric by 8.91%, 4.64% and 5.97% respectively.
25

---

Published in Transactions on Machine Learning Research (08/2024)
Table 10: Performance metrics (%) evaluated on the Fiqa dataset across various ranges of visited documents. The best value for each metric, corresponding to specific number of visited documents, is indicated in bold font.
Metric
Method 10%
20%
30%
40%
50%
60%
70%
80%
90%
100%
Fxact Search
30. 13
ScaNN
23.95
29.08
29.69
29.98
29.83
30.13
30.08
30.08
30.13
30.13
R@10
Faiss-IVF
26.8
29.67
29.9
29.93
29.93
29.93
29.98
30.13
30.13
30.13
EHI
27.43
32.02
32.12
32.19
32.42
32.42
32.42
32.42
32.47
32.47
Exact Search
36.0
ScaNN
28.54
34.45
35.34
35.67
35.7
36.0
35.88
35.95
36.0
36.0
R@20
Faiss-IVF
32.17
35.51
35.77
35.8
35.8
35.8
35.85
36.0
36.0
36.0
EHI
33.47
39.01
39.36
39.45
39.68
39.68
39.68
39.68
39.73
39.73
Exact Search
45.4
ScaNN
35.93
43.38
44.51
45.11
45.02
45.12
45.07
45.48
45.5
45.4
R@50
Faiss-IVF
39.78
44.45
44.86
45.19
45.19
45.19
45.24
45.4
45.4
45.4
EHI
41.03
49.77
50.35
50.62
50.75
50.83
50.99
50.99
51.04
51.04
Exact Search
53.78
SeaNN
51.0
51.9
52.33
52.69
52.63
53.39
52.49
52.97
53.64
53.78
R@100
Faiss-IVF
52.75
53.27
53.65
53.65
53.65
53.63
53.78
53.78
53.78
53.78
EHI
57.36
58.42
58.97
59.18
59.13
59.29
59.29
59.34
59.34
59.39
Table 11: Performance metrics (%) evaluated on the MS Marco dataset across various ranges of visited documents. The best value for each metric, corresponding to specific number of visited documents, is indicated in bold font.
Metric
Method
10%
20%
30%
40%
50%
60%
80%
90%
100%
Exact Search
22.41
ScaNN
22.06
22.38
22.36
22.36
22.11
22.39
22.11
22.41
22.41
22.41
nDCG@1
Faiss-IVF
22.16
22.32
22.38
22.41
22.41
22.41
22.39
22.39
22.39
22.41
EHI
22.39
22.48
22.46
22.46
22.46
22.46
22.46
22.46
22.46
22.46
Exact Search
32.5
ScaNN
31.96
32.36
32.42
32.43
32.49
32.48
32.5
32.5
32.5
32.5
nDCG@3
Faiss-IVF
32.05
32.25
32.39
32.44
32.48
32.48
32.48
32.5
32.5
32.5
EHI
32.45
32.61
32.61
32.6
32.6
32.6
32.6
32.6
32.6
32.6
Exact Search
36.03
ScaNN
35.38
35.82
35.91
35.94
36.01
36.01
36.03
36.02
36.03
36.03
nDCG@5
Faiss-IVF
35.56
35.75
35.89
35.96
36.0
36.0
36.02
36.04
36.01
36.03
EHI
35.9
36.08
36.08
36.08
36.08
36.08
36.08
36.08
36.08
36.08
Exact Search
39.39
ScaNN
38.61
39.1
39.25
39.3
39.37
39.36
39.38
39.38
39.39
39.39
nDCG@10
Faiss-IVF
38.86
39.09
39.23
39.31
39.35
39.36
39.36
39.38
39.38
39.39
EHI
39.19
39.42
39.42
39.42
39.43
39.43
39.43
39.43
39.43
39.43
Exact Search
33.77
ScaNN
33.16
33.58
33.67
33.7
33.76
33.75
33.77
33.77
33.77
33.77
MRR@10
Faiss-IVe IVF
33.34
33.54
33.66
33.72
33.75
33.75
33.74
33.76
33.76
33.77
EHT
33.64
33.81
33.81
33.8
33.81
33.8
33.8
33.8
33.8
33.8
26

---

Published in Transactions on Machine Learning Research (08/2024)
0.85
0.85
0.90
0.80
Recail @ 20
0.80
@80
00.85
Recal
0.75
0.75
Search
0.80
Search
DE ScaNN
<xDE ScaNN
DE ScaNN
0.70
* B-il (Ouns)
0.70
*+HH (Ours)
0.75
*HHI (Ours)
100
100
50
% Documents visited
% Documents visited
% Documents visited
(a) Scifact R@10
(b) Scifact R820
(c) Scifact RØ50
Figure 11: Additional Results for other Recallåk metrics of EHI trained on the Scifact dataset.
Performance metrics (%) evaluated on the various datasets showcasing the effect of negative mining when finetuned from distilbert-cos model trained on MS Marco benchmark.
Table 12:
Dataset
Metric
EHI's indexer negative mining
ANCE
Scifact
R@100
95.27
95.27
MS Marco (distilbert-cos)
MRR010
33.8
33.79
MS Marco (distilbert-cos)
nDCG@10
39.43
39.4
Furthermore, note that most prior works, including ScaNN, DSI, and NCI, all report their numbers for single seed. To show the robustness of our approach to initializations, we also report our numbers on multiple seeds and showcase that we significantly out perforrmm our baselines across all the boards as further described in Section 4.4.
C.2 Negative Mining
In this section, we discuss additional details on the exact working of our possible negative mining approach. The idea is that given positive document (d+), you would like to sample other documents which reached the same bucket as di but are not considered relevant to the query (g) per the ground truth data. Note that the dynamic negative mining something EHI achieves for almost free, as the built indexer is also used for the ANNS. Prior approaches including Monath al. (2023); Dahiya et al. (2023); Hofstätter et al. (2021) do study this dynamic negative sampling approaches well. However, the difference lies in the fact that EHI does not need to build new index for clustering every few steps, and the hicrarchical tree index learned by EHI is used for the downstream ANNS task. This contrasts with works such as Monath et al. (2023); Dahiya et al. (2023); Hofstätter et al. (2021), where the index is used during training and hard negative mining, but finally discarded. We leave this as an hyperparameter left to the practitioner. Note that we fine-tune a distilbert model which was trained on the MS Marco dataset with mined hard negatives. Thus, addition of this negative mining feature in experiments does not. lead to any significant drops as further shown in Table 12, and leads to efficient re-using of the same indexer learnt during inference.
C.3 Comparisons against ELIAS
In this section, we compare against other end-to-end training paradigms such as ELIAS (Gupta et al., 2022) which requires warm starting of knn index. It should be noted that ELIAS and BLISS are tailored for the domain of extreme classification while focus the dense information literature domain (note that label classifiers, and pre-computing KNN structure for datascts the size of MSMARCO is indeed a significant overhead). Regardless, we report our findings on the L.-AAmazrTiilee 131K (XC dataset) below (see Table 13). Note that below comparison is still unfair, as ELIAS utilizes additional million auxiliary parameters in their label classifiers, while EHT does not. Furthermore, EHT uses label features for encoding, while ELIAS does not. Although EHI outperforms ELIAS by 4.16% in P@1, it must be recognized that comparing against extreme classification based methods is unfair.
27

---

Published in Transactions on Machine Learning Research (08/2024)
Table 13:
Performance metrics (%) evaluated on the TFFAAmazorTitlee1131111 dataset showcasing the efficacy
of EHI comparison ELIAS.
LF- TT--AazzonTillss 131K
P01
P@3
P85
ELIAS
40.13
27.11
19.54
EII
41.8
28.64
20.79
MSMarco
MSMarco
0.39
0.335
MRR @ 10
nDCG @
0.330
0.38
d2/ =
0.325
d2/ =2
d2/ 2
0.37
0.320
d21
0.315
0
25
50
75
100
25
50
100
% Documents visited
% Documents visited
(a) MS MARCO: nDCG@10
(b) MS MARCO: MRR@10
Figure 12: Ablation study of docs to leaves on the MS Marco benchmark. Overall, we notice that as long as computational power available, increasing number leaves indexed for documents aids performance retrieval.
C.4 Allowing documents to have multiple semantics
Note that throughout training, model works with beam size for both query and document and tries to send relevant documents and queries to the same leaves. While testing, we have only presented results where beam size is used to send queries to multiple leaves, limiting the number of leaves per document to 1. We could use a number of leaves per document as hyperparameter and allow documents to reach multiple leaves as long as our computational power allows.
Accuracy: this section, we extend our results on increasing this factor on the MS Marco dataset as presented Figure Intuitively, each document could have multiple views, and simply sending them only one leaf might not be ideal, but efficient indexing. Furthermore, increasing this hyperparameter from the default d21 setting shows strictly superior performance and improves our metrics the MRR@10 well nDCG@10 metric.
Advantages of Multiple Leaves per Document: While assigning document to multiple leaf nodes incurs additional memory costs, uncover intriguing semantic properties from this approach. For the analysis, we assign multiple leaves for a document (Doc) ranked based on decreasing order of similarity, as shown in Figure 13. For Figure 13, we showcase three distinct examples of diverse domains showcasing not only the semantic correlation of the documents which reach similar leaves but also the queries which reach the same lcaf as our example. can be seen that cach Icaf captures slightly different semantics for the document. On further observations, queries mapped to the same leaf also share the semantic similarity with the Doc. Note that semantic similarity decreases the document's relevance to label decreases (left to right). This shows that EHl's training objection (Part 2 and Part 3) ensures queries with similar semantic meaning documents and similar documents are clustered together.
Diverse Concepts Across Leaf Nodes: Our analysis, as showcased in Figure 14, reveals a wide range of diverse concepts captured by the word clouds derived from the leaf nodes ofEHI. For Figure 14, we visualize leaf nodes clarity, which was obtained through
only a subset of randomly chosen leaves (54) instead random sampling. Below, we highlight some intriguing findings from our analysis. These concepts span various domains such as phones, plants, healthcare, judiciary, and more. Moreover, within each leaf node, we observe cohesive theme among the words. For instance, the word cloud at coordinates (1, 1) reflects
28

---

Published in Transactions on Machine Learning Research (08/2024)
Doc: Depending the particular form, the govemmeett generally carries
principles
positive
Law SraEsStisee
government
Doc. words
S g eéces
PORC THKEE UnLLed B
B provide
Dgreece
atatee
9 Eity yew lg eh
Government
Dy student
dam Vee
anci ent g Feek
wa te aedera government
source Sa Jrjeet
WE t nnn form a
state, tg perypeople
gover nment
music
Baaee me Tica
thaepoond keaning ieat wigh witt
smelll dirwe Ten
Query words
gover remennt
comic
taxes
dennark
allbpp
meatl
year
qut let
begin
publ
cay projec ject wid 1n
cycle
schsteu main
gover dovernmenta
dig SE60/30
state
ditrerance Treeeu wr ing
play lostiiitation
government
rule
good
necessary
qui
power public
people oold
first rel elesse
level SN Canada
gunitary &democr system
considered
grading prolinirary
skills) Turra carief
Dec: Iron (III) oxide is Fe203 therefore there are 3 oxygen atoms, which Is
answer #1, three exygens and two irons.
b yco
1odds
tneplas Seseessseeee fneer
el let Orozttnn TURER
tW LE cang
Doc. words
b
blood
Lood
Ce H
200
SW SWOTI 10 pop
ron
OX
Oxygen hest Seear
hemo
Lobin
water 000
ing
make
mo
number
YOU
eTement
Oxygen
deepima
ke
many
Query words
Spend t transport
phase
life e co
formula
tt
county
phone
ee Sig
time najority
tish
oh
Formul a
sal
preces
Adutres
ron
pood
on
mon thoron
AnersyC
quiziet
cycle
oxygen en numbe
E
mo Molluul
Jack mesh
new meeeiin
a tom
Doc: Like phono cartridges headphones and loudspeakers, the microphone is
.eeeegy (sound) and transiates It into equivalent electrical energy
PC
device
ervice STrE SEE
type
dama
Doc. words
Computer
type
deoics
Frver
mellory usB
Thent
Laptop yye First
per son
sdatar guus d
Cour
power
eeirrou Vo
device
hard drive
device
system SP
O
clutch
common
st.cpu
OxY
vvygen gen
Query words
Computer
type death
heart
Server
Connect
Ct
C
wo K
fireecccc
partcounty
cause e
e elettrical
Figure 13: We deploy the concept of word cloud to summarize the semantics of all queries and documents in leaf. Here, every document and query can hold multiple semantics therefore for the analysis, each document (Doc) was assigned to multiple leafs with decreasing order of similarity. It can be seen that for each leaf captures slightly different semantic for the document. On further observations, queries that were mapped to the same leaf also shares the semantic similarity with the Doc. Note that semantic similarity decreases as the relevance of document to a label decreases (left to right). This shows that EHT's training objection (Part 2 and Part 3) ensure queries with similar semantic meaning as document as as well as similar documents are clustered together.
29

---

Published in Transactions on Machine Learning Research (08/2024)
madey oo foodneed
u eeur O
Kurgenc 22 cause
Teg. cra sramp
hay may. dog
eat
meaning
ymbo
buy
vehicle
buy-aayneed
ca
tapioca
Foot
Computer
ance
hee Chelper eraap
goo
plant
time
Coin
feed
muscle
Toan
Loese
nane Find
known
anem
d
Glne
fire
nay
water
m
vine
bedspersd Spersom
Likaline
blood cell
India
autorrune diseass
noun
unn ted states
souse nane
so1l+
language
measure
remperatur
e
calor Morree
number
mall
poddy
narr
ause
spouse
beernay
FRE an lutal
year
tife
BO
zone
Chew veare
Falle
biane
en mineraa
word
aoll ther
divorce
value
nai
parpent
growth
state
ga
cal
ium
women
Tat
lesophagus
TWO
KON n
car
shou
der
Lade
b Lttonnn
muss le
speed
cause
acount
COSt
mass
wo conc Casti
fee
eanan OD
new
mph
pain
Si
twitter
cause b
EO
mO
S.eeeeed
Speed limit
mbps
shoulder
new
new
tweet
Chapter
volume
book
uSi
texas
sentence
make
Sirgson son
i ladd
conkaed
aaa
university ty example
O
indiana
in or on
sea
Function S student
UO program
Acoltiege
verb
wo
stress
year
pron
with
Coconut oll
mexico : Riee
people wothh
puencosrico
artro
schoo
peam
united state
Seee
Fient
key
dy mininum Wage
blue
Two change
us ing
Usine too work state
net wor
ato
tes enntin
Ehuit fire
shalk
name
emp loyee
hade
test
TO time
ii Teus
work
Sys
may
souhd
may
Species
Sfunnd
EnpLover
form
Emmllaon
YPR
électron
Ther of decay
neutron
new
letter
rother hear power
Finght
eye
arug fruit
part Die
Sofware
time
dog
&
otthe
son
phone: number
part
rengine ne
order
may
born
adeno
de
weites
mango
age
price
songgr
father
Docaned va
computer
LocK
Arre most
cause
name
necica center
eeaapp
born carch
may
pain
pregabalin
si gng
may
labana
Junee
people
1yrica
mar
diabetes
iphone
years
cause
year
Symptom
medication
nerve pain
wat ch
age Natch
OD first
Sideee
treatment
drug
Tie
Work
phonp
pa payment
cet
Sme
bee rabos
part
function plant
Lationaion
medical
type
gargen
amino acid
dna
vote
M f fowe Bossanee
protein
number
S
using
ssevvice
calorie
Figure 14: Figure shows that EHT's learning objective encourages leaf nodes to have different semantics. Here
semantic information in each leaf is shown in the form of word cloud. The scemingly similar leaves such (1,1), (5, 5) or (8,2) shares the similar words like "year" however, they capture completely different meaning such as financial, criminal and sports respectively.
30

---

Published in Transactions on Machine Learning Research (08/2024)
international finance theme, while the word cloud at (1, pertains to fitness. This observation highlights the importance of grouping similar documents into the same leaf nodes and excluding irrelevant documents using our objective function. Note that seemingly similar leaves such as (1,1), (5, 5), or (8, 2) share similar words like "year"; however, they capture completely different meanings such as financial, criminal, and sports, respectively.
D Extending to billion-scale corpus
While this work focused on datasets up to 10M in size (MS MARCO) showing no training bottlenecks ( seee Section 4), EHT holds strong theoretical advantages for scaling:
Sublinear Retrieval Complexity The hierarchical tree structure results in sublinear time complexity for retrieval, outperforming flat indexes the dataset grows. This is supported by our findings in Section 4.9, where deeper trees show benefits for larger datasets.
Efficient Search Space Partitioning The tree efficiently partitions the search space, allowing retrieval to focus on smaller candidate subsets and significantly reducing distance computations.
Regarding the retrieval process itself, determining relevant Icaf nodes for a query requires just a single forward pass through the indexer which is a standard IVF style data-structure, Final document retrieval from those leaves leverage efficient, parallelizable matrix multiplications, similar to techniques used in billion-scale systems like ScaNN (Guo et al., 2020).
The primary computational cost lies in the initial indexing task, demanding forward pass through the corpus. However, once indexed, retrieval becomes analogous to ScaNN, reducible to a parallelizable matrix multiplication with a (D, 768) matrix (D being the candidate document set, controlled by the beam size) and the (1, 768) query vector.
For inference, the learned index structure is equivalent to ScaNN-like data structures so the same scaling will hold (Note that even the ScaNN paper showcased its efficacy the Glove-1.2M benchmark, while parallelizing to billions of documents in production). Purthermore, ScaNN is deployable at a billion-scale so EHI would also hold.
Training EHI is analogous to standard contrastive-stylee learning and poses no thcoretical bottlencck in terms of document size. While our primary focus this paper is to showcase the potential our novel paradigm shift for dense retrieval and demonstrate its efficacy on well-tested gold benchmarks such as MS MARCO, running EHI on billion-scale benchmark would be interesting direction for future work. Large-scale retrieval systems like Differentiable Search Index (DSI; Tay el al. (2022)) and Neural Corpus Indexer (NCI; Wang et al. (2022)) operate in similar regimes of MS MARCO/NO NNQ-2,,,, and exploring EHI's scalability in such contexts could unlock further benefits.
E
Additional information of EHI's indexer
In this section, we hope to shed more light on the indexing process of EHI at both training and inference time. We expand upon our Figure 3, and only discuss the indexer part or the tree segment of the image in Figure 15. For a given query/ /document, the root node uses the embedding information (Vis in this case) alone, and we use a linear layer (U1) to transform the embeddings received by the encoder to another 768-dimensional embedding. Having these alfine transformations different aspects of the same query/ query/docummn document at each height to classify which children node the item needs to go to. This is followed by another linear layer (W1) 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 lowards (p² = [P1, Pa)). 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 embodding input (Vas) and the one hot version of the probabilities the previous height to make the predictions. During training, we have at the index with the highest probability and 0 otherwise (assume P; > P2. and hence one-hot version is [1, 0)) However, once the conditional probability at the next.
each height is essential as we would like to focus on
31

---

Published in Transactions on Machine Learning Research (08/2024)
EHI Indexer
[Pa, Pal Pi
Query 98
iphone se walmart cost
Vs8
TVsa.
H=0
[PL Pz]
[Vss, 1,01
TV
W2 softmax
U1
W1
U2
say P1 >P
Path Embedding:
[P1. Pz. PiPs. P1
Figure 15: EHI indexer a simple tree-structure where each height. uses a neural network decide which path the given query or document needs to go. Following the similar paradigm Figure 3, variables Vgs is the dense representations (embeddings) of the text of query document.
height is computed, we multiply by this Pi to have the actual probability of reaching nodes 3 and 4 (note these numberings follow the same logic as used in Figure 3). The final path embedding is the concatenation of these probability distributions leading to [Pi, P2, Pi P3, P1, Pa] as depicted Figure 15. During inference, the indexer follows similar setup where, instead of only doing one forward pass with the highest probability, do b forward passes given height with top-b highest. probabilities.
32

---

Published in Transactions on Machine Learning Research (08/2024)
Algorithm 1 Training step for EHI. u(g) & v(d) denote the query (q) & document (d) representation from encoder Eg. Similarly path embedding by the indexer (I₂) is denoted by T(4) & Tdd). Please refer to Algorithm 2 in appendix for the definition of TOPK-INDEXER and INDEXER. Note that the Update (.)
function updates the encoder and indexer parameters through the back-propagation of the loss through AdamW (Loshchilov & Hutter, 2017) in our case, while other optimizers could also be used for the same (e.g. SGD, Adam, etc.).
Require: Φ < Indexer Parameters, 0 < Encoder parameters, T + similarity threshold, B < indexer branching factor, H + indexer height, B + beam size, M + Mapping of leaf documents, Ó {qr 92...., {s} Q (randomly sampled mini-batch), INDEXER Returns path embedding, TOPK-INDEXER Returns most relevant leaf
procedure RETRIEVER(U, 6, B) Retrieves most relevant doc Leafs = INDE-INDEX o, B,H, return UteLeateM(I) end procedure
b Refer to Alg 2
procedure TRAINING(() jointly update encoder and indexer parameters 0, Φ
loss
for qi E Q do
d, ~ (klya 1 h; ~ [d|RETRIEVER(& 6, B)} (Aly 1}
Sample relevant document,
> Sample a hard neg. from retriever
loss +- ((Ea(q)),
b Triplet loss against hard negative
for E and q, + qj do
dj (klysk = -1 =)
> Sample an in-batch negative
loss + ((E))
Uses 0 to encode doc, and query Uses INDEXER in Alg for T
loss
loss + ((SIM(v(d)) < (ol,
end for end for return Update(0, Φ, loss) end procedure
> Update encoder and indexer parameters with gradient of loss
33

---

Published in Transactions on Machine Learning Research (08/2024)
Algorithm Training step for EHI. u(q) & v(d) denote the query (q) document (d) representation from encoder Eg. Similarly path embedding by the indexer (Ip) is denoted by T(g) & T(d).
Require: Φ <(WH,... W1, U H....: U1} < Indexer Parameters, 0 Encoder parameters, T + Inter- document similarity threshold, B f- branching factor of the Indexer, H height of the indexer, 8 +- beam size, . <- Mapping of leaf to documents, Q = (4,,2, ................................ -ss ~ Q (randomly sampled mini-batch)
procedure INDEXER [u, o, B, H) Computes indexer path emb. for a query/documner
= IOPK-INDEXER u, 6, B, H. = 1)
To
Pi =
for e.l... H do
S([o(i) h- o th h- :((]:)))))) W.. U.)
To [Ta: P pi]
Pt
end for return T end procedure
procedure JOPK-INDEXER( while H do
Indexes query/documents in the 3 most relevant leaves > Tuple of score, path and node ID
P = Max-Heap(ssize = 3) for s, P.pop() do
S=S(p: Wt, Un) s; For all i 1.. B: p.push (f)) p. B+i>)
end for P= end while return P end procedure
34


Reviews:
[{"review_id": "DsZcqkxa4g", "reviewer": "Reviewer_Qqpy", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Several embedding ideas seem useful and interesting, and the combined method does appear to have strong empirical performance.\n\nThe main weakness of the paper is poor presentation. Many components are not formally defined. The description of the encoder (BERT) is dismissed with a simple citation without any further details on architecture and training. The motivation for design are very superficial, e.g., the path is computationally intractable (why would a simple path be intractable?). The training and updating is treated very superficially (e.g. by having a call to Update() in Algorithm 2).\n\nWhile the experiments do appear to show improvement, there is no indication of statistical significance.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a hierarchical indexing algorithm for document retrieval that trains the embedding and indexing components simultaneously. Empirical evaluations show the merit of individual components of the architecture. Moreover, the the proposed method is shown the outperform several baselines on some standard benchmarks.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The description and motivation of the algorithm should be considerably improved.\n\nThere should be a detailed description of the experimental set-up, and clear indication of the significance of the results. The current presentation of the experiment might be sufficient for a conference paper that highlights the promise of an algorithm, but it is insufficient for TMLR that has a strong emphasis on the trustworthiness of the presented claims.", "broader_impact_concerns": "none", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "QiNLDejwpY", "reviewer": "Reviewer_y44s", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\nThis paper is well-written, making it easy to understand the proposed method. The motivation for the proposed approach is clear, and the method of jointly learning the encoder and ANNS is highly intriguing. The novelty of the proposed method is emphasized through a wide range of experimental results. The paper holds significant value by offering a new paradigm in information retrieval.\n\n## Weaknesses\n- The introduction of a new paradigm is highly commendable. However, several questions remain unanswered upon reading this paper:\n    - Is this method applicable to \"real\" scenarios at a billion-scale? (While the cost of extracting dense features may be comparable to other methods, I am curious about the costs incurred during the retrieval process.)\n    - If applicable, what are the computational costs?\n    - When adding multiple new search documents, what procedures must be followed with the trained indexer?\n- A simple question remains: What are the costs involved in hard-negative mining?\n- I am curious about the reason for using distilBERT, pretrained on MS MARCO. According to the main paper and the appendix, the main contribution is the introduction of a new paradigm, but the authors would also agree that performance is crucial. Therefore, for fair comparison, it should be clear if other comparative models also used distilBERT. According to Table 4 in the appendix, only the parameters of the encoder are presented, without a detailed description of the pretrained model used.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a novel method for dense retrieval that integrates embedding generation and approximate nearest neighbor search (ANNS) into a single, end-to-end learning framework. The proposed End-to-end Hierarchical Indexing (EHI) addresses the misalignment issues found in conventional two-stage approaches by jointly optimizing the embedding encoder and the search index structure. The authors highlight the method's capability to learn balanced and uniform clusters without needing proxy information or additional clustering algorithms. Extensive empirical evaluations demonstrate EHI's superiority over state-of-the-art techniques, showing significant improvements in retrieval accuracy on benchmarks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It would be beneficial if the paper included content that addresses my opinions on its weaknesses. If such content is already included, please indicate where it can be found.", "broader_impact_concerns": "There are no broader impact concerns of this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "q00KSCUtYa", "reviewer": "Reviewer_Nywg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros)\nThis paper presents the first end-to-end approach for training discrete indices-based ANNS models for text query-document retrieval. The novel tree-based architecture successfully overcomes the issues of previous two-stage approaches efficiently. Sufficient experiments are conducted to prove the advantages of the proposals. Detailed discussions on ablation results and method components are provided.\n\nCons)\nIn general, end-to-end approaches for ANNS already exist in visual domains, as mentioned in references [1-2]. It would be beneficial to clarify that this work represents the first approach specifically for the text domain.\n\n[1] Yu, Tan, et al. \"Product quantization network for fast image retrieval.\" Proceedings of the European Conference on Computer Vision (ECCV). 2018.\n[2] Jang, Young Kyun, and Nam Ik Cho. \"Generalized product quantization network for semi-supervised image retrieval.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes an end-to-end hierarchical indexing (EHI) method, which overcomes the limitations of previous approximate nearest neighbor search (ANNS) methods that require separate two-stage processes for semantic search tasks. The single-stage approach of the EHI method has the advantage of reducing potential mismatches between the embedding space and the requirements of the ANNS structure. It also allows for consideration of specific characteristics of the query distribution. EHI employs a tree-structured inverted index without any initialization of cluster points, utilizing significantly compressed path embedding and a parameterized indexer. EHI is a model-agnostic method that can be integrated into various models with different encoder architectures or training schemes. Experiments on various benchmarks show that EHI achieves state-of-the-art results.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Need explanation of what DE+ANNS stands for in section 1.2.\nIn Figure 1, it would be better to change e2e to end-to-end for clarification.\n\nOverall, paper is well written and straightforward.", "broader_impact_concerns": "I suppose nothing special for this paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper has mixed but mostly positive opinions. The main concerns include

- The end-to-end learning framework is not first proposed by this work (Reviewer Nywg)
  - This concern is addressed by revising Section 2, and focusing on "dense text retrieval"
- Comparison with SOTA -- because it uses distilBert (Reviewer y44s)
  - Table 4 was revised to include the backbone details to resolve the concern
- The presentation could be improved
  - The reason why the path is computationally intractable is now mentioned in Section 3.3.
  - In my opinion, the presentation quality still can be improved, but the current version looks okay. But, if the authors can update the manuscript, I slightly recommend revising the paper to make the motivation and experimental details more visible to novice readers. See my comment below.
- The lack of statistical significance for the experiments
  - The authors refuted this argument by citing Table 11
  - In my opinion, the current empirical evidence is somewhat okay, but the visibility of the results can be improved. See my comment below.

There were some additional comments, such as

- Practicalness of the proposed method to the billion-scale DB (Reviewer y44s)
  - As far as I understood, the discussion is not explicitly included in the revised paper. If my understanding is correct, I strongly recommend adding the related contents in the revision.

In my opinion, considering that BERT is now a very famous architecture whose detail can be omitted, the current manuscript has a weak problem with reproducibility (reproducibility can still be a problem because this submission does not include any implementation, but it is acceptable, IMO).

Overall, I think this paper can satisfy the TMLR evaluation criteria with a minor revision.

I request a minor revision by moving major contents in the Appendix to the main page. Please note that there is no page limitation for TMLR papers; the page limitation is only used to set the initial review period for the reviewers (2 weeks for less than 12 pages, 4 weeks for more than 12 pages). 

In my opinion, the appendix should be an "appendix," namely, the TMRL criteria (claim and evidence) should be supported whenever the appendix's contents are ignored. For example, dataset explanations or model hyperparameter details can be in the Appendix, but many experiments that support the main claim should be in the main paper (unless an experiment has a duplicated message with any previous experiment in the main paper).

I strongly recommend moving Appendix B (Motivation for EHI) and many important experiments in Appendix D and E to the main paper. After moving the contents to the main paper, I think this paper will need a polishing stage (e.g., enhancing clarity of the manuscript), which would be a minor modification.
Decision: Accept (poster)