# Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

arXiv: https://arxiv.org/abs/2104.06967 · HTML: https://mlsanity.com/paper/2104.06967/

## Key facts

- Cited by 238 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 49 references that are in the map.
- First posted on arXiv 14 Apr 2021.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [MS MARCO: A Human Generated MAchine Reading COmprehension Dataset](https://mlsanity.com/paper/1611.09268/) (2016)
- [ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT](https://mlsanity.com/paper/2004.12832/) (2020)
- [TinyBERT: Distilling BERT for Natural Language Understanding](https://mlsanity.com/paper/1909.10351/) (2019)
- [Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval](https://mlsanity.com/paper/2007.00808/) (2020)
- [Passage Re-ranking with BERT](https://mlsanity.com/paper/1901.04085/) (2019)
- [RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2010.08191/) (2020)
- [Hard Negative Mixing for Contrastive Learning](https://mlsanity.com/paper/2010.01028/) (2020)
- [Distilling Knowledge from Reader to Retriever for Question Answering](https://mlsanity.com/paper/2012.04584/) (2020)
- [Sparse, Dense, and Attentional Representations for Text Retrieval](https://mlsanity.com/paper/2005.00181/) (2020)
- [Overview of the TREC 2020 deep learning track](https://mlsanity.com/paper/2102.07662/) (2021)
- [Improving Efficient Neural Ranking Models with Cross-Architecture Knowledge Distillation](https://mlsanity.com/paper/2010.02666/) (2020)
- [The Expando-Mono-Duo Design Pattern for Text Ranking with Pretrained Sequence-to-Sequence Models](https://mlsanity.com/paper/2101.05667/) (2021)
- [CEDR: Contextualized Embeddings for Document Ranking](https://mlsanity.com/paper/1904.07094/) (2019)
- [Context-Aware Sentence/Passage Term Importance Estimation For First Stage Retrieval](https://mlsanity.com/paper/1910.10687/) (2019)
- [Rethink Training of BERT Rerankers in Multi-Stage Retrieval Pipeline](https://mlsanity.com/paper/2101.08751/) (2021)
- [Ranking Distillation: Learning Compact Ranking Models With High Performance for Recommender System](https://mlsanity.com/paper/1809.07428/) (2018)
- [Distilling Dense Representations for Ranking using Tightly-Coupled Teachers](https://mlsanity.com/paper/2010.11386/) (2020)
- [TwinBERT: Distilling Knowledge to Twin-Structured BERT Models for Efficient Retrieval](https://mlsanity.com/paper/2002.06275/) (2020)
- [Understanding BERT Rankers Under Distillation](https://mlsanity.com/paper/2007.11088/) (2020)
- [Learning To Retrieve: How to Train a Dense Retrieval Model Effectively and Efficiently](https://mlsanity.com/paper/2010.10469/) (2020)
- [Let's measure run time! Extending the IR replicability infrastructure to include performance aspects](https://mlsanity.com/paper/1907.04614/) (2019)
- [Distilling Knowledge for Fast Retrieval-based Chat-bots](https://mlsanity.com/paper/2004.11045/) (2020)

## Built on by (most cited first, 30 of 238)

- [BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models](https://mlsanity.com/paper/2104.08663/) (2021)
- [Unsupervised Dense Information Retrieval with Contrastive Learning](https://mlsanity.com/paper/2112.09118/) (2021)
- [M3-Embedding: Multi-Linguality, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation](https://mlsanity.com/paper/2402.03216/) (2024)
- [Precise Zero-Shot Dense Retrieval without Relevance Labels](https://mlsanity.com/paper/2212.10496/) (2022)
- [ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction](https://mlsanity.com/paper/2112.01488/) (2021)
- [Large Dual Encoders Are Generalizable Retrievers](https://mlsanity.com/paper/2112.07899/) (2021)
- [Text and Code Embeddings by Contrastive Pre-Training](https://mlsanity.com/paper/2201.10005/) (2022)
- [Tool Learning with Large Language Models: A Survey](https://mlsanity.com/paper/2405.17935/) (2024)
- [SGPT: GPT Sentence Embeddings for Semantic Search](https://mlsanity.com/paper/2202.08904/) (2022)
- [MedCPT: Contrastive Pre-trained Transformers with Large-scale PubMed Search Logs for Zero-shot Biomedical Information Retrieval](https://mlsanity.com/paper/2307.00589/) (2023)
- [RocketQAv2: A Joint Training Method for Dense Passage Retrieval and Passage Re-ranking](https://mlsanity.com/paper/2110.07367/) (2021)
- [SPLADE v2: Sparse Lexical and Expansion Model for Information Retrieval](https://mlsanity.com/paper/2109.10086/) (2021)
- [Mr. TyDi: A Multi-lingual Benchmark for Dense Retrieval](https://mlsanity.com/paper/2108.08787/) (2021)
- [GPL: Generative Pseudo Labeling for Unsupervised Domain Adaptation of Dense Retrieval](https://mlsanity.com/paper/2112.07577/) (2021)
- [Knowledge-Augmented Language Model Prompting for Zero-Shot Knowledge Graph Question Answering](https://mlsanity.com/paper/2306.04136/) (2023)
- [How to Train Your DRAGON: Diverse Augmentation Towards Generalizable Dense Retrieval](https://mlsanity.com/paper/2302.07452/) (2023)
- [Task-aware Retrieval with Instructions](https://mlsanity.com/paper/2211.09260/) (2022)
- [From Distillation to Hard Negative Sampling: Making Sparse Neural IR Models More Effective](https://mlsanity.com/paper/2205.04733/) (2022)
- [Dense X Retrieval: What Retrieval Granularity Should We Use?](https://mlsanity.com/paper/2312.06648/) (2023)
- [Llama2Vec: Unsupervised Adaptation of Large Language Models for Dense Retrieval](https://mlsanity.com/paper/2312.15503/) (2023)
- [What Evidence Do Language Models Find Convincing?](https://mlsanity.com/paper/2402.11782/) (2024)
- [RankVicuna: Zero-Shot Listwise Document Reranking with Open-Source Large Language Models](https://mlsanity.com/paper/2309.15088/) (2023)
- [A Few Brief Notes on DeepImpact, COIL, and a Conceptual Framework for Information Retrieval Techniques](https://mlsanity.com/paper/2106.14807/) (2021)
- [Retrieval Augmented Generation (RAG) and Beyond: A Comprehensive Survey on How to Make your LLMs use External Data More Wisely](https://mlsanity.com/paper/2409.14924/) (2024)
- [A Multi-Task Embedder For Retrieval Augmented LLMs](https://mlsanity.com/paper/2310.07554/) (2023)
- [FiD-Light: Efficient and Effective Retrieval-Augmented Text Generation](https://mlsanity.com/paper/2209.14290/) (2022)
- [Leveraging Large Language Models in Conversational Recommender Systems](https://mlsanity.com/paper/2305.07961/) (2023)
- [Salient Phrase Aware Dense Retrieval: Can a Dense Retriever Imitate a Sparse One?](https://mlsanity.com/paper/2110.06918/) (2021)
- [A Review of Modern Recommender Systems Using Generative Models (Gen-RecSys)](https://mlsanity.com/paper/2404.00579/) (2024)
- [Improving Query Representations for Dense Retrieval with Pseudo Relevance Feedback](https://mlsanity.com/paper/2108.13454/) (2021)

## Abstract (arXiv)

A vital step towards the widespread adoption of neural retrieval models is their resource efficiency throughout the training, indexing and query workflows. The neural IR community made great advancements in training effective dual-encoder dense retrieval (DR) models recently. A dense text retrieval model uses a single vector representation per query and passage to score a match, which enables low-latency first stage retrieval with a nearest neighbor search. Increasingly common, training approaches require enormous compute power, as they either conduct negative passage sampling out of a continuously updating refreshing index or require very large batch sizes for in-batch negative sampling. Instead of relying on more compute capability, we introduce an efficient topic-aware query and balanced margin sampling technique, called TAS-Balanced. We cluster queries once before training and sample queries out of a cluster per batch. We train our lightweight 6-layer DR model with a novel dual-teacher supervision that combines pairwise and in-batch negative teachers. Our method is trainable on a single consumer-grade GPU in under 48 hours (as opposed to a common configuration of 8x V100s). We show that our TAS-Balanced training method achieves state-of-the-art low-latency (64ms per query) results on two TREC Deep Learning Track query sets. Evaluated on NDCG@10, we outperform BM25 by 44%, a plainly trained DR by 19%, docT5query by 11%, and the previous best DR model by 5%. Additionally, TAS-Balanced produces the first dense retriever that outperforms every other method on recall at any cutoff on TREC-DL and allows more resource intensive re-ranking models to operate on fewer passages to improve results further.
