Elamur Research
☰
home / papers / 2021 / Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

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/2104.06967

Summary

Dense retrievers such as the dual-encoder BERT DOT give fast first-stage search. Recent methods for training them cost a lot of compute: ANCE keeps refreshing a negative-sampling index, and RocketQA needs very large batches. Hofstätter et al. ask whether a strong dense retriever can be trained on modest hardware instead.

Their answer is TAS-Balanced, a way to compose training batches. Before training, MSMARCO training queries are clustered once by semantic similarity, and each batch draws its queries from a single cluster. Passage pairs are then picked so that the teacher score margins are evenly spread. A 6-layer BERT DOT student learns from two teachers at once: BERT CAT supplies pairwise signals and ColBERT supplies in-batch negative signals.

Training fits on one consumer GPU in under two days. On TREC-DL '19 and '20 the model reaches state-of-the-art quality for low-latency retrieval while using small batches. This matters because competitive neural first-stage retrieval becomes reachable for groups without large server clusters.

Key points

  • Cost: under 48 hours on a single 11GB GPU, compared with the usual 8x V100 setup. Clustering all ~400K MSMARCO training queries takes under 10 minutes.
  • Effectiveness: at 64 ms per query, NDCG@10 is 44% higher than that of BM25, with further gains of 19% over a dense retriever trained the standard way, 11% over docT5query, and 5% over the strongest earlier dense retriever. On TREC-DL'19, nDCG@10 is .722.
  • Recall: this is the first BERT DOT model to beat BM25 and docT5query at every recall cutoff on TREC-DL. Fusing it with docT5query raises recall further.
  • Training choices and robustness: Margin-MSE was the most reliable loss for the dual-teacher setup. Five runs with different random orderings varied by less than .01 nDCG on TREC-DL.
  • Re-ranking and code: with shallow re-ranking (duo-T5 on the top-10), the pipeline beats systems that are 2–6x slower. However, mono-duo-T5 at depth 1K does not benefit from the extra recall. Code is available at github.com/sebastian-hofstaetter/tas-balanced-dense-retrieval.

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

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

Cite

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

@misc{hofstatter2021efficiently,
  title = {{Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling}},
  author = {Hofstatter, Sebastian and Lin, Sheng-Chieh and Yang, Jheng-Hong and Lin, Jimmy and Hanbury, Allan},
  year = {2021},
  doi = {10.48550/arxiv.2104.06967},
  eprint = {2104.06967},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2104.06967}
}

References

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

  1. Zeynep Akkalyoncu Yilmaz, Wei Yang, Haotian Zhang, and Jimmy Lin. 2019. Cross-Domain Modeling of Sentence-Level Evidence for Document Retrieval. In Proc. of EMNLP-IJCNLP. find
  2. Payal Bajaj, Daniel Campos, Nick Craswell, Li Deng, Jianfeng Gao, Xiaodong Liu, Rangan Majumder, Andrew McNamara, Bhaskar Mitra, and Tri Nguyen. 2016. MS MARCO : A Human Generated MAchine Reading COmprehension Dataset. In Proc. of NIPS. find our card
  3. Zhe Cao, Tao Qin, Tie-Yan Liu, Ming-Feng Tsai, and Hang Li. 2007. Learning to Rank: From Pairwise Approach to Listwise Approach. In Proc. of ICML. find
  4. Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. 2021. Unsupervised Learning of Visual Features by Contrasting Cluster Assignments. arXiv:2006.09882 (2021). — [arXiv:2006.09882](https://arxiv.org/abs/2006.09882) find
  5. Xuanang Chen, Ben He, Kai Hui, Le Sun, and Yingfei Sun. 2020. Simplified TinyBERT: Knowledge Distillation for Document Retrieval. arXiv:2009.07531 (2020). — [arXiv:2009.07531](https://arxiv.org/abs/2009.07531) find
  6. Daniel Cohen, Scott M. Jordan, and W. Bruce Croft. 2019. Learning a Better Negative Sampling Policy with Deep Neural Networks for Search. In Proc. of ICTIR. find
  7. Nick Craswell, Bhaskar Mitra, Emine Yilmaz, and Daniel Campos. 2019. Overview of the TREC 2019 Deep Learning Track. In TREC. find our card
  8. Nick Craswell, Bhaskar Mitra, Emine Yilmaz, and Daniel Campos. 2020. Overview of the TREC 2020 Deep Learning Track. In TREC. find
  9. Zhuyun Dai and Jamie Callan. 2019. Context-Aware Sentence/Passage Term Importance Estimation for First Stage Retrieval. arXiv:1910.10687 (2019). — [arXiv:1910.10687](https://arxiv.org/abs/1910.10687) find
  10. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proc. of NAACL. find
  11. Yingqi Qu Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. 2020. RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering. arXiv:2010.08191 (2020). find our card
  12. Luyu Gao, Zhuyun Dai, and Jamie Callan. 2020. Understanding BERT Rankers Under Distillation. arXiv:2007.11088 (2020). — [arXiv:2007.11088](https://arxiv.org/abs/2007.11088) find
  13. Luyu Gao, Zhuyun Dai, and Jamie Callan. 2021. Rethink Training of BERT Rerankers in Multi-Stage Retrieval Pipeline. arXiv:2101.08751 (2021). — [arXiv:2101.08751](https://arxiv.org/abs/2101.08751) find
  14. Sebastian Hofstätter, Sophia Althammer, Michael Schröder, Mete Sertkan, and Allan Hanbury. 2020. Improving Efficient Neural Ranking Models with Cross-Architecture Knowledge Distillation. arXiv:2010.02666 (2020). — [arXiv:2010.02666](https://arxiv.org/abs/2010.02666) find
  15. Sebastian Hofstätter and Allan Hanbury. 2019. Let’s Measure Run Time! Extending the IR Replicability Infrastructure to Include Performance Aspects. In Proc. of OSIRRC. find
  16. Gautier Izacard and Edouard Grave. 2020. Distilling Knowledge from Reader to Retriever for Question Answering. arXiv:2012.04584 (2020). — [arXiv:2012.04584](https://arxiv.org/abs/2012.04584) find
  17. Nick Jardine and Cornelis Joost van Rijsbergen. 1971. The Use of Hierarchic Clustering in Information Retrieval. Information Storage and Retrieval 7, 5 (1971), 217–240. find
  18. Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2019. TinyBERT: Distilling BERT for Natural Language Understanding. arXiv:1909.10351 (2019). find
  19. Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2017. Billion-Scale Similarity Search with GPUs. arXiv:1702.08734 (2017). find our card
  20. Yannis Kalantidis, Mert Bulent Sariyildiz, Noe Pion, Philippe Weinzaepfel, and Diane Larlus. 2020. Hard Negative Mixing for Contrastive Learning. Advances in Neural Information Processing Systems 33 (2020). find
  21. Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen tau Yih. 2020. Dense Passage Retrieval for Open-Domain Question Answering. arXiv:2004.04906 (2020). find our card
  22. Omar Khattab and Matei Zaharia. 2020. ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT. In Proc. of SIGIR. find our card
  23. Diederik P. Kingma and Jimmy Ba. 2014. Adam: A Method for Stochastic Optimization. arXiv:1412.6980 (2014). find
  24. Sheng-Chieh Lin, Jheng-Hong Yang, and Jimmy Lin. 2020. Distilling Dense Representations for Ranking using Tightly-Coupled Teachers. arXiv:2010.11386 (2020). — [arXiv:2010.11386](https://arxiv.org/abs/2010.11386) find
  25. Wenhao Lu, Jian Jiao, and Ruofei Zhang. 2020. TwinBERT: Distilling Knowledge to Twin-Structured BERT Models for Efficient Retrieval. arXiv:2002.06275 (2020). — [arXiv:2002.06275](https://arxiv.org/abs/2002.06275) find
  26. Yi Luan, Jacob Eisenstein, Kristina Toutanova, and Michael Collins. 2020. Sparse, Dense, and Attentional Representations for Text Retrieval. arXiv:2005.00181 (2020). find our card
  27. Sean MacAvaney, Franco Maria Nardini, Raffaele Perego, Nicola Tonellotto, Nazli Goharian, and Ophir Frieder. 2020. Training Curricula for Open Domain Answer Re-Ranking. In Proc. of SIGIR. find
  28. Sean MacAvaney, Andrew Yates, Arman Cohan, and Nazli Goharian. 2019. CEDR: Contextualized Embeddings for Document Ranking. In Proc. of SIGIR. find
  29. James MacQueen. 1967. Some Methods for Classification and Analysis of Multivariate observations. In Proc. of the Fifth Berkeley Symposium on Mathematical Statistics and Probability, Vol. 1. 281–297. find
  30. Rodrigo Nogueira and Kyunghyun Cho. 2019. Passage Re-ranking with BERT. arXiv:1901.04085 (2019). find our card
  31. Rodrigo Nogueira and Jimmy Lin. 2019. From doc2query to docTTTTTquery. Online preprint (2019). find
  32. Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. 2017. Automatic Differentiation in PyTorch. In Proc. of NIPS-W. find
  33. Ronak Pradeep, Rodrigo Nogueira, and Jimmy Lin. 2021. The Expando-Mono-Duo Design Pattern for Text Ranking with Pretrained Sequence-to-Sequence Models. arXiv:2101.05667 (2021). — [arXiv:2101.05667](https://arxiv.org/abs/2101.05667) find
  34. Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. DistilBERT, A Distilled Version of BERT: Smaller, Faster, Cheaper and Lighter. arXiv:1910.01108 (2019). find
  35. Zhiqiang Shen, Zechun Liu, Zhuang Liu, Marios Savvides, and Trevor Darrell. 2020. Rethinking Image Mixture for Unsupervised Visual Representation Learning. arXiv:2003.05438 (2020). — [arXiv:2003.05438](https://arxiv.org/abs/2003.05438) find
  36. Jiaxi Tang and Ke Wang. 2018. Ranking Distillation: Learning Compact Ranking Models with High Performance for Recommender System. In Proc. of SIGKDD. find
  37. Amir Vakili Tahami, Kamyar Ghajar, and Azadeh Shakery. 2020. Distilling Knowledge for Fast Retrieval-based Chat-bots. In Proc. of SIGIR. find
  38. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention Is All You Need. In Proc. of NIPS. find
  39. Ellen M. Voorhees. 1985. The Cluster Hypothesis Revisited. In Proc. of SIGIR. find
  40. Xuanhui Wang, Cheng Li, Nadav Golbandi, Michael Bendersky, and Marc Najork. 2018. The LambdaLoss Framework for Ranking Metric Optimization. In Proc. of CIKM. find
  41. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoe… find our card
  42. Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Arnold Overwijk. 2020. Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval. arXiv:2007.00808 (2020). find our card
  43. Peilin Yang, Hui Fang, and Jimmy Lin. 2017. Anserini: Enabling the Use of Lucene for Information Retrieval Research. In Proc. of SIGIR. find
  44. Jingtao Zhan, Jiaxin Mao, Yiqun Liu, Min Zhang, and Shaoping Ma. 2020. Learning To Retrieve: How to Train a Dense Retrieval Model Effectively and Efficiently. arXiv:2010.10469 (2020). — [arXiv:2010.10469](https://arxiv.org/abs/2010.10469) find
  45. Justin Zobel. 1998. How Reliable are the Results of Large-Scale Information Retrieval Experiments?. In Proc. of SIGIR. find