Contextual Document Embeddings
Authors: John Xavier Morris
Alexander M Rush
Venue: ICLR 2025 Poster (2025)
Keywords: text
embeddings
retrieval
context
contrastive
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: A new way to do contrastive learning for text embedding models

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.

Published as a conference paper at ICLR 2025
CONTEXTUAL DOCUMENT EMBEDDINGS
John X. Morris Cornell University jxm3 cornellleduu
Alexander M. Rush
Comell University mosh@comeelledu
ABSTRACT
Dense document embeddings are central 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 document embedding should take into account both the document and neighboring documents in context analogous to contextualized word embeddings. We propose two com- plementary methods for contextualized document embeddings: first, an alternative contrastive learning objective that explicitly incorporates document neighbors into the intra-batch contextual loss; second, 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, -ataset- instructions, intra-GPU example-sharing or extremely large batch sizes. Our method can be applied to improve performance any contrastive learning dataset and any biencoder.
INTRODUCTION
Machine learning approaches text retrieval aim learn an embedded representation for indexing documents, Classically, this area was dominated by statistical approaches using sparse lexical matching methods based on n-gram frequencies such as BM25 (Robertson & Zaragoza 2009). Only recently have neural networks become competitive with state-of-the-art models on retrieval tasks (Karpukhin et al., 2020; Thakur et al. 202T). The primary neural method is dual encoder architecture that independently encodes both document and query to dense latent space for retrieval lookup. This document embedding space can improve upon a statistical model since learned end-to-end for retrieval.
However, there is at least one notable benefit of statistical approaches that is lost by neural models. Statistical models can easily incorporate prior corpus statistics such as inverse document frequency (IDF), into their representation. This prior term imparts context-dependence onto the model, since be updated based information specific retrieval in given domain test time. We contrast this contextual formulation with neural document encoders that are by definition function of the document itself. For example consider the following document:
The National Football League Draft is an annual event in which the National Football League (NFL) teams select eligible college football players...
Depending on the retrieval domain, e.g. Wikipedia search, sports articles, or televised events, IDF may weight terms such as NFL, draft or annual higher; a neural document embedding model would need to select global weighting for this document.
In this work, we explore contextualization of document embeddings produced by dense encoders. The goal is produce embeddings that better able to handle retrieval tasks in specific challenging contexts. We propose two complementary changes to document encoders: a contextual training procedure and architecture.
For contextual training. we aim to build a notion of neighboring documents directly into the contrastive learning process. We propose method that uses fast query-document clustering produce group

---

Published as a conference paper at ICLR 2025
First stage (Gather context)
Second stage (Embed context)
Text netr invell dataset
Contextual batching partitions dotaset of documents and queries into batches that share similar context.
Contextual embedding produces embedding for text that incorporates corpus-level information.
Figure I: Overview of our system for contextual document embeddings (CDE). Our model operates two stages: a first stage used characterize the dataset from samples, and second stage used to embed the final document.
of neighbors for each training batch. Each update for training is constructed purely from neighboring documents to ensure that embeddings can distinguish documents in the most challenging contexts.
For the architecture, propose encoder that injects information about the contextual docu- ments during embedding. The proposed architecture augments the standard BERT-style encoder with additional conditioning that provides aggregated document-level information about neighboring doc- uments. We call our method Contextual Document Embedding (CDE). Analogously to pre-computed corpus-level statisties, this method provides a manner for the embedding take into account the relative frequency of terms in context. The final output is still an embedding of the same size, so this does not require any additional storage or other changes to the retrieval process. When indexing, we utilize information from the corpus to produce document and query embeddings that are specific to particular domain.
Experiments compare these two extensions standard approaches for training document embeddings. Our results show that contextual contrastive training improves standard text embedding model training, and can be without other approaches such additional hard negatives. With the contextual encoder architecture, we see additional improvements over baseline model in all settings tested, with larger improvements in highly specific domains such small datasets of financial and medical documents. When trained at industry-scale, our model achieves state-of-the-aa results for small (<250M parameter) models the MTEB benchmark.
RELATED WORK
Text retrieval. Our work is related to the general field of text retrieval; propose specific improvements to the training of "biencoder" text embedding models such as DPR (Karpukhin et al.
2020), GTR (i al. 2021). Contriever (Izacard al. 2022). LaPraDoR [Xu al. 2022). Instructor
Suu et al. 2023), Nomic-Embet (Nussbaum et al.) 2024), E5 (Wang et al. 2024). and GTE (Li et al.
2023). We focus on the problem adapting these text retrieval models new corpora test time:
some prior work has noted this problem (Dai et al., 2022 Sciavolino 2021) and proposed solutions such as unsupervised span apan-sampling and training on test corpora (Gao & Callan 202T) and distillation on the test corpus from reranker (Sung et al. 2023). Late interaction methods (Khattab & Zaharia 2020, Santhanam et al., 2022) also offer one way to improve oft-of-domaan retrieval performance, but increase the runtime and complexity of search. We propose better sampling scheme that can be used train any biencoder or late interaction model as well as a training freee method for test tist adaptation.
2

---

Published as a conference paper at ICLR 2025
Contrastive learning. Much research has focused on the effect of hard negatives on the performance of contrastive learning methods Chen et al. (020)) Qu et al. (202 Robinson et al. (021); Wang
al. (023). Zhang & Stratos, 2021b observe that harder negatives provide a better approximation
of the overall cross-entropy loss, but do not consider batch-level optimizations for negative selection.
Hotstätter et al. (202T) cluster queries before training and show that this improves performance.
Sachidananda et al. (2023) also consider contrastive batch sampling as global optimization problem. text embedding models. (Ma et al. 2024) use clustering algorithm to partition a dataset into several sub-datasets, but train different model on each sub-dataset, Solatorio (0224) also use a pre-trained model address the problem of in-batch false negatives from randomly sampled batches. Our training algorithm aims to find the hardest possible high-quality batches to train text embedding models.
but do not apply their technique to state-of-the-an transformer-bant
Test-time adaptation. Our method can be compared to other solutions test-time adaptation, problem that has been well-studied across variety of domains (Jang et al. 2023). 3. In retrieval, one form of test-time adaptation reeevo-relevance feedback (PRF) (Rocchio 1971 et al. 2018) Wang et al. [021), where documents relevant to the query are used to construct a final, enhanced query representation. The query side of our model can be seen as a form of pseudo-relevance feedback; however, we train from scratch to support a more general form of PRF natively, on the document representation as well as the query.
Non- Non-parametric modeling. Our contextual document model can be seen a form of non- parametric modeling. This shows connections with the a large body of deep learning research such as the non-parametric transformer (NPT) (Kossen et al. (022) and the subfield of Neural Processes (Garnelo et al. 2018 Kim et al. 2019 Nguyen & Grover 2023). Semi-parametri models have been recently applied in NLP, specifically to the task of language modeling (Borgeaud et al.
2022 Khandelwal et al. (020). Instead of using retrieval model to build semi-parametric langauge model, we build a semi -armmeareetticc model specifically for the task of retrieval.
3 BACKGROUND
We can view text retrieval methods probabilistically as computing a distribution over potential documents based on a scalar score function ((d. q) matching documents and queries:
exp fd. q)
p(d (dd1)) EveD exp f(d'.
where D is a finite set of documents in a dataset. There is a wide variety of different definitions for f including full pairwise neural parameterizations (Nogueira Cho 2020). In this work, we focus on efficient retrieval methods using vector- methods, also known as embedding models.
Vector retrieval methods assume that f(d, q) can be factored into two embedding terms, ((d) (g))- the document and query embedding respectively. This factorization allows precomputation of the E D. This is critical for facilitating fast computation of arg max, p(d q) or top-k variants Douze et al. 2024).
document embeddings ((d) for
In statistical retrieval, á and vú are closed-form functions of the data, often representing unigram or bigram counts by the relative frequency of word types. Notably for this work, these methods can also utilize distributional properties of the test dataset as a prior, for example through inverse document frequency (IDF). We represent this integration of dataset-level information by writing the vector product (dd; D)- D).
In neural retrieval, we instead learn the representation training corpus of document and query pairs (these may be supervised, ie, goad-standard annotations, or unsupervised, i.e. noised synthetic examples), Dr ((d), g²)), (-), ,q") }}, with the aim of learning the embedding function & and vo.
dense vector. We assume access
Training can be motivated maximizing likelihood of the document corresponding to each query, i.e. Σn log p(d q²). Unfortunately, since retrieval datasets can have (D) exceed millions of documents, computing the normalizer in at each training step is not an option. Instead contrastive learning is used where the likelihood is replaced with biased approximation calculated from negative samples:
3

---

Published as a conference paper at ICLR 2025
exp exp f(d ,0)
Σ log p(d 11)) 20 Elog Loog
TIRX O
E Eaennlg)) exp f(d', qP)
where H. is a set of examples used to approximate the normalizing constant. In implementation, in addition to these hard negative examples, other examples from the mini-batch are also used to compute the normalizer since it requires additional compute for calculating d(d).
METHODS
In our work, we are interested in integrating contextual information into our embedding functions and db. The standard neural Ø is purely function of the document o(d) and does not take into account any notion of context. This contrasts with the statistical model <b D) and s D). Arguably this is not an issue if retrieval is completely in domain, é is capable of learning statistics such as IDF and average document length on the training set through gradient descent.
However, in many retrieval benchmarks, models are trained over single set of documents D and then tested many other domains ID that differs significantly from Dr In this setting, training on Dr alone may not be able to provide robust embeddings when used in contexts such as D.
4.1 CONTEXTUAL TRAINING WITH ADVERSARIAL CONTRASTIVE LEARNING
Returning to the example from the introduction, we assume that in general purpose training corpus Dr. the term NFL is a rare word appearing in relatively few documents and useful signal. However, if at test time D is corpus of sports articles, this word would be exceedingly common. Evaluation in this domain is, in statistical sense, adversarial to the original dataset. To handle this issue, lett-leamiig--ttlee objectives have shown be effective for training document embedders. In these approaches, instead of sampling documents-query pairs iid, the objective first sample a domain and then sample batch of examples. This ensures that the model mostly sees related training points in each domain.
We propose training objective that synthesizes large set fine-grained domains train the model on. Formally, our aim is to partition the training dataset Dr into groups (B3 ......B.)) such that each group represents self-similar pseudo-domain:
exp f(d, q) E(d exp [d',
max 0.10
log ()) = max
Σ
Σ (d.g)e
log
0.0.
(d.g)EB*
Computationally, the inner term can be implemented single batch and computed efficiently without the need for separate hard negatives (H)) Ideally we want groups that are as challenging as possible. Zhang & Stratos (021) show that increasing the partition term improves the contrastive approximation to the maximum likelihood the gradient. We can formalize this search for the most difficult configuration of batches optimization problem:
max Ba
Σ (d.))ER* (d' "')""
+ = max (8 Bn.
Σ
Σ (d-q)Es* (d' (d',g")E8*
(o(d) -0d') + o(d') ((g))
(8)
Solving this combinatorial objective exactly is intractable, but we can approximate solution using clustering. We first move from maximization to minimization by replacing the two dot products with L2 distance M ((d, q), (d', 42) = (lĐ(d) <>(g')l ||| [lo('d) rb(g)) (which is equivalent for normalized embeddings). We then note when that treated as symmetric pairs, this term obeys the triangle inequality for any other centroid pair
M((dg)),c) M(c, (')) 2 (((d.g),dd
4

---

Published as a conference paper at ICLR 2025
This implies that the following centroid- based objective represents upper-bound original objective:
Σ
Σ
M((d,)).c
(d.))EB*
(m
For known B, this search defines an asymmetric K-Means clustering problem. A solution can be efficiently computed using extremely fast Euclidean K-Means packages be treating each data point as two separate vectors ((d) @ (q) and v)()) @ o(d), where @ is concatenation.
Cluster Embeddings. Since clustering is performed before training, we do not have dense encoders é and s when constructing the groups. Borrowing methods from bard-negative mining (Robinson al. 2021) can replace the & and & with simpler embedding model when constructing groups. We experiment with sparse vector representation and with pretrained dense representations, settling on GTR et. 202T). popular and generic text embedding model.
Filtering False Negatives. Our method is especially sensitive to false negatives, they will be more likely to be included in given batch. Unfortunately, traditional retrieval datasets are not designed with this type of global objective in mind: false negatives common in most retrieval datasets and their prevalence increases with dataset scale. As one datapoint, Qu et al. (201 found that over 70% of top-retrieved passages in MS Marco are false negatives.
To avoid situation where each batch contains a large number of false negatives, we compute an equivalence class: S(g, d) = {d' D f(q, d') 2 fqg d) + +c for some surrogate scoring function and boundary term e. At training time, we alter the partition function for so that it no longer includes the elements of S(g, d), which are not definitively negative examples:
exp f(d
log p(d q) f(d. 9) + Ea
((g.a) exp f(d'
For simplicity, we again select f to be a simple pre-trained embedding model. This method likely over-prunes some potential true negatives found by the surrogate model; however we found to be critical to model accuracy.
Packing. Clusters found by our algorithm will be of varying sizes, and need to be packed into equal-sized batches. We apply a post-hoc procedure. We consider both random partitioning and grouping via greedy cluster-level traveling salesman, similar to Shh al. 2024). In both cases, we split large group into into smaller batches, and merge close small batches from within the same domain into evenly-sized batches. This has an added benefit of introducing randomness into the groups when training for multiple epochs. We leave it future work to analyze the full effects of different packing strategies such as expensive Balanced K-Means or heuristic approaches such as Equal K-Means (Gurraaggan et al. 2023)
4.2 CONTEXTUAL DOCUMENT EMBEDDING (CDE)
Contextualization can also be added directly to the architecture. Taking inspiration from sparse vector retrieval which uses corpus statistics to determine the form of the embedding, we modify the encoders to have access to the corpus itself, i.c. (dd; D) and (dd; D). This effectively augments the biencoder model give the ability to contextualize documents directly.
The main challenge is how to design a neural architecture that can take into account dataset contex- tualization. On one extreme, we could follow methods like BM25 and precompute a fixed set of corpus statistics that could be fed to the document encoder. On the other extreme, we could allow the encoder full access to the entire corpus, through some form of cross attention. The latter approach has been explored on a small scale in methods like neural processes (Gamelo al. 2018): however, it has been shown to be difficult to scale to larger datasets Borgeaud et al. (2022).
We opt for middleground that allows the model to learn corpus statistics, but is also relatively efficient to compute, shown in Figure [ Specifically, we note that document embeddings retain surprising amount of lexical information even after embedding (Morris et al. 2023). Therefore, we pre-embed a subset of the corpus, we believe we can still dynamically calculate key dataset information during encoding.
5

---

Published as a conference paper at ICLR 2025
We produce contextualized embeddings via two-stage process:
First stage:
Gather and embed context. Given context documents d' ...., ...,d's E D, we embed each using a unique embedding model and concatenate embeddings into sequence M1 )..d) M (d').
Second stage:
Embed document with additional context tokens. To compute ó for document we integrate contextual embedding sequence the input of second-stage embedding model M2:
d(d'; ( Mi
,(())..,,((((
Here Mj is the first-stage encoder model, My is second-stage encoder model, and E is the token embedding matrix of M2 applied to each token d'. In practice, we parameterize both M1 and M2 using traditional bidirectional transformers, so our model is comprised of two biencoder-like backbones called sequence.
There is a similar contextualized model for the query encoder vp which is also given document context (as we do not have query context at test time):
((q:)) =
E(gr))
(6)
We note several implementation properties of this architecture. During training, computing contextual embeddings for each contextual document for each training instance would naively increase training by a computational factor proportional to J, the number of documents in context. This time increase would not be tractable, since contrastive training can already take many days. We overcome this difficulty by sharing context d', ,..,, d' within a batch of documents; this allows us to compute representations just once per training step and reuse them between documents via computational graph.
When indexing new corpus DD first stage representations M (dd (d') can be computed once and cached, so M1 does not add parameters or runtime to the search process; only M2 needs to be stored and used at search time. Query representations can also use the cached context, which only require additional inputs to the encoder. (Our model does not include contextualized queries, only documents, typically do not assume access example queries test-time.)
Embedding without context. Individual corpora during training may not have sufficient or available context. To improve our model's generalization, we use sequence dropout, where we randomly replace context embeeddigs M, (d") with some null token vy according to some uniform probability
At test time, if no corpus information is available, our model can now function as non-contextual biencoder simply by replacing all sequence token inputs with tig.
Position-agnostic embedding. Since documents of DD are unordered, we remove all positionality from the neural encodings. When parameterizing 0 with traditional transformer, this can be achieved by omitting positional embeddings at the positions corresponding to D. In practice, we use transformers implementations dependent on FlashAttention with rotary positional embeddings at each self-attention layer. Full details of how we disable positionality are available in Section T0.4
Two-stage gradient caching. To improve training we employ a gradient-caching technique analo- gous to two-stage version of GradCache (Gao al. 2021). This technique allows to fit larger batches, longer sequences with more contextual samples without running out of memory. Essentially, compute first-stage and second-stage representations independently without gradients. We then use these frozen representations to compute the loss, and gradients with respect to the second-stage representations. We then re-run the second stage with gradients enabled and use the output gradients to backpropagate through the second -tage model, and obtain gradients for the first-stage representa- tions. We repeat this process for the first-stage representations. This allows us to tradeoff computation (running each transformer forward pass twice) for memory.
'Context reuse is only feasihle hecause documents within the same batch typically share large amount of context anyway, since they are clustered.

---

Published as a conference paper at ICLR 2025
5 EXPERIMENTAL SETUP
We consider range of retrieval experiments across different scales. To run experiments across suit- able number of settings, we devise a small setting: six-layer transformer, maximum sequence length of 64, and maximum number of 64 additional contextual tokens. In this scenario, truncated version of the BEIR benchmark (Thakur et al. (02T). Given the low cost of each experiment, are able to pre-train and fine-tune both biencoder and contextual models across a variety of batch sizes in {256, 512, 1024, 2048, 4096) and cluster sizes {64, 256, 1024, 4096, .... 2097152 4194304). As typical state-of-the-an text embedding models are trained in two phases, large weakly-supervised pre-training phase and a short supervised phase, we run all experiments for both phases.
evaluate
For the large setting, we use the best settings found via small experiments. We train a single model on sequences of length 512 with 512 contextual documents, evaluating on the full MTEB benchmark (Muennighoff et al., 2022). This includes tasks from retrieval as well as tasks like classification, clustering, and reranking. We compare our model's performance to the top small-size (under 250M parameters) models on MTEB (Nussbaum et al. 2024 Xiao et al.
2024 Solatorio 2024 et al.
2023
Training Data and Metries We train on the meta-datasets collected in Nussbaum al. 2024) for training text embedding models. This collection of datasets includes data from 24 datasets scraped from web sources such as Wikipedia and Reddit. Our unsupervised training phase trains on 200M weakly-suueerised datapoints scraped from large internet sources such as Reddit and Wikipedia. The supervised training phase includes 1.8M human-written query-document pairs intended for text retrieval, and is aggregated from popular retrieval datasets such as HotpotQA and MS MARCO YYngg) et al. 2018 Bajaj et al., 2018). For our full model, we also consider supervised training on the BGE meta-datasets (Xiao et al. (024).. We evaluate models using NDCG@ 10, conventional retrieval metric that enables comparison across many disparate datasets.
Implementation When partitioning our dataset into batches, we encode documents and queries using GTR (Ni al. 202T) and implement our clustering algorithm on top of FAISS (Douze et al. 2024.. We cluster per-domain for 100 steps and take the best clustering out of 3 attempts. We select NoRT as our pre-trained model backbone (Nussbaum et al. 2024), which has 137M parameters. We prepend all texts with short task-specific prefixes to identify each task; prefixes are
listed in Section 10.7 When pooling, we pool over text tokens only, never contextual tokens.
Training We initialize both Mj and M2 using the BERT-base model from Nussbaum et al. 2024 that includes flash attention. Weights are shared between Φ and vb. but notably not between M, and M2. For all experiments, we train with the Adam optimizer with 1000 steps of warmup to learning rate
10-5 and linearly decay to throughout training. For the filtering model we select nomi cemiied-medyy which was trained on the same datasets (Nussbaum et al. 2024). We train for three epochs unless otherwise specified. We set the maximum sequence length for all inputs to 512 and the number of contextual inputs to 512 (so the second-stage model has input length of 1024). When computing contrastive loss, we use a fixed temperature r = 0.02. When sequence dropout is set contextual input tokens null vectors with uniform
enabled in our contextual architecture, probability p 0.005. If the batch size exceeds the number of contextual documents, we randomly sample produce contextual inputs.
6
RESULTS
The main results are highlighted in Table I and Section the smaller setting, we observe that both adversarial contrastive learning and our contextual architecture improve performance compared to vanilla biencoder training. We observe the largest improvement when we combine these techniques.
Contextual batching After controlling for batch size and filtering for false negatives, we observe a strong correlation (visualized in Figure 13) between batch difficulty and downstream performance: reordering datapoints to make batches harder definitively enhances overall learning. This corrob-
orates prior findings (Xiong et al. 2020 Qu al. 202T) and theory (Zhang & Stratos 2021) that
7

---

Published as a conference paper at ICLR 2025
Contextual
Batch
Arch
Batch Size
Cluster Size
Train loss
Train acc.
NDCG@10
16384
0.39
90.3
59.9
512
512
0.81
77.7
61.7
16384
0.37
90.7
62.4
512
512
0.68
80.9
63.1
Table Performance our small models with and without the two improvements proposed this paper, measured on a shortened version of the BEIR benchmark. Numbers are NDCG@ 10.
more difficult batches in contrastive learning form better overall gradient approximation and learn more effectively.
Section TO.12 showcases model performance across batch and cluster sizes after both phases of training. We observe that although large batch and cluster size are useful when filtering is not enacted, when including filtering, smaller cluster (and harder) are clearly better, and large batches do not add much. When comparing filtered to non-filtered models (Figure 15p. filtering false negatives clearly improves performance.
Contextual architecture addition adversarial batching, we compare our contextual architec- ture to a biencoder across the datasets of BEIR in Table | (full results in appendix). Our architecture generally matches or improves performance on downstream datasets, with largest improvements in ArguAna and ClimateFEVER, two of the smaller and more out-of-domain datasets.
Full-seale training Figure 16 shows our models' performance when trained for multiple epochs on the supervised datasets, relative to the best similar-sized embedding model (dashed line). We find best performance when training for four epochs on the BGE meta-datasets. Although our best model does use a single hard negative per query, we are still able to achieve state-of-the-an performance without using any hard negatives.
For our final model (cde-small-vi) we select the best of the supervised models" which comes from finetuning the BGE dataset. On MTEB, cde-small-vl obtains flate-o-the-ee results compared to models of the same size. Although inspired by problems in the specific domain text retrieval, we observe that our approach improves embedding performance in all domains, including clustering, classification, and semantic similarity. We also evaluate a "random documents" bascline, where we sample random documents from the training dataset to simulate scenario where lack access to the test corpus. In this setting, we drop around 1.2 points on average across all tasks; the STS tasks in particular appear to produce representations that are close to context-agnostic.
Mathed
Arg
CQA
CREVER
DRP
REVER
FIQA
HPQA
MSMRC
NFC
NQ
QUORA
SCID
SCIF
TRFC
TOUCHE
Mean
Unsmpervised
Boseline
548
41.4
24.7
40.2
74,4
39.9
63.8
35.0
35.7
48.6
88.2
20.2
72.0
62.2
19.2
48.0
Contestual
54.
88.9
77.1
Supervised
59.4
79.9
Contextual
53.8
41.2
38.8
73.9
35.9
61.6
87.1
82.6
Table 2: Results (NDCG@ 10) on the retrieval setting of the MTEB benchmark.
ANALYSIS
How hard are our clusters? To analysis the relationship between cluster size in our clustering algorithm and the overall average difficulty of in-batch negatives, we measure the average difficulty of 1000 batches across a variety of batch and cluster sizes and plot the data Figure 17 We observe that larger batches bring casier non-negative examples, and decreasing cluster size clearly increases the average hardness of negative examples in a given cluster.
We select models by choosing the hest batch size and cluster size according to average NDCG BEIR.
8

---

Published as a conference paper at ICLR 2025
Clssfetn Cluster
PairCIs
Rerank
Retrvl
STS
Summ.
Mean
nomic-embed-vl
74.1
43,9
85,2
55.7
52.8
82.1
30.1
62.39
stella-baseenn-22
75.3
44.9
86.5
58.8
50.1
83.0
32.5
62.61
bge-base-en-vl.5
75.5
45.8
86.6
58.9
53.3
82.4
31.1
63.56
GIST-Embedding--V
76.0
46.2
86.3
59.4
52.3
83.5
30.9
63.71
gie-base-en-v1.5
77.2
46.8
85.3
57.7
54.1
82.0
31.2
64.11
cde ode-small-vi
[Random] [Contextual]
81.3
46.6
84.1
55.3
51.1
81.4
31.6
63.81
81.7
48.3
84.7
56.7
53.3
81.6
31.2
65.00
Table 3: Performance of models with 250M or fewer parameters on the MTEB benchmark for text embedding models. "Random" indicates the performance of our model with random training documents included instead of per-domain contextual documents.
Rand.
Cont.
Diff
ClimateFEVER
19.2
25.7
+6.5
ArguAna
66.5
72.0
+5.5
FiQA2018
40.7
45.4
+4.7
FEVER
81.6
86.1
+4.5
ChoxsivCuusteringP2P
41.5
44.9
+3.4
TRECCOVID
76.1
79.4
+3.3
McdrxivClusteringP2
34.6
37.7
+3.1
HotpotQA
63.8
66.7
+2.9
NQ
57.1
60.0
+2.9
SciDocsRR
80.1
82.9
+2.8
Table 4: Tasks with most improvement when contextual documents are provided. Scores are NDCG@ 10.
Which contextual documents help? To confirm that the CDE model is utilizing contextual information from we consider how different contextual documents help for given docuent d. Figure 18 measures results on CQADupstack, collection of Stack Exchange forum posts. We randomly sample inputs from D from domain (x-axis) and use them as input to the downstream task d marked along the y-axis. We mark square as red if its score comes within I point of NDCG of top score for its domain. Generally utilizing in-domain works best, but there are some crossover interactions.
CONCLUSION
We propose two improvements to traditional biencoder models for generating embeddings. The first improvement involves an algorithm for reordering training datapoints to make batches harder and improves vanilla training with minimal changes. Our second improvement involves new corpus-aware architecture for retrieval and allows us to train a state-of-the-an text embedding model.
ACKNOWLEDDGEMENTS
Thanks to Orion Weller, Vin Sachidananda, and Zach Nussbaum for valuable feedback on this research. We would also like to acknowledge Nomic and Hyperbolic for providing the compute necessary to conduct this research. This work was partially supported by Intelligence Advanced Research Projects Activity (IARPA), via the HIATUS Program #2022-2207220000 JM is supported by an NSF GFRP fellowship.
9

---

Published as a conference paper at ICLR 2025
REFERENCES
Payal Bajaj, Daniel Campos, Nick Craswell, Li Deng, Jianfeng Gao, Xiaodong Liu, Rangan Ma- jumder, Andrew McNamara, Bhaskar Mitra, Tri Nguyen, Mir Rosenberg, Xia Song, Alina Stoica, Saurabh Tiwary, and Tong Wang. Ms marco: A human generated machine reading comprehension dataset, 2018 URL https://axxiv.ormmmm wwtp://rrggabbb/bb/ao/oo
09268
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George van den Driessche, Jean- Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, and Laurent Sifre. Improving language models by retrieving from trillions of tokens, 2022.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations, 2020.
William Coster and David Kauchak. Simple English Wikipedia: A new text simplification task. In Dekang Lin, Yuji Matsumoto, and Rada Mihalcea (eds.), Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL ogy. org/P11-2117
https://aclaathor logy
Zhuyun Dai, Vincent Y. Zhao, Ji Ma, Yi Luan, Jianmo Ni, Jing Lu, Anton Bakalov, Kelvin Guu, Keith B. Hall, and Ming-Wei Chang. Promptagator: Few-shot dense retrieval from 8 examples, 2022.
Matthijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Pierre-Emmanuel Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The faiss library, 2024.
Anthony Fader, Luke Zettlemoyer, and Oren Etzioni. Open Question Answering Over Curated and Extracted Knowledge Bases, In KDD, 2014.
Angela Fan, Yacine Jernite, Ethan Perez, David Grangier, Jason Weston, and Michael Auli. ELI5: long form question answering. In Anna Korhonen, David R. Traum, and Lluís Mårquez (eds.), Proceedings of the 57th Conference of the Association for Computational Linguistics, 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pp. 3558 3567. Association for Computational Linguistics, 2019, doi: 10. 10.18663/11111-113446
| htttps://doi.orrg
18653/v1/p19 1346
Katja Filippova and Yasemin Altun. Overcoming the lack of parallel data in sentence compression. In David Yarowsky, Timothy Baldwin, Anna Korhonen, Karen Livescu, and Steven Bethard (eds.), Proceedings of the 2013 Conference on Empirical Methods Natural Language Processing, pp. 1481-1491, Seattle, Washington, USA, October 2013. Association for Computational Linguistics. URL https: http:s////aLaatttoooooooo//////
Wikimedia Foundation. Wikimedia downloads, 2024. URL htlttttt///uupppo.
org
Luyu Gao and Jamie Callan. Unsupervised corpus aware language model pre-training for dense passage retrieval, 2021.
Luyu Gao, Yunyi Zhang, Jiawei Han, and Jamie Callan. Scaling deep contrastive learning batch size under memory limited setup, 2021.
Marta Garnelo, Dan Rosenbaum, Chris Maddison, Tiago Ramalho, David Saxton, Murray Shana- han, Yee Whye Teh, Danilo J. Rezende, and S. M. Ali Eslami. Conditional neural processes, 2018.
Mansi Gupta, Nitish Kulkarni, Raghuveer Chanda, Anirudha Rayasam, and Zachary C Lipton. Amazonqa: review-based question answering task, 2019.
Suchin Gururangan, Margaret Li, Mike Lewis, Weijia Shi, Tim Althoff, Noah A. Smith, and Luke Zettlemoyer. Scaling expert language models with unsupervised domain discovery, 2023.
10

---

Published as a conference paper at ICLR 2025
Felix Hamborg, Norman Meuschke, Corinna Breitinger, and Bela Gipp. news-please: A generic news crawler and extractor. In Proceedings of the 15th International Symposium of Information Science, pp. 218-223, March 2017. doi: 10.528
Christopher Hidey and Kathy McKeown. Identifying causal relations using parallel Wikipedia articles. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 1424-1433, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: 10. 111111861/1111---// 1135. htt.gpss//aaaaattooooo.ooo////L/ 1135
Sebastian Hofstätter, Sheng-Chien Chich Lin, Jheng-Hong Yang. Jimmy Lin, and Allan Hanbury. Efficiently
teaching an effective dense retriever with balanced topic aware sampling, 2021. URI. https
V/arxiv. //axxi..rg/abbs// 2444. 06967
Hamel Husain, Ho-Hsiang Wu, Tiferet Gazit, Miltiadis Allamanis, and Mare Brockschmidt. CodeSearchNet challenge: Evaluating the state of semantic code search. arXiv preprint urXiv:1909.09453 09436, 2019.
Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning, 2022.
Minguk Jang, Sae-Young Chung, and Hye Won Chung. Test-time adaptation via self-training with nearest neighbor information, 2023.
Vladimir Karpukhin, Barlas Oğuz. Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen tau Yih. Dense passage retrieval for open open-domain question answering, 2020.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models, 2020.
Daniel Khashabi, Amos Ng. Tushar Khot, Ashish Sabharwal, Hannaneh Hajishirzi, and Chris
Callison-Bured -ursch. Gooaq: Open question answering with diverse answer types, 2021.
Omar Khattab and Matei Zaharia. Colbert: Efficient and effective passage search via contextualized late interaction over bert, 2020. UtLptttps://arxiv wrg. 2004.. 12832
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Gamelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes, 2019.
Jannik Kossen, Neil Band, Clare Lyle, Aidan N. Gomez, Tom Rainforth, and Yarin Gal. Self-attention between datapoints: Going beyond individual input-output pairs in deep learning, 2022.
Mahnaz Koupaee and William Yang Wang. Wikihow: large scale text summarization dataset, 2018.
Patrick Lewis, Yuxiang Wu, Linging Liu, Pasquale Minervini, Heinrich Küttler, Aleksandra Piktus, Pontus Stenetorp, and Sebastian Riedel. Paq: 65 million probably-aste aaseed questions and what you can do with them, 2021.
Canjia Li, Yingfei Sun, Ben He, Le Wang, Kai Hui, Andrew Yates, Le Sun, and Jungang Xu. NPRF: neural pseudo relevance feedback framework for ad-hoc information retrieval. Ellen Riloff, David Chiang, Julia Hockenmaier, and Jun ichi Tsujii (eds.), Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing. pp. 4482-4491, Brussels, Belgium, October-November 2018. Association for Computational Linguisties. doi: 10. 18553/11DD1--1444 URL
Zehan Li, Xin Zhang, Yanzhao Zhang, Dingkun Long, Pengjun Xie, and Meishan Zhang. Towards general text embeddings with multi-stage contrastive learning, 2023.
Kyle Lo, Lucy Lu Wang, Mark Neumann, Rodney Kinney, and Dan S. Weld. S2ore: The semantic scholar open research corpus, 2020.
Jiawei Ma. Po- Yao Huang, Saining Xie, Shang- Wen Li, Luke Zettlemoyer, Shih-Fu Chang, Wen- Tau Yih, and Hu Xu. Mode: Clip data experts via clustering, 2024.
11

---

Published as a conference paper at ICLR 2025
John X. Morris, Volodymyr Kuleshov, Vitaly Shmatikov, and Alexander M. Rush. Text embeddings reveal (almost) as much as text, 2023.
Niklas Muennighoff, Nouamane Tazi, Loie Magne, and Nils Reimers. Mteb: Massive text embedding 2022. doi: 10.4650/ARXX12210077166 URL
benchmark. arXiv preprint
https::/arxtt org/ httgg/////// 07316
Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning
sequence modeling, 2023.
Jianmo Ni. Jiacheng Li, and Julian McAuley. Justifying recommendations using distandl-labelldd reviews and fne-grained aspects. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.). Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 188-197. Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/1//D1-100111 1/D19-1018
URL htttp:://aelaattioooooooo/////// 1018
Jianmo Ni, Chen Qu, Jing Lu, Zhuyun Dai, Gustavo Hernández Ábrego. Ji Ma, Vincent Y. Zhao, Yi Luan, Keith Hall, Ming-We Beng-Wei Chang, and Yinfei Yang. Large dual encoders are generalizable retrievers, 2021.
Rodrigo Nogueira and Kyunghyun Cho. Passage re-ranking with bert, 2020.
Zach Nussbaum, John X. Morris, Brandon Duderstadt, and Andriy Mulyar. Nomic embed: Training a reproducible long context text embedder, 2024.
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 dense passage retrieval for open-domain question answering, 2021.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ Questions for Machine Comprehension of Text. arXiv e-prints, art. arXiv:1606.052500 2016
Nils Reimers, Elliot Choi, Amr Kayid, Alekhya Nandula, Manoj Govindassamy, and Ab-
dullah Elkady. Introducing embed v3, Nov 2023. URL. https::/ttt..cheeeee com/
introduc ingrroducing---beee -aneeeid-neembee-yeeeee
Stephen Robertson and Hugo Zaragoza. The Probabilistic Relevance Framework: BM25 and Beyond.
Now Publishers Inc., 2009.
Joshua Robinson, Ching- Yao Chuang, Suvrit Sra, and Stefanie Jegelka. Contrastive learning with hard negative samples, 2021.
J. J. Rocchio. Relevance feedback in information retrieval. 1971.
URL https://api
Demanttccscolar..oor wtt.ccho.......geooooooptlllaaaa
Vin Sachidananda, Ziyi Yang, and Chenguang Zhu. Global selection of contrastive batches via opti- mization on sample permutations. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 29542-29562. PMLR, 23-29 Jul 2023. URL httpt::t/ httpt:s:/ppoceeddnngeeee www.iinnng............. pww.sssss...r..e2222..2 sachi Rachidananda23. ht.l
Keshav Santhanam, Omar Khattab, Jon Saad-Faloon Christopher Potts, and Matei Zaharia. Colbertv2: Effective and efficient retrieval via lightweight late interaction, 2022. URI. httth:tt/t//att
brg/aba/ 2112.111111 01488
Christopher Sciavolino. Towards universal dense retrieval for open-domain question answering. 2021.
Abigail See, Peter J. Liu, and Christopher D. Manning. Get to the point: Summarization with pointer-generator networks. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1073-1083, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 18.186 1118653/v1/P1--/// 1099. URL https wab.orrg/amthoo.og//11111- httpchhooooyyy11111/1
www. ac web. org. antho
12

---

Published as a conference paper at ICLR 2025
Weijia Shi, Sewon Min, Maria Lomeli, Chunting Zhou, Margaret Li, Gergely Szilvasy, Rich James, Xi Victoria Lin, Noah A. Smith, Luke Zettlemoyer, Scott Yih, and Mike Lewis. In-context pretraining: Language modeling beyond document boundaries, 2024.
Aivin V. Solatorio. Gistembed: Guided in-sammll selection of training negatives for text embedding abs 2402. 16829
fine-tuning, 2024. URL
Hongjin Su, Weijia Shi, Jungo Kasai, Yizhong Wang, Yushi Hu, Mari Ostendorf, Wen tau Yih, Noah A. Smith, Luke Zettlemoyer, and Tao Yu. One embedder, any task: text embeddings, 2023.
Mujeen Sung, Jungsoo Park, Jaewoo Kang. Danqi Chen, and Jinhyuk Lee. Optimizing test-time query representations for dense retrieval, 2023.
Nandan Thakur, Nils Reimers, Andreas Rücklé, Abhishek Srivastava, and Iryna Gurevych. Beir: A heterogenous benchmark for zero-shot evaluation of information retrieval models, 2021.
Liang Wang, Nan Yang. Xiaolong Huang, Binxing Jiao, Linjun Yang. Daxin Jiang, Rangan Majumder, and Furu Wei. Simlm: Pre-training with representation bottleneck for dense passage retrieval, 2023.
Liang Wang, Nan Yang, Xiaolong Huang, Binxing Jiao, Linjun Yang, Daxin Jiang, Rangan Majumder, and Furu Wei. Text embeddings by weakly-supervine contrastive pre-training, 2024.
Xiao Wang, Craig Macdonald, Nicola Tonellotto, and Iadh Ounis. Pseudo-relevance feedback for multiple representation dense retrieval. In Proceedings of the 2021 ACM SIGIR International Conference on Theory of Information Retrieval, ICTIR '21. ACM, July 2021. doi: 10.1145/
3471158.347220 URL http: //cx: doi. org. 10. 1145 2444/34411555 3472250,
Shitao Xiao, Zheng Liu, Peitian Zhang, Niklas Muennighoff, Defu Lian, and Jian- Yun Nie. C-pack: Packaged resources to advance general chinese embedding, 2024. URL /ttpp::///rxxxv//
abs 220//339.000 07597
Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Amold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval, 2020. URL htttt:::/aaai...rrtttt
abs /207 /ttps///00000000 .00000000088
Canwen Xu, Daya Guo, Nan Duan, and Julian McAuley. Laprador: Unsupervised pretrained dense retriever for zero-shot text retrieval, 2022.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. Hotpotqa: dataset for diverse, explainable multi-hop question answering, 2018. hhttpe://aaxii.err//ab///BB 09600
Wenzheng Zhang and Karl Stratos. Understanding hard negatives in noise contrastive estimation, 2021.
13


Reviews:
[{"review_id": "Km35v86mxa", "reviewer": "Reviewer_2sbA", "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.", "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?\n2. When evaluating on MTEB, how the in-context documents are selected for each task in the second-stage encoding?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "9Av249R7VB", "reviewer": "Reviewer_zMkE", "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.\n\nThey 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.", "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.\n- I also don't fully get how the batch of documents can capture the domain shift at inference time.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- A new batch sampling technique.\n  - 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "vMNo3XtDqS", "reviewer": "Reviewer_2RGZ", "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.", "questions": "See above.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "(1) Overall, the paper works on an important problem and provides two clean improvements that seem well-motivated and effective.\n\n(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.\n\n(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?\n\n(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?\n\n(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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "TMSFeLbbVD", "reviewer": "Reviewer_Ja4L", "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.\n\nThe 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.\n\nThe 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.", "questions": "See weaknesses section for questions and concerns.\n\nFew suggestions below:\n\n1. 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.\n\n2. 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.\n\n3. 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "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.\n\n2. 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.\n\n3. 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.\n\n4. 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?\n\n2. 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.\n\n3. 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?\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)