[Paper Review] RAFT: Adapting Language Model to Domain Specific RAG 논문 리뷰

Abstract

• 특정 도메인에서 LLM을 사용할 때 새로운 지식을 RAG-based-prompting이나 fine-tuning을 통해 추가하는 것이 일반적임
• RAFT(Retrieval Augmented Fine Tuning)는 특정 도메인 RAG에 대해 언어 모델의 성능을 향상시키기 위한 훈련 기법
• RAFT는 질문에 대답하는데 도움이 되지 않는 문서(distractor documents)를 무시하도록 모델을 훈련시킴
• RAFT는 관련 문서에서 질문에 대답하는데 도움이 되는 올바른 시퀀스를 인용하고, chain-of-thought 스타일 응답을 통해 모델의 추론 능력을 향상시킴
• RAFT는 PubMed, HotpotQA, Gorilla 데이터셋에서 지속적으로 성능을 향상시킴

1. Introduction

• How to adapt pre-trained LLMs for Retrieval Augmented Generation (RAG) in specialized domains?
• 이 논문에서는 특정 도메인에서 “open-book” 설정으로 질문에 답하는 모델의 능력을 향상시키는 RAFT를 제안함: RAG + SFT
• RAFT는 fine-tuning을 통해 특정 domain의 답변을 잘하게 할 뿐만 아니라 리트리벌이 잘못 뽑혔을때, robustness하게 답변하도록 학습 (retrieval된 output들 중 도움이 되지 않는 distractor 문서를 무시하도록 훈련함)

2. LLMs for Open-Book Exam

• Closed-Book Exam에서 LLM은 질문에 답하기 위해 추가 문서나 참조 자료에 접근할 수 없음
• Open Book Exam에서는 LLM이 질문에 답하기 위해 외부 정보 출처를 참조할 수 있음
• Domain Specific Open-Book Exam
  • 이 논문에서는 general open book exam이 아닌, domain specific open book exam에 집중. → pre-trained LLM을 특정 도메인에 adaptation 시키고, retrieved documents & distractors에 대해 robust하게 만듦

3. RAFT

• Supervised Fine-tuning(SFT)은 Dataset(D)에서 유도된 Question(Q)과 Answer(A) 쌍을 통해 모델을 훈련시킴
  • classical SFT에서 모델은 pre-training phase에서 학습한 knowledge의 Q에 대답하는 능력을 향상시키도록 학습
  • Inference 시에 RAG를 활용하여 학습된 지식에 대해 더 대답을 잘 하도록 하기도 함.
• RAFT에서는 모델이 RAG된 결과를 잘 활용하도록 fine-tuning data를 준비
  • question(Q), documents(D_k), Chain-of-thought 스타일 answer(A*)
  • 2가지 타입의 documents
    • ‘oracle’ documents (D*): 질문에 대한 답을 할 수 있는 문서
    • ‘distractor’ documents (D_i): 답과 관련 없는 정보의 문서
• P% 비율의 데이터는 distractor 문서와 함께 oracle 문서(D∗)를 포함하고, (1-P)% 비율은 oracle 문서 없이 distractor 문서만 포함함
  • P % of data: Q + D∗ + D2 + . . . + Dk → A∗
  • (1 − P) % of data: Q + D1 + D2 + . . . + Dk → A∗
• training quality를 높이는데 가장 중요한 factor는 CoT와 같은 reasoning process를 도입한 것

4. Evaluation

• RAFT의 성능을 여러 baseline과 비교 실험함
• Baseline
  1. LLaMA2-7B-chat model (w/ 0-shot prompting) : reference doc 없이 instruction을 통한 QA
  2. LLaMA2-7B-chat model (w/ RAG): ref doc 있음 (나머진 1 세팅과 비슷)
  3. domain specific finetuning (DSF) (w/ 0-shot prompting): doc 없이 supervised finetuning
  4. domain specific finetuning (DSF) (w/ RAG): RAG를 통해 뽑힌 doc과 함께 SFT (RAFT 세팅과 다르게 distractor 문서가 없음)
• Datasets
  • Natural Questions (NQ), Trivia QA, HotpotQA (Wikipedia 기반) - general domain
  • HuggingFace/Torch Hub/TensorFlow Hub (APIBench 데이터셋) - specific domain
  • PubMed QA (bio관련 QA) - specific domain

4.1 Results

• RAFT가 모든 도메인 특화 데이터셋 (PubMed, HotpotQA, HuggingFace, Torch Hub, TensorFlow Hub)에서 baselines보다 훨씬 좋은 성능을 보임
• GPT-3.5와 비교해도 RAFT가 유의미한 성능 향상을 보임
• DSF만으로는 QA 스타일에 align되는 것에는 도움이 되지만 context에서 정보를 추출하는 데에는 한계가 있음
• RAFT는 QA 스타일 align과 document processing을 동시에 잘하도록 학습시킬 수 있음

4.2 Effect of CoT

• QA만 제공하는 것보다 reasoning chain을 함께 제공하는 것이 모델 성능 향상에 큰 도움이 됨
• Chain-of-thought를 활용하면 학습 안정성도 높일 수 있음

4.3 Qualitative Analysis

• DSF는 질문에 대한 답을 직접적으로 제공하지 못하고 엉뚱한 답변을 하는 경우가 있음
• RAFT는 주어진 context에서 정확한 답을 찾아낼 수 있음
• 이는 QA pair로만 학습하는 것보다 context를 함께 고려하며 학습하는 것이 중요함을 시사함

4.4 Should we train the LLM always with the oracle context for RAG?

• 실험을 통해 학습 데이터의 p%에만 oracle document를 포함시키고 (1-p)%에는 포함시키지 않는 것이 성능 향상에 도움이 된다는 것을 발견함
• 이는 oracle document만으로 학습하는 것보다 일부 학습 데이터에서는 oracle document 없이 학습하는 것이 모델의 RAG 성능을 높이는 데 도움이 된다는 것을 의미함

5. RAFT Generalizes to Top-K RAG

• Retriever가 제공하는 문서의 개수가 몇 개인지에 따라 RAFT의 robutness가 어떻게 달라지는지 실험함
• LLM은 irrelevant text에 취약하다고 알려져 있는데, 이는 LLM+RAG에서 특히 중요한 문제임 (top-k RAG가 높은 recall을 위해 자주 사용되기 때문)
• 이러한 문제를 다루기 위해 RAFT에서는 학습 시 distractor document를 함께 사용함

5.1 Making Model Robust to top-K RAG

• Oracle document와 함께 다양한 개수의 distractor document로 학습시키는 것이 모델이 test time에 다양한 개수의 문서가 주어졌을 때에도 robust하게 만드는 데 도움이 됨
• 실험 결과, oracle document만으로 학습시키는 것보다 distractor document와 함께 학습시키는 것이 더 좋은 성능을 보임
• RAFT에서는 일반적으로 1개의 oracle document와 4개의 distractor document로 학습시키는 것이 가장 좋은 밸런스를 보임

5.2 Generalization to a variable number of test-time documents

• 학습 시 사용한 distractor document의 개수에 따라 모델이 test time에 다양한 개수의 문서가 주어졌을 때 얼마나 잘 generalize되는지 실험함
• 실험 결과, 학습 시 distractor document를 사용하는 것이 test time에 다양한 개수의 문서가 주어졌을 때에도 일관된 성능을 보일 수 있도록 해줌
• 이는 실제 응용에서 마주칠 수 있는 다양한 상황에 대한 robustness를 높여줌

6. Conclusion

• RAFT는 특정 도메인에서 “open-book” 설정으로 질문에 답하는 모델의 성능을 향상시키기 위한 훈련 전략
• ⭐️ Distractor 문서와 함께 모델을 훈련시키고, 데이터셋의 일부가 oracle 문서 없이 구성되도록 하며 (1-p)%, chain-of-thought 방식으로 관련 텍스트에서 직접 인용하여 답변을 생성하게 함
• domain specific QA task에서 작은 모델 fine-tuning을 잘 하면 LLM만큼 성능이 좋을 수도 있다