DSI ++: Updating Transformer Memory with New Documents (2023)

Abstract: Differentiable Search Indices (DSIs) encode a corpus of documents in the parameters of a model and use the same model to map queries directly to relevant document identifiers. Despite the solid performance of DSI models, successfully deploying them in scenarios where document corpora change with time is an open problem. In this work, we introduce DSI++, a continual learning challenge for DSI with the goal of continuously indexing new documents while being able to answer queries related to both previously and newly indexed documents. Across different model scales and document identifier representations, we show that continual indexing of new documents leads to considerable forgetting of previously indexed documents. We also hypothesize and verify that the model experiences forgetting events during training, leading to unstable learning. To mitigate these issues, we investigate two approaches. The first focuses on modifying the training dynamics. Flatter minima implicitly alleviates forgetting, so we explicitly optimize for flatter loss basins and show that the model stably memorizes more documents (+12%). Next, we introduce a parametric memory to generate pseudo-queries for documents and supplement them during incremental indexing to prevent forgetting for the retrieval task. Extensive experiments on a novel continual indexing benchmark based on Natural Questions demonstrate that our proposed solution mitigates the forgetting in DSI++ by a significant margin and improves the average Hits@10 by +21.1% over competitive baselines.

## introduction differentiable search indices (dsis; @xcite ) represent a new modeling paradigm for information retrieval tasks using sequence-tosequence learning. specifically, dsis leverage transformer memory @xcite to encode all of the information in a corpus of documents and then use that memory to answer user queries directly, thereby simplifying the retrieval process. dsis achieve this functionality by jointly optimizing for indexing (or memorization) and retrieval tasks. the indexing task requires learning a mapping from document content to its identifier, typically represented by integers or short strings (document identifiers, abbreviated docids). then, the retrieval task necessitates mapping user queries to relevant docids. besides its simplicity and endto-end differentiable nature, dsi significantly outperforms state-of-the-art "retrieve-and-rank" methods based on dual-encoders @xcite . despite the remarkable performance of dsi models, there remain open questions about their applicability in the practical setting of dynamic corpora. consider the realistic scenario wherein new documents are continually added to the indexed corpus. updating the index in dual-encoder-based methods requires computing embeddings for new documents, followed by re-indexing all document embeddings @xcite . in contrast, index construction using a dsi involves training a transformer model. therefore, the model must be re-trained from scratch every time the underlying corpus is updated, thus incurring prohibitively high computational costs compared to dual-encoders. in this work, we aim to address this issue by devising methods for effective incremental indexing using transformer memory without re-training the dsi model from scratch. lifelong (or continual) learning @xcite is a biologically-inspired machine learning paradigm that deals with continuous learning of new tasks by preserving past knowledge and using it to learn new concepts efficiently. based on this paradigm, we propose dsi++ figure 1 : indexing accuracy of d 0 , d 1 , and d 2 document corpora visualized as we continuously index new documents (averaged over 3 runs). we observe that continual indexing of new documents leads to severe forgetting of the previously memorized documents. to reduce forgetting during memorization, we propose explicitly optimizing for flatter loss basins using sharpness-aware minimization (sam; @xcite ). recent works have shown that geometrical properties of the minima play a vital role in forgetting, especially models in flatter loss basins tend to undergo less forgetting while lifelong learning from task sequences @xcite . next, we introduce a generative memory to sample pseudo-queries for already indexed documents and use them to alleviate forgetting of the retrieval task during incremental indexing of the new documents. also, the generative memory enables continual semi-supervised learning of the retrieval task by generating pseudo-queries for an incoming batch of new documents. our main contributions can be summarized as follows: ## explicit forgetting: generative memory the dsi paradigm consists of two tasks -memorization and retrieval. the previous section showcases that sam alleviates implicit forgetting by stably memorizing documents. in this section, we focus on the forgetting phenomenon that arises from the continual indexing of new documents, specifically in the context of the retrieval task. through our systematic study (in §2.4), we show that irrespective of the model scale and docid representations, dsi models undergo severe forgetting. moreover, we observe that the learning performance la n keeps declining for the retrieval task (see figures 2 and 6 for hits@1 and hits@10, respectively). this observation suggests that as we continuously update the dsi model with the indexing objective, the model forgets the retrieval task. in dsi, both memorization and retrieval tasks return docid for input. by setup, we can assume access to previous documents and continue indexing old and new documents to reduce forgetting of the retrieval task. however, in figure 4 , we see that the model still undergoes forgetting (more in §5.2). ## experimentation in this section, the models are initialized with the pre-trained t5-base model, while the additional parameters for atomic docid tokens are randomly initialized. see §a.1 for implementation details. ## conclusion dsi++ introduces a new approach to address a crucial requirement of dsi models for practical use in production setups, where continuous addition of new documents to the corpus is necessary. through experiments, we demonstrate the effectiveness of our proposed solutions: sharpness-aware minimization and generative memory, which significantly reduce catastrophic forgetting. this work establishes a foundation for further research, benefiting both dsi models and the broader community of continual (semi-supervised) learning.