논문 리뷰 | NLP | Query2doc: Query Expansion with Large Language Models

지난 포스트에서 읽어본 Query Expansion by Prompting LLMs 논문에 이어 Query Expansion 관련 논문을 다시 읽어보았습니다. 이 논문에서는 LLM을 활용하여 쿼리를 전체 문서 형식으로 변환하는 방법론인 Query2doc을 다루고 있는데, Few-shot 프롬프트 실험으로만 이루어져 있다는 차이점이 있습니다.

0. Abstract

• 본 논문에서는 Query2doc이라는 간단하면서도 효과적인 쿼리 확장(Query Expansion) 기법을 제안
  • Few-shot prompting을 활용한 LLM 기반 Pseudo-document 생성 후, 이를 사용하여 쿼리를 확장
• LLM은 웹 스케일 텍스트 데이터로 학습되어 있으며, 방대한 지식을 기억하고 있음 → 이를 활용해 생성된 Pseudo-document는 쿼리의 의미를 명확히 하고 검색 성능을 향상
• 실험 결과, Query2doc은 MS MARCO, TREC DL 등 Ad-hoc IR 데이터셋에서 BM25 성능을 3~15% 향상
  • 별도의 모델 파인튜닝 없이도 성능 개선 가능
• 또한, 최신 Dense Retriever에도 긍정적인 영향을 미치며, In-domain 및 Out-of-domain 검색 성능 모두 향상됨

1. Introduction

• Information Retrieval (IR)
  • 사용자의 쿼리에 대해 관련 문서를 대규모 코퍼스에서 찾아주는 기술로, 현대 검색 엔진의 핵심 구성 요소
  • IR에는 두 가지 주요 패러다임이 존재
    1. Sparse retrieval (예: BM25) – 단어 기반 검색
       • 도메인 외 데이터(OOD)에서 경쟁력 유지
    2. Dense retrieval – 임베딩을 활용한 검색
       • 라벨링된 데이터가 많을 때 우수한 성능
• Query Expansion (QE)
  • 기존 쿼리를 보강하여 검색 성능을 높이는 오래된 기술
  • Sparse retrieval에서는 쿼리와 문서 간의 어휘적 차이를 줄이는 역할 수행
  • 그러나 기존 QE 기법 (예: RM3)은 한계가 있으며, 최신 Dense Retriever에서는 잘 사용되지 않음
  • 반면, Doc2Query(doc expansion) 기법은 Sparse Retrieval에서 효과적임
• Query2doc
  • LLM을 활용한 Query Expansion 기법
    • Few-shot Prompting을 사용하여 Pseudo-document를 생성
    • 생성된 Pseudo-document를 원래 쿼리와 결합하여 새로운 확장 쿼리를 생성
  • 간단한 구현 가능 → 훈련 파이프라인이나 모델 구조 변경 없이 적용 가능
  • 향후 연구에서 쉽게 확장 가능
• 실험 및 결과
  • MS MARCO, TREC DL 2019/2020 데이터셋에서 평가
  • OpenAI GPT-3 (text-davinci-003)를 활용해 Pseudo-document 생성
  • BM25 성능을 크게 향상시키며, 특히 TREC DL의 어려운 쿼리에서 성능 개선이 두드러짐
  • DPR, SimLM, E5 등 최신 Dense Retriever도 Query2doc으로 성능 향상
  • 하지만, 강력한 Cross-Encoder 기반 Re-ranker를 사용할 경우 성능 개선 폭이 줄어드는 경향

📍 왜 Cross-Encoder 기반 Re-ranker를 사용하면 성능 개선 폭이 줄어들까?

1. Re-ranking이란?

• IR 시스템에서는 초기 검색(Retrieval) 을 수행한 후, 검색된 문서들을 더 정교한 기준으로 다시 정렬(Re-ranking) 하여 최종 결과를 제공함.
• 초기 검색 단계에서는 BM25나 Dense Retriever를 활용하여 수천 개의 후보 문서를 찾고,
• Re-ranker는 이 후보 문서들을 보다 정밀하게 평가하여 가장 관련성이 높은 문서를 상위에 배치함.

2. Cross-Encoder란? (↔ Bi-Encoder)

• 쿼리와 문서를 하나의 Transformer 모델에 함께 입력하여 유사도 계산 → 강력한 문맥 이해 가능 (대신 느림)

3. 그렇다면 왜 성능 개선 폭 줄어드는지?

• Cross-Encoder 자체가 강력한 문맥 이해 능력을 가지고 있어, Query Expansion이 제공하는 추가 정보를 활용할 여지가 줄어들기 때문

• Zero-shot OOD(도메인 외) 실험에서도 강력한 베이스라인보다 우수한 성능 달성
• LLM 크기가 클수록 Query2doc의 효과가 커짐 (소형 모델은 제한적인 성능 향상만 제공)
• 재현성을 위해 text-davinci-003을 활용한 생성 데이터 공개

2. Method

Query2doc은 주어진 쿼리 $q$ 에 대해 Few-shot prompting을 활용해 pseudo-document $d'$ 를 생성한 후, 이를 원래 쿼리와 결합하여 확장된 쿼리 $q^+$ 를 만드는 방식이다.

Pseudo-document 생성

• 프롬프트 구성:
  • 짧은 지시문: "Write a passage that answers the given query:"
  • k개의 레이블된 쿼리-문서 샘플 (논문에서는 $k=4$)
• 확장된 쿼리 생성:
  • $q^+=\text{concat}(q, d')$

Sparse Retrieval에서의 적용 (BM25)

• 일반적으로 쿼리 $q$ 가 pseudo-document $d'$ 보다 짧기 때문에, 쿼리 단어의 가중치를 증가시키기 위해 $q$ 를 n번 반복하여 pseudo-document 와 결합
• 확장된 쿼리 $q^+$ :
  • $$ q^+ = \text{concat}(\{q\} \times n, d') $$
  • 여기서는 $n=5$ 가 기본값으로 설정됨

Dense Retrieval에서의 적용

• 쿼리와 pseudo-document를 단순히 [SEP] 토큰으로 구분하여 결합
• 확장된 쿼리 $q^+$ :
• $$ q^+=\text{concat}(q, [SEP], d') $$
• 학습 방법

1) BM25 Hard Negative 샘플을 이용한 학습

📍 Hard Negative?

• Information Retrieval (IR) 시스템에서는 모델을 학습할 때 양성 문서 (Positive document)와 음성 문서 (Negative document) 를 구분해야 함
• Hard Negative란, 진짜 정답 문서는 아니지만 정답과 매우 유사해서 모델이 혼동할 가능성이 높은 문서
  • 너무 무관한 문서는 학습에 도움이 되지 않으므로, 적절히 어려운 (Hard) 문서를 Negative 샘플로 사용하면 모델 학습이 효과적

1. BM25로 초기 검색 수행

• 쿼리에 대해 BM25를 사용하여 가장 높은 점수를 가진 문서들을 찾음

2. Positive & Hard Negative 샘플 선택

• 정답 문서 (Positive document): BM25가 반환한 문서 중 실제 정답
• Hard Negative 문서: BM25가 높은 점수를 줬지만 실제 정답이 아닌 문서

3. Dense Retriever 학습

• Contrastive loss를 활용해 정답 문서는 가깝게, Hard Negative 문서는 멀어지도록 학습
  • 정답 문서 $h_d$ 와 쿼리 $h_q$ 의 코사인 유사도를 최대화
  • Hard Negative 문서들 $h_{d_i}$ 와 쿼리의 유사도를 최소화
  • 즉, 모델이 정답과 Hard Negative를 더 명확하게 구분할 수 있도록 학습

요약하면...

• BM25 Hard Negative = BM25가 높은 점수를 줬지만 실제 정답이 아닌 문서
• 왜 필요한가?
  • 너무 쉬운 Negative 샘플은 Dense Retriever 학습에 도움이 안 됨
  • Hard Negative를 활용하면 모델이 더 정밀하게 학습됨
• 어떻게 사용되는가?
  • BM25로 검색 → 정답 문서(Positive)와 Hard Negative 샘플 선택 → Contrastive loss 학습
  • 정답 문서는 더 가깝게, Hard Negative 문서는 더 멀게 학습

• Dense retriever (DPR) 모델을 BERT-base로 초기화
• 대조 학습 (Contrastive Loss) 적용:
  • $$ L_\text{cont} = -\log \frac{e^{h_q \cdot h_d}}{e^{{h_q \cdot h_d}} + \sum_{d_i \in N} e^{h_q \cdot h_d}} $$
  • $h_q$ : 쿼리 임베딩
  • $h_d$ : 정답 문서 (positive document) 임베딩
  • $h_{d_i}$ : Hard negative 문서 임베딩
  • $N$ : Hard negative 문서 집합

2) Cross-Encoder 기반 Re-ranker를 활용한 지식 증류 (Knowledge Distillation)

• 강력한 Dense retriever에 적용
• Cross-Encoder의 확률 분포 $p_{ce}$ 를 Student 모델의 확률 분포 $p_{stu}$ 로 근사
• 최종 Loss:
  • $$ \min D_{KL}(p_{ce}, p_{stu}) + \alpha L_\text{cont} $$
  • $D_{KL}$ : Kullback-Leibler divergence
  • $\alpha$ : Distillation loss 와 Contrastive loss의 균형을 조정하는 계수

기존 Pseudo-relevance Feedback (PRF) 방법과 비교

• 기존 PRF 방법은 초기 검색 결과에서 top-k 문서를 선택하여 query expansion 수행
• Query2doc은 초기 검색 결과의 품질에 의존하지 않고, LLM을 활용해 더 관련성 높은 pseudo-document를 생성
• 따라서, 노이즈가 많은 초기 검색 결과의 영향을 받지 않고, 더욱 효과적인 query expansion이 가능

3. Experiments

3.1 Setup

Evaluation Datasets

• In-domain evaluation:
  • MS-MARCO passage ranking
  • TREC DL 2019
  • TREC DL 2020
• Zero-shot out-of-domain evaluation: BEIR benchmark에서 5개의 low-resource 데이터셋을 선택.
• 평가 지표: MRR@10, R@k (k∈{50, 1k}), nDCG@10.

Hyperparameters

• Sparse retrieval: BM25 및 RM3은 Pyserini의 기본 구현 사용.
• Dense retrieval: SimLM과 동일한 하이퍼파라미터 사용하지만, pseudo-document 포함을 위해 query 최대 길이를 144로 증가.
• LLM prompting:
  • 4개의 in-context 예시 포함.
  • 샘플링 시 temperature=1 설정.
  • 최대 128개의 토큰 생성.

3.2 Main Results

In-domain Retrieval Results

• BM25+query2doc:
  • BM25 대비 TREC DL 2019 및 2020에서 15% 이상의 성능 향상.
  • TREC DL 데이터셋에는 long-tail entity-centric queries가 많아, 정확한 lexical match가 중요한 경우 성능 개선이 두드러짐.
  • 기존 query expansion 기법인 RM3는 R@1k 지표에서만 소폭 향상.
• 문서 확장(Document Expansion) vs. query2doc
  • docT5query는 MS-MARCO dev set에서 더 좋은 성능을 보이지만, T5 기반 query generator를 모든 라벨링된 데이터로 학습해야 함.
  • 반면, query2doc은 모델 fine-tuning 없이 사용 가능.
• Dense Retrieval 모델과 query2doc 결합
  • query2doc을 적용한 모든 dense retriever 모델이 기존 baseline 대비 향상.
  • 다만, intermediate pre-training이나 cross-encoder 기반 knowledge distillation을 사용할 경우 query2doc의 성능 향상 효과가 감소
    • SimLM+query2doc, E5+query2doc 실험에서 확인됨.

📍 왜 성능 저하?

(1) Intermediate Pre-training이 영향을 미치는 이유

• Intermediate pre-training(중간 사전 학습)은 일반적으로 도메인 적응(domain adaptation)이나 특정 retrieval 작업에 적합한 방식으로 모델을 학습하는 단계이다.
• 예를 들어, MS-MARCO와 같은 데이터셋에서 미리 학습된 dense retriever는 해당 데이터의 패턴을 학습하며, 적절한 쿼리와 문서를 매칭하는 능력이 향상됨.
• 이러한 retriever는 기본적으로 입력된 쿼리만으로도 충분한 정보를 얻을 수 있도록 학습되므로, query2doc이 추가적인 정보를 제공하더라도 성능 향상이 크지 않을 수 있음.

(2) Cross-Encoder 기반 Knowledge Distillation이 영향을 미치는 이유

• Cross-encoder 기반 re-ranker는 쿼리와 문서를 함께 입력하여 더 정밀한 relevance score를 계산하는 방식이다.
• Cross-encoder를 teacher model로 사용하여 dense retriever를 학습시키면,
  • teacher model이 문맥적 의미를 더 잘 이해하도록 지도 학습을 수행
  • dense retriever는 원래 쿼리만 가지고도 좋은 embedding을 생성
  • 즉, query2doc이 추가적인 정보를 줘도 teacher model로부터 이미 충분히 학습했기 때문에 효과가 줄어듦.
• 특히, strong cross-encoder 기반 re-ranker를 사용할 경우 query2doc의 영향이 더 미미해지는 경향이 있음.

Zero-shot Out-of-Domain Retrieval

• Entity-centric 데이터셋(DBpedia 등)에서는 가장 큰 성능 향상.
• NFCorpus, SciFact 데이터셋에서는 소폭 성능 저하 발생.
  • 원인: 학습 데이터와 평가 데이터 간 분포 차이(distribution mismatch).

(1) Distribution Mismatch란?

• Query2doc이 효과적으로 작동하려면, LLM이 생성한 pseudo-document가 실제 검색해야 할 문서와 유사한 패턴을 가져야 함.
• 하지만 학습 시 사용한 데이터와 평가할 데이터의 분포가 다르면, LLM이 생성한 pseudo-document가 실제 평가 데이터셋에 적합하지 않을 가능성이 커짐.
• 이러한 데이터 분포 차이(distribution mismatch)가 발생하면, query2doc이 오히려 성능을 저하시킬 수 있음.

(2) NFCorpus와 SciFact에서의 문제

• NFCorpus:
  • 건강 및 의학 관련 정보가 포함된 데이터셋으로, 일반적인 웹 문서보다 전문 용어 및 도메인 특화된 용어가 많음.
  • Query2doc에서 생성한 pseudo-document가 일반적인 웹 검색 패턴을 반영한 문장이 많다면, 실제 평가 데이터의 문서들과 일치하지 않아 retrieval 성능이 떨어질 가능성이 있음.
• SciFact:
  • 과학 논문 및 연구 문헌을 기반으로 하며, 정확한 팩트와 인용 정보가 중요함.
  • LLM이 생성한 pseudo-document가 일반적인 언어 패턴을 따른다면, 과학적 근거가 부족하거나 부정확한 정보가 포함될 가능성이 있음.
  • 그 결과, 실제 문서를 검색하는 데 방해가 되어 성능이 저하될 수 있음.

(3) 해결 방안?

• LLM의 prompting 방식을 개선하여 해당 도메인에 맞는 pseudo-document를 생성하도록 유도
  • 예: NFCorpus에서는 의학 논문 스타일의 문장을 생성하도록 few-shot 예제를 조정
  • SciFact에서는 팩트 기반의 근거가 포함된 문장을 생성하도록 유도
• 또는, 도메인별로 학습된 smaller LLM을 활용하여 query expansion 수행
  • 예: 일반적인 GPT-4가 아니라, 의료나 과학 도메인에 특화된 LLM을 활용하는 방식.

 

4. Analysis

Scaling up LLMs is Critical (LLM 규모가 클수록 성능 향상)

• Query expansion 성능은 모델 크기가 증가할수록 꾸준히 향상됨.
• 1.3B 모델부터 175B 모델까지 성능이 증가하며, 특히 davinci-003 모델이 davinci-001보다 더 좋은 성능을 보임 → 더 나은 훈련 데이터와 instruction tuning 때문.
• GPT-4(OpenAI, 2023)가 가장 높은 성능을 기록.
• 작은 언어 모델은 짧고 부정확한 텍스트를 생성하는 경향이 있어 성능이 떨어짐.

Performance Gains are Consistent across Data Scales (데이터 크기와 관계없이 일관된 성능 향상)

• DPR+query2doc 변형 모델이 DPR baseline보다 항상 약 1% 향상된 성능을 보임.
• Fine-tuning에 사용된 데이터 양과 관계없이 성능 개선 효과가 유지됨 → query2doc 방식이 데이터 크기와 독립적으로 성능 향상을 제공한다는 점을 시사.

How to Use Pseudo-documents (Psuedo-document 활용 방식 비교)

• 세 가지 활용 방식 비교
  • BM25 + query2doc
  • w/ query only
  • w/ pseudo-doc only
• Query와 pseudo-document를 단순히 연결(concatenation)하는 방식이 가장 좋은 성능을 보임.
  → 이는 pseudo-document가 원래 쿼리를 보완해 lexical mismatch를 줄여주기 때문.

Case Analysis

• Pseudo-document가 retrieval 성능을 향상시키는 이유는 추가 정보를 제공하여 query와 문서 간 lexical mismatch를 줄이기 때문.
• 어떤 경우에는 pseudo-document 자체만으로도 충분한 정보를 제공하여, retrieval 과정 없이도 사용자의 정보 요구를 충족할 수 있음.
• 하지만 LLM이 생성하는 pseudo-document에는 팩트 오류(factual error)가 포함될 가능성이 있음.
  • 예시:
    • "It’s a Jungle Out There"라는 노래는 2002년이 아니라 2003년에 사용되었음.
    • 이러한 미세한 오류도 신뢰성 문제를 유발할 수 있음.

5. Related Work

Query Expansion and Document Expansion

• Sparse retrieval 시스템의 검색 품질을 향상시키기 위한 대표적인 두 가지 기법.
• 쿼리 확장(Query Expansion):
  • 쿼리를 관련 피드백(relevance feedback)(Lavrenko & Croft, 2001; Rocchio, 1971) 또는 어휘 자원(WordNet, Miller, 1992)을 기반으로 다시 작성.
  • 라벨이 없는 경우, 상위 k개의 검색된 문서를 의사 관련 피드백(pseudo-relevance feedback)으로 활용 가능 (Lv & Zhai, 2009).
  • Liu et al.: encoder-decoder 모델을 미세 조정하여 문맥적 단서를 생성.
• 문서 확장(Document Expansion):
  • 문서에 관련 키워드를 추가하여 표현을 풍부하게 만듦.
  • Doc2query (Nogueira et al., 2019): 문서에서 seq2seq 모델을 활용해 가상의 쿼리를 생성 후, 이를 문서 인덱스에 추가.
  • Sparse retrieval 모델 (SPLADE, uniCOIL)은 문서의 단어 가중치를 end-to-end 학습.
  • 그러나 Dense Retriever(Ren et al., 2021; Wang et al., 2023)들은 일반적으로 확장 기법을 사용하지 않음.
  • (↔ 본 논문에서는 strong dense retriever의 경우에도 LLM을 사용한 query expansion의 이점 있다고 밝힘)

Large Language Models (LLMs)

• GPT-3, PaLM, LLaMA 같은 LLM은 수조 개의 토큰과 수십억 개의 파라미터로 훈련되어 강력한 일반화 성능을 보임.
• LLM은 Zero-shot learning 또는 Few-shot prompting을 활용하여 새로운 쿼리를 생성할 수 있음.
• 본 연구에서는 Few-shot prompting을 사용하여 pseudo-documents 생성.
• HyDE (Gao et al., 2022):
  • Zero-shot 설정에서 pseudo-document 임베딩을 사용한 검색(similarity search) 기법.
  • 그러나 pseudo-doc과 실제 문서가 의미적으로 같다고 가정, 이는 항상 성립하지 않을 수 있음.
• Question Answering 관련 연구:
  • RECITE (Sun et al., 2022), GENREAD (Yu et al., 2022):
    • LLM을 강력한 문맥 생성 모델로 활용하여 방대한 사실 정보를 인코딩.
  • 하지만 LLM이 사실과 다른 정보(false claims)를 생성할 수 있어 신뢰성이 문제가 될 수 있음.

6. Conclusion

• 본 연구에서는 LLM을 활용한 Query Expansion 기법 "query2doc"을 제안함.
  • Few-shot prompting을 통해 LLM이 pseudo-documents를 생성하도록 유도한 후, 이를 기존 Sparse 또는 Dense Retriever에 통합하여 검색 성능을 향상시킴.
  • 본 기법은 LLM을 직접 학습시키지 않고도 Prompting을 통해 LLM의 지식을 활용하는 효과를 가짐.
  • 실험 결과, 다양한 검색 모델과 데이터셋에서 일관된 성능 향상을 확인함.
• 한계점: 검색의 효율성

 

 

1. LLM 기반 추론 속도 문제
   • LLM의 토큰 단위 Autoregressive Decoding으로 인해 추론 속도가 느려질 수 있음.
2. 검색 속도 저하
   • Query Expansion으로 인해 쿼리 길이가 증가하면 Inverted Index 검색 속도도 느려질 가능성이 있음.

⇒ 실제 응용 시 성능뿐만 아니라 효율성도 고려해야 함.

 

Appendix

A. Implementation Details

• Dense Retrieval 실험 (Table 1) 관련 하이퍼파라미터 (Table 7)
  • DPR (BERTbase) & SimLM/E5base-unsup 모델 사용
  • Learning Rate: 2×10⁻⁵ (DPR) / 3×10⁻⁵ (SimLM/E5)
  • Batch Size: 64
  • Warmup 스텝: 1000
  • 네거티브 샘플 개수: 15 (DPR) / 23 (SimLM/E5)
  • 쿼리 및 문서 길이: 144
  • 학습에 4개의 GPU 사용, 최대 10시간 소요
• Knowledge Distillation을 활용한 Dense Retriever 학습
  • Wang et al.이 공개한 Cross-Encoder Teacher Score 사용.
  • SimLM & E5 체크포인트는 Hugging Face에서 제공됨.
  • SimLM: [CLS] 벡터 사용, E5: Mean Pooling 사용 → 훈련 방식과 일관성 유지.
• LLM Prompting 설정
  • MS-MARCO 데이터셋에서 4개의 in-context 예제 샘플링하여 Prompting 수행.
  • 총 55만 회 API 호출 → 약 5000달러 비용 발생.
  • GPT-4를 활용할 경우, 명확한 지시를 위해 system message 설정:
    • "You are asked to write a passage that answers the given query. Do not ask the user for further clarification."
• Out-of-Domain 평가
  • BEIR 벤치마크의 DBpedia, NFCorpus, SciFact, TREC-COVID, Touche2020 데이터셋 사용.
  • SimLM의 결과는 Wang et al.이 공개한 체크포인트를 활용함.

B. Exploring Other Prompting Strategies

• Iterative Prompting 실험
  • GPT-4를 활용하여 pseudo-document 생성 후 자체 검증(Self-Verification) 수행.
  • 추가 프롬프트를 통해 LLM이 스스로 생성한 내용을 검토하고 수정하도록 요청:

You are asked to rewrite the passage that potentially answers the given query. 
You should only correct the factual errors in the passage, do not ask for clarification or make unnecessary changes.

• 하지만 GPT-4는 대부분의 경우 기존 내용을 거의 수정하지 않음 → 기존 pseudo-document의 품질이 이미 충분히 높거나, LLM이 스스로 오류를 교정할 능력이 제한적일 가능성.
• 실험 결과, Iterative Prompting을 적용해도 성능 향상 효과는 미미 (DL 2019: 69.2 → 68.6, DL 2020: 64.5 → 64.8).

C. Results Across Multiple Runs

• Query2doc의 랜덤성 평가
  • Few-shot 예제 선택 & Auto-regressive Sampling 방식이 결과에 영향을 미칠 가능성 존재.
  • 3회 실행 후 평균 및 표준 편차 계산:

• 성능 변동이 크지 않음 → Query2doc 방식이 일관된 성능을 유지하는 것으로 보임.