[논문 리뷰] Ultron: An Ultimate Retriever on Corpus with a Model-based Indexer

Yujia Zhou, Jing Yao외 4명|arXiv (Cornell University)|2022. 08. 19.

Topic Modeling인용 수 15

한 줄 요약

Ultron은 문서 지식을 생성형 트랜스포머 모델에 직접 인코딩하는 모델 중심의 문서 검색 시스템을 제안한다. 별도의 색인을 필요로 하지 않으며, 의미적 문서 ID(예: URL/제목 또는 제품 양자화 코드)와 삼단계 훈련 전략을 사용함으로써, 이중 인코더 기반 모델(DPR 등)에 비해 메모리 소비와 지연 시간을 크게 줄이고 최신 기술 수준의 검색 성능을 달성한다.

ABSTRACT

Document retrieval has been extensively studied within the index-retrieve framework for decades, which has withstood the test of time. Unfortunately, such a pipelined framework limits the optimization of the final retrieval quality, because indexing and retrieving are separated stages that can not be jointly optimized in an end-to-end manner. In order to unify these two stages, we explore a model-based indexer for document retrieval. Concretely, we propose Ultron, which encodes the knowledge of all documents into the model and aims to directly retrieve relevant documents end-to-end. For the model-based indexer, how to represent docids and how to train the model are two main issues to be explored. Existing solutions suffer from semantically deficient docids and limited supervised data. To tackle these two problems, first, we devise two types of docids that are richer in semantics and easier for model inference. In addition, we propose a three-stage training workflow to capture more knowledge contained in the corpus and associations between queries and docids. Experiments on two public datasets demonstrate the superiority of Ultron over advanced baselines for document retrieval.

연구 동기 및 목표

• 문서 검색에서 인덱싱과 검색을 분리하는 파이프라인 기반 색인-검색 프레임워크의 한계를 극복하기 위해, 인덱싱과 검색을 통합한 엔드 투 엔드 최적화를 가능하게 한다.
• 모델 학습과 추론 효율성을 향상시키기 위해 더 풍부하고 의미적으로 유의미한 문서 ID를 설계한다.
• 검색 분야에서의 지도 학습 데이터 부족 문제를 해결하기 위해, 일반 데이터와 검색 중심 데이터로 사전 훈련한 후 지도 미세조정을 수행하는 삼단계 훈련 워크플로우를 도입한다.
• 외부 색인 없이 쿼리에서 직접 순서 기반 문서 ID 순위를 생성하는 시퀀스-투-시퀀스 생성 방식을 통해 엔드 투 엔드 문서 검색을 가능하게 한다.

제안 방법

• 쿼리에서 문서 ID 목록으로의 시퀀스-투-시퀀스 생성 작업으로 간주하여, T5 기반의 인코더-디코더 모델을 사용해 문서 검색을 수행한다.
• 키워드 기반(예: URL 및 제목)과 의미 기반(문서 임베딩의 제품 양자화 코드)의 두 가지 유형의 의미적 문서 ID를 도입한다.
• 삼단계 훈련 파이프라인을 적용: 일반적인 사전 훈련(문서 내용과 문서 ID 쌍), 검색 중심 사전 훈련(가짜 쿼리 사용), 그리고 관련성 레이블이 부여된 데이터에 대한 지도 미세조정.
• 추론 시 제약 조건이 부여된 베이즈 서치를 적용하여 모델의 자동 회귀 확률에서 직접 상위-k 문서 ID 순위를 생성한다.
• 문서 ID 시퀀스의 접두사 트리만 저장함으로써, 벡터 기반 색인에 비해 메모리 오버헤드를 크게 감소시킨다.
• 제품 양자화(PQ)를 활용해 유사한 문서들이 공통된 접두사를 공유하도록 컴 pact하고 의미적으로 일관된 식별자를 생성한다.

실험 결과

연구 질문

• RQ1 모델 기반 색인기는 전통적인 인덱스-검색 파이프라인에 비해 문서 검색 정확도에서 뛰어나게 성능을 높일 수 있는가?
• RQ2 다양한 문서 ID 표현 방식(키워드 대비 의미 기반)이 검색 성능과 모델 수렴에 어떤 영향을 미치는가?
• RQ3 제한된 지도 학습 데이터에서 삼단계 훈련 전략이 모델 일반화 능력과 검색 성능 향상에 얼마나 기여하는가?
• RQ4 쿼리에서 문서 ID 순위를 직접 생성하는 엔드 투 엔드 생성 방식이 이중 인코더 기반 방법(DPR 등)에 비해 더 낮은 메모리 소비와 지연 시간을 달성할 수 있는가?

주요 결과

• Ultron-PQ는 MS MARCO에서 MRR@10 0.3155를 기록하여, DPR(0.2908) 및 기타 기준 모델들을 능가했으며, 메모리 사용량은 69.6MB에 불과했다.
• 320K개 문서 코퍼스에서 추론 지연 시간을 DPR의 18.87ms에서 Ultron은 8.90ms로 줄여, 더 뛰어난 효율성을 입증했다.
• Ultron은 DPR의 벡터 기반 색인 대비 메모리 비용을 90% 감소시켰다. 이는 문서 ID의 접두사 트리만 저장하기 때문이다.
• 삼단계 훈련 전략은 성능 향상에 크게 기여했으며, 특히 지도 미세조정 단계에서 최종 성능 향상의 대부분 기여를 했다.
• 의미 기반 문서 ID(PQ)는 공통된 접두사를 공유하는 문서들 간에 더 높은 유사도 일관성을 보였으며, 이는 모델 수렴에 기여했다.
• 제거 실험(ablation study) 결과, 지도 미세조정을 제거할 경우 검색 성능에 심각한 하락이 발생함을 확인하여, 이 단계의 필수성을 입증했다.