[论文解读] 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 Techniques被引用 4

一句话总结

EmbedDistill 提出了一种用于信息检索的几何知识蒸馏方法，通过在教师模型与学生模型之间对查询和文档嵌入的相对几何结构进行对齐，提升了学生模型的性能，方法结合了嵌入匹配与查询生成。在 MSMARCO 基准测试中，使用仅为教师模型十分之一大小的非对称学生模型，该方法实现了教师模型 95%-97% 的性能表现。

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.

研究动机与目标

• 通过利用嵌入空间中的几何关系，减少教师模型与学生模型在信息检索任务中的泛化差距。
• 通过直接对齐查询和文档嵌入，超越基于得分的知识蒸馏，提升蒸馏效果。
• 通过构建代理嵌入，实现从交叉编码器教师模型到双编码器学生模型的有效蒸馏。
• 设计非对称学生模型架构，在不增加推理成本的前提下保持高性能。
• 在不同教师训练范式（包括单阶段与多阶段框架）下实现泛化能力。

提出的方法

• 提出一种新颖的蒸馏目标，通过最小化教师模型与学生模型在共享嵌入空间中查询和文档嵌入之间的距离，实现对齐。
• 引入查询生成以探索数据流形，减少低数据区域中的分布差异。
• 采用双池化（dual-pooling）与重建损失，从交叉编码器教师模型中提取有意义且语义丰富的查询嵌入。
• 设计非对称学生模型，使查询编码器与文档编码器采用不同架构，在不增加在线推理成本的前提下提升嵌入对齐效果。
• 使用兼容的投影层处理教师模型与学生模型之间的维度不匹配问题。
• 结合嵌入匹配与教师模型得分的知识蒸馏，提供更强的多信号训练信号。

(a) DE to DE distillation

实验结果

研究问题

• RQ1 教师模型与学生模型之间的嵌入几何结构的直接对齐，是否能减少信息检索任务中的泛化差距？
• RQ2 基于嵌入匹配的蒸馏方法是否在双编码器与交叉编码器设置下均优于传统的基于得分的蒸馏？
• RQ3 通过构建代理查询嵌入，能否实现从交叉编码器教师模型到双编码器学生模型的有效蒸馏？
• RQ4 查询生成在数据稀疏区域中如何提升蒸馏性能？
• RQ5 非对称学生模型架构是否能在不增加在线推理成本的前提下实现更好的嵌入对齐？

主要发现

• EmbedDistill 在 MSMARCO 基准测试中，使用仅为教师模型十分之一大小的学生模型，实现了教师模型 95%-97% 的性能表现。
• 在 NQ dev 数据集上，该方法将 Recall@100 从直接训练的 82.0% 提升至 90.4%（使用 1600 万参数的非对称学生模型）。
• 在 MSMARCO 重排序任务中，该方法将 MRR@10 从直接训练的 29.7% 提升至 33.0%（使用 1600 万参数的学生模型）。
• 该方法在不同教师训练范式下均表现出泛化能力，在单阶段与多阶段训练的教师模型上均取得一致性能提升。
• 嵌入匹配显著减少了教师模型与学生模型之间查询内距离的差异，如距离差异的直方图分析所示。
• 双池化交叉编码器模型生成的查询嵌入具有语义意义，能将相似查询分组聚集，而 [CLS]-池化的交叉编码器模型则退化为两个簇。

(b) CE to DE distillation