[논문 리뷰] Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

Sebastian Hofstätter, Sheng-Chieh Lin외 3명|arXiv (Cornell University)|2021. 04. 14.

Topic Modeling인용 수 4

한 줄 요약

이 논문은 이중 지도 학습(쌍별 및 인-배치 교사 모델)을 통한 효율적인 밀도 검색기 훈련을 위한 주제 인식형 및 마진 균형 임의 추출 방법인 TAS-Balanced를 제안한다. 단일 소비자용 GPU에서 48시간 이내에 훈련이 가능하며, 두 개의 TREC-DL 데이터셋에서 NDCG@10 기준으로 44% 향상된 BM25보다 뛰어난 성능(1 query당 64ms의 저지연)을 달성한다.

ABSTRACT

A vital step towards the widespread adoption of neural retrieval models is their resource efficiency throughout the training, indexing and query workflows. The neural IR community made great advancements in training effective dual-encoder dense retrieval (DR) models recently. A dense text retrieval model uses a single vector representation per query and passage to score a match, which enables low-latency first stage retrieval with a nearest neighbor search. Increasingly common, training approaches require enormous compute power, as they either conduct negative passage sampling out of a continuously updating refreshing index or require very large batch sizes for in-batch negative sampling. Instead of relying on more compute capability, we introduce an efficient topic-aware query and balanced margin sampling technique, called TAS-Balanced. We cluster queries once before training and sample queries out of a cluster per batch. We train our lightweight 6-layer DR model with a novel dual-teacher supervision that combines pairwise and in-batch negative teachers. Our method is trainable on a single consumer-grade GPU in under 48 hours (as opposed to a common configuration of 8x V100s). We show that our TAS-Balanced training method achieves state-of-the-art low-latency (64ms per query) results on two TREC Deep Learning Track query sets. Evaluated on NDCG@10, we outperform BM25 by 44%, a plainly trained DR by 19%, docT5query by 11%, and the previous best DR model by 5%. Additionally, TAS-Balanced produces the first dense retriever that outperforms every other method on recall at any cutoff on TREC-DL and allows more resource intensive re-ranking models to operate on fewer passages to improve results further.

연구 동기 및 목표

• 일반적으로 대규모 GPU 클러스터나 막대한 배치 크기를 요구하는 효과적인 밀도 검색 모델 훈련의 높은 계산 비용을 해결한다.
• 각 훈련 배치당 정보 수확을 극대화하는 추출 전략을 도입하여 고비용 인프라 의존도를 줄인다.
• 하드웨어 요구 사항을 늘리지 않고도 훈련 효율성과 모델 성능을 향상시켜 강력한 밀도 검색 모델의 광범위한 접근성을 보장한다.
• 1단계 리콜 향상 및 재랭커와의 호환성 향상을 통해 저지연 추론을 가능하게 하여 고성능 검색 파이프라인을 구현한다.
• 다양한 무작위 시드 및 교사 지도 설정에 대해 안정적이고 자원 효율적인 훈련 방법을 설계한다.

제안 방법

• 훈련 이전에 의미적 도트 곱 유사도를 사용해 쿼리를 주제 그룹으로 군집화하여 주제 인식형 배치 구성 가능하게 한다.
• 각 배치에서 하나의 주제 군집에서 쿼리를 추출하여 의미 정보를 집중하고 인-배치 음성 학습을 향상시킨다.
• 쌍별 교사 점수 마진을 균형 조절하여 유사한 마진 값이 있는 문단 쌍을 선택함으로써 훈련 안정성 향상 및 일반화 능력 향상.
• 더 높은 신호 품질을 확보하기 위해 쌍별 연결 BERT(BERT-CAT)와 ColBERT 기반의 인-배치 음성 교사 모델을 모두 활용한 이중 지도 학습을 적용한다.
• 양쪽 교사 신호를 모두 사용하여 Margin-MSE 손실을 적용한 경량 6층 이중 인코더 BERT-DOT 모델을 훈련시켜 순위 정렬 품질 향상.
• 대규모 배치나 지속적으로 업데이트되는 색인 전략이 필요 없도록 단일 소비자용 GPU(11GB VRAM)를 사용해 엔드 투 엔드 훈련을 수행한다.

실험 결과

연구 질문

• RQ1 무작위 추출과 비교해 주제 인식형 추출(TAS)이 훈련 효율성과 검색 성능 향상에 얼마나 효과적인가?
• RQ2 쌍별 교사 점수 마진을 균형 조절하는 것(TAS-Balanced)이 다양한 훈련 설정에서 모델 성능을 추가로 향상시키는가?
• RQ3 쌍별 및 인-배치 음성 교사에서 유도된 이중 지도 학습이 모델의 효과성과 내성에 어떤 영향을 미치는가?
• RQ4 대규모 GPU 클러스터에서 훈련된 모델과 비교해도 TAS-Balanced가 단일 소비자용 GPU에서 뛰어난 성능을 내는 밀도 검색기를 훈련시킬 수 있는가?
• RQ5 특히 저지연 및 중간 지연 조건에서 재랭커와 함께 사용할 경우 TAS-Balanced 검색기가 엔드 투 엔드 검색 파이프라인에서 어떻게 성능을 내는가?

주요 결과

• TAS-Balanced는 두 개의 TREC-DL 데이터셋에서 NDCG@10 기준으로 SOTA 성능을 달성하여, BM25보다 44% 향상되었고, 일반 훈련된 DR 모델보다 19% 향상되었으며, 이전 최고의 DR 모델보다 5% 높은 성능을 기록했다.
• 모델은 1쿼리당 64ms의 저지연 추론 시간을 기록하여 실시간 응용에 적합하다.
• TAS-Balanced는 TREC-DL에서 모든 컷오프 기준으로 모든 베이스라인을 초월하는 첫 번째 밀도 검색기를 생성하여 1단계 리콜을 크게 향상시켰다.
• TAS-Balanced를 docT5query와 융합하면 Recall@1K가 향상되며, 2~6배 높은 지연을 가진 시스템을 능가하는 것으로 나타나 강력한 파이프라인 호환성을 입증했다.
• 단일 소비자용 GPU(11GB VRAM)에서 48시간 이내에 훈련이 가능하여 다중 GPU나 대규모 배치 훈련 설정이 필요 없어졌다.
• Margin-MSE 손실을 사용한 이중 지도 학습이 가장 높은 성능을 내며, 특히 리콜 측면에서 일관된 성과 향상을 보였고, 다양한 무작위 시드 및 교사 설정에서도 유사한 개선 효과를 보였다.