[논문 리뷰] Pre-training Tasks for Embedding-based Large-scale Retrieval

Wei-Cheng Chang|arXiv (Cornell University)|2020. 02. 10.

Topic Modeling인용 수 101

한 줄 요약

이 논문은 두-타워 Transformer 검색 모델의 사전 학습 과제를 분석하고, 단락 수준의 과제(Inverse Cloze Task, Body First Selection, Wiki Link Prediction)가 BM-25 대비 검색 성능을 크게 향상시키는 반면 토큰 수준 MLM은 이득이 제한적임을 보인다.

ABSTRACT

We consider the large-scale query-document retrieval problem: given a query (e.g., a question), return the set of relevant documents (e.g., paragraphs containing the answer) from a large document corpus. This problem is often solved in two steps. The retrieval phase first reduces the solution space, returning a subset of candidate documents. The scoring phase then re-ranks the documents. Critically, the retrieval algorithm not only desires high recall but also requires to be highly efficient, returning candidates in time sublinear to the number of documents. Unlike the scoring phase witnessing significant advances recently due to the BERT-style pre-training tasks on cross-attention models, the retrieval phase remains less well studied. Most previous works rely on classic Information Retrieval (IR) methods such as BM-25 (token matching + TF-IDF weights). These models only accept sparse handcrafted features and can not be optimized for different downstream tasks of interest. In this paper, we conduct a comprehensive study on the embedding-based retrieval models. We show that the key ingredient of learning a strong embedding-based Transformer model is the set of pre-training tasks. With adequately designed paragraph-level pre-training tasks, the Transformer models can remarkably improve over the widely-used BM-25 as well as embedding models without Transformers. The paragraph-level pre-training tasks we studied are Inverse Cloze Task (ICT), Body First Selection (BFS), Wiki Link Prediction (WLP), and the combination of all three.

연구 동기 및 목표

• 대규모 질의-문서 검색 문제와 2단계 시스템에서 효율적인 검색의 필요성을 동기화한다.
• 사전 학습 과제가 두-타워 Transformer 검색기 성능에 미치는 영향을 조사한다.
• 단락 수준 사전 학습 과제 ICT, BFS, WLP 및 그 조합을 토큰 수준 MLM 및 BM-25와 비교 평가한다.
• 적절히 설계된 사전 학습 과제가 두-타워 모델이 BM-25 및 BoW 기초선보다 검색 설정에서 더 나은 성능을 발휘하도록 함을 보여준다.

제안 방법

• 쿼리 인코더와 문서 인코더가 Transformer 아키텍처를 통해 임베딩을 생성하는 두-타워 검색 모델을 정의한다.
• 전체 소프트맥스를 근사하기 위해 Sampled Softmax를 사용한 후보 문서에 대한 Softmax로 학습한다.
• 세 가지 단락 수준 사전 학습 과제 ICT(Inverse Cloze Task), BFS(Body First Selection), WLP(Wiki Link Prediction) 및 ICT+BFS+WLP의 조합을 토큰 수준 MLM과 비교하여 제안하고 평가한다.
• 사전 학습 데이터를 위키피디아에서 추출해 양의 (q, d) 쌍을 구성하고 다운스트림 검색 QA 데이터셋(SQuAD, Natural Questions) 및 오픈 도메인 설정에서 미세 조정한다.
• 효과를 평가하기 위해 BM-25 및 BoW-MLP 기초선과 비교한다.
• 실험 설정은 512차원 임베딩을 갖는 두-타워 인코더, 64토큰 질의, 288토큰 문서, 32 TPU v3에서의 100K 사전 학습 스텝을 사용한다.

실험 결과

연구 질문

• RQ1 대규모 검색에서 다양한 사전 학습 과제가 두-타워 Transformer 검색 모델의 효과성에 어떤 영향을 미치는가?
• RQ2 단락 수준 사전 학습 과제가 BM-25와 같은 토큰 수준 MLM 및 전통적인 IR 기초선보다 검색 작업에서 성능을 향상시키는가?
• RQ3 ICT, BFS, WLP를 결합하면 특히 데이터가 적거나 오픈 도메인 설정에서 추가 이득이 생기는가?
• RQ4 모델의 깊이와 임베딩 차원이 사전 학습 과제와 상호 작용하여 검색 재현(RRecall)에 어떤 영향을 미치는가?

주요 결과

Train/Test Ratio	Encoder	Pre-training task	R@1	R@5	R@10	R@50	R@100
1%/99%	BM-25	No Pretraining	4.99	11.91	15.41	24.00	27.97
1%/99%	BoW-MLP	No Pretraining	0.14	0.35	0.49	1.13	1.72
1%/99%	BoW-MLP	ICT+BFS+WLP	22.55	41.03	49.93	69.70	77.01
1%/99%	Transformer	No Pretraining	0.02	0.06	0.08	0.31	0.54
1%/99%	Transformer	MLM	0.18	0.51	0.82	2.46	3.93
1%/99%	Transformer	ICT+BFS+WLP	37.43	61.48	70.18	85.37	89.85
5%/95%	BM-25	No Pretraining	41.87	57.98	63.63	74.17	77.91
5%/95%	BoW-MLP	ICT+BFS+WLP	26.23	46.49	55.68	75.28	81.89
5%/95%	Transformer	ICT+BFS+WLP	45.90	70.89	78.47	90.49	93.64
80%/20%	BM-25	No Pretraining	41.77	57.95	63.55	73.94	77.49
80%/20%	BoW-MLP	ICT+BFS+WLP	32.24	55.26	65.49	83.37	88.50
80%/20%	Transformer	ICT+BFS+WLP	58.35	82.76	88.44	95.87	97.49

• 적절히 설계된 단락 수준 사전 학습 과제를 가진 두-타워 Transformer 모델은 검색 작업에서 BM-25 및 BoW 기초선보다 크게 우수하다.
• 단락 수준 사전 학습 ICT, BFS, WLP는 상당한 이득을 주지만 토큰 수준 MLM은 지엽적 개선에 그친다.
• ICT+BFS+WLP의 조합은 SQuAD 및 Natural Questions에서 단일 과제보다 일관되게 우수하며 특히 저자원 데이터 및 오픈 도메인 설정에서 이익이 크다.
• Transformer 인코더는 BoW-MLP 인코더보다 단락 수준 사전 학습으로부터 더 큰 이점을 얻고, 더 큰 임베딩 차원이 성능을 향상시킨다.
• 오픈 도메인 검색 실험에서 ICT+BFS+WLP 및 ICT가 큰 후보 집합에서도 견고한 이득을 제공한다.