Soft Prompt Tuning for Augmenting Dense Retrieval with Large Language Models, Peng et al., arXiv 2023

Motivation

Dense retrieval 모델들은 query와 문서를 dense vector로 임베딩해 검색 성능을 극대화하지만, 도메인 이동 시 성능이 저하되는 문제가 있다.
Hard prompts는 수작업 설계가 필요하며, 다양한 도메인에 걸쳐 효과적으로 작동하기 어려움.
Soft Prompt Tuning은 LLM의 매개변수를 고정한 상태에서 soft prompt만 최적화함으로써 비용을 절감하며 성능 향상을 도모한다.

작동방식

Data Preparation
- 예시: MS MARCO 데이터셋에서 일부 문서를 training set으로 사용하고 나머지는 validation 및 test로 나눔.
- unlabeled data: 뉴스 아티클이나 웹 페이지 같은 다양한 도메인의 문서를 사용해 학습에 활용할 weak queries를 생성.
Soft Prompt Tuning
- 매개변수 고정: LLM(GPT-3 또는 Flan-T5)의 내부 매개변수는 고정하고, soft prompt만 학습함.
- → 이 부분이 조금 헷갈렸는데, (B, S, D)가 기존 input 이라면 **(B, S+K, D)**를 input으로 넣음. soft prompt는 토큰이 아니라 K의 길이를 가진 D 차원 벡터 (K, D). 그래서 이 (K,D)만 학습하는 것이 이 논문의 핵심
- 예시: "Generate a query for this document."이라는 hard prompt를 soft prompt로 변환.
  - 초기 hard prompt: "Generate a query for this document."
  - Soft prompt(최적화 후): "[TASK:QUERY] Generate question: What purpose does the document serve?"
- 최적화 과정: 학습 중 다양한 문서와 쿼리 쌍을 조합해 최적의 soft prompt를 찾음. Adafactor optimizer를 사용해 손실을 최소화.
Soft Prompt Filtering
- 예시: 생성된 쿼리 중 "What is the main argument?"와 같은 질문이 포함된 경우, BM25 기반 평가를 통해 해당 query가 관련된 문서 상위 20개 내에 포함되지 않으면 폐기.
- 목적: 불필요한 노이즈를 줄이고 고품질 데이터만 남김.
Soft Prompt Augmentation
- 학습된 soft prompts를 사용해 대규모 unlabeled 데이터에 대한 weak queries를 생성.
- 예시: 뉴스 기사에서 "What event does this article describe?"와 같은 쿼리를 자동 생성.
- 다양한 도메인에 걸친 queries를 생성해 retriever가 여러 환경에서 잘 동작하도록 학습함.
Weak Data Filtering
- BM25 필터를 사용해 쿼리가 Top-k 문서 내에 포함되지 않으면 폐기.
- 예시: "What is the economic impact of this event?"와 같은 질문이 적절한 문서와 일치하지 않으면 해당 쿼리를 제거.
Dense Retriever 학습 및 평가
- 학습: 생성된 query-document pairs를 사용해 ColBERT와 같은 dense retriever 모델을 학습.
- 예시: MS MARCO 벤치마크에서 retriever를 평가하며, MRR@10, nDCG@10 등의 지표를 사용해 성능을 측정.
- 결과: 학습된 dense retriever는 다양한 도메인에서 높은 검색 성능을 유지하며 latency도 최소화됨.

기대 효과 및 활용 결과

Labeled 데이터 없이도 도메인 적응이 가능.
다양한 retriever 모델의 성능을 향상시키며, latency를 줄여 실시간 검색에 적합.
여러 벤치마크에서 zero-shot 성능을 크게 향상.

성능 비교

Dataset	Baseline (InPars)	SPTAR
MS MARCO MRR@10	0.1519	0.2114 (+0.0595)
FiQA-2018 R@100	0.5289	0.5747 (+0.0458)
DL2019 nDCG@10	0.4460	0.5253 (+0.0793)
DL2020 MAP	0.2407	0.3042 (+0.0635)

장단점

장점:
- Soft prompt tuning을 통해 비용을 절감하면서도 성능을 향상.
- 노이즈 필터링으로 high-quality weak queries를 유지.
- 다양한 retriever 모델에 적용 가능하며 reproducibility 보장.
단점:
- Soft prompts의 성능은 사용된 example pairs에 따라 달라질 수 있음.
- 대규모 weak query 생성에는 시간과 자원이 필요함.

한계 및 Future Work

다양한 retriever 모델과의 통합 가능성 연구 필요.
Soft prompts의 다국어 확장성을 평가하여 적용 범위를 확대해야 함.
실시간 검색 시스템에 맞춘 latency 최소화 전략 개발 필요.

기타 디테일

리소스: NVIDIA A100 GPU 사용.
데이터셋: MS MARCO, FiQA-2018, DL2019, DL2020 사용.
코드 공개: [GitHub 링크 제공 예정]