Query Drift Compensation: Enabling Compatibility in Continual Learning of Retrieval Embedding Models

TL;DR

This work addresses embedding drift in continual dense retrieval, where updating a retriever causes old corpus embeddings to become incompatible with new query embeddings. It introduces a re-indexing-free framework that combines embedding distillation with Query Drift Compensation (QDC) to project new query embeddings back to the old embedding space, preserving compatibility without re-indexing. The approach is instantiated on a nomic dense retriever with MLM pretraining and supervised contrastive finetuning, augmented by a cosine-distance distillation loss and a drift-based query projection. Across five BEIR datasets in a task-incremental setting, QDC (especially with embedding distillation) significantly reduces forgetting and nearly matches joint training performance, while outperforming re-indexing baselines. The method offers practical impact for real-world retrieval systems needing continual knowledge updates without expensive re-indexing.

Abstract

Text embedding models enable semantic search, powering several NLP applications like Retrieval Augmented Generation by efficient information retrieval (IR). However, text embedding models are commonly studied in scenarios where the training data is static, thus limiting its applications to dynamic scenarios where new training data emerges over time. IR methods generally encode a huge corpus of documents to low-dimensional embeddings and store them in a database index. During retrieval, a semantic search over the corpus is performed and the document whose embedding is most similar to the query embedding is returned. When updating an embedding model with new training data, using the already indexed corpus is suboptimal due to the non-compatibility issue, since the model which was used to obtain the embeddings of the corpus has changed. While re-indexing of old corpus documents using the updated model enables compatibility, it requires much higher computation and time. Thus, it is critical to study how the already indexed corpus can still be effectively used without the need of re-indexing. In this work, we establish a continual learning benchmark with large-scale datasets and continually train dense retrieval embedding models on query-document pairs from new datasets in each task and observe forgetting on old tasks due to significant drift of embeddings. We employ embedding distillation on both query and document embeddings to maintain stability and propose a novel query drift compensation method during retrieval to project new model query embeddings to the old embedding space. This enables compatibility with previously indexed corpus embeddings extracted using the old model and thus reduces the forgetting. We show that the proposed method significantly improves performance without any re-indexing. Code is available at https://github.com/dipamgoswami/QDC.

Figures (3)

• Figure 1: Continual Document Retrieval (CDR). Here, we consider a two-task continual setting and illustrate two different approaches to tackle the embedding drift issue for old task retrieval in CDR which arises due to non-compatibility between query embeddings from the updated model $Q_1^{f_2}$ (in red) and corpus embeddings indexed using the old model $I_1^{f_1}$ (in green). (left) A naive approach to make the query and corpus embeddings compatible is by re-indexing the corpus documents using the updated model. However, re-indexing large amounts of documents from all old tasks after updating the model on every new task is time-consuming and computationally expensive. (right) To avoid re-indexing, we propose to estimate embedding drift from old to new model $\Delta^{t-1\rightarrow t}$ and compensate the drift from $Q_1^{f_2}$ during retrieval. The proposed query drift compensation approach enables compatibility by projecting the query embedding back to the embedding space of $f_1$ without any need for expensive re-indexing.
• Figure 2: Illustration of Query Drift Compensation. (left) We show how to estimate the query drift vectors $\Delta^{t-1\rightarrow t}$ for each task transition (from $t-1$ to $t$) after training on task $t$. We store the drift vectors of all old tasks. For an old task $t'$ ($t' < t$), we obtain the drift vector $\Delta^{t'\rightarrow t}$ by addition of all drift vectors from task $t'$ to $t$. (right) During retrieval of old task $t'$, we compensate the query embeddings with drift vector $\Delta^{t'\rightarrow t}$ to project them from embedding space of task $t$ to that of $t'$. As a result, we compare the query and previously indexed document embeddings in the same embedding space of task $t'$ (in purple), thus avoiding the non-compatibility issue.
• Figure 3: (left) UMAP visualization of the drift in query and corpus embeddings of MS MARCO after fine-tuning on NQ and Hotpot QA (using Model after task 3). (right) Analysis of the performance when using multiple drift vectors to represent drift between embedding spaces of old and new task.