LLM-Augmented Retrieval: Enhancing Retrieval Models Through Language Models and Doc-Level Embedding

## Abstract

Abstract Recently embedding-based retrieval or dense retrieval have shown state of the art results, compared with traditional sparse or bag-of-words based approaches. This paper introduces a model-agnostic doc-level embedding framework through large language model (LLM) augmentation. In addition, it also improves some important components in the retrieval model training process, such as negative sampling, loss function, etc. By implementing this LLM-augmented retrieval framework, we have been able to significantly improve the effectiveness of widely-used retriever models such as Bi-encoders (Contriever, DRAGON) and late-interaction models (ColBERTv2), thereby achieving state-of-the-art results on LoTTE datasets and BEIR datasets.

## 1 Introduction

The Bi-encoder (Karpukhin et al., ) is a type of neural network architecture that is widely used in information retrieval. It consists of two encoders, typically in the form of transformer models (Vaswani et al., ), which encode an vector representation for user queries and potential documents or passages respectively. These two encoders can be shared or using two separate models. The similarity between these two embedding vectors can then be computed, often using dot product or cosine similarity, to determine the relevance of the document or passage to the user’s query.

Figure: Figure 1: Overall view on LLM-augmented retrieval framework. Synthetic relevant queries and synthetic titles are generated from LLM and then assembled into doc-level embedding together with chunks (passages) split from the original document. The final retrieval is based on the similarity between user query and the doc-level embedding.
Refer to caption: extracted/5517725/relevance_0.png

Cross-encoders (Nogueira and Cho, ), unlike bi-encoders, amalgamate the inputs at an early stage, allowing for a more intricate interaction between user queries and documents. Here the user query and the document are concatenated, based on which a joint embedding vector is computed. The joint embedding vector is then used to make predictions, such as the relevance of a document to a query in an information retrieval task. Cross-encoders often outperform bi-encoders in tasks requiring a nuanced understanding of the interplay between inputs.

Late-interaction models, such as ColBERT (Khattab and Zaharia, ), ColBERTv2 (Santhanam et al., ) or SPALDE++ (Formal et al., ), are model architectures that hybrids cross-encoder models and bi-encoder models. Queries and documents are independently encoded into token-level vector representations. So in some sense, this is a bag of embedding vectors model. The interaction between these representations, which constitutes the “late interaction”, involves computing the cosine similarity or dot product scores over the token-level vector embedding.

All the model architectures require informative embedding of user queries and target documents. While we cannot control the user queries during retrieval tasks, we hypothesize that enriching the embedding of documents can improve the quality and robustness of text retrieval.
This work makes the following contributions:

- 1.
We propose LLM-augmented retrieval, a model-agnostic framework that enriches the contextual information in the vector embedding of documents to improve the quality and robustness of existing retrievers.
- 2.
We propose doc-level embedding, which combines more contextual information in the context embedding.
- 3.
We evaluate this framework across different models and a wide range of datasets, establishing state-of-art quality beyond original models.
- 4.
We propose improvements to some key components in retrieval model training, such as negative sampling, loss function, etc.

## 2 Background & Related Work

To the best of our knowledge, there’s no existing work on constructing a unified LLM-augmented framework for doc-level embedding. Therefore, we summarized some related knowledge that’s important to our work.

### 2.1 Data Augmentation in Information Retrieval

Data augmentation is a widely used technique in information retrieval training. In Contrastive Learning (Izacard et al., ), inverse cloze task, independent cropping, as well as random word deletion, replacement or masking are introduced to enrich the diversity of training data. While in training the DRAGON model, (Lin et al., ), not only query augmentation by query generation models are studied, but also label augmentation methods with diverse supervision are discussed and compared.

Large pre-trained language models are good at generating text data of high quality (Anaby-Tavor et al., ; Papanikolaou and Pierleoni, ; Yang et al., ; Kumar et al., ; Schick and Schütze, ; Meng et al., ). Some past work has been done to utilize the generation capability of language models to create synthetic training data for retriever models (Bonifacio et al., ; Wang et al., ).

### 2.2 Pseudo Queries Generation

Pre-generated pseudo queries are proved to be effective in improving retrieval performance. Previously the similarity of pseudo-queries and user-queries are calculated through BM25 or BERT model in determining the final relevance score of the query to document through relevance score fusion (Chen et al., ; Wen et al., ). An alternative method to generate pseudo queries is to generate pseudo query embedding through K-means clustering algorithm (Tang et al., ).

### 2.3 Retrieval Augmented Generation (RAG)

In a Retrieval Augmented Generation (RAG) system (Lewis et al., ), a dedicated retriever module is introduced to retrieve relevant documents from a set of corpus based on the input query, and the retrieved documents are then integrated by a language model as part of the context, in order to refine its final response generation. This approach can be applied to not only decoder-only models but also encoder-decoder architectures (Yu, ).

### 2.4 Similarity Scores

Having the embedding vectors of input query and target document, we need to compute the similarity score between these two vectors to determine the relevance between them (Jones and Furnas, ). Dense retrievers like Contriever (Izacard et al., ) and DRAGON (Lin et al., ) use dot products to measure the similarity of query and document in the text embedding space. SPLADE (Formal et al., ) and SPLADEv2 (Formal et al., ) use dot products as well while ColBERT (Khattab and Zaharia, ) and ColBERTv2 (Santhanam et al., ) use cosine similarity scores.
The main difference between cosine similarity and dot product is the sensitivity to the magnitude of the embedding vectors. And dot product is computationally cheaper and more stable.

### 2.5 Training

The contrastive InfoNCE loss (Izacard et al., ; Lin et al., ) is widely adopted in dense retrieval training and is defined as:

$$ $L(q,d_{+})=-\frac{exp(s(q,d_{+})/\tau)}{exp(s(q,d_{+})/\tau)+\sum_{i=1}^{K}exp (s(q,d_{i})/\tau)}$ (1) $$

where $\tau$ is the temperature, $s$ is the similarity score function, $q$ is the input query, $d_{i}$ is any candidate document and $d_{+}$ is the relevant document. Other popular loss functions include point-wise Binary Cross Entropy (Sun et al., ), list-wise Cross Entropy (Bruch, ), RankNet (Burges et al., ) and LambdaLoss (Wang et al., ). Pointwise Binary Cross Entropy is calculated based on each of the query-document pairs independently. List-wise Cross-Entropy loss is widely used in passage ranking and minimizes list-wise softmax cross entropy on all passages. RankNet is used to calculate the pair-wise relevant order of passages.

## 3 LLM-augmented Retrieval

In this section, we first discuss the components of the LLM-augmented retrieval framework. After that, we explain how this framework can be adapted to different retriever model architectures. In particular, we propose doc-level embedding for bi-encoders and late-interaction encoders under the LLM-augmented retrieval framework and demonstrate how it is applied to enhance the end-to-end retrieval quality.

### 3.1 LLM-Augmented Retrieval Framework

#### 3.1.1 Synthetic Relevant Queries

The inspiration for this concept was drawn from web search techniques (Xue et al., ; Guo et al., , ; Chuklin et al., ). To illustrate the idea, let us consider an example on a user query of "MIT". Without prior knowledge, it is challenging to figure out that "Massachusetts Institute of Technology" and "MIT" are equivalent. Nevertheless, in web search, we can observe that the home page of "Massachusetts Institute of Technology" has received numerous clicks from the query "MIT", allowing us to infer that the home page of "Massachusetts Institute of Technology" must be closely associated with the query "MIT". On the other side, we typically don’t have click data for each user query in the scenario of contextual retrieval. However, large language models are good at generating synthetic queries (Anaby-Tavor et al., ; Papanikolaou and Pierleoni, ; Yang et al., ; Kumar et al., ; Schick and Schütze, ; Meng et al., ) so we can use the synthetic queries as proxy “click data” to direct user queries to related documents.

An important point is that in traditional retrieval tasks, we are using similarity to express relevance (Jones and Furnas, ). The similarity score is mathematically defined as either dot product or cosine of the encoded vectors of user query and documents. However, sometimes this similarity score might not reflect semantic relevance (Rao et al., ). For example, “Who is the first president of the United States?” might be very close, in terms of similarity score, to “Who became the first president of America?”. But our target answer might be a Wiki page or autobiography about “George Washington”, whose similarity score to the query is not that high. While if we use Washington’s autobiography to create synthetic queries, “Who became the first president of America?” may be one of them. The user query “Who is the first president of the United States?” can easily match to the relevant query through similarity scores. And the latter points to the target document (Washington’s autobiography). Therefore the generated synthetic queries express the semantic of the original document from different angles which are helpful to match the relevant queries.

Figure: Figure 2: Through synthetic relevant queries, the relevance relationship is not solely expressed by the similarity now but also expressed by the augmentation steps of the large language models
Refer to caption: extracted/5517725/relevance_2.png

#### 3.1.2 Title

The title of a document plays a crucial role in determining its relevance and usefulness to a user’s query. The title is often the first piece of information that a user sees when searching for documents, and it can greatly influence their decision to click on a particular result. A well-crafted title can provide important context and keywords that help users quickly understand the content and purpose of a document. If the original document has a title, we can use them directly. If it does not, we can leverage the large language models to generate a synthetic title for a that document.

#### 3.1.3 Chunks (Passages)

Chunking refers to the process of dividing a large document or piece of text into smaller, more manageable units. These units, known as "chunks," or "passages", are typically created by grouping together related pieces of information. Due to the limitation of the context window of retriever models, (in other words, max length of the model input), we typically divide a long document into several chunks whose number of tokens is below the context window limit. The chunk data is from original documents and is not from LLM-augmentation. The optimal chunking size is different for various retriever models. For bi-encoders like Contriever and DRAGON, we found the optimal chunking size to be 64 after empirical studies. For token-level late-interaction models such as ColBERT and ColBERTv2, since it’s already calculating the similarity score at token-level, it’s not necessary to chunk the original documents unless the context window limit is reached.

### 3.2 Doc-level embedding

In this section, we will first introduce the high level idea of doc-level embedding for information retrieval, and then use bi-encoders and token-level late-interaction models to illustrate how the doc-level embedding can be adaptive to different retriever model structures.

Document fields. For convenience, we call the above mentioned information source, synthetic queries, title and chunk, fields of a document. These fields express the semantic of the original document from different angles and will be composed into the doc-level embedding of a document which is static and can be pre-computed and cached for information retrieval. Embedding indexes can be pre-built to speed up the retrieval inference, and each doc-level embedding points to the original document.

Figure: Figure 3: The graphic representation of "relevance" in doc-level embedding
Refer to caption: extracted/5517725/relevance_3.png

#### 3.2.1 For Bi-encoders

Bi-encoders are typically “Two-Tower” model structures. Given a query and a document, a query encoder and a doc encoder are applied to compute the embedding vectors for the query and the document respectively. Then these two embedding vectors are fed into the dot products (or cosine similarity) to compute the similarity scores between the query and the document. As we want to enrich the document embedding vectors by injecting the synthetic queries and titles, we propose to compute the similarity as the following:

$$ $sim(q,d)=max_{i}s(q,c_{i})+\sum_{f}w_{f}\times s(q,e_{f})$ (2) $$

The first term on the right hand side calculates the maximum similarity score of query chunk embedding pairs, where $s$ is the similarity score function, $q$ is the embedding vector of the input query, $c_{i}$ is the embedding vector of the $i$-th chunk in the document. This term is commonly used in current embedding based retrievals, which determines the similarity between a query and a document based on the query and the most relevant chunk in the document.

The second term is innovative and considers more information, where $e_{f}$ is the embedding vector of each document field. The similarity scores between the query embedding and each field embedding are computed and combined together, each with a field weight parameter $w_{f}$. As described before, those document fields include the synthetic query, title and chunk fields.

Now let’s consider how to compute the query field embedding $e_{f}$ for each field. For the title field, which just contains one phase or sentence, it is straightforward. We just apply the doc encoder to compute the embedding vector of the title as the title field embedding.

For the chunk field which may contain multiple doc chunks, we can compute the embedding vector of each chunk with the doc encoder. The question is how to combine these embedding vectors to represent the whole document? Actually people have already considered a similar problem for sentence embedding: having the embedding vectors of all the tokens in one sentence, how to come up with a representation of the whole sentence? A simple but effective idea was proposed in (Arora et al., ) for this problem. Here we adapt this idea to the doc embedding problem. Namely we compute the average of all the chunk embedding vectors as the chunk field embedding. Similarly, for the synthetic query field, we compute the embedding vector of each query with the query encoder, and then compute the average of these embedding vectors as the query field embedding. This simple approach works very well in our experiments, while clearly more advanced approaches can be explored here in the future.

Furthermore, as the similarity function is linear(^1^11Dot product is linear. Cosine similarity is also linear if we normalize the embedding vectors to unit length.), the above equation can be simplified as the following:

$$ $sim(q,d)=max_{i}s(q,c_{i}+\sum_{f}w_{f}\times e_{f})$ (3) $$

Therefore, we can treat $c_{i}+\sum_{f}w_{f}\times e_{f}$ as the chunk embedding vector of each chunk $c_{i}$ of the original document, and still apply the algorithms like approximate nearest neighbors (Indyk and Motwani, ) to retrieve the most relevant documents.

#### 3.2.2 For Token-Level Late-Interaction Models

Instead of using a single embedding vector for query and a single embedding vector for each document, late-interaction models such as ColBERT and ColBERTv2 use token-level embedding and the embedding vectors of all the tokens are kept and will participate in computing the similarity score between the query and document.

$$ ${sim(q,d)=\sum_{i}max_{j}s(q_{i},t_{j})}$ (4) $$

where $q_{i}$ and $t_{j}$ are token-level embedding vectors for the input query and document respectively. Therefore, for each query token, the most similar token from the document is identified, and their similarity score is recorded. All these scores are summed up over all the query tokens to obtain the overall similarity between the query and the document. Since the similarity score calculation is done at token-level, we can concatenate the synthetic queries and titles to the original document passages. After that, we decide whether to chunk the concatenated documents if the number of tokens reaches the context window limit.

## 4 Experiments

In this section we will introduce our experiment set up, the datasets and the models used in our experiments.

### 4.1 Datasets

BEIR Data The BEIR (Benchmark for Evaluating Information Retrieval) dataset (Thakur et al., ), provides a comprehensive benchmark for evaluating and comparing different IR models, particularly in the context of out-of-domain tests. BEIR is designed to address the limitations of other previous datasets by offering a more diverse and extensive range of queries and passages. It encompasses a wide array of topics, thereby facilitating a more robust and comprehensive evaluation of IR models.

LoTTE Data LoTTE data (Santhanam et al., ) is a dataset specifically designed for Long-Tail Topic-stratified Evaluation. The primary focus of LoTTE is on natural user queries that are associated with long-tail topics, which may not be adequately covered by an entity-centric knowledge base such as Wikipedia. LoTTE is composed of 10 distinct test sets, each containing 500 – 2,000 queries and 100,000 – 2,000,000 passages. These test sets are divided by topic, and each is accompanied by a validation set comprising related yet disjoint queries and passages. In this experiment we are using the test split sorely for evaluation.

### 4.2 Models

Contriever For Contriever we use the Roberta-base (Liu et al., ) model architecture, the checkpoint is trained on Wiki passages (Karpukhin et al., ) and CC100 (Conneau et al., ) data through Contrastive Learning. It has 125M parameters, a context window of 512 tokens, with 12 layers, 768 hidden dimensions, and 12 attention heads. We use the single Roberta-base model for query encoder and context encoder, namely the “Two Tower” of the bi-encoder is shared in model architecture.

DRAGON For DRAGON we also use the Roberta-base model architecture, the checkpoint is trained by the author and released publicly. Different from the Contriever model, DRAGON has separate Roberta-base models for query encoder and context encoder.

ColBERTv2 For ColBERTv2, we use bert-base-uncased model architecture, the same as the default settings in the original paper. It has 110M parameters and a context window of 256 tokens, with 12 layers, 768 hidden dimensions, and 12 heads. The checkpoint is trained on the MSMARCO dataset (Nguyen et al., ) from the author’s checkpoint.

### 4.3 Implementation Details

We choose open source Llama-70B (Touvron et al., , ) for synthetic query generation and title generation. The prompt templates used for generating synthetic queries and titles are in Table 9 and 10.

For Bi-encoders, we implemented the doc-level embedding as above mentioned and chose $w_{query}$=1.0, $w_{title}$=0.5, $w_{chunk}$=0.1 for the Contriever model and $w_{query}$=0.6, $w_{title}$=0.3, $w_{chunk}$=0.3 for the DRAGON model. We chose $chunk_size$=64 after empirical experiments and found 64 normally performed the best in retrieval results for Bi-encoders. Note these hyper-parameters of field weights are not fully optimized. They are chosen according to the performance of Bi-encoder models on the single LoTTE-lifestyle-forum dataset, and then fixed for the evaluations on all the other datasets.

For ColBERTv2, as mentioned previously, we concatenate the title with all the synthetic queries for each document and make it an additional “passage” of the original document. Thus there’s no field weights hyper-parameters in these experiments. There could be other better assembling methods for composing the doc-level embedding under a late-interaction model architecture. We set index_bits=8 when building the ColBERT index.

## 5 Results

The result on LoTTE and BEIR data for all three models can be seen in Table 1, 2 and 3. We can observe that the LLM augmented retrieval and doc-level embeddings boost the recall@3 and recall@10 of Bi-encoders (Contriever and DRAGON) significantly. For token-level late-interaction models (ColBERTv2), the increase on the LoTTE and BEIR Dataset is still clear, although not as much as on Bi-encoders. We hypothesis this due to that the baselines for token-level late-interaction models are much higher than those of Bi-encoders.

Moreover, the performance of the LLM-augmented Contriver has exceeded the performance of the vanilla DRAGON in most datasets. Similarly, the LLM-augmented DRAGON even exceeds vanilla ColBERTv2 on BEIR-ArguAna, BEIR-SciDocs and BEIR-CQADupstack-English datasets and greatly reduce the performance gap in the remaining datasets, although ColBERTv2 introduces a more complex late-interaction architecture than DRAGON. Therefore, we can see that, after enriching the embedding of documents with LLM augmentation, we can greatly improve the retriever models recall performance without further fine-tuning.

**Table 1: Results on Contriever: The performance of LLM-augmented Contriver has greatly exceeded the vanilla Contriever on both LoTTE and BEIR dataset, and even exceeds the performance of the vanilla DRAGON in most datasets. * means base model plus the doc-level embedding ($w_{query}$=1.0, $w_{title}$=0.5, $w_{chunk}$=0.1).**
| LoTTE Dataset |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Recall | Lifestyle Forum | Lifestyle Search | Recreation Forum | Recreation Search | Science Forum | Science Search | Technology Forum | Technology Search | Writing Forum | Writing Search |
| Contriever | recall@3 | 0.4366 | 0.3358 | 0.3486 | 0.1948 | 0.1046 | 0.1005 | 0.1826 | 0.1242 | 0.3950 | 0.2745 |
| recall@10 | 0.6149 | 0.4690 | 0.4895 | 0.2857 | 0.1706 | 0.1637 | 0.3174 | 0.1896 | 0.5390 | 0.3950 |  |
| Contriever* | recall@3 | 0.6244 | 0.6021 | 0.5455 | 0.4610 | 0.2395 | 0.2901 | 0.3663 | 0.3557 | 0.5970 | 0.5724 |
| recall@10 | 0.7622 | 0.7821 | 0.6948 | 0.6320 | 0.3570 | 0.4684 | 0.5494 | 0.5017 | 0.7365 | 0.6919 |  |

**Table 2: Results on DRAGON: The performance of LLM-augmented DRAGON has greatly exceeded the vanilla DRAGON on both LoTTE and BEIR dataset, and even exceeds vanilla ColBERTv2 on BEIR-ArguAna, BEIR-SciDocs and BEIR-CQADupstack-English datasets, as well as greatly reduces the performance gap in the remaining datasets. * means base model plus the doc-level embedding ($w_{query}$=0.6, $w_{title}$=0.3, $w_{chunk}$=0.3).**
| LoTTE Dataset |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Recall | Lifestyle Forum | Lifestyle Search | Recreation Forum | Recreation Search | Science Forum | Science Search | Technology Forum | Technology Search | Writing Forum | Writing Search |
| DRAGON | recall@3 | 0.5270 | 0.5598 | 0.4560 | 0.4253 | 0.2578 | 0.2601 | 0.2854 | 0.3591 | 0.5300 | 0.5798 |
| recall@10 | 0.6798 | 0.7035 | 0.5949 | 0.5325 | 0.3704 | 0.3938 | 0.4232 | 0.5101 | 0.6675 | 0.7311 |  |
| DRAGON* | recall@3 | 0.6883 | 0.7625 | 0.6079 | 0.6472 | 0.3099 | 0.4498 | 0.4192 | 0.5285 | 0.6520 | 0.7031 |
| recall@10 | 0.8172 | 0.8911 | 0.7468 | 0.7944 | 0.4427 | 0.6062 | 0.6038 | 0.7097 | 0.7725 | 0.8170 |  |

**Table 3: Results on ColBERTv2: The performance of LLM-augmented ColBERTv2 has greatly exceeded the performance of vanilla ColBERTv2 on both LoTTE and BEIR dataset. * means base model plus the doc-level embedding.**
| LoTTE Dataset |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Recall | Lifestyle Forum | Lifestyle Search | Recreation Forum | Recreation Search | Science Forum | Science Search | Technology Forum | Technology Search | Writing Forum | Writing Search |
| ColBERTv2 | recall@3 | 0.6988 | 0.7927 | 0.6344 | 0.6677 | 0.3932 | 0.5073 | 0.4496 | 0.5940 | 0.6960 | 0.7423 |
| recall@10 | 0.8087 | 0.8911 | 0.7498 | 0.7868 | 0.5285 | 0.6613 | 0.6292 | 0.7315 | 0.8050 | 0.8366 |  |
| ColBERTv2* | recall@3 | 0.7308 | 0.8003 | 0.6753 | 0.7100 | 0.4026 | 0.5024 | 0.4626 | 0.5956 | 0.7145 | 0.7544 |
| recall@10 | 0.8447 | 0.9107 | 0.7862 | 0.8268 | 0.5558 | 0.6726 | 0.6517 | 0.7383 | 0.8260 | 0.8571 |  |

## 6 Ablation Studies

In this section, we want to see how the chunk, query and title fields impact the retrieval quality of different retriever models. For Bi-encoders (Contriever and DRAGON), we further control the field weights of chunk, synthetic query and title to see how those parameters can affect the performance. For the token-level late-interaction model (ColBERTv2), we just control the model to use one of the chunk, query or title field only to see how the they affect the end-to-end retrieval quality.

For the Contriever model (Table 4), we have observed that most of the time the synthetic queries play the most critical role in boosting the recall performance, comparing to the other two fields, while in BEIR-SciDocs and BEIR-Scifact, synthetic queries’ importance is smaller. Therefore, a weighted sum of multiple fields in doc-level embedding yields better performance in most cases. And these weights can be further tuned as hyper-parameters.

For the DRAGON model (Table5), there’s less strong pattern that which field plays a more important role in doc-level embedding. In the LoTTE dataset it’s more driven by the title field. However, the chunk field matters more in datasets of BEIR-ArguAna, BEIR-Quora and BEIR-SciFact. Once again, a weighted sum of multiple document fields in doc-level embedding yields better performance in most cases. For the different patterns we have observed in DRAGON vs Contriever, one reason might be that DRAGON uses separate query and context encoders, while Contriever uses shared query and context encoders in our setup. Thus Contriever is better at identifying similarity instead of relevance and that’s why the synthetic query field has a more significant impact in the Contriever model, since it better transforms similarity to relevance as explained before.

For ColBERTv2 (Table 6), we have observed from the LoTTE dataset that further chunking the passages actually hurts the performance. This may be because the similarity calculation is done at token-level so chunking the passages (meanwhile increasing the number of chunks) will not help the ColBERTv2 model digest the granular contextual information. Therefore we did not evaluate the chunk-only scenario on the BEIR dataset. The synthetic queries are more critical than titles for ColBERTv2 across the datasets, while combining them all often yields even better recall results. Note again there’s no field weights hyper-parameters for token-level late-interaction models.

## 7 Supervised Fine-tuning

The above studies are all based on zero-shot evaluation, namely the checkpoints in (Lin et al., ; Santhanam et al., ) are directly used in evaluation. Actually our LLM-augmented retrieval and doc-level embedding also support supervised fine-tuning on domain-specific datasets. Some popular training methods include picking hard negatives, constructing in-batch or cross-batch negatives, and calculating InfoNCE loss as training loss (Izacard et al., ; Lin et al., ). In this section, we propose some effective training techniques for fine-tuning which we used internally on proprietary dataset and performed better than the common methods mentioned above.

### 7.1 Adaptive Negative Sampling

The process of acquiring negative docs for each query can be approached in several ways. A straightforward method involves the random sampling of doc contexts, excluding the ground truth. Alternatively, a model such as BM25 (Robertson et al., , ) could be employed to select “strong” negative doc contexts that have been assigned high ranking scores by the model. However, for more effective training, it is advantageous to obtain “difficult” negative doc contexts. To this end, an adaptive negative sampling approach was developed. The retrieval model under training is used to rank all doc contexts for each query. The doc contexts ranked in the top positions, excluding the ground truth doc contexts, are selected and then represented as the most challenging negative examples that could potentially confound our retrieval model. Consequently, it is imperative to train the retrieval model to effectively handle these samples.

A significant challenge of this approach is the time and space complexity. For each query (or batch of queries), it is necessary to re-compute the embedding for all doc contexts to identify the top negative examples. These doc context embedding cannot be pre-calculated due to the continuous updating of the model parameters during training, which correspondingly alters the doc context embedding. To address this issue, we can periodically pre-calculate and update the top negative samples for multiple training batches simultaneously.

### 7.2 Loss Function

In supervised training, for each training query, each doc has a binary label: relevant or not relevant.

In the context of binary label problems, the cross-entropy loss function is a prevalent choice. However, the cross-entropy is primarily for classification problems. This is due to its requirement for model output scores to be maximized for positive examples and minimized for negative examples. Within the realm of information retrieval, the primary challenge lies in the task of ranking of target documents. The absolute numerical values of output scores are not the primary focus. Instead, the focus is on the relative order and gaps between the model outputs of positive and negative doc contexts. Consequently, we adopted the margin ranking loss function (Nayyeri et al., ), which aims to maximize the "margin" between the scores of positive and negative examples.

$$ $Loss(o_{1},o_{2},y)=max(0,-y(o_{1}-o_{2})+margin)$ (5) $$

In the above equation, $o_{1}$, $o_{2}$ are model output scores, the value of $y$ equals 1 or -1, indicating whether $o_{1}$ should be larger or smaller than $o_{2}$ respectively. For example, when $y$=1, and the gap between $o_{1}$ and $o_{2}$ is larger than the “margin” parameter, the loss value is 0. Otherwise, the smaller the gap between $o_{1}$ and $o_{2}$ , the larger the loss value. Our experiments showed that this margin loss was better than the cross entropy loss on proprietary domain data fine-tuning.

## 8 Conclusion

This paper presents a novel framework, LLM-augmented retrieval, which significantly improves the performance of existing retriever models by enriching the embedding of documents through large language model augmentation. The proposed framework includes a doc-level embedding that encodes contextual information from synthetic queries, titles, and chunks, which can be adapted to various retriever model architectures. The proposed approach has achieved state-of-the-art results across different models and datasets, demonstrating its effectiveness in enhancing the quality and robustness of neural information retrieval. Future research could explore further enhancements to the LLM-augmented retrieval framework, such as the integration of additional contextual information into doc-level embedding, the application of more advanced similarity score measures, more complicated approaches to combine the embedding of multiple chunks/queries into one chunk/query field embedding, etc.

## 9 Limitations

One limitation of study include the extra computation resources it required in augmenting relevant queries and titles for original documents and sometimes to size of augmented texts can be comparable to the size of the original documents. This computational limitation may restrict the usage of this approach where computational resource is limited.

Another limitation or risk is that the hallucination in large language models may pose extra inaccuracy in augmented corpus to the original documents. Hallucination remains an unsolved problem in the field of large language model’s study.

## Appendix

**Table 4: Ablation study on doc-level embedding with Contriever. In most cases the ensemble of relevant queries, title and chunks gives the best results. * means base model plus the doc-level embedding (chunk:0.1, query:1.0, title:0.5).**
| LoTTE Dataset |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Recall | Lifestyle Forum | Lifestyle Search | Recreation Forum | Recreation Search | Science Forum | Science Search | Technology Forum | Technology Search | Writing Forum | Writing Search |
| Contriever<br>$w_{chunk}$=1.0 | recall@3 | 0.4476 | 0.4342 | 0.3871 | 0.2695 | 0.1358 | 0.1896 | 0.2116 | 0.1879 | 0.4425 | 0.3968 |
| recall@10 | 0.6459 | 0.6172 | 0.5415 | 0.4145 | 0.2196 | 0.3063 | 0.3693 | 0.3003 | 0.600 | 0.5369 |  |
| Contriever<br>$w_{query}$=1.0 | recall@3 | 0.6194 | 0.6967 | 0.5355 | 0.4437 | 0.2335 | 0.2901 | 0.3523 | 0.3305 | 0.5860 | 0.5472 |
| recall@10 | 0.7762 | 0.7837 | 0.6863 | 0.6115 | 0.3461 | 0.4295 | 0.5180 | 0.4883 | 0.7410 | 0.6910 |  |
| Contriever<br>$w_{title}$=1.0 | recall@3 | 0.5310 | 0.4902 | 0.4975 | 0.3789 | 0.2345 | 0.1896 | 0.3468 | 0.2668 | 0.5315 | 0.4809 |
| recall@10 | 0.6958 | 0.6641 | 0.6404 | 0.5314 | 0.3421 | 0.3241 | 0.5200 | 0.4077 | 0.6725 | 0.6153 |  |
| Contriever* | recall@3 | 0.6244 | 0.6021 | 0.5455 | 0.4610 | 0.2395 | 0.2901 | 0.3663 | 0.3557 | 0.5970 | 0.5724 |
| recall@10 | 0.7622 | 0.7821 | 0.6948 | 0.6320 | 0.3570 | 0.4684 | 0.5494 | 0.5017 | 0.7365 | 0.6919 |  |

**Table 5: Ablation study on doc-level embedding with DRAGON. In most cases the ensemble of relevant queries, title and chunks gives the best results. * means base model plus the doc-level embedding (chunk:0.3, query:0.6, title:0.3).**
| LoTTE Dataset |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Recall | Lifestyle Forum | Lifestyle Search | Recreation Forum | Recreation Search | Science Forum | Science Search | Technology Forum | Technology Search | Writing Forum | Writing Search |
| DRAGON<br>$w_{chunk}$=1.0 | recall@3 | 0.6244 | 0.7126 | 0.5375 | 0.5779 | 0.2385 | 0.3695 | 0.3239 | 0.4362 | 0.6190 | 0.6471 |
| recall@10 | 0.7627 | 0.8636 | 0.6813 | 0.7359 | 0.3738 | 0.5462 | 0.5010 | 0.6023 | 0.7620 | 0.7871 |  |
| DRAGON<br>$w_{query}$=1.0 | recall@3 | 0.6583 | 0.7247 | 0.5839 | 0.6071 | 0.2707 | 0.3647 | 0.3892 | 0.4866 | 0.6235 | 0.6583 |
| recall@10 | 0.8017 | 0.8654 | 0.7108 | 0.7478 | 0.3991 | 0.5219 | 0.5704 | 0.6812 | 0.7420 | 0.7656 |  |
| DRAGON<br>$w_{title}$=1.0 | recall@3 | 0.6913 | 0.7610 | 0.6294 | 0.6472 | 0.3565 | 0.4408 | 0.4616 | 0.5436 | 0.6550 | 0.6928 |
| recall@10 | 0.8167 | 0.8790 | 0.7458 | 0.7879 | 0.4834 | 0.5948 | 0.6477 | 0.7064 | 0.7690 | 0.8011 |  |
| DRAGON* | recall@3 | 0.6883 | 0.7625 | 0.6079 | 0.6472 | 0.3099 | 0.4498 | 0.4192 | 0.5285 | 0.6520 | 0.7031 |
| recall@10 | 0.8172 | 0.8911 | 0.7468 | 0.7944 | 0.4427 | 0.6062 | 0.6038 | 0.7097 | 0.7725 | 0.8170 |  |

**Table 6: Ablation study on doc-level embedding with ColBERTv2. In all cases the ensemble of relevant queries, title and chunks gives the best results. * means base model plus the doc-level embedding.**
| LoTTE Dataset |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Recall | Lifestyle Forum | Lifestyle Search | Recreation Forum | Recreation Search | Science Forum | Science Search | Technology Forum | Technology Search | Writing Forum | Writing Search |
| ColBERTv2<br>chunk64 only | recall@3 | 0.6184 | 0.7474 | 0.5949 | 0.6245 | 0.3535 | 0.4797 | 0.3927 | 0.5168 | 0.6680 | 0.7199 |
| recall@10 | 0.7537 | 0.8759 | 0.7258 | 0.7586 | 0.4819 | 0.6353 | 0.5758 | 0.6846 | 0.7945 | 0.8273 |  |
| ColBERTv2<br>query only | recall@3 | 0.7088 | 0.7413 | 0.6479 | 0.6580 | 0.3634 | 0.4327 | 0.3643 | 0.4530 | 0.6835 | 0.7274 |
| recall@10 | 0.8222 | 0.8759 | 0.7642 | 0.7727 | 0.4948 | 0.5997 | 0.5259 | 0.5419 | 0.7890 | 0.8254 |  |
| ColBERTv2<br>title only | recall@3 | 0.6004 | 0.6218 | 0.5210 | 0.5487 | 0.3128 | 0.3695 | 0.4336 | 0.4715 | 0.5425 | 0.5780 |
| recall@10 | 0.7368 | 0.7458 | 0.6479 | 0.6937 | 0.4378 | 0.5024 | 0.5968 | 0.6141 | 0.6505 | 0.6853 |  |
| ColBERTv2* | recall@3 | 0.7308 | 0.8003 | 0.6753 | 0.7100 | 0.4026 | 0.5024 | 0.4626 | 0.5956 | 0.7145 | 0.7544 |
| recall@10 | 0.8447 | 0.9107 | 0.7862 | 0.8268 | 0.5558 | 0.6726 | 0.6517 | 0.7383 | 0.8260 | 0.8571 |  |

**Table 7: Statistical description of BEIR dataset**
| Question Set | #Questions | #Passages |
| --- | --- | --- |
| ArguAna | 1406 | 8674 |
| FIQA | 6648 | 57600 |
| Quora | 15000 | 522929 |
| SciDocs | 1000 | 25313 |
| SciFact | 1109 | 5183 |
| CQADupstack English | 1570 | 40221 |
| CQADupstack Physics | 1039 | 38316 |

**Table 8: Statistical description of LoTTE Test dataset**
| Question Set | #Questions | #Passages | Subtopics |
| --- | --- | --- | --- |
| Lifestyle Search | 661 | 219k | Cooking, Sports, Travel |
| Lifestyle Forum | 2002 |  |  |
| Recreation Search | 924 | 167k | Gaming, Anime, Movies |
| Recreation Forum | 2002 |  |  |
| Science Search | 617 | 1.694M | Math, Physics, Biology |
| Science Forum | 2017 |  |  |
| Technology Search | 596 | 639k | Apple, Android, UNIX, Security |
| Technology Forum | 2004 |  |  |
| Writing Search | 1071 | 200k | English |
| Writing Forum | 2000 |  |  |

Here we share the ablation studies’ results and the prompts we used in generating synthetic queries and titles using Llama-70B.

**Table 9: Prompt for generating relevant queries for documents**
| I will give you an article below. What are some search queries or questions that are relevant for this article or this article can answer? |
| --- |
| Separate each query in a new line. |
| This is the article: {document} |
| Only provide the user queries without any additional text. Format every query as ’query:’ followed by the question. Don’t write empty queries. |

**Table 10: Prompt for generating titles for documents.**
| I will give you an article below. Create a title for the below article. |
| --- |
| This is the article: {document} |
| Only provide the title without any additional text. Format the reply starting with ’title:’ followed by the question. Don’t write empty title. |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: - How was the chunk size determined?
- I suggest a retouch on line 348-350 “augmentation of document embeddings with LLMs can substantially elevate the recall performance of retrieval models without necessitating additional fine-tuning”. I don’t think one model is simply fine-tuned relative to another, but it’s their complexity that differs.
- Could you cite some prior work on chunk-level embeddings (line 246-248)? Using a single document-level embedding, such as truncating, seems to be a common practice. In such cases, the proposed method will increase the document index size by |doc_len| / |chunk_size|.
soundness: 2
strengths: - I like the idea of enhancing document embeddings prior to the inference step. Specifically, this can be pre-computed, offering practical advantages over query expansion methods which have been studied extensively.
- Proposed method can be applied to both bi-encoder and token-level late-interaction models, and it achieves consistently strong results across most datasets.
summary: This paper introduces a novel, model-agnostic approach for document embedding, leveraging an augmented document field comprising synthetically generated queries, titles, and document chunks. With a new similarity function, the method computes a similarity score between the query and document field as a proxy for query-document relevance. The method outperforms using naive document embeddings from three existing retrievers (Contriever, DRAGON, ColBERTv2) on two benchmarks (BeIR, LoTTE).
weaknesses: - There lacks detailed analysis of how (synthetically generated) document fields help the retrieval.  For example, how many queries were generated per document on average? The limitations section mentions that the augmented text may be as long as the original text, which suggests a substantial number of queries—clarification here would be valuable. Also, could you include some qualitative examples? 
- In the ablation study, weights for each field are exclusively set to 1.0. However, the “left term” at equation (1) remains active, meaning the query-chunks still affect the results. This raises concerns about whether the ablation study truly isolates the impact of each field, especially for the query and title.
- The results are evaluated only on the subset of BeIR benchmark without any explanations.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Clarified on implementation details: according to Figure1, for each passage coming from the same document, it seems you use the same synthetic queries and titles? However, in the experiments, you mention you use the original chunk from the datasets; then, how do you know which chunks are coming from the same document and use the document to generate synthetic queries and titles for those chunks.
2. Which dev set are used to tune the hyperparameters, $w_{query}, w_{title}, w_{chunk}$?
3. From the prompts, it seems that the models are instructed to output multiple relevant queries; however, only one generated query is used for document expansion. How do you choose the query among all the generated ones? Or if you use multiple generated queries, how do you combine them?
soundness: 2
strengths: 1. The proposed method is straightforward and easy to implement and shows significant improvements for existing state-of-the-art retrievers.
2. The proposed approach is useful, which combines the document embedding from different sources and forms a single document embedding. Thus, the index size is still the same for single-vector dense retrieval.
summary: The paper presents an approach to document expansion using LLMs to enhance the zero-shot retrieval effectiveness of existing state-of-the-art bi-encoder retrievers (Contriever, DRAGON and ColBERTv2).
weaknesses: 1. Although the proposed approach is simple and effective, I think the major concern for me is the limited discussion in ablations and comparisons, which are important for the readers who want to adopt this approach. I'm willing to raise score if the authors can address the concern below: (1) As a reader, one may want to know the quality and cost tradeoff of the generated titles and queries. For instance, the paper prompts Llama-70B to make it work; how about using smaller size of Llama; 7B, 13B etc.? and the additional cost of document indexing with LLM document expansion should be reported. (2) There are no ablations on the three combinations: chunk + query, chunk + title and query + title, which helps readers to choose the combinations with a balanced tradeoff between effectiveness and cost  (3) A comparison with LLM-based retrievers, such as RepLlama[1] or MistralE5[2], in terms of retrieval effectiveness, query latency and indexing time, can provide more useful information for readers to choose which approach to use. (4) Can we apply this approach to multilingual retrieval tasks?
2. Although the proposed method show improvements over existing SoTA retrievers, the proposed approach is a bit incremental. I think the approach to augment the document is very similar to the previous work[3][4]. The authors should make a comparison with the previous work in Related Work to make the contribution more clear.
3. Some experimental details are missing (see Questions).

[1] Xueguang Ma, Liang Wang, Nan Yang, Furu Wei, and Jimmy Lin. 2023. Fine-tuning llama for multi-stage text retrieval.
[2] Liang Wang, Nan Yang, Xiaolong Huang, Linjun Yang, Rangan Majumder, and Furu Wei. 2023. Improving text embeddings with large language models.
[3] Rodrigo Nogueira and Jimmy Lin. 2019. From doc2query to docTTTTTquery.
[4] Yongqi Li, Nan Yang, Liang Wang, Furu Wei, and Wenjie Li. 2023. Multiview identifiers enhanced generative retrieval.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: All clear to me.
soundness: 3
strengths: The paper is clear and the results are good.
summary: This paper proposes a novel retrieval framework called to improve the performance of current retrieval models. The framework achieves this by enriching document embeddings with information from large language models (LLMs), such as synthetic queries and titles. The proposed document-level embedding approach integrates this augmented information with existing document chunks, and it's shown to be effective for both bi-encoder and late-interaction retrieval models. The authors demonstrate the framework's effectiveness by achieving state-of-the-art results on benchmark datasets LoTTE and BEIR.
weaknesses: It's a rather straightforward model. The novelty is simply calling some LLM to generate titles and queries for each doc, which the retriever will later use to retrieve the document. I think ICLR requires more model novelty.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. I am curious whether this method could be effective on larger document benchmarks, such as MSMARCO and FEVER within the BEIR benchmark. Could the authors provide results on these datasets?

2. Could the authors provide an analysis of the costs associated with using an LLM to enhance these tasks, such as the number of tokens or text checks required for generation?

3. It would be interesting to see the performance when using the original query to search against the pseudo queries, to quantify any potential leakage from the LLM (i.e., the ability to recover the query based on the document).
soundness: 2
strengths: 1. The motivation is solid, with a straightforward idea: using pseudo-queries to augment document representation, an important direction given the short context windows of current embedding models that limits their effectiveness.

2. The observed improvement appears substantial on the presented datasets.

3. The authors provide some valuable insights that could be helpful, such as the impact of titles and approaches to handle cases where titles are missing, which aligns well with real-world scenarios.
summary: This paper proposes an LLM-augmented retrieval framework that uses a large language model (LLM), specifically LLaMA-70B, to generate pseudo-queries for document retrieval. For each document, the framework includes synthetic queries generated by the LLM, existing or LLM-generated titles, and document chunks. These text fields are embedded and combined with the query embedding using inner product, followed by weighted sum and max pooling. Evaluations are conducted with Contriever, Dragon, and ColBERT on the LoTTE dataset and a subset of the BEIR benchmark, with considerable improvements.
weaknesses: 1. The cost of this solution appears too high, given that it relies on LLaMA-70B, and all documents require query generation. Additionally, I have concerns regarding potential data leakage in the LLM, as it may have previously encountered these query-document pairs, potentially generating similar queries based on the document. This issue could be addressed if the authors provided examples and analysis.

2. The inference process requires weighting strategies to combine the relevance between the query and different fields of the document. I have concerns about the scalability of this solution, as the presented dataset contains a relatively small number of documents (see Q1).

3. Some presentation issues: (i) The inconsistent capitalization of "Bi-encoder" is concerning, with both "Bi-encoder" and "bi-encoder" used interchangeably throughout the paper, which appears unprofessional. Additionally, on line 19, to my knowledge, ColBERTv2 falls under the category of bi-encoders but employs a more complex relevance scoring process. (ii) Certain findings in sections 3.1.2, 3.1.3, and others could be more effectively presented in an analysis section following the experiments section, providing a clearer flow and focused discussion on results.

Meta-review: This paper proposes an LLM-based method for document expansion in document retrieval. It enhances each document with synthetic queries, titles, and chunks generated by the LLM. The proposed method is training-free and model-agnostic because of the use of LLM. The method demonstrates substantial performance improvements across various retrievers (Contriever, DRAGON, ColBERTv2) and benchmarks (LoTTE, BEIR). Reviewers found the approach straightforward (z7Jy, ZRAX) and practical to implement (nf7F, z7Jy), with substantial empirical gains (fyxe, ZRAX).

The reviewers also raised several concerns:
1. The high inference cost of LLM-generated augmentations and the unfair comparison with previous methods in terms of inference efficiency (z7Jy, ZRAX).
2. The limited novelty, as the method is incremental with respect to existing document augmentation techniques (z7Jy, fyxe).
3. Concerns about the evaluation metrics (nf7F, z7Jy) and datasets (nf7F, ZRAX).

During the rebuttal, the authors provided additional results to address concern 3; however, concerns 1 and 2 were not sufficiently addressed. A rejection is recommended. We encourage the authors to improve the current version based on reviewers’ feedback for future publication.

Decision: Reject