Transformer Memory as a Differentiable Search Index | Summary

This article explains the key points of Transformer Memory as a Differentiable Search Index.

• 2022-02-14 (arXiv), 36th Conference on Neural Information Processing Systems (NeurIPS 2022)
• Tay, Yi, Tran, Vinh Q., Dehghani, Mostafa, Ni, Jianmo, Bahri, Dara, Mehta, Harsh, Qin, Zhen, Hui, Kai, Zhao, Zhe, Gupta, Jai, et al.
• Google Research
• Paper (v3, revised 2022-10-21)

Read this article in Korean

Summary

• This paper studies how to retrieve relevant documents by storing a search index in a Transformer’s parameters. On NQ320K, a Natural Questions dataset of questions paired with Wikipedia articles, an 11B-parameter Differentiable Search Index (DSI) using identifiers derived from document clusters achieved 40.4% Hits@1, the fraction of queries whose correct document ranks first, versus 24.3% for an 11B-parameter T5 dual encoder that searches separately encoded queries and documents.
• DSI learns both to associate document text with document identifiers and to generate identifiers from queries. Its results depend on identifier design, indexing strategy, model size, and the balance between indexing and retrieval training.
• The evaluation uses one question-answering retrieval dataset and moderate corpus sizes. The appendix reports substantial fluctuations caused by forgetting indexed documents, making the validation checkpoint choice consequential.

1 Introduction

Information retrieval maps a query to relevant documents, typically through separate retrieval and ranking stages. DSI instead uses a sequence-to-sequence model, which generates an output sequence from an input sequence, to map a query directly to a document identifier (docid).

The upper path embeds queries and documents before a separate similarity search; the lower path trains an encoder-decoder to emit document identifiers and uses beam search to form a ranking.

The table locates the architectural change in indexing and retrieval: DSI learns document-to-identifier mappings and predicts identifier probabilities, while the dual encoder stores document vectors and searches by similarity.

2 Related Work

Autoregressive entity linking generates meaningful entity names, whereas DSI can generate arbitrary identifiers for documents containing an answer. The paper also distinguishes DSI from retrieval-augmented generation, which retrieves documents to help a language model generate text, and closed-book question answering, which generates answers from model parameters.

3 Differentiable Search Index

DSI has two operations: indexing associates each document’s content with its docid, and retrieval generates candidate docids from a query. Beam search, which keeps multiple candidate output sequences during decoding, can produce a ranked list.

3.1 Indexing Strategies

The final indexing method, Inputs2Targets, trains doc_tokens → docid; alternatives generate document tokens from docids, combine both directions, or use span corruption, which masks spans for reconstruction. Direct indexing uses the first L document tokens in order, while set indexing removes repeated terms and common stopwords and the inverted-index variant samples contiguous token chunks.

3.2 Representing Docids for Retrieval

Atomic identifiers assign a distinct output token to each document; naive string identifiers tokenize arbitrary integers and decode them sequentially. Semantic string identifiers use hierarchical clustering so similar documents share prefixes: the authors recursively cluster document embeddings from an 8-layer BERT model into 10 groups and stop recursion at clusters of at most 100 documents.

Cluster choices become successive identifier digits, so documents in the same branch share a prefix and decoding can navigate the hierarchy.

3.3 Training and Optimization

Models use sequence-to-sequence cross-entropy and teacher forcing, in which training supplies previous target tokens to the decoder. Joint training of indexing and retrieval with task prompts performed better than indexing followed by retrieval fine-tuning, so the authors adopted multi-task training.

4 Experiments

Natural Questions (NQ) supplies natural-language questions and Wikipedia articles; NQ10K, NQ100K, and NQ320K denote the sizes of the combined query-document pair sets. The authors report Hits@1 and Hits@10, use pretrained T5 models from Base to XXL, and train for at most 1M steps with batch size 128; the main configuration uses direct indexing with L = 32.

4.1 Baselines

Baselines include BM25, a term-based retrieval method, and T5 dual encoders trained contrastively on NQ pairs. Zero-shot experiments, which provide no labeled query-to-docid training pairs, also compare SentenceT5 and raw T5 embeddings.

4.2 Experimental Results

With supervised training on NQ10K, 11B semantic-string DSI reaches 48.5% Hits@1 versus 22.1% for the 11B dual encoder. On NQ320K, its 40.4% Hits@1 and 70.3% Hits@10 compare with 24.3% and 67.3% for that baseline.

The strongest NQ320K supervised result uses 11B semantic-string DSI, but identifier choices vary substantially: the 11B naive-string variant scores 23.8% Hits@1 on the same set.

In zero-shot NQ320K retrieval, 11B atomic-docid DSI scores 25.1% Hits@1 versus 11.6% for BM25 and 16.9% for SentenceT5. Atomic identifiers perform best on NQ100K and NQ320K in this setting, but their supervised results are variable and the authors report intermittent optimization failures.

The best identifier type changes with dataset size: semantic strings lead on NQ10K, while atomic identifiers lead on NQ100K and NQ320K.

On NQ100K with naive docids, Inputs2Targets and bidirectional indexing achieve 13.5 and 13.2 Hits@1, while Targets2Inputs and span corruption achieve 0%. On NQ320K, direct indexing performs best among document representations; an indexing-to-retrieval example ratio of 32 generally works well in the reported co-training study.

On NQ320K, semantic-string DSI gains more Hits@1 as model size grows than the dual encoder, while naive-string DSI starts much lower and improves with scale.

The plotted retrieval score peaks at an indexing-to-retrieval ratio of 32 and falls at higher ratios, showing sensitivity to the training mixture.

Direct indexing with 32 or 64 tokens outperforms the plotted set and inverted-index alternatives; the 128-token direct variant also performs worse than shorter direct inputs.

5 Conclusion

The authors identify larger corpora, alternative document and identifier representations, and updates to dynamic collections as future research directions. The experiments establish a proof of concept on NQ rather than a broad comparison across retrieval tasks.

Appendix

• Dataset Statistics: NQ10K, NQ100K, and NQ320K contain 10K, 86K, and 228K unique documents, respectively, with 8K/2K, 80K/20K, and 307K/8K training/validation pairs. Extended Results report indexing Hits@1 from 76.3% to 92.7% across the listed model and identifier configurations.
• Discussion of DSI Training Dynamics: indexing includes documents from both training and validation splits, while retrieval training uses only training queries. Validation documents were concatenated after training documents and shuffled with a 5,000-example buffer; forgetting produced validation peaks and troughs, so selecting the best checkpoint favored minimum forgetting. For 11B semantic docids on NQ320K, Hits@1 was 40.4% at minimum forgetting, 12.2% at maximum forgetting, and 26.3% in the reported average.

Brief Thoughts

The comparison with similarly sized T5 dual encoders supports the proposed retrieval formulation, but checkpoint sensitivity materially narrows the headline result: the reported average for 11B semantic DSI is 26.3% Hits@1, compared with 24.3% for the 11B dual encoder. The DSI average assumes validation documents are distributed uniformly through indexing, while the dual-encoder value is the main-table result. Performance on larger or changing corpora remains untested in this paper, and the experiments used 128–256 TPUv4 chips to train models above 1B parameters.