Source-linked AI summary
Transformer Memory as a Differentiable Search Index
Yi Tay, Vinh Q. Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, Tal Schuster, William W. Cohen, Donald Metzler
TL;DR
Information retrieval traditionally uses separate retrieval machinery, while the paper asks whether a single Transformer can encode a corpus and retrieve documents directly. It introduces DSI, which maps text queries to docids through model parameters and studies representation, training, and scale choices. DSI outperforms dual-encoder baselines across the reported settings and also beats BM25 in zero-shot evaluation, while its scaling to larger corpora remains open.
Problem
The paper examines whether information retrieval can be performed by encoding corpus information entirely in a Transformer instead of relying on conventional retrieval procedures.
Method
DSI trains a seq2seq Transformer to map queries directly to relevant docids while exploring document and identifier representations, indexing, training, and scale.
Results
DSI significantly outperforms dual-encoder baselines in reported fine-tuned settings and outperforms BM25 by 14 Hits@1 points in a zero-shot setting.
Takeaways & Limitations
DSI shows that generative indexing can outperform strong retrieval baselines while mapping retrieval into standard model inference and training.
Takeaways & Limitations
The experiments use moderate-sized corpora from 10k to 320k documents derived from one retrieval task, leaving scaling to larger corpora for future work.
Abstract
from arXiv · showhide
In this paper, we demonstrate that information retrieval can be accomplished with a single Transformer, in which all information about the corpus is encoded in the parameters of the model. To this end, we introduce the Differentiable Search Index (DSI), a new paradigm that learns a text-to-text model that maps string queries directly to relevant docids; in other words, a DSI model answers queries directly using only its parameters, dramatically simplifying the whole retrieval process. We study variations in how documents and their identifiers are represented, variations in training procedures, and the interplay between models and corpus sizes. Experiments demonstrate that given appropriate design choices, DSI significantly outperforms strong baselines such as dual encoder models. Moreover, DSI demonstrates strong generalization capabilities, outperforming a BM25 baseline in a zero-shot setup.
1 Introduction
The paper introduces DSI, a Transformer-based architecture that maps queries directly to relevant docids, storing corpus information in model parameters. Experiments examine representations, training procedures, and scale, finding strong performance against dual encoders and BM25.
- 1 Introduction: DSI uses a seq2seq system to map a query directly to a relevant docid.At inference, the trained model can generate ranked lists with beam search.
- 1 Introduction: All corpus information in DSI is encoded within the parameters of a large pre-trained Transformer.The model serves as the learned index rather than relying on a separate fixed search procedure.
- 1 Introduction: For a base-sized T5 model, Hits@1 rises from 12.4% for a DE to 33.9% for DSI on the smallest corpus.On a corpus 30× larger, performance improves by nearly 7 points; larger models produce still larger gains.
- 1 Introduction: DSI significantly and consistently improves with model scale and can outperform BM25 by 14 Hits@1 points in zero-shot evaluation.The summary reports favorable results against strong baselines and improved docid representations over naive choices.
- 1 Introduction: Structured docids can scale better to large corpora, while performance varies substantially across indexing and training strategies.The paper considers both semantically structured identifiers and tokenized integer identifiers.
2 Related Work
The paper distinguishes DSI from related retrieval paradigms including autoregressive entity linking, retrieval-augmented generation, dual encoders, and closed-book language-model retrieval. Its distinctive target is a document identified by an arbitrary docid rather than a semantically meaningful answer or title.
- 2 Related Work: Unlike autoregressive entity linking, DSI retrieves a document containing an answer rather than a document whose title is the answer.DSI permits arbitrary docids, whereas entity linking generates semantically meaningful canonical names.
- 2 Related Work: Retrieval-augmented generation retrieves auxiliary documents to enhance language models, whereas DSI uses generation to replace retrieval.The related methods typically rely on traditional retrieval systems such as dual encoders.
- 2 Related Work: Dual encoders independently encode queries and documents, then retrieve by similarity across embedding pairs using contrastive learning.This contrasts with DSI's direct query-to-docid generation.
- 2 Related Work: Prior Transformer work retrieves facts encoded in model parameters, while DSI retrieves full documents based on docids.The paper frames this as a distinction from closed-book question answering and language models used as knowledge bases.
3 Differentiable Search Index
DSI folds indexing and retrieval into one sequence-to-sequence model that stores the index in its parameters and generates document identifiers from queries. The section compares indexing strategies and docid representations, including structured identifiers designed to make decoding more efficient.
- Core DSI operations: DSI parameterizes the retrieve-then-rank pipeline within a single neural model, supporting document-to-docid indexing and query-to-ranked-docid retrieval.Indexing uses document tokens as input and docids as outputs; retrieval uses autoregressive generation.
- Indexing strategies: The indexing task trains the model to associate each document with its docid, with Inputs2Targets using doc_tokens → docid as the final strategy.Inputs2Targets places the identifier directly in the denoising target, while alternative formulations reverse or combine the mapping and include span corruption.
- Document representations: Direct indexing preserves the first L document tokens, while set and inverted indexing respectively filter terms or map sampled contiguous chunks to docids.Inverted indexing allows the model to look beyond the first k tokens by subsampling a contiguous chunk.
- Docid representations: Unstructured atomic identifiers assign arbitrary unique integers, whereas string-based identifiers decode sequentially and avoid a large output softmax over individual docids.Beam search approximates top-k retrieval scores for sequentially decoded identifiers.
- Semantically structured identifiers: Semantically structured identifiers encode document semantics and share prefixes among similar documents, reducing the search space after each decoding step.They are constructed through hierarchical clustering of document embeddings into a decimal tree or trie; this preprocessing is fully unsupervised.
4 Experiments
Experiments evaluate DSI on Natural Questions across corpus and model scales, comparing supervised and zero-shot retrieval with BM25, dual encoders, and other unsupervised baselines. Results show strong performance gains, sensitivity to training and representation choices, and more optimistic scaling than dual encoders.
- Supervised Finetuning Results: DSI outperforms DE across all NQ dataset sizes, with the largest reported relative Hits@1 gain reaching +66% on NQ320K.On NQ10K, the best DSI variant achieves twice DE performance; on NQ100K, it gains +5% on Hits@1 and Hits@10.
- Zero-Shot Results: DSI outperforms BM25 and other unsupervised retrieval baselines in zero-shot NQ retrieval, without labeled query-document pairs.The best zero-shot result is generally obtained with unstructured atomic identifiers on NQ100K and NQ320K.
- Document Identifiers: Semantic structured identifiers generally improve over unstructured identifiers, although atomic identifiers show instability, high variance, and occasional non-convergence.The paper attributes these issues to a possible optimization quirk and defers further investigation.
- Indexing Strategies: Indexing strategy produces substantial performance variance, and the Inputs2Targets plus bidirectional formulation performs best among the tested strategies.Without indexing, the model achieves 0% Hits@1 because the docids are not meaningful without the indexing task.
- Document Representations: Direct indexing performs best among document representations, while performance substantially dips beyond 64 tokens and inverted-index training is difficult.The inverted-index method repeatedly exposes the same docid to different tokens, complicating optimization.
- Scaling Laws: For an 11B-parameter T5 model, DSI improves Hits@1 by more than 25 points over DE on the small corpus and more than 15 points on the large corpus.DSI with semantic IDs improves substantially with scale, while DE performance is relatively plateaued at smaller parameterizations.
5 Conclusion
DSI is presented as an end-to-end search paradigm that encodes indexing and retrieval within Transformer parameters and performs favorably against BM25 and dual encoders. The paper identifies scaling, dynamic updates, alternative representations, and broader uses as future directions.
- DSI encodes the relationship between terms and docids completely within Transformer parameters, creating a unified end-to-end search system.
- Experiments on Natural Questions show that DSI performs favorably against BM25 and dual encoders in fine-tuning and zero-shot settings.
- Future work includes alternative document and docid representations, mixture-of-experts models for scaling memory capacity, and updates for dynamic corpora.
- The paper also proposes investigating DSI for unsupervised representation learning and as a memory store for other language models.
7 Appendix
The appendix provides dataset and experiment context, indexing-performance observations, training-data details, and an analysis of forgetting cycles that affect validation scores. Despite maximum forgetting, semantic-docid DSI remains competitive with BM25 and averages above the dual-encoder baseline.
- Appendix materials: Additional appendix figures and tables document the NQ datasets, indexing memorization objective, and results at different forgetting checkpoints.Table 4 describes dataset statistics, Table 5 indexing performance, and Table 6 minimum-forgetting, maximum-forgetting, and average results.
- Indexing/Memorization Performance: Indexing performance is relatively strong across methods and model sizes, and increasing model size improves indexing performance.
- Experimental setup: Indexing uses documents from NQ train and validation splits, while retrieval trains on NQ train queries and evaluates on NQ validation.
- DSI Training Dynamics: Buffered-shuffle training produces cycles of minimum and maximum forgetting, creating regular peaks and valleys in validation performance.Selecting the checkpoint with maximum validation performance implicitly selects the checkpoint with minimum forgetting.
- DSI Training Dynamics: Even under maximum forgetting, semantic-docid DSI remains competitive with BM25, while its average score still exceeds the Dual Encoder baseline.