[논문 리뷰] EmbedDistill: A Geometric Knowledge Distillation for Information Retrieval

Seungyeon Kim, Ankit Singh Rawat외 7명|arXiv (Cornell University)|2023. 01. 27.

Advanced Image and Video Retrieval Techniques인용 수 4

한 줄 요약

EmbedDistill는 정보 검색을 위한 기하학적 지식 정련 방법을 제안하며, 교사 및 학생 모델 간 쿼리 및 문서 임베딩의 상대적 기하학적 구조를 일치시켜 학생 모델 성능을 향상시킵니다. 임베딩 매칭과 쿼리 생성을 사용하여, MSMARCO 벤치마크에서 1/10 크기의 비대칭 학생 모델로도 교사 성능의 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.

연구 동기 및 목표

• 임베딩 공간 내 기하관계를 활용하여 교사 및 학생 IR 모델 간 일반화 갭을 줄이기.
• 점수 기반 정련을 초월하여 쿼리 및 문서 임베딩을 직접 일치시킴으로써 정련 성능을 향상시키기.
• 프록시 임베딩 구축을 통해 크로스-엔코더 교사에서 듀얼-엔코더 학생으로 효과적인 정련을 가능하게 하기.
• 추론 비용을 증가시키지 않으면서도 높은 성능을 유지하는 비대칭 학생 아키텍처 설계하기.
• 단일 단계 및 다단계 프레임워크를 포함한 다양한 교사 훈련 제도에서의 일반화 성능 입증하기.

제안 방법

• 공유 임베딩 공간에서 교사 및 학생의 쿼리 및 문서 임베딩 간 거리 최소화를 목표로 하는 새로운 정련 목적함수 제안.
• 데이터 다양체를 탐색하고 낮은 데이터 영역의 분포 차이를 줄이기 위해 쿼리 생성 도입.
• 이중 풀링과 재구성 손실을 활용해 크로스-엔코더 교사로부터 의미 있는, 의미적으로 풍부한 쿼리 임베딩을 추출.
• 쿼리 및 문서 엔코더가 서로 다른 아키텍처를 가지는 비대칭 학생 모델 설계로, 추론 비용을 증가시키지 않으면서도 임베딩 일치도 향상시키기.
• 교사 및 학생 모델 간 차원 불일치를 처리하기 위해 호환성 있는 투영 레이어 사용.
• 임베딩 매칭과 교사 점수로부터의 지식 정련을 결합하여 더 강력한 다중 신호 학습 신호 제공하기.

(a) DE to DE distillation

실험 결과

연구 질문

• RQ1 교사 및 학생 모델 간 임베딩 기하학의 직접 일치가 IR에서 일반화 갭을 줄일 수 있는가?
• RQ2 임베딩 매칭 기반 정련이 듀얼-엔코더 및 크로스-엔코더 설정 모두에서 전통적인 점수 기반 정련을 초월하는가?
• RQ3 프록시 쿼리 임베딩을 구성함으로써 크로스-엔코더 교사에서 듀얼-엔코더 학생으로 효과적인 정련을 달성할 수 있는가?
• RQ4 쿼리 생성은 데이터가 적은 영역에서 정련 성능을 어떻게 향상시키는가?
• RQ5 비대칭 학생 아키텍처는 추론 비용을 증가시키지 않으면서도 더 나은 임베딩 일치도를 달성할 수 있는가?

주요 결과

• EmbedDistill은 교사 모델 크기의 1/10인 학생 모델을 사용하여 MSMARCO 벤치마크에서 교사 성능의 95-97%를 달성합니다.
• NQ 개발 세트에서, 16M 파라미터 비대칭 학생을 사용할 경우 Recall@100이 직접 훈련 시 82.0%에서 90.4%로 향상됩니다.
• MSMARCO 재정렬 작업에서, 16M 파라미터 학생을 사용할 경우 MRR@10이 직접 훈련 시 29.7%에서 33.0%로 상승합니다.
• 이 방법은 교사 훈련 제도에 관계없이 일반화 가능하며, 단일 단계 및 다단계로 훈련된 교사 모두에서 일관된 성능 향상을 보입니다.
• 임베딩 매칭은 교사 및 학생 모델 간 상호 쿼리 거리의 격차를 크게 줄이며, 거리 격차 히스토그램 분석을 통해 이를 입증합니다.
• 이중 풀링된 크로스-엔코더 모델은 유사한 쿼리를 함께 묶는 의미 있는 쿼리 임베딩을 생성하는 반면, [CLS]-풀링된 CE 모델은 두 개의 클러스터로 분리되어 기능을 상실합니다.

(b) CE to DE distillation