Optimizing Test-Time Query Representations for Dense Retrieval (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 (Test-Time Optimization of Query Representations), which further optimizes instance-level query representations guided by signals from test-time retrieval results. We leverage a cross-encoder 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 re-ranker. 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.4x faster with an efficient implementation.

## introduction 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) @xcite @xcite and information retrieval @xcite . dense retrieval often uses a dual encoder architecture, which enables the pre-computation of congiven the initial representation of a test query q 0 , tour iteratively optimizes its representation (e.g., q 0 → q 1 → q 2 → q 3 ) based on top-k retrieval results. the figure shows how each query vector retrieves new context vectors and updates its representation to find the gold answer (e.g., 1983). our cross-encoder re-ranker provides a relevance score for each top retrieval result making the query representation closer to the final answer. 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 @xcite , which finetunes the query encoder over training queries in a new domain. instead, we fine-tune query rep-resentations for each test query. cross-encoders are known to exhibit better generalization ability in unseen distributions compared to dual encoders @xcite . accordingly, we leverage cross-encoder re-rankers @xcite 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; @xcite , 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 @xcite ) and passage retrieval @xcite 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. ## methodology 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). ## experiments we test tour on multiple open-domain qa datasets. specifically, we evaluate its performance on phrase retrieval and passage retrieval. ## conclusion 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.