Elamur Research
☰
home / papers / 2022 / Transformer Memory as a Differentiable Search Index

Transformer Memory as a Differentiable Search Index

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "arXiv — non-exclusive license to distribute" does not allow redistributing the text itself. Original — arxiv.org/pdf/2202.06991

Summary

Conventional retrieval relies on multi-stage pipelines, such as inverted indexes or dual encoders with nearest-neighbour search. This paper asks whether one Transformer can do retrieval by storing the entire corpus in its weights. The authors propose the Differentiable Search Index (DSI), a seq2seq model based on T5. During indexing it learns to generate a document's identifier (docid) from the document's text. At query time it produces docids for a question, and beam search yields a ranked list.

The authors compare several design choices:

  • how documents are represented
  • naive integer docids, atomic single-token docids, and semantically structured docids built by hierarchical clustering
  • different indexing objectives
  • how indexing and retrieval training are mixed

Experiments use Natural Questions subsets of 10k–320k documents. With the right choices, DSI clearly beats dual encoders and also beats BM25 without labelled queries. Its advantage also grows with model size.

The work matters because it turns indexing into ordinary model training. This opens retrieval to standard ML techniques and makes the search component differentiable.

Key points

  • On the smallest NQ corpus with T5-Base, DSI raises Hits@1 from 12.4% (dual encoder) to 33.9%. With the 11B-parameter T5, the gain over DE exceeds 25 points on the small corpus and 15 points on NQ320K.
  • In the zero-shot setting, DSI improves Hits@1 by about 14 points over BM25. Atomic docids work best there, often doubling the results of the other docid variants.
  • Indexing objectives differ sharply in how well they work:
    • Training without any indexing gives 0% Hits@1.
    • Inputs2Targets (13.5) and Bidirectional (13.2) perform best on NQ100K.
    • Targets2Inputs and span corruption with docids fail entirely.
  • Truncating documents to short inputs works best, and performance drops beyond roughly 64 tokens. Stopword removal and set processing bring no benefit.
  • Training indexing first and retrieval afterwards performs poorly; multi-task co-training works better. An indexing-to-retrieval example ratio of about 32 performs well. Unlike DSI, dual encoders largely plateau as model size increases.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Other

Cited by

117 papers of our library cite it (of the 140673 whose bibliographies we read); 11 of them have a page here, 106 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2202.06991 · arXiv 2202.06991

@misc{tay2022transformer,
  title = {{Transformer Memory as a Differentiable Search Index}},
  author = {Tay, Yi and Tran, Vinh Q. and Dehghani, Mostafa and Ni, Jianmo and Bahri, Dara and Mehta, Harsh and Qin, Zhen and Hui, Kai and Zhao, Zhe and Gupta, Jai and Schuster, Tal and Cohen, William W. and Metzler, Donald},
  year = {2022},
  doi = {10.48550/arxiv.2202.06991},
  eprint = {2202.06991},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2202.06991}
}

References

29 references, 7 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Samira Abnar, Mostafa Dehghani, Behnam Neyshabur, and Hanie Sedghi. Exploring the limits of large scale pre-training. arXiv preprint arXiv:2110.02095, 2021. — [arXiv:2110.02095](https://arxiv.org/abs/2110.02095) find
  2. Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. arXiv preprint arXiv:2112.04426, 2021. find
  3. Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020. find
  4. Nicola De Cao, Gautier Izacard, Sebastian Riedel, and Fabio Petroni. Autoregressive entity retrieval. arXiv preprint arXiv:2010.00904, 2020. find
  5. Mostafa Dehghani, Hamed Zamani, Aliaksei Severyn, Jaap Kamps, and W Bruce Croft. Neural ranking models with weak supervision. In Proceedings of the 40th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 65–74, 2017. find
  6. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018. find
  7. Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. arXiv preprint arXiv:2112.06905, 2021. find our card
  8. William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961, 2021. find
  9. Tianyu Gao, Xingcheng Yao, and Danqi Chen. SimCSE: Simple contrastive learning of sentence embeddings. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 6894–6910, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational L… find our card
  10. Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. arXiv preprint arXiv:2012.14913, 2020. find
  11. Daniel Gillick, Alessandro Presta, and Gaurav Singh Tomar. End-to-end retrieval in continuous space. arXiv preprint arXiv:1811.08008, 2018. — [arXiv:1811.08008](https://arxiv.org/abs/1811.08008) find
  12. Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning, 2020. URL https://arxiv.org/abs/1908.10396. — [arXiv:1908.10396](https://arxiv.org/abs/1908.10396) ↗
  13. Kelvin Guu, Kenton Lee, Zora Tung, and Panupong Pasupat. REALM: Retrieval-Augmented Language Model Pre-Training. In Proceedings of ICML 2020, 2020. find
  14. Martin Josifoski, Nicola De Cao, Maxime Peyrard, and Robert West. Genie: Generative information extraction. arXiv preprint arXiv:2112.08340, 2021. find
  15. Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. find
  16. Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprint arXiv:2004.04906, 2020. find our card
  17. Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin Kenton Lee, Kristina Toutanova, Llion Jones Matthew Kelcey, Ming-Wei Chang, Andrew M Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural Questions… find our card
  18. Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668, 2020. find
  19. Donald Metzler, Yi Tay, Dara Bahri, and Marc Najork. Rethinking search: making domain experts out of dilettantes. In ACM SIGIR Forum, volume 55, pages 1–27. ACM New York, NY, USA, 2021. find
  20. Jianmo Ni, Gustavo Hernández Ábrego, Noah Constant, Ji Ma, Keith B Hall, Daniel Cer, and Yinfei Yang. Sentence-t5: Scalable sentence encoders from pre-trained text-to-text models. arXiv preprint arXiv:2108.08877, 2021. find our card
  21. Fabio Petroni, Tim Rocktäschel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, Alexander H Miller, and Sebastian Riedel. Language models as knowledge bases? arXiv preprint arXiv:1909.01066, 2019. find our card
  22. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv:1910.10683, 2019. find
  23. Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.02217, 2020. find
  24. Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge can you pack into the parameters of a language model? arXiv preprint arXiv:2002.08910, 2020. find
  25. Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Procee… find
  26. Ilya Sutskever, Oriol Vinyals, and Quoc V Le. Sequence to sequence learning with neural networks. arXiv preprint arXiv:1409.3215, 2014. find
  27. Yi Tay, Mostafa Dehghani, Jinfeng Rao, William Fedus, Samira Abnar, Hyung Won Chung, Sharan Narang, Dani Yogatama, Ashish Vaswani, and Donald Metzler. Scale efficiently: Insights from pre-training and fine-tuning transformers. arXiv preprint arXiv:2109.10686, 2021. — [arXiv:2109.10686](https://arxiv.org/abs/2109.10686) find
  28. Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239, 2022. find our card
  29. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998–6008, 2017. find