[Paper Review] EmbedDistill: A Geometric Knowledge Distillation for Information Retrieval

Seungyeon Kim, Ankit Singh Rawat+7 more|arXiv (Cornell University)|Jan 27, 2023

Advanced Image and Video Retrieval Techniques4 citations

TL;DR

EmbedDistill proposes a geometric knowledge distillation method for information retrieval that improves student model performance by aligning the relative geometry of query and document embeddings between teacher and student models, using embedding matching and query generation. It achieves 95-97% of teacher performance with 1/10th-sized asymmetric student models on MSMARCO benchmarks.

ABSTRACT

Large neural models (such as Transformers) achieve state-of-the-art performance for information retrieval (IR). In this paper, we aim to improve distillation methods that pave the way for the resource-efficient deployment of such models in practice. Inspired by our theoretical analysis of the teacher-student generalization gap for IR models, we propose a novel distillation approach that leverages the relative geometry among queries and documents learned by the large teacher model. Unlike existing teacher score-based distillation methods, our proposed approach employs embedding matching tasks to provide a stronger signal to align the representations of the teacher and student models. In addition, it utilizes query generation to explore the data manifold to reduce the discrepancies between the student and the teacher where training data is sparse. Furthermore, our analysis also motivates novel asymmetric architectures for student models which realizes better embedding alignment without increasing online inference cost. On standard benchmarks like MSMARCO, we show that our approach successfully distills from both dual-encoder (DE) and cross-encoder (CE) teacher models to 1/10th size asymmetric students that can retain 95-97% of the teacher performance.

Motivation & Objective

• To reduce the generalization gap between teacher and student IR models by leveraging geometric relationships in embedding space.
• To improve distillation beyond score-based distillation by directly aligning query and document embeddings.
• To enable effective distillation from cross-encoder teachers to dual-encoder students via proxy embedding construction.
• To design asymmetric student architectures that maintain high performance without increasing inference cost.
• To demonstrate generalization across different teacher training regimes, including single-stage and multi-stage frameworks.

Proposed method

• Proposes a novel distillation objective based on minimizing the distance between teacher and student query and document embeddings in the shared embedding space.
• Introduces query generation to explore the data manifold and reduce distributional discrepancies in low-data regions.
• Employs dual-pooling and reconstruction losses to extract meaningful, semantically rich query embeddings from cross-encoder teachers.
• Designs asymmetric student models where query and document encoders have different architectures to improve embedding alignment without increasing online inference cost.
• Uses compatible projection layers to handle dimension mismatches between teacher and student models.
• Combines embedding matching with knowledge distillation from teacher scores to provide a stronger, multi-signal training signal.

(a) DE to DE distillation

Experimental results

Research questions

• RQ1 Can direct alignment of embedding geometry between teacher and student models reduce the generalization gap in IR?
• RQ2 Does embedding matching-based distillation outperform traditional score-based distillation in both dual-encoder and cross-encoder settings?
• RQ3 Can effective distillation be achieved from cross-encoder teachers to dual-encoder students by constructing proxy query embeddings?
• RQ4 How does query generation improve distillation performance in data-sparse regions?
• RQ5 Can asymmetric student architectures achieve better embedding alignment without increasing online inference cost?

Key findings

• EmbedDistill achieves 95-97% of the teacher model's performance on MSMARCO benchmarks using student models that are 1/10th the size of the teacher.
• On the NQ dev set, the method improves Recall@100 from 82.0% (direct training) to 90.4% with a 16M-parameter asymmetric student.
• For the MSMARCO re-ranking task, the method boosts MRR@10 from 29.7% (direct training) to 33.0% with a 16M-parameter student.
• The method generalizes across teacher training regimes, showing consistent gains on both single-stage and multi-stage trained teachers.
• Embedding matching significantly reduces the discrepancy in inter-query distances between teacher and student models, as shown in the histogram analysis of distance discrepancies.
• Dual-pooled cross-encoder models produce semantically meaningful query embeddings that group similar queries together, unlike [CLS]-pooled CE models which degenerate into two clusters.

(b) CE to DE distillation