Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

TL;DR

This paper introduces TAS-Balanced, a resource-efficient training method for dense retrieval models that achieves state-of-the-art low-latency results using only a single GPU, significantly reducing training costs while improving retrieval performance.

Contribution

The paper presents a novel topic-aware sampling technique and dual-teacher supervision for training dense retrievers efficiently on limited hardware, outperforming existing methods.

Findings

Achieves state-of-the-art low-latency retrieval results

Outperforms BM25 and previous dense models on TREC-DL datasets

Operates effectively on a single consumer-grade GPU

Abstract

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.