In-Batch Negatives for Knowledge Distillation with Tightly-Coupled Teachers for Dense Retrieval (2021)

Abstract: We present an efficient training approach to text retrieval with dense representations that applies knowledge distillation using the ColBERT late-interaction ranking model. Specifically, we propose to transfer the knowledge from a bi-encoder teacher to a student by distilling knowledge from ColBERT’s expressive MaxSim operator into a simple dot product. The advantage of the bi-encoder teacher–student setup is that we can efficiently add in-batch negatives during knowledge distillation, enabling richer interactions between teacher and student models. In addition, using ColBERT as the teacher reduces training cost compared to a full cross-encoder. Experiments on the MS MARCO passage and document ranking tasks and data from the TREC 2019 Deep Learning Track demonstrate that our approach helps models learn robust representations for dense retrieval effectively and efficiently.

## introduction for well over half a century, solutions to the ad hoc retrieval problem-where the system's task is return a list of top k texts from an arbitrarily large corpus d that maximizes some metric of quality such as average precision or ndcg-has been dominated by sparse vector representations, for example, bag-of-words bm25. even in modern multi-stage ranking architectures, which take advantage of large pretrained transformers such as bert @xcite , the models are deployed as rerankers over initial candidates retrieved based on sparse vector representations; this is sometimes called "first-stage retrieval". one well-known example of this design is the bert-based reranker of @xcite ; see @xcite for a recent survey. the standard reranker architecture, while effective, exhibits high query latency, on the order of seconds per query @xcite because expensive neural inference must be applied at query time on query-passage pairs. this design is known as a cross-encoder @xcite , which exploits query-passage attention interactions across all transformer layers. as an alternative, a biencoder design provides an approach to ranking with dense representations that is far more efficient than cross-encoders @xcite @xcite @xcite @xcite . prior to retrieval, the vector representations can be precomputed for each of the texts in a corpus. when retrieving texts in response to a given query, computationally expensive transformer inference is replaced by much faster approximate nearest neighbor (ann) search @xcite . recently, researchers have proposed bi-encoders that produce multiple vectors to represent a query (or a passage) @xcite @xcite , which have proven to be effective both theoretically and empirically. however, the main disadvantage of these designs is their high storage requirements. for example, colbert @xcite ) requires storing all the wordpiece token vectors of each text (passage) in the corpus. on the ms marco passage corpus comprising 8.8m passages, for example, this requires 154 gib. of course, a common alternative is to produce single vectors for queries and passages @xcite . although this design is less storage-demanding, it sacrifices ranking effectiveness since its structure breaks rich interactions between queries and passages compared to multi-vector bi-encoders or cross-encoders. hence, improving the effectiveness of single-vector biencoders represents an important problem. one approach to improving the effectiveness of single-vector bi-encoders is hard negative mining, by training with carefully selected negative examples that emphasize discrimination between relevant and non-relevant texts. there are several approaches to accomplish this. @xcite and @xcite leverage large in-batch negatives to enrich training signals. @xcite and @xcite propose to mine hard negatives using the trained bi-encoder itself. by searching for global negative samples from an asynchronously updated ann index, the bi-encoder can learn information not present in the training data produced by sparse representations @xcite . however, both large in-batch negative sampling and asynchronous ann index updates are computationally demanding. the later is especially impractical for large corpora since it requires periodic inference over all texts in the corpus to ensure that the best negative examples are retrieved. there is also work that explores knowledge distillation (kd) @xcite to enhance retrieval effectiveness and efficiency. most related to our study is @xcite , who demonstrate that kd using a cross-encoder teacher significantly improves the effectiveness of bi-encoders for dense retrieval. similarly, @xcite investigate the effectiveness of distilling a trained cross-encoder into a bi-encoder for sentence similarity tasks. @xcite explore kd combinations of different objectives such as language modeling and ranking. however, the above papers use computationally expensive cross-encoder teacher models; thus, combining them for kd with more advanced negative sampling techniques can be impractical. in light of existing work on hard negative mining and knowledge distillation, we propose to improve the effectiveness of single-vector bi-encoders with a more efficient kd approach: in-batch kd using a bi-encoder teacher. the advantage of our design is that, during distillation, it enables the efficient exploitation of all possible query-passage pairs within a minibatch, which we call tight coupling (illustrated in figure 1 ). this is a key difference between our kd approach and previous methods for dense retrieval, where only the scores of given query-passage triplets (not all combinations) are bi-encoder ## background we focus on improving the training efficiency and retrieval effectiveness of dense retrieval and begin by formalizing it as a dense representation learning problem. to be more specific, we propose to use knowledge distillation to enrich training signals and stabilize the representation learning procedure of bi-encoder models in the context of the well-known noise-contrastive estimation (nce) framework. ## experiments in this section, we conduct experiments on the ms marco passage and document corpora. for passage ranking, we first train models on bm25 negatives as warm-up and compare different kd methods. we then further train models on the hard negatives retrieved by the bm25 warmed-up checkpoint. for document ranking, following previous work @xcite @xcite , we start with our bm25 warmed-up checkpoint for passage ranking and conduct additional hard negative training. ## conclusions improving the effectiveness of single-vector biencoders is an important research direction in dense retrieval because of lower latency and storage requirements compared to multi-vector approaches. we propose a teacher-student knowledge distillation approach using tightly coupled bi-encoders that enables exhaustive use of query-passage combinations in each minibatch. more importantly, a bi-encoder teacher requires less computation than a cross-encoder teacher. finally, our approach leads to robust learned representations. overall, our hard negative sampling strategy leads to an effective and efficient dense retrieval technique, which can be further combined with sparse retrieval techniques in dense-sparse hybrids. together, these designs provide a promising solution for end-to-end text retrieval that balances quality, query latency, and storage requirements.