Contextual Document Embeddings

Authors: John Xavier Morris, Alexander M Rush
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: text, embeddings, retrieval, context, contrastive

Abstract:
Dense document embeddings are central to neural retrieval. The dominant paradigm is to train and construct embeddings by running encoders directly on individual documents. In this work, we argue that these embeddings, while effective, are implicitly out-of-context for targeted use cases of retrieval, and that a contextualized document embedding should take into account both the document and neighboring documents in context - analogous to contextualized word embeddings. We propose two complementary methods for contextualized document embeddings: first, an alternative contrastive learning objective that explicitly incorporates the document neighbors into the intra-batch contextual loss; second, a new contextual architecture that explicitly encodes neighbor document information into the encoded representation. Results show that both methods achieve better performance than biencoders in several settings, with differences especially pronounced out-of-domain. We achieve state-of-the-art results on the MTEB benchmark with no hard negative mining, score distillation, dataset-specific instructions, intra-GPU example-sharing, or extremely large batch sizes.  Our method can be applied to improve performance on any contrastive learning dataset and any biencoder.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1730210682547

thread_type:
official_review

<thread_object>
creator: reviewer_ja4l

<content>
summary:
The paper explores contextual document embeddings by factoring in the information from surrounding documents in context in order to effectively handle retrieval tasks in out-of-domain scenarios that can be extremely challenging for existing dense neural encoders.

The authors propose two complementary methods to improve the dense neural encoders in such settings: a) contrastive learning objective to explicitly incorporate neighboring documents into the contextual batch loss b) architectural modification to explicitly encode information from neighboring documents.

The results on MTEB benchmark demonstrate SOTA performance without the need of sophisticated strategies like hard negative mining, intra-GPU example sharing or large batch sizes which are often employed by prior works. The proposed method is agnostic and can be applied to any contrastive learning dataset or dual encoder models.

strengths:
1. The proposed method incorporates corpus-level statistics through statistical retrieval method into dense neural biencoders to better contextualize the embeddings by conditioning on the neighboring documents. This effectively address the lack of contextual information not captured by existing neural biencoders and makes the contextualized document embeddings to be particularly beneficial in domain-specific retrieval tasks.

2. Most of the prior works on training dense neural biencoders often employ strategies like hard negatives or large batch sizes to achieve effective performance. It is interesting that the proposed method achieves competitive or superior performance without relying on such strategies.

3. The authors conduct extensive ablation studies to understand the impact of different components on the downstream performance, and these findings will be valuable for future research in this area.

4. The paper mostly list all the hyperparameters, datasets and models used in this work, ensuring the reproducibility of work.

weaknesses:
1. The proposed method relies on the availability of relevant document neighbors for contextualization in the second stage of proposed architecture. The authors provide some evidence (Figure 7) on how the model performance varies in scenarios with limited context by simulating using random documents. However, it is not clear how the model trained with the proposed method would perform in absence of context (using null tokens in place of contextual tokens). Does the method makes the model overly reliant on context?

2. Following on the previous point, the contextual tokens might need to be augmented with the null tokens and this may potentially lead to unnecessary computational overhead.

3. The effectiveness of the proposed method heavily depends on accurately clustering documents into fine-grained pseudo-domains using GTR model, as these clusters are later used for contextual training. I am uncertain if using a simpler embedding model would result in obtaining difficult batch configurations, as it might fail to account for nuances. Do you have any ablations with alternative embedding models?

4. The authors mention that they set the dataset input tokens to zero during the self-attention phase in the section 9.4 in the appendix ("position-agnostic embedding" para in the manuscript). However, self-attention plays a critical role in contextualizing and routing necessary information to the subsequent modules. I am not sure how contextualization will be achieved if this is true. I would appreciate clarification from the authors.

questions:
See weaknesses section for questions and concerns.

Few suggestions below:

1. Equation on page 4 in the manuscript uses symbol "m" in two different contexts which could make the notation a bit confusing. It might be better to use distinct symbols to avoid ambiguity.

2. The contextual architecture resembles to be somewhat similar to RETRO in terms of inspiration. It would be useful to add some discussion in the section 4.2.

3. Section 6 mentions that using single hard negative per query achieves better performance compared to without using any hard negatives. The authors should also report the results for both variants in Figure 5.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730722976352

thread_type:
official_review

<thread_object>
creator: reviewer_2rgz

<content>
summary:
The paper proposes two improvements for document embedding models: (1) using clustering to group similar documents together to make contrastive batches more difficult, and (2) making document embedding contextual through an alternate architecture that takes in not only the document itself, but also neighboring documents. Both improvements lead to better performance on document retrieval benchmarks.

strengths:
(1) Overall, the paper works on an important problem and provides two clean improvements that seem well-motivated and effective.

(2) The paper is well-written and clear: each section is well-organized, the methods are well-motivated, and each aspect of the method is explained clearly.

(3) The experiments are thorough and include useful ablations and analysis.

weaknesses:
(1) The contextual architecture seems more expensive than its non-contextual counterpart, so the experiments would be clearer if they also reported wall clock training time. For example, from what I understand, all four methods in table 1 were trained for the same number of steps; would the results change if they were trained for the same amount of time instead?

(2) The paper states that hyperparameters were chosen based on the small-scale experiments, but later also states "For our final model (anon-model-v1), we select the best of the supervised models, which comes from finetuning on the BGE dataset." What is the set of models you are selecting from, and what metric are you using to perform selection?

(3) The paper mentions that the "largest improvements [were] in ArguAna and SciFact, two of the smaller and more out-of-domain datasets," but I wasn't able to find a breakdown of the BEIR results for individual datasets in Table 1 or the supplementary material.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730752036324

thread_type:
official_review

<thread_object>
creator: reviewer_zmke

<content>
summary:
The paper proposed a batch sampling and architecture change to allow for contextual document embeddings. They achieve state-of-the-art results on the MTEB benchmark.

They embed a document in a two-stage approach. They first use clustering on embeddings generated by existing text embedding model to form a batch. Then they generate a single embedding by taking both the input document and the context documents in the same batch.

strengths:
- A new batch sampling technique.
  - State-of-the-art on the MTEB benchmark.

weaknesses:
Even though the paper is motivated by adapts the model to out-of-domain corpus, it's not evaluated on the domain-shift paradigm.

questions:
- The paper motivates the work by introducing context into document embeddings. However, the context is collected by clustering and therefore it is more similar to increase the difficulty of the in-batch negatives.
- I also don't fully get how the batch of documents can capture the domain shift at inference time.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730782491037

thread_type:
official_review

<thread_object>
creator: reviewer_2sba

<content>
summary:
The paper hypothesizes that dense embeddings lose the benefit of statistical embedding approaches in using corpus statistics, and proposes a two-step method to leverage the corpus information (e.g., documents in the same domain). The method first uses an extra embedding model to rearrange the training set by using their embeddings for clustering, such that similar documents can be in same batches for contrastive learning; in inference time, it further uses similar documents to get embeddings contextualized on these given documents.

strengths:
1. The method is well-motivated, although it is unclear whether without the method, current state-of-the-art embedding models are not able to provide good nuanced representations. 
2. The authors interpret the training of dense embedding methods and the method itself from a statistical perspective which is convincing.

weaknesses:
The authors claim that no hard negative mining is required to achieve state-of-the-art. However, the first step of the method (grouping similar documents) is essentially hard negative mining and is shown to be a key contribution to the performance. At the end, it is mentioned that an extra hard negative per query is used to achieve the best performance.

questions:
1. Can the authors provide more information about how the model is trained in the two-stage pass? Also, is the final model M1 or M2, or the both of them?
2. When evaluating on MTEB, how the in-context documents are selected for each task in the second-stage encoding?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732162676218

thread_type:
official_comment

<thread_object>
creator: area_chair_gse8

<content>
title:
Reminder: Please respond and update the score if necessary

comment:
Dear Reviewers,

Kindly ensure that you respond proactively to the authors' replies so we can foster a productive discussion. If necessary, please update your score accordingly. We greatly appreciate the time and effort you’ve dedicated to the review process, and your contributions are key to making this process run smoothly.

Thank you,

AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces a two-step method aimed at leveraging corpus information, particularly focusing on documents within the same domain. The first step involves using an auxiliary embedding model to rearrange the training data by clustering documents based on their embeddings. This allows for similar documents to be placed in the same batches during contrastive learning. In the inference phase, the method employs these related documents to generate embeddings that are contextualized by their similar counterparts.

This well-conceived approach significantly enhances dense neural biencoders by incorporating corpus-level statistics, thereby improving the contextualization of embeddings through a statistical retrieval strategy. It effectively overcomes the shortcomings of existing neural biencoders in capturing contextual nuances and is especially beneficial for domain-specific retrieval tasks. Impressively, the method achieves competitive or superior performance without resorting to traditional strategies such as using hard negatives or large batch sizes.

The paper is written and well-organized, clearly explaining its methodologies and findings. It is complemented by extensive experimental evaluations and ablation studies that provide valuable insights into the method's effectiveness. Additionally, the paper ensures reproducibility by meticulously detailing the hyperparameters, datasets, and models utilized in the research. Considering these strengths, I recommend the paper for acceptance.