Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

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