Source-linked AI summary
Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling
Sebastian Hofstätter, Sheng-Chieh Lin, Jheng-Hong Yang, Jimmy Lin, Allan Hanbury
TL;DR
Dense retrieval training can demand substantial compute despite its low-latency inference benefits. TAS-Balanced addresses this with topic-aware query and balanced-margin passage sampling plus dual-teacher supervision, achieving strong TREC-DL results with modest hardware requirements and supporting more efficient reranking pipelines.
Problem
Dense retrieval models offer low-latency first-stage retrieval, but improving their quality can require substantial hardware resources for training and indexing.
Method
TAS-Balanced clusters queries by topic, balances pairwise teacher-score margins when selecting passage pairs, and trains BERTDOT with pairwise and in-batch teacher supervision.
Results
TAS-Balanced achieves state-of-the-art low-latency results on two TREC-DL query sets and improves NDCG@10 over BM25 by 44%, a plainly trained DR by 19%, docT5query by 11%, and the previous best DR model by 5%.
Takeaways & Limitations
TAS-Balanced is the first dense retriever reported to outperform every other method at every recall cutoff on densely judged TREC-DL queries and can support reranking fewer passages.
Takeaways & Limitations
Retraining the mono-duo-T5 rerankers on the dense retriever's shifted candidate distribution is outside this work's scope.
Abstract
from arXiv · showhide
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.
KEYWORDS
The paper concerns dense retrieval, knowledge distillation, and batch sampling, and illustrates Topic Aware Sampling with clustered queries.
- The paper's keywords are Dense Retrieval, Knowledge Distillation, and Batch Sampling.
- The work was published at SIGIR 2021.
- Figure 1 visualizes eight randomly sampled topic clusters and example queries for Topic Aware Sampling.
1 INTRODUCTION
The introduction frames resource-efficient dense retrieval as a central adoption challenge and proposes TAS-Balanced with dual-teacher supervision to improve training batches and retrieval quality. It evaluates effectiveness, robustness, comparative performance, and use in a larger retrieval pipeline.
- Motivation: Dense retrieval offers low-latency retrieval but can require substantial hardware resources for training and indexing.
- Method: The method combines pairwise and in-batch negative knowledge-distillation signals through dual supervision.
- Method: TAS-Balanced combines topic-clustered query batches with passage-pair sampling that balances teacher score margins.
- Method: Selecting queries from one topic cluster concentrates in-batch information that random query batches largely fail to provide.
- Efficiency: TAS-Balanced can train an effective BERTDOT model on one consumer-grade GPU in under 48 hours rather than a common eight-V100 configuration.
- Results: TAS-Balanced improves pairwise training, in-batch negatives, and dual-supervision training, with Margin-MSE consistently outperforming other tested losses.
- Robustness: Across five randomization instances, metric variation remained below .01 nDCG on TREC-DL and .001 MRR on MSMARCO-DEV.
- Results: On two TREC-DL query sets, BERTDOT improves NDCG@10 over BM25 by 44%, plainly trained DR by 19%, docT5query by 11%, and the previous best DR by 5%.
2 RETRIEVAL MODEL BACKGROUND
The paper contrasts BERT-based teacher architectures with BERTDOT dense retrieval, whose independent encoding and dot-product scoring support indexed, low-latency retrieval. It then motivates dual-teacher supervision as a compromise between BERTCAT’s effectiveness and ColBERT’s efficiency.
- BERT Teacher Models: BERTCAT concatenates query and passage sequences, pools the BERT CLS representation, and scores the pair with a linear layer.
- BERT Teacher Models: BERTCAT is effective but requires candidate selection and cannot pre-compute indexed passage representations, making it slow in practice.
- BERT Teacher Models: ColBERT delays query-document interactions until after encoding and aggregates them with per-query-term max pooling followed by summation.
- BERT Teacher Models: ColBERT can theoretically index passage representations, but storage scales with collection terms; therefore, the paper uses it as an in-batch-negative teacher rather than a dense retriever.
- BERTDOT Dense Retrieval Model: BERTDOT independently encodes queries and passages into single vectors, then uses dot-product scoring for nearest-neighbor retrieval after passages are indexed.
- BERTDOT Dense Retrieval Model: 64ms is the measured single-query latency for Top-1000 retrieval with DistilBERT and a brute-force Faiss FlatIP index on one TITAN RTX GPU.
3 TOPIC AWARE SAMPLING
TAS-Balanced improves dense-retriever batch composition by grouping queries into topical clusters and balancing passage-pair teacher margins, addressing the weak information gain of random in-batch negatives.
- Motivation: Random batches draw queries from unrelated topics, limiting the information gained from in-batch negative interactions.In-batch negatives reuse already computed representations, but random query combinations offer little additional training information.
- TAS: Topic Aware Sampling: TAS clusters training queries once with k-means using baseline representation vectors, then samples queries from selected clusters for each batch.The procedure is more efficient than clustering passages because there are fewer queries and they are faster to encode.
- TAS-Balanced: TAS-Balanced filters passage pairs into uniformly covering margin ranges between each query’s minimum and maximum teacher-model margins.The method targets the skew caused by queries having fewer relevant than non-relevant passages and defines distant negatives as easy.
- TAS-Balanced: Sampling a margin range first and then a filtered passage pair unskews the passage-pair distribution while retaining topic-aware batches.Together, topic clustering and margin balancing form the TAS-Balanced batch-sampling strategy.
- Training procedure: Batch composition runs concurrently in a subprocess, continuously generates new batches, and uses early stopping instead of a fixed epoch count.The same batch is not repeated across multiple epochs.
4 EXPERIMENT DESIGN
The experiments evaluate dense retrieval on MSMARCO and TREC-DL query sets using lightweight DistilBERT-initialized models, specified clustering and margin settings, and an approximated early-stopping set.
- Infrastructure: The main training and inference stack uses PyTorch, HuggingFace Transformers, and Faiss for query clustering and brute-force nearest-neighbor retrieval.
- Datasets: MSMARCO-Passage contains 8.8 million passages, while evaluation uses 6,980 MSMARCO-DEV queries and TREC-DL ’19 and ’20 sets of 43 and 54 queries.TREC graded relevance is binarized at 2 for MRR, MAP, and recall.
- Configuration: TAS clustering creates 2K clusters from 400K training queries, samples n=1 cluster with batch size b=32, and balances margins into 10 bins.The dual-teacher combination hyperparameter is set to α=0.75.
- Model initialization: All BERTDOT and ColBERT instances start from a fresh 6-layer DistilBERT checkpoint rather than retrieval-trained generational checkpoints.ColBERT is trained with pairwise teacher signals and implemented dynamically for in-batch teaching.
- Training control: An approximated early-stopping set indexes the top 100 passages for 3,200 uniformly sampled DEV-49K queries using a pairwise-trained baseline.The stopping procedure is designed to compare methods without disadvantaging models that take longer to train.
5 RESULTS
TAS-Balanced and dual-teacher supervision consistently improve dense-retrieval effectiveness, while maintaining strong low-latency performance across TREC-DL and MSMARCO-DEV. The method also supports effective first-stage retrieval and multi-stage pipelines with limited latency overhead.
- Source of Effectiveness: Margin-MSE generally outperforms list-based losses, with especially noticeable recall gains on both TREC-DL query sets.It encourages the model to follow teacher score distributions rather than only the ordering of in-batch negatives.
- Source of Effectiveness: TAS-Balanced improves pairwise, in-batch-negative, and dual-supervision training, producing the strongest overall configuration across three query sets.TAS alone provides limited gains for pairwise supervision, whereas balanced margin sampling improves most pairwise-teaching metrics by at least 1 percentage point.
- Source of Effectiveness: Dual-teacher supervision substantially improves recall over a single in-batch teacher on all three query sets, while nDCG and MRR improve on two and tie on one.The authors therefore recommend combining dual supervision with TAS-Balanced sampling.
- Comparing to Baselines: TAS-Balanced models outperform all other dense-retrieval methods on both TREC-DL sets, improving nDCG@10 by at least 4%, MRR@10 by 3%, and Recall@1K by at least 9%.The evaluation uses a lightweight 6-layer encoder with batch sizes of 32, 96, and 256; larger batches show a clear improvement only on MSMARCO-DEV.
- Comparing to Baselines: TAS-Balanced outperforms BM25 by 44% on nDCG@10 and by 9-14% on Recall@1K on TREC’19 and TREC’20 at comparable query latency.On MSMARCO-DEV it also outperforms other baselines, especially on Recall@1K, while RocketQA’s 4,000-sample batch remains stronger than its 128-sample setting.
- TAS-Balanced Retrieval in a Pipeline: Fusing TAS-Balanced with docT5query yields the highest Recall@1K across every query set at virtually no latency cost beyond merging result lists.The fused low-latency system often matches or exceeds medium-latency methods that are 2-6x slower.
6 CONCLUSION
TAS-Balanced improves dense retrieval training through cost-neutral topic-aware and balanced-margin sampling with dual teacher supervision, while remaining feasible on modest hardware. It performs consistently across random orderings, integrates effectively into larger search pipelines, and benefits shallow reranking.
- TAS-Balanced combines topic-aware query sampling with balanced passage-pair margins and pairwise plus in-batch teacher supervision.
- Under 48 hours of training on a single consumer-grade GPU enables strong dense retrieval without large compute servers.
- TAS-Balanced remains consistent across different random orderings and teacher supervisions.
- Fusing TAS-Balanced with docT5query outperforms many systems with 2-6x higher latency.
- TAS-Balanced also improves retrieval pipelines operating at low reranking depths.