Optimizing Test-Time Query Representations for Dense Retrieval
Year: 2023

Abstract:
Recent developments of dense retrieval rely on quality representations of queries and contexts from pre-trained query and context encoders. In this paper, we introduce TOUR (TestTime Optimization of Query Representations), which further optimizes instance-level query representations guided by signals from testtime retrieval results. We leverage a crossencoder re-ranker to provide fine-grained pseudo labels over retrieval results and iteratively optimize query representations with gradient descent. Our theoretical analysis reveals that TOUR can be viewed as a generalization of the classical Rocchio algorithm for pseudo relevance feedback, and we present two variants that leverage pseudo-labels as hard binary or soft continuous labels. We first apply TOUR on phrase retrieval with our proposed phrase re-ranker, and also evaluate its effectiveness on passage retrieval with an off-the-shelf reranker. TOUR greatly improves end-to-end open-domain question answering accuracy, as well as passage retrieval performance. TOUR also consistently improves direct re-ranking by up to 2.0% while running 1.3–2.4× faster with an efficient implementation.1

heading: 1 Introduction
text: Recent progress in pre-trained language models gave birth to dense retrieval, which typically learns dense representations of queries and contexts in a contrastive learning framework. By overcoming the term mismatch problem, dense retrieval has been shown to be more effective than sparse retrieval in open-domain question answering (QA) (Lee et al., 2019; Karpukhin et al., 2020; Lee et al., 2021a) and information retrieval (Khattab and Zaharia, 2020; Xiong et al., 2020).
Dense retrieval often uses a dual encoder architecture, which enables the pre-computation of con-
∗Work partly done while visiting Princeton University. 1Our code is available at https://github.com/
dmis-lab/TouR.
text representations while the query representations are directly computed from the trained encoder during inference. However, directly using trained query encoders often fails to retrieve the relevant context (Thakur et al., 2021; Sciavolino et al., 2021) as many test queries are unseen during training.
In this paper, we introduce TOUR, which further optimizes instance-level query representations at test time for dense retrieval. Specifically, we treat each test query as a single data point and iteratively optimize its representation. This resembles the query-side fine-tuning proposed for phrase retrieval (Lee et al., 2021a), which finetunes the query encoder over training queries in a new domain. Instead, we fine-tune query rep-
5731
resentations for each test query. Cross-encoders are known to exhibit better generalization ability in unseen distributions compared to dual encoders (Rosa et al., 2022). Accordingly, we leverage cross-encoder re-rankers (Nogueira and Cho, 2019; Fajcik et al., 2021) to provide pseudo relevance labels on intermediate retrieval results and then iteratively optimize query representations using gradient descent. For phrase retrieval, we also develop a cross-encoder phrase re-ranker, which has not been explored in previous studies.
We theoretically show that our framework can be viewed as a generalized version of the Rocchio algorithm for pseudo relevance feedback (PRF; Rocchio, 1971), which is commonly used in information retrieval to improve query representations. While most PRF techniques assume that the topranked results are equally pseudo-relevant, our method dynamically labels the top results and updates the query representations accordingly. We leverage our pseudo labels as either hard binary or soft continuous labels in two instantiations of our method, respectively. Lastly, to reduce computational overhead, we present an efficient implementation of TOUR, which significantly improves its runtime efficiency.
We apply TOUR on phrase (Lee et al., 2021a) and passage retrieval (Karpukhin et al., 2020) for open-domain QA. Experiments show that TOUR consistently improves performance in both tasks, even when the query distribution changes greatly. Specifically, TOUR improves the end-to-end opendomain QA accuracy by up to 10.7%, while also improving the accuracy of the top-20 passage retrieval by up to 8.3% compared to baseline retrievers. TOUR requires only a handful of top-k candidates to perform well, which enables TOUR to run up to 1.3–2.4× faster than the direct application of re-ranker with our efficient implementation while consistently improving the performance by up to 2.0%. The ablation study further shows the effectiveness of each component, highlighting the importance of fine-grained relevance signals.

heading: 2 Background

heading: 2.1 Dense Retrieval
text: Dense retrieval typically uses query and context encoders—Eq(·) and Ec(·)—for representing queries and contexts, respectively (Lee et al., 2019; Karpukhin et al., 2020). In this work, we focus on improving phrase or passage retrievers for open-
domain QA. The similarity of a query q and a context c is computed based on the inner product between their dense representations:
sim(q, c) = Eq(q)⊺Ec(c) = q⊺c. (1)
Dense retrievers often use the contrastive learning framework to train encoders Eq and Ec. After training the encoders, top-k results are retrieved from a set of contexts C: Cq1:k = [c1, . . . , ck] = top-kc∈Csim(q, c), (2)
where the top-k operator returns a sorted list of contexts by their similarity score sim(q, c) in descending order, i.e., sim(q, c1) ≥ · · · ≥ sim(q, ck). Dense retrievers aim to maximize the probability that a relevant context c∗ exists (or is highly ranked) in the top results.

heading: 2.2 Query-side Fine-tuning
text: After training the query and context encoders, the context representations {c | c ∈ C} are typically pre-computed for efficient retrieval while the query representations q are directly computed from the query encoder during inference. However, using the dense representations of queries as is often fails to retrieve relevant contexts, especially when the test query distribution is different from the one seen during training.
To mitigate the problem, Lee et al. (2021a) propose to fine-tune the query encoder on the retrieval results of training queries {q | q ∈ Qtrain} over the entire corpus C. For phrase retrieval (i.e., c denotes a phrase), they maximize the marginal likelihood of relevant phrases in the top-k results:
Lquery = − ∑
q∈Qtrain log
∑
c∈Cq1:k,c=c∗ Pk(c|q), (3)
where Pk(c|q) = exp(sim(q,c))∑k i=1 exp(sim(q,ci)) and c = c∗ checks whether each context matches the gold context c∗ or not. Note that c∗ is always given for training queries. The query-side fine-tuning significantly improves performance and provides a means of efficient transfer learning when there is a query distribution shift. In this work, compared to training on entire training queries as in Eq. (3), we treat each test query q ∈ Qtest as a single data point to train on and optimize instance-level query representations at test time. This is in contrast to distillation-based passage retrievers (Izacard and Grave, 2020; Ren et al., 2021), which fine-tune the parameters of the retrievers directly on all training data by leveraging signals from cross-encoders.

heading: 2.3 Pseudo Relevance Feedback
text: Pseudo relevance feedback (PRF) techniques in information retrieval (Rocchio, 1971; Lavrenko and Croft, 2001) share a similar motivation to ours in that they refine query representations for a single test query. Unlike using the true relevance feedback provided by users (Baumgärtner et al., 2022), PRF relies on heuristic or model-based relevance feedback, which can be easily automated. Although most previous work uses PRF for sparse retrieval (Croft et al., 2010; Zamani et al., 2018; Li et al., 2018; Mao et al., 2021), recent work has begun to apply PRF for dense retrieval (Yu et al., 2021; Wang et al., 2021; Li et al., 2021).
PRF aims to improve the quality of the retrieval by updating the initial query representation from the query encoder (i.e., Eq(q) = q0):
qt+1 ← g(qt, Cqt1:k), (4)
where g is an update function and qt denotes the query representation after t-th updates over q0.
The classical Rocchio algorithm for PRF (Rocchio, 1971) updates the query representation as:
g(qt, Cqt1:k) = αqt + β 1
|Cr| ∑
cr∈Cr cr − γ
1 |Cnr| ∑
cnr∈Cnr cnr,
(5)
where Cr and Cnr denote relevant and non-relevant sets of contexts, respectively. α, β, and γ determine the relative contribution of the current query representation qt, relevant context representations cr, and non-relevant context representations cnr, respectively, when updating to qt+1. A common practice is to choose top-k′ contexts as pseudorelevant among top-k (k′ < k), i.e., Cr = Cqt1:k′ :
g(qt, Cqt1:k) =
αqt + β 1
k′
k′∑
i=1
ci − γ 1 k − k′ k∑
i=k′+1
ci. (6)
In this work, we theoretically show that our testtime query optimization is a generalization of the Rocchio algorithm. While Eq. (6) treats the positive (or negative) contexts equally, we use crossencoder re-rankers (Nogueira and Cho, 2019) to provide fine-grained pseudo labels and optimize the query representations with gradient descent.

heading: 3 Methodology
text: In this section, we provide an overview of our method (§3.1) and its two instantiations (§3.2, §3.3). We also introduce a relevance labeler for phrase retreival (§3.4) and simple techniques to improve efficiency of TOUR (§3.5).

heading: 3.1 Optimizing Test-time Query Representations
text: We propose TOUR (Test-Time Optimization of Query Representations), which optimizes query representations at the instance level. In our setting, the query and context encoders are fixed after training, and we optimize the query representations solely based on their retrieval results. Figure 1 illustrates an overview of TOUR.
First, given a single test query q ∈ Qtest, we use a cross-encoder re-ranker ϕ(·) to provide a score of how relevant each of the top-k contexts c ∈ Cq1:k is with respect to a query:
s = ϕ(q, c), (7)
where ϕ(·) is often parameterized with a pretrained language model, which we detail in §3.4. Compared to simply setting top-k′ results as pseudo-positive in PRF, using cross-encoders enables more fine-grained judgments of relevance over the top results. In addition, it allows us to label results for test queries as well without access to the gold label c∗.

heading: 3.2 TOUR with Hard Labels : TOURhard
text: First, we explore using the scores from the crossencoder labeler ϕ and selecting a set of pseudopositive contexts Cqhard ⊂ C q 1:k defined as the smallest set such that:
Pk(c̃ =c ∗|q, ϕ) = exp(ϕ(q, c̃)/τ)∑k
i=1 exp(ϕ(q, ci)/τ)∑
c̃∈Cqhard
Pk(c̃ = c ∗|q, ϕ) ≥ p,
(8)
where τ is a temperature parameter and c̃ ∈ Cqhard denotes a pseudo-positive context selected by ϕ. Intuitively, we choose the smallest set of contexts as Cqhard whose marginal relevance with respect to a query under ϕ is larger than the threshold p. This is similar to Nucleus Sampling for stochastic decoding (Holtzman et al., 2020).
Then, TOUR optimizes the query representation with the gradient descent algorithm based on the
relevance judgment Cqhard made by ϕ:
Lhard(q, Cq1:k) = − log ∑
c̃∈Cqhard
Pk(c̃|q), (9)
where Pk(c̃|q) = exp(sim(q,c̃))∑k i=1 exp(sim(q,ci)) . Similar to the query-side fine-tuning in Eq. (3), we maximize the marginal likelihood of (pseudo) positive contexts Cqhard. We denote this version as TOURhard. Unlike query-side fine-tuning that updates the model parameters of Eq(·), we directly optimize the query representation q itself. TOURhard is also an instance-level optimization over a single test query q ∈ Qtest without access to the gold label c∗.
For optimization, we use gradient descent:
qt+1 ← qt − η ∂Lhard(qt, Cqt1:k)
∂qt , (10)
where η denotes the learning rate for gradient descent and the initial query representation is used as q0. Applying gradient descent over the test queries shares the motivation with dynamic evaluation for language modeling (Krause et al., 2019), but we treat each test query independently unlike the series of tokens for the evaluation corpus of language modeling. For each iteration, we perform a single step of gradient descent followed by another retrieval with qt+1 to update Cqt1:k into C qt+1 1:k .
Relation to the Rocchio algorithm Eq. (10) could be viewed as performing PRF by setting the update function g(qt, Cqt1:k) = qt − η ∂Lhard(qt,Cqt1:k) ∂qt . In fact, our update rule Eq. (10) is a generalized version of the Rocchio algorithm as shown below:
g(qt, Cqt1:k) = qt + η ∑
c̃
P (c̃|qt)(1− Pk(c̃|qt))c̃
− η ∑
c̃
[ P (c̃|qt) ∑
c∈Cqt1:k,c ̸=c̃ Pk(c|qt)c
] ,
(11)
where c̃ ∈ Cqthard and P (c̃|qt) = exp(sim(qt,c̃))∑ c̃′ exp(sim(qt,c̃′)) (proof in Appendix A). Although our update rule seems to fix α in Rocchio to 1, it can be dynamically changed by applying weight decay during gradient descent, which sets α = 1−ηλdecay multiplied by qt. Then, the equality between Eq. (6) and Eq. (11) holds when Cqthard = C qt 1:k′ with Pk(c|qt) being equal for all c ∈ Cqt1:k, namely Pk(c|qt) = 1/k. This reflects that the Rocchio algorithm treats all top-k′ results equally (i.e., P (c̃|qt) = 1/k′). Then, β = γ = η k−k
′ k holds (Appendix C).
In practice, Cqthard would be different from C qt 1:k′ if some re-ranking happens by ϕ. Also, each pseudopositive context vector c̃ in the second term of the RHS of Eq. (11) has a different weight. The contribution of c̃ is maximized when it has a larger probability mass P (c̃|qt) among the pseudo-positive contexts, but a smaller probability mass Pk(c̃|qt) among the top-k contexts; this is desirable since we want to update qt a lot when the initial ranking of pseudo-positive context in top-k is low. For instance, if there is a single pseudo-positive context c̃ (i.e., P (c̃|qt)) = 1) ranked at the bottom of top-k with a large margin with top-1 (i.e., Pk(c̃|qt) = 0), then P (c̃|qt)(1− Pk(c̃|qt)) = 1 is maximized.

heading: 3.3 TOUR with Soft Labels : TOURsoft
text: From Eq. (11), we observe that it uses pseudopositive contexts Cqthard sampled by the crossencoder labeler ϕ, but the contribution of c̃ (the second term in RHS) does not directly depend on the scores from ϕ. The scores are only used to make a hard decision in pseudo-positive contexts. Another version of TOUR uses the normalized scores of a cross-encoder over the retrieved results as soft labels. We can simply change the maximum marginal likelihood objective in Eq. (9) to reflect the scores from ϕ in g. Specifically, we change Eq. (9) to minimize Kullback-Leibler (KL) divergence loss as follows:
Lsoft(qt, Cqt1:k) =
− k∑
i=1
P (ci|qt, ϕ) log Pk(ci|qt) P (ci|qt, ϕ) , (12)
where P (ci|qt, ϕ) = P (ci = c∗|qt, ϕ) defined in Eq. (8). We call this version TOURsoft. The update rule g for TOURsoft changes as follows:
g(qt, Cqt1:k) = qt + η k∑
i=1
P (ci|qt, ϕ)ci − η k∑
i=1
Pk(ci|qt)ci.
(13)
Eq. (13) shows that qt+1 reflects ci weightaveraged by the cross-encoder (i.e., P (ci|qt, ϕ)) while removing ci weight-averaged by the current retrieval result (i.e., Pk(ci|qt)) (proof in Appendix B).

heading: 3.4 Relevance Labeler for Phrase Retrieval
text: In the previous section, we used a cross-encoder reranker ϕ to provide a relevance score si over
a pair of a query q and a context c. While it is possible to use an off-the-shelf re-ranker (Fajcik et al., 2021) for passage retrieval, no prior work has introduced a re-ranker for phrase retrieval (Lee et al., 2021b). In this section, we introduce a simple and accurate phrase re-ranker for TOUR.
Inputs for re-rankers For phrase retrieval, sentences containing each retrieved phrase are considered as contexts, following Lee et al. (2021b). For each context, we also prepend the title of its document and use it as our context for re-rankers. To train our re-rankers, we first construct a training set from the retrieved contexts of the phrase retriever given a set of training queriesQtrain. Specifically, from the top retrieved contexts C1:k for every q ∈ Qtrain, we sample one positive context c+q and one negative context c−q . In open domain QA, it is assumed that a context that contains a correct answer to each q is relevant (positive). Our re-ranker is trained on a dataset Dtrain = {(q, c+q , c−q )|q ∈ Qtrain}.
Architecture We use the RoBERTa-large model (Liu et al., 2019) as the base model for our re-ranker. Given a pre-trained LM M, the cross-encoder re-ranker ϕ outputs a score of a context being relevant:
s = ϕ(q, c) = w⊤M(q ⊕ c)[CLS] (14)
where {M,w} are the trainable parameters and ⊕ denotes a concatenation of q and c using a [SEP] token. Since phrase retrievers return both phrases and their contexts, we use special tokens [S] and [E] to mark the retrieved phrases within the contexts.
Re-rankers are trained to maximize the probability of a positive context c+q for every (q, c + q , c − q ) ∈ Dtrain. We use the binary cross-entropy loss defined over the probability P+ = exp(h
+) exp(h+)+exp(h−) where
h+ = ϕ(q, c+q ) and h − = ϕ(q, c−q ). We pre-train ϕ on reading comprehension datasets (Rajpurkar et al., 2016; Joshi et al., 2017; Kwiatkowski et al., 2019), which helped improve the quality of ϕ. For the ablation study of our phrase re-rankers, see Appendix D for details.
Score aggregation After running TOUR, aggregating the reranking scores with the retreival scores provides consistent improvement. Specifically, we linearly interpolate the similarity score sim(q, ci) with the re-ranking score si and use this to obtain the final results: λsi + (1− λ)sim(q, ci).

heading: 3.5 Efficient Implementation of TOUR
text: TOUR aims to improve the recall of gold candidates by iteratively searching with updated query representations. However, it has high computational complexity, since it needs to label top-k retrieval results with a cross-encoder and perform additional retrieval. To minimize the additional time complexity, we perform up to t = 3 iterations with early stopping conditions. Specifically, at every iteration of TOURhard, we stop when the top-1 retrieval result is pseudo-positive, i.e., c1 ∈ Cqthard. When using TOURsoft, we stop iterating when the top-1 retrieval result has the highest relevance score. Additionally, we cache ϕ(q, ci) for each query to skip redundant computation.

heading: 4 Experiments
text: We test TOUR on multiple open-domain QA datasets. Specifically, we evaluate its performance on phrase retrieval and passage retrieval.

heading: 4.1 Datasets
text: We mainly use six open-domain QA datasets: Natural Questions (Kwiatkowski et al., 2019), TriviaQA (Joshi et al., 2017), WebQuestions (Berant et al., 2013), CuratedTrec (Baudiš and Šedivỳ, 2015), SQuAD (Rajpurkar et al., 2016), and EntityQuestions (Sciavolino et al., 2021). Following previous works, Entity Questions is only used for testing. See statistics in Appendix E.

heading: 4.2 Open-domain Question Answering
text: For end-to-end open-domain QA, we use phrase retrieval (Seo et al., 2019; Lee et al., 2021a) for TOUR, which directly retrieves phrases from the entire Wikipedia using a phrase index. Since a singlestage retrieval is the only component in phrase retrieval, it is easy to show how its open-domain QA performance can be directly improved with TOUR. We use DensePhrases (Lee et al., 2021a) for our base phrase retrieval model and train a crossencoder labeler as described in §3.4. We report exact match (EM) for end-to-end open-domain QA. We use k = {10, 40} for our phrase re-ranker and k = {10, 20} for TOUR on open-domain QA while k = 10 is used for both whenever it is omitted. For the implementation details of TOUR, see Appendix D.
Baselines Many open-domain QA models take the retriever-reader approach (Chen et al., 2017; Lee et al., 2019; Izacard and Grave, 2021; Singh
methods, we denote the improvement from their base retrievers in ↑x.x . Best performance is denoted in bold. *: We could not measure the exact latency of DPRmulti (large) as its checkpoint has not been released.
et al., 2021). As our baselines, we report extractive open-domain QA models, which is a fair comparison with retriever-only (+ re-ranker) models whose answers are always extractive. For re-ranking baselines of retriever-reader models, we report ReConsider (Iyer et al., 2021), which re-ranks the outputs of DPR + BERT. For a PRF baseline, GAR (Mao et al., 2021), which uses context generation models for augmenting queries in BM25, is reported.
Results Table 1 shows the results on the five open-domain QA datasets in the in-domain evaluation setting where all models use the training sets of each dataset they are evaluated on. First, we observe that using our phrase re-ranker largely improves the performance of DensePhrasesmulti. Compared to adding a re-ranker on the retriever-reader model (DPRmulti + Re-ranker by Iyer et al., 2021), our phrase re-ranking approach performs 5× faster with a larger top-k due to the efficient retriever-only method. Furthermore, the performance gain is significantly larger possibly due to the high top-k accuracy of phrase retrievers. Unlike using the Rocchio algorithm, using TOURhard or TOURsoft greatly improves the performance of the base retriever. Compared to our phrase re-rankerk=40, TOURhard,k=20 runs 1.3× faster as well as outperforming it by up to 2.0%. Even TOURhard,k=10 often outperforms re-rankerk=40 with 2.4× faster inference. For this task, TOURhard and TOURsoft work similarly with
exceptions on NQ and TREC.
Latency vs. performance Figure 2 compares the query latency and performance of TOUR and other baselines on the NQ development set. We vary the top-k value from 10 to 50 by 10 (left to right) to visualize the trade-off between latency and performance. The result shows that TOUR with only top-10 is better and faster than the re-ranker with the best top-k. Specifically, TOURhard,k=10 outperforms re-rankerk=40 by 1.0% while being 2.5× faster. This shows that TOUR requires a less number of retrieval results to perform well, compared to a re-ranker model that often requires a larger k.
from DensePhrasesmulti or DPRmulti in ↑x.x . †: unseen query distribution.
Query distribution shift In Table 2, we show open-domain QA results under query distribution shift from the training distribution. Compared to DensePhrasesmulti in Table 1, which was trained on all five open-domain QA datasets, we observe huge performance drops on unseen query distributions when using DPRNQ and DensePhrasesNQ. DPRNQ seems to suffer more (e.g., 0.1 on SQuAD) since both of its retriever and reader were trained on NQ, which exacerbates the problem when combined.
On the other hand, using TOUR largely improves the performance of DensePhrasesNQ on many unseen query distributions even though all of its component were still trained on NQ. Specifically, TOURhard,k=20 gives 6.5% improvement on average across different query distributions, which easily outperforms our phrase re-rankerk=40. Interestingly, TOURhard consistently performs better than TOURsoft in this setting, which requires more investigation in the future.

heading: 4.3 Passage Retrieval
text: We test TOUR on the passage retrieval task for open-domain QA. We use DPR as a passage retriever and DensePhrases as a phrase-based passage retriever (Lee et al., 2021b). In this experiment, we use an off-the-shelf passage re-ranker (Fajcik et al., 2021) to show how existing re-rankers can serve as a pseudo labeler for TOUR. We report the top-k retrieval accuracy, which is 1 when the answers exist in top-k retrieval results. For passage retrieval, we use k = 100 for both the re-ranker and TOUR due to the limited resource budget.
Results Table 3 shows the results of passage retrieval for open-domain QA. We find that using TOUR consistently improves the passage retrieval accuracy. Under the query distribution shift similar to Table 2, DPRmulti + TOURsoft improves the original DPR by 8.3% and advances the off-the-shelf re-ranker by 1.8% on EntityQuestions (Acc@20). Notably, Acc@100 always improves with TOUR,
which is not possible for re-rankers since they do not update the top retrieval results. Unlike the phrase retrieval task, we observe that TOURsoft is a slightly better option than TOURhard on this task.

heading: 5 Analysis

heading: 5.1 Train-Test Overlap Analysis
text: Open-domain QA datasets often contain semantically overlapping queries and answers between training and test sets (Lewis et al., 2021), which overestimates the generalizability of QA models. Hence, we test our models on train-test overlap
2While the passage retrieval accuracy is mostly reported for Entity Questions (Ram et al., 2022; Lewis et al., 2022), we also report EM for open-domain QA.
splits provided by Lewis et al. (2021). Table 4 shows that TOUR consistently improves the performance of test queries that do not overlap with training data (i.e., None). Notably, on WebQuestions, while the performance on the none overlap split has been improved by 1.5% from the re-ranker, the performance on query overlap is worse than the re-ranker since unnecessary exploration is often performed on overlapping queries. Our finding on the effectiveness of query optimization is similar to that of Mao et al. (2021), while our approach often improves performance on query overlap cases.

heading: 5.2 Ablation Study
text: Table 5 shows an ablation study of TOURhard on end-to-end open-domain QA. We observe that using fine-grained relevance signals generated by our phrase re-ranker (i.e., Cqhard) is significantly more effective than simply choosing top-k′ as relevance
signals (i.e., C1:k′). Using SGD or aggregating the final scores between the retriever and the re-ranker gives additional improvement.
Figure 3 shows the effect of multiple iterations in TOURhard compared to the Rocchio algorithm. While PRFRocchio with t = 1 achieves slightly better performance than DensePhrases, it shows a diminishing gain with a larger number of iterations. In contrast, the performance of TOURhard benefits from multiple iterations until t = 3. Removing the score aggregation between the retriever and the re-ranker (i.e., λ = 0) causes a performance drop, but it quickly recovers with a larger t.
Efficient implementation Simple techniques introduced in §3.5 such as early stopping and caching significantly reduce the run-time of TOUR. Figure 4 summarizes the effect of optimization techniques to improve efficiency of TOUR. Without each technique, the latency increases linearly with the number of iterations. By adding the caching mechanism for ϕ and the stop condition of c1 ∈ Cqthard, the latency is greatly reduced. Prediction sample Figure 5 shows a sample prediction of TOUR. We use DensePhrasesmulti + TOURhard with k = 10, from which the top-5 results are shown. While the initial result at t = 0 failed to retrieve correct answers in the top-10, the
next round of TOURhard gives new results including the correct answer, which were not retrieved before. As the iteration continues, the correct answer starts to appear in the top retrieval results, and becomes the top-1 at t = 3.

heading: 6 Conclusion
text: In this paper, we propose TOUR, which iteratively optimizes test query representations for dense retrieval. Specifically, we optimize instance-level query representations at test time using the gradientbased optimization method over the top retrieval results. We use cross-encoder re-rankers to provide pseudo labels where our simple re-ranker or off-theshelf re-rankers can be used. We theoretically show that gradient-based optimization provides a generalized version of the Rocchio algorithm for pseudo relevance feedback, which leads us to develop different variants of TOUR. Experiments show that our test-time query optimization largely improves the retrieval accuracy on multiple open-domain QA datasets in various settings while being more efficient than traditional re-ranking methods.
Limitations
In this paper, we focus on the end-to-end accuracy and passage retrieval accuracy for opendomain QA. We have also experimented on the BEIR benchmark (Thakur et al., 2021) to evaluate our method in the zero-shot document retrieval task. Overall, we obtained 48.1% macro-averaged NDCG@10 compared to 47.8% by the re-ranking method. For some tasks, TOUR obtains significant improvements with a pre-trained document retriever (Hofstätter et al., 2021). For example, TOUR improves the baseline retriever by 11.6% and 23.8% NDCG@10 on BioASQ and TRECCOVID, respectively, while also outperforming the re-ranker by 2.1% and 2.4% NDCG@10. We plan to better understand why TOUR performs better specifically on these tasks and further improve it.
TOUR also requires a set of validation examples for hyperparameter selection. While we only used in-domain validation examples for TOUR, which were also adopted when training re-rankers, we observed some performance variances depending on the hyperparameters. We hope to tackle this issue with better optimization in the future.

heading: Acknowledgements
text: We thank Zexuan Zhong, Mengzhou Xia, Howard Yen, and the anonymous reviewers for their helpful feedback. This work was supported in part by the ICT Creative Consilience program (IITP-20232020-0-01819) supervised by the IITP (Institute for Information & communications Technology Planning & Evaluation), National Research Foundation of Korea (NRF-2023R1A2C3004176), and the Hyundai Motor Chung Mong-Koo Foundation.

heading: A Derivation of the Gradient for
text: TOURhard
Proof. We compute the gradient of Lhard(qt, Cqt1:k) in Eq. (10) with respect to the query representation qt. Denoting ∑ c̃ Pk(c̃|qt) as Z, the gradient is:
∂Lhard(qt, Cqt1:k) ∂qt = ∂Lhard(qt, Cqt1:k) ∂Z ∂Z ∂qt = − 1 Z ∑
c̃
∂Pk(c̃|qt) ∂qt
= − 1 Z
∑
c̃
k∑
i=1
∂Pk(c̃|qt) ∂q⊤t ci ∂q⊤t ci ∂qt
= − 1 Z
∑
c̃
k∑
i=1
(δ[ci = c̃]− Pk(ci|qt))Pk(c̃|qt)ci
= − ∑
c̃
[ P (c̃|qt) k∑
i=1
(δ[ci = c̃]− Pk(ci|qt))ci ]
= − ∑
c̃
P (c̃|qt) [ (1− Pk(c̃|qt))c̃− ∑
c∈Cqt 1:k ,c ̸=c̃
Pk(c|qt)c ]
= − ∑
c̃
P (c̃|qt)(1− Pk(c̃|qt))c̃
+ ∑
c̃
[ P (c̃|qt) ∑
c∈Cqt 1:k ,c ̸=c̃
Pk(c|qt)c ]
Then, we have:
g(qt, Cqt1:k) = qt − η ∂Lhard(qt, Cqt1:k)
∂qt
= qt + η ∑
c̃
P (c̃|qt)(1− Pk(c̃|qt))c̃
− η ∑
c̃
[ P (c̃|qt) ∑
c∈Cqt1:k,c ̸=c̃ Pk(c|qt)c
] .

heading: B Derivation of the Gradient for TOURsoft
text: Proof. We compute the gradient of Lsoft(qt, Cqt1:k) in Eq. (12) with respect to qt. Denoting P (ci = c∗|qt, ϕ) as Pi, we expand the loss term as:
Lsoft(qt, Cqt1:k) = − k∑
i=1
Pi log Pk(ci|qt)
Pi
= − k∑
i=1
Pi(q ⊤ t ci − log
k∑
j=1
exp(q⊤t cj)− logPi).
Then, the gradient is:
∂Lsoft(qt, Cqt1:k) ∂qt
= − k∑
i=1
Pi ∂ ∂qt (q⊤t ci − log
k∑
j=1
exp(q⊤t cj)− logPi)
= − k∑
i=1 Pi(ci − 1∑k j=1 exp(q ⊤ t cj)
k∑
j=1
cj exp(q ⊤ t cj))
= − k∑
i=1
Pi(ci − k∑
j=1 cj exp(q⊤t cj)∑k l=1 exp(q ⊤ t cl) )
= − k∑
i=1
Pi(ci − k∑
j=1
Pk(cj |qt)cj)
= − k∑
i=1
Pici + k∑
i=1
Pk(ci|qt)ci.
Putting it all together:
g(qt, Cqt1:k) = qt − η ∂Lsoft(qt, Cqt1:k)
∂qt
= qt + η k∑
i=1
P (ci|qt, ϕ)ci − η k∑
i=1
Pk(ci|qt)ci.

heading: C Relation to the Rocchio Algorithm
text: Proof. We derive Eq. (6) from Eq. (11).
g(qt, Cqt1:k) = qt + η ∑
c̃
P (c̃|qt)(1− Pk(c̃|qt))c̃
− η ∑
c̃
[ P (c̃|qt) ∑
c∈Cqt1:k,c ̸=c̃ Pk(c|qt)c
]
= qt + η k′∑
i=1
1 k′ (1− 1 k )ci
− η k′∑
i=1
[ 1 k′ k∑
j=1,i ̸=j
1 k cj ]
= qt + η k − 1 k′k
k′∑
i=1
ci
− η 1 k′k
k′∑
i=1
[ k′∑ j=1,i ̸=j cj + k∑ j=k′+1 cj ]
= qt + η k − 1 k′k
k′∑
i=1
ci
− η 1 k′k
[ (k′ − 1)
k′∑
i=1
ci + k ′
k∑
i=k′+1
ci ]
= qt + η k − k′ k′k
k′∑
i=1
ci − η 1
k
k∑
j=k′+1
cj .
(15)
Then, the equality holds when α = 1, β = η k−k
′ k , and γ = η k−k′ k .
D Implementation Details
Phrase re-ranker To train a cross-encoder reranker for phrase retrieval (§3.4), we first annotate the top 100 retrieved results from DensePhrases. We use three sentences as our context, one that contains a retrieved phrase and the other two that surround it. This leads to faster inference than using the whole paragraph as input while preserving the performance. During the 20 epochs of training, we sample positive and negative contexts for every epoch while selecting the best re-reanker based on the validation accuracy of the re-ranker. We modified the code provided by the Transformers library3 (Wolf et al., 2020) and used the same hyperparameters as specified in their documentation except for the number of training epochs. The ablation study in Table 6 shows that we can achieve stronger performance by prepending titles to inputs, using larger language models, using three sentences as our context, and pre-training over reading comprehension datasets. Using entire paragraphs as input contexts only slightly increases performance compared to using three sentences, but it doubles the query latencies of re-ranking.
Dense retriever We modified the official code of DensePhrases4 (Lee et al., 2021a) and DPR5 (Karpukhin et al., 2020) to implement TOUR on dense retrievers. While pre-trained models and indexes of DensePhrasesmulti and DPRNQ are publicly available, the indexes of
3https://github.com/huggingface/ transformers/blob/v4.13.0/examples/pytorch/ text-classification/run_glue.py
4https://github.com/princeton-nlp/ DensePhrases
5https://github.com/facebookresearch/DPR
DensePhrasesNQ and DPRmulti have not been released as of May 25th, 2022. When necessary, we reimplemented them to experiment with opendomain QA and passage retrieval in the query distribution shift setting.
Hyperparamter When running TOUR , we use gradient descent with momentum set to 0.99 and use weight decay λdecay = 0.01. We also perform a linear learning rate scheduling per iteration. Both the threshold p and temperature τ for pseudo labels are set to 0.5. Table 7 lists the hyperparameters that are used differently for each task. All hyperparameters of TOUR were tuned using the in-domain development set.

heading: E Data Statistics
text: Table 8 shows the statistics of the datasets used for end-to-end open-domain QA and passage retrieval tasks. For EntityQuestions, we only use its test set for the query distribution shift evaluation.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:

heading: 3 A1. Did you describe the limitations of your work?
text: Yes, please see the limitations section.
A2. Did you discuss any potential risks of your work? Not applicable. Left blank.

heading: 3 A3. Do the abstract and introduction summarize the paper’s main claims?
text: Yes, please see the abstract and introduction sections.

heading: 3 A4. Have you used AI writing assistants when working on this paper?
text: Grammarly for grammar checking. The contents are original work from the human authors.
B 3 Did you use or create scientific artifacts? Yes, please see section 4.1 Datasets.

heading: 3 B1. Did you cite the creators of artifacts you used?
text: Yes, please see section 4.1 Datasets.

heading: 3 B2. Did you discuss the license or terms for use and / or distribution of any artifacts?
text: Yes, please see section 4.1 Datasets.

heading: 3 B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided
text: that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Yes, please see section 4.1 Datasets.
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Not applicable. Left blank.

heading: 3 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and
text: linguistic phenomena, demographic groups represented, etc.? Yes, please see section 4.1 Datasets.
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Yes, please see Appendix E Data Statistics.
C 3 Did you run computational experiments? Yes, please see section 4 Experiments.

heading: 3 C1. Did you report the number of parameters in the models used, the total computational budget
text: (e.g., GPU hours), and computing infrastructure used? Yes, please see section 4 Experiments.
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.

heading: 3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? Yes, please see appendix F Implementation Details.

heading: 3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary
text: statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Yes, please see section 4 Experiments.

heading: 3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did
text: you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Yes, please see appendix F Implementation Details.
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? Not applicable. Left blank.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? Not applicable. Left blank.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? Not applicable. Left blank.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? Not applicable. Left blank.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? Not applicable. Left blank.

heading: Future Work
text: We hope to tackle this issue with better optimization in the future