Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

By Sebastian Hofstätter, Sheng-Chieh Lin, Jheng-Hong Yang, Jimmy J. Lin, Allan Hanbury. 2021

521 citations recorded by the source corpus.

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. 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. 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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Information Retrieval & Search

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Information Retrieval & Search

Citations in the corpus

Unsupervised Dense Information Retrieval with Contrastive Learning (2021, neutral): “We also compare to TAS-B (Hofstätter et al., 2021), which performs distillation from a cross-encoder to a bi-encoder, and GenQ, which creates synthetic query-document pairs with a generative model.”
You Only Need One Model for Open-domain Question Answering (2021, methodology): “TAS (Hofstätter et al., 2021) and Xiong et al. (2021b) used clustering of embeddings for the same.”
Pretrained Transformers for Text Ranking: BERT and Beyond (2020, background): “In the simplest approach, ranking becomes the problem of approximate nearest neighbor (ANN) search based on some simple metric such as cosine similarity [10, 12, 13, 17, 19, 22–24, 35, 36, 40, 42].”
M3-Embedding: Multi-Linguality, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation (2024, unknown): “…of L dense , L lex , L mul and L inter as the loss without self-knowledge distillation: (4) In previous studies, the training quality of embedding model can benefit from knowledge distillation, which takes advantage of fine-grained soft labels from another ranking model (Hofst¨atter et al., 2021).”
Introducing Neural Bag of Whole-Words with ColBERTer: Contextualized Late Interactions using Enhanced Reduction (2022, methodology): “Compared to a TAS-Balanced trained dense retriever [16] ColBERTer is not statistically significantly worse on any single collection.”
Salient Phrase Aware Dense Retrieval: Can a Dense Retriever Imitate a Sparse One? (2021, methodology): “BM25, in comparison, has a latency of 55ms using the Anserini toolkit (Hofstätter et al., 2021), and a hybrid DPR+BM25 retriever has a latency of 58ms assuming DPR and BM25 retrieval can be done in parallel.”
ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction (2021, methodology): “While ColBERT outperforms single-vector systems like RepBERT, ANCE, and even TAS-B, improvements in supervision such as distillation from cross-encoders (see §2) enable systems like SPLADEv2, PAIR, and RocketQAv2 to achieve higher quality than vanilla ColBERT.”
Towards Unsupervised Dense Information Retrieval with Contrastive Learning (2021, result): “We also compare to TAS-B (Hofstätter et al., 2021), which performs distillation from a cross-encoder to a bi-encoder, and GenQ, which creates synthetic query-document pairs with a generative model.”
SPLADE v2: Sparse Lexical and Expansion Model for Information Retrieval (2021, methodology): “…providing results that are still competitive; (3) finally, we use distillation techniques [12] to boost SPLADE performance, leading to close to state-of-the-art results on the MS MARCO passage ranking task as well as the BEIR zero-shot evaluation benchmark [26]. and IR [10, 11, 13, 16, 24, 30].”
TU Wien at TREC DL and Podcast 2021: Simple Compression for Dense Retrieval (2021, methodology): “At the DL track, we tested our TAS-Balanced [7] training approach against a standalone dense retrieval baseline, and a compressed version of TAS-Balanced dense retriever, trained with a simple dimensionality reduction technique, which we present in this paper.”
BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models (2021, methodology): “TAS-B model prefers to retrieve documents with shorter lengths.”
A Few Brief Notes on DeepImpact, COIL, and a Conceptual Framework for Information Retrieval Techniques (2021, unknown): “…al., 2020) and ANCE (Xiong et al., 2021), but also encompassing many other techniques (Gao et al., 2021b; Hofstätter et al., 2020; Qu et al., 2021; Hofstätter et al., 2021; Lin Table 1: Our conceptual framework for organizing recent developments in information retrieval. et al., 2021), can be…”
Jointly Optimizing Query Encoder and Product Quantization to Improve Retrieval Performance (2021, methodology): “To generate result candidates based on embedded queries and documents in the space, most existing DR models [16, 22, 43, 48] rely on brute-force search.”
SGPT: GPT Sentence Embeddings for Semantic Search (2022, methodology): “1B from §3 (c) BM25+CE [44], the current overall state-of-the-art on BEIR (d) TAS-B [17], the original Bi-Encoder state-of-the-art on BEIR (e) Contriever [20], a similar training scheme as [27] but using an encoder transformer (f ) GTR-XXL [29], the current Bi-Encoder state-of-the-art on BEIR with 4”
Text and Code Embeddings by Contrastive Pre-Training (2022, methodology): “First, we observe that our unsupervised model performs competitively even with some previous embedding methods that leverage supervised MSMARCO data (Xiong et al., 2020; Hofstätter et al., 2021).”
Task-aware Retrieval with Instructions (2022, background): “The bi-encoder architecture is known to be less expressive since it enables only limited interactions between queries and documents (Khattab and Zaharia, 2020), especially when the training data is limited (Hofstätter et al., 2021).”
Large Dual Encoders Are Generalizable Retrievers (2021, methodology): “We follow prior work (Xiong et al., 2020; Hofstätter et al., 2021) to initialize dual encoders from pre-trained language models.”
Tool learning with large language models: a survey (2024, unknown): “Conversely, semantic-based methods ( i.e., dense retrieval) utilize neu-Changle ral networks to learn the semantic relationship be-tween queries and tool descriptions [116–120], and then calculate the semantic similarity using meth-ods such as cosine similarity.”
Mr. TyDi: A Multi-lingual Benchmark for Dense Retrieval (2021, background): “, 2021), and ColBERT (Khattab and Zaharia, 2020), but there is much recent work along these lines (Gao et al., 2020; Hofstätter et al., 2020; Hofstätter et al., 2021; Lin et al., 2021b), just to list a few papers.”