[Paper Review] Distilling Dense Representations for Ranking using Tightly-Coupled Teachers

Sheng-Chieh Lin, Jheng-Hong Yang+1 more|arXiv (Cornell University)|Oct 22, 2020

Machine Learning and Algorithms24 references56 citations

TL;DR

The paper distills ColBERT’s MaxSim-based ranking into a tight-coupled dot-product bi-encoder, enabling single-pass ANN search and a dense-sparse hybrid retrieval with reduced storage and latency.

ABSTRACT

We present an approach to ranking with dense representations that applies knowledge distillation to improve the recently proposed late-interaction ColBERT model. Specifically, we distill the knowledge from ColBERT's expressive MaxSim operator for computing relevance scores into a simple dot product, thus enabling single-step ANN search. Our key insight is that during distillation, tight coupling between the teacher model and the student model enables more flexible distillation strategies and yields better learned representations. We empirically show that our approach improves query latency and greatly reduces the onerous storage requirements of ColBERT, while only making modest sacrifices in terms of effectiveness. By combining our dense representations with sparse representations derived from document expansion, we are able to approach the effectiveness of a standard cross-encoder reranker using BERT that is orders of magnitude slower.

Motivation & Objective

• Motivate efficient ad hoc retrieval with dense representations as an alternative to costly cross-encoders.
• Propose a tightly-coupled knowledge distillation setup from ColBERT to a dot-product bi-encoder.
• Reduce storage and latency while preserving retrieval effectiveness through distillation.
• Integrate sparse signals with dense representations to approach cross-encoder performance.
• Evaluate end-to-end retrieval effectiveness and latency on MS MARCO and TREC 2019 DL.

Proposed method

• Formulate dense retrieval with a bi-encoder using PoolAvg for embeddings.
• Distill knowledge from ColBERT MaxSim into a dot-product student via a tightly-coupled training loop.
• Use a batch-based soft label distillation with KL divergence between teacher and student distributions.
• Train in two stages: fine-tune teacher with MaxSim, then freeze teacher and distill into a PoolAvg student.
• Combine dense and sparse signals end-to-end via a hybrid scoring rule (Eq. 7).
• Evaluate using MS MARCO dev and TREC 2019 DL with latency measurements.

Experimental results

Research questions

• RQ1 Can tightly-coupled knowledge distillation from ColBERT improve a dot-product bi-encoder sufficiently for single-stage ANN search?
• RQ2 Does combining dense representations with sparse signals yield competitive end-to-end retrieval performance and latency?
• RQ3 How does distillation influence reranking vs retrieval setups in MS MARCO and TREC benchmarks?
• RQ4 What are the practical latency costs of the proposed dense–sparse hybrid approach?

Key findings

• Bi-encoder with Tight Coupled Distillation (TCT-ColBERT) outperforms the non-distilled PoolAvg baseline in MS MARCO and TREC2019 DL.
• Dense representations with distillation plus sparse signals surpass several single-stage dense methods and approach cross-encoder performance when using doc2query-T5.
• Hybrid dense–sparse retrieval yields state-of-the-art recall in both MS MARCO and TREC2019 DL tasks while being substantially faster than multi-stage methods.
• Ablation shows in-batch distillation improves retrieval performance, and triplet-plus-in-batch sampling offers gains over triplet-only distillation.
• End-to-end latency breakdown shows dense retrieval plus small score-combination overhead remains efficient (roughly 3 ms for query encoding and 100 ms for dot-product search on GPU).
• The proposed method is four times faster than ColBERT and roughly thirty times faster than cross-encoder reranking, with competitive effectiveness.