[ICLR 2023] PROMPTGATOR: Few-shot dense retrieval from 8 example

1. Introduction

• dual encoder 기반의 모델을 통해, 딥러닝 기반의 검색 모델이 발달
• 하지만, 검색기가 특정 도메인에서 학습이 되었더라도, 다른 도메인에서는 좋은 성능을 보이지 않을 수 있음
• BEIR 벤치마크는 이러한 문제점을 지적하며 학습 데이터가 없는 상황(zero-shot)에서 검색 성능을 측정하는 벤치마크
• 본 논문에서는 학습 데이터가 적은 상황에서 학습 데이터를 augmentation 하기 위해 LLM을 사용
• LLM을 통해 생성된 query-document pair에 대한 검증을 위해 round-trip consistency 적용

 

2. Few-shot Retrieval Task

2.1 Retrieval Task

• Retrieval은 각 도메인이 가지고 있는 documents D, 사용자 질의인 query Q, 도데인의 search intent L로 정의
• query가 주어졌을 때, 가장 연관성이 있는 document를 찾는 것이 목표
• 연관성이 있다는 것의 기준이 각 도메인의 search intent
• 이 논문에서는 특히 document는 있지만, document와 짝을 이루는 query가 없거나 적은 상황을 가정

2.2 Few-shot BEIR Setting

• BEIR는 zero-shot setting에서의 검색기의 성능을 측정하기 위한 벤치마크로 9개의 도메인이 있는 18개의 데이터셋
• 본 논문에서는 BEIR의 evaluation 데이터의 document-query 쌍을 기반으로 BEIR few-shot setting을 구축

 

3. PROMPTGATOR

3.1 Prompt-Base Query Generation

• document는 있지만, document와 쌍을 이루는 query가 충분하지 않아 query를 생성하는 과정
• LLM에 query-document description과 query-document example을 프롬프트로 입력 

• e(d), e(q)는 document와 query에 대한 description
• 예를 들어, ArguAna 데이터 셋에서 query가 논증이고, document가 query와 반대되는 반대 논증이라면, e(d) = "Argument: {d}", e(q) = "Counter Argument: {q}"
• k개의 document-query example이 입력되고, query를 생성하고 싶은 document가 마지막에 들어감
• LLM은 FLAN을 사용

 

3.2 Consistency Filtering using only Generated Data

• Round-Trip Consistency를 이용하여 유효한 document-query 쌍을 검증

1. 생성된 document-query 쌍을 통해 초기 검색 모델 학습
2. 학습된 초기 검색 모델을 통해 생성된 query를 넣었을 때, 해당 document가 top-k안에 있을때만 유효하다고 판단

 

3.3 Few-shot PROMPTGATOR Retriever

• T5 모델의 Transformer encoder 부분을 사용하여 초기화
• Contriever의 C4 방식으로 만들어진 데이터셋 사용하여 pre-training
• 마지막으로 생성/필터링이 끝난 document-query 쌍을 통해 duel-encoder 학습 (PROMPTGATOR)
• reranker로 사용하기 위한 cross-encoder 모델도 학습 (PROMPTGATOR++)

 

3.4 Zero-shot PROMPTGATOR Retriever

• query 생성시 example 없이 생성하여 document-query 쌍 생성
• LLM에 prompt로 "{d} Read the passage and generate a query" 입력

 

4. Experiments

4.1 Main Results

• zero-shot PROMPTGATOR만 사용해도 좋았지만, few-shot PROMPTGAOTR을 사용하면 더 좋은 성능
• Reranker로 PROMPTGATOR++을 사용하는 경우 심지어 더 향상
• 향상 폭이 컸던 touche와 ArguAna 데이터셋의 경우, 전통적인 retrieval과는 다른 search intent를 가지고 있었음

4.2 Ablation Study

• Round Trip Filtering을 적용한 경우, 대부분의 데이터셋에서 성능 향상
• 성능이 떨어지는 데이터셋의 경우 document의 수 자체가 적어 학습되는 초기 검색기가 overfitting이 일어남
• Round Trip Filtering을 통해 너무 general한 query, hallucination을 포함한 query를 효과적으로 제거할 수 있었음

• query를 생성하는 LLM이었던 FLAN의 학습 데이터가 이미 QA 데이터인 NQ와 Quora 데이터셋을 포함하고 있음
• 만약 NQ와 Quora 없이 학습된 FLAN을 이용한다면 어떨까에 대한 실험
• NQ와 Quora를 학습 데이터에서 제외하더라도 PROMPTGATOR는 잘 작동