RAFT : Adapting Large Language Model to Domain Specific RAG (2024.05.)
Tianjun Zhang, Shishir G. Patil, Naman Jain, Sheng Shen 외 3인 (UC Berkeley)
https://arxiv.org/abs/2403.10131
오늘은 LLM을 특수 도메인에 적용하기 위해 주로 활용되어지는 RAG, 파인튜닝 방법론의 각 단점을 보완하고 두 방식을 결합함으로써 효율적인 도메인 최적화 방식을 제안한 RAFT : Adapting Large Language Model to Domain Specific RAG 논문을 리뷰하도록 하겠습니다! 😁
📌 Abstract & Introduction

일반 지식 추론 작업에서 LLM은 방대한 양의 Corpus를 사전학습 함으로써 눈부신 발전을 이루었고, 이후 위의 그림처럼 의료, 법률과 같은 특수 도메인에서 활용될 수 있도록 점차 발전해 나아가고 있습니다.
특수 도메인에 LLM을 적용할 때에는 일반적인 지식 추론 보다 주어진 문서를 기반으로 답변 정확도를 최대화하는 것이 더 중요한데, 효과적으로 LLM을 적용하기 위한 크게 두 갈래의 접근 방향성이 있었습니다.
- RAG를 통한 In-Context Learning

RAG - Supervised Fine-Tuning

Supervised Fine-Tuning
각각의 방식을 아래 그림과 같이 시험 응시에 빗대어 설명할 수 있습니다. Fine-Tuning은 왼쪽처럼 모든 지식을 학습을 통해 기억하고 있다가 질문에 대해 학습된 지식을 통해 답변합니다. 이 경우, 도메인 특수성은 고려할 수 있으나 사전 학습된 지식 외의 내용을 물어보는 질문을 마주하면 LLM은 좋은 답변을 내기 어렵습니다.
반대로 RAG는 오른쪽과 같이 외부 문서를 참조하여 답변을 보강할 수 있기 때문에 일종의 오픈북 시험을 보는 것과 같지만, LLM이 학습되지 않는다면 오픈북이더라도 어디가 중요한지 모른채로 오픈북 시험을 보는 것과 동일하다는 단점이 있습니다.

본 논문에서는 각 방법론의 단점을 보완하기 위해 RAG와 Fine-Tuning을 결합한 RAFT 방법론을 제안합니다.
LLM은 문서 D*에서 질문 Q에 대해 답변 A*를 생성하도록 훈련됨으로써 도메인 지식을 통합할 수 있고, 더불어 불필요한 검색 정보를 활용하는 것에 강건해질수 있습니다. 이 과정에서, 문서 D*, 즉 Context에는 Query와 관련 없는 문서도 일부 포함되어 학습 과정에서 관련성 높은 문서와 낮은 문서를 필터링 하는 방식을 익힐 수 있고, 답변은 관련 있는 문서의 시퀀스를 직접 인용하여 CoT(Chain of Thought) 방식으로 작성됩니다.

CoT는 여러 단계의 중간 reasoning 과정들을 통해 답을 더 쉽게 도출할 수 있게 하는 프레임워크입니다. 사람이 복잡한 문제를 해결할 때처럼 step-by-step reasoning을 추가해줌으로써 더욱 정확한 답을 추출할 수 있고, 모델의 해석력을 높일 수 있습니다.
결론적으로 앞선 오픈북 시험 예시에 다시 비추어 보면, RAFT는 마치 오픈북 시험을 위해 책 안에서 어떤 부분이 시험과 관련있는지를 공부하는 것과 같다고 봐주시면 될 것 같습니다.
📌 RAFT

기존 파인튜닝과 RAG의 흐름을 다시 한번 살펴보면 아래와 같습니다. LLM은 QA Pair를 통해 학습되며, Inference Time의 RAG 환경 내에서는 LLM의 답변 품질 향상을 위해 프롬프트에 D와 같은 추가 문서가 도입될 수 있습니다.

RAFT에서는 RAG와 유사하게 문서 집합(Dₖ) 및 문서 중 하나 또는 여러 개의 정답 기반 Chain-of-Thought 스타일 답변 A∗를 포함하는 학습 데이터를 사용합니다. 아래 그림은 데이터 예시입니다.

이 때 문서의 종류를 두 가지로 구분하는데, 질문에 답을 추론하는 데 핵심적인 정보가 포함된 문서인 Golden 문서 D*와 답변과 관련 없는 정보를 담고 있어, 모델이 해당 문서를 무시하도록 학습하는 역할을 하는 Distractor 문서 Di가 그 두 가지 종류에 해당합니다.
학습 데이터는 전체의 일정 비율 P에 대해 질문, golden 문서 및 distractor 문서를 포함하고, 나머지 (1 − P) 비율에는 distractor 문서만 포함하여 구성됩니다. golden 문서가 반드시 단일 문서일 필요는 없습니다.

데이터셋의 구성 이후, 모델은 표준 Supervised Fine-Tuning 방법론에 따라 전반적인 학습을 진행하고, 테스트에서는 Top-K개의 문서만 Q와 함께 프롬프트를 통해 모델에 전달되고 이후 순차적으로 답변을 생성합니다.
결국 핵심은 golden 문서가 사용되었을 때의 답변과 그 외 답변을 비교함으로써 모델이 학습을 통해 자연스럽게 golden 문서를 찾아내는 능력을 기르는 것입니다.
📌 Evaluation
본 논문에서는 RAFT가 다양한 Baseline과 비교하여 얼마나 잘 수행되는지 연구하기 위해 다양한 실험들을 설계하여 수행하였습니다.
[Dataset]
먼저 본 논문에서는 베이스라인들과 RAFT의 평가를 위해 아래와 같은 평가용 데이터셋을 활용하였습니다.
- Natural Questions (NQ), Trivia QA, Hotpot QA: Wikipedia 기반 오픈 도메인 질문, 주로 일반 지식(영화, 스포츠 등)
- HuggingFace, Torch Hub, TensorFlow Hub: Gorilla 논문에서 제안된 APIBench에 포함된 데이터셋, 문서에 따라 정확하고 기능적인 실행 가능한 API 호출을 생성하는 방법
- PubMed QA: 생물 의학 연구 질의응답에 맞춤화된 데이터셋
[Baseline]
다음으로 RAFT와 비교할 4가지의 베이스라인은 아래와 같습니다.
- LlaMA2-7B-chat model with 0-shot prompting : QA 작업을 위해 일반적으로 사용되는 instruction-finetuned 모델이며, 명확하게 작성된 instructions을 제공하지만 참조 문서는 제공하지 않는 모델
- LlaMA2-7B-chat model with RAG (Llama2 + RAG) : 이전 설정과 유사하지만 도메인 특정 QA 작업을 처리하기 위해 참조 문서를 포함하는 형식
- Domain-Specific Finetuning with 0-shot prompting (DSF) : Context 없이 표준 지도학습 미세 조정하는 방식
- Domain-Specific Finetuning with RAG (DSF + RAG): Domain-specific finetuned 모델에 RAG를 사용하여 외부 지식을 통합하는 방식
[Result]

대체적으로 모든 데이터셋에서 RAFT가 베이스라인들에 비해 월등히 좋은 성능을 보입니다만, PubMed의 경우 예/아니오의 이진 답변 형식이기 때문에 DSF+RAG와 큰 차이가 없는 것을 관찰할 수 있습니다.

또한 CoT 답변 형식에 대한 Ablation Study를 진행했을 때, 해당 방식의 적용은 모델의 이해를 풍부하게 하는 추론 체인을 통합함으로써 전체 정확도를 향상시키고 간결한 답변에 대한 과적합을 방지할 수 있음을 시사합니다. 실제로 위의 표를 보더라도, 대부분의 데이터셋에서 성능 향상이 있습니다.

마지막으로, 본 논문의 연구자들은 컨텍스트에 golden document가 없는 훈련 데이터의 일부를 통합하면 RAG 작업에서 모델의 성능이 향상되는 것을 파악하였습니다. 위 그래프에 golden document를 포함해야 하는 훈련 instance의 백분율을 나타내는 하이퍼파라미터 P에 대한 연구 결과가 나타나있는데, 최적 P는 데이터셋에 다르지만 확실한 지점은 golden document가 없는 특정 비율의 훈련 데이터를 포함하는 것이 모델의 컨텍스트 판별력을 향상시켜, 모델의 노이즈 강건성 높이 성능 향상에 긍정적으로 기여한다는 것입니다.
📌 RAFT Generalizes to Top-K RAG

P%만큼 golden document를 포함하는 QA 인스턴스를 생성할 때, 위의 추가 실험을 통해 golden document와 distractor document의 비율을 어떻게 설정할 지에 대한 인사이트를 제공해줍니다.
논문에서는 실험적으로 NQ 데이터셋에서는 3개의 distractor를 섞어주는 것이, Hotpot QA에서는 1개의 distractor를 섞어주는 것이 가장 성능이 좋아 해당 비율을 활용하였다고 합니다만, 결론적으로 아예 golden 문서만 활용하는 것보다 일정 비율로 데이터셋에 따라 distractor를 섞어 FT를 진행하는 것이 테스트시 다양한 양의 문서가 주어질 때, 모델이 관련 있는 정보를 효과적으로 추출하는 능력을 향상시킴을 알 수 있습니다.
📌 Conclusion
RAFT는 일종의 오픈북 설정, 문서 Context를 참조할 수 있는 환경에서 특정 도메인 내 질문에 대한 모델의 성능을 향상시키도록 설계된 학습 전략입니다. 방해 문서와 함께 모델을 학습시키고, 데이터셋을 구성하여 일부는 해당 Context에 핵심 문서가 부족하도록 하고, 관련 텍스트에서 직접 인용한 CoT 방식으로 답변을 구성하는 것과 같은 효과적인 방식들을 설계하였습니다. 여러 데이터셋에 대한 평가지표를 통해 RAFT의 상당한 잠재력을 확인할 수 있습니다.
📌 느낀점(논문을 읽으면서..)
현재 LLM 관련 프로젝트 수행에 있어서 RAG 프레임워크들을 하나씩 공부해나가고 있는데, Naive RAG 이외에도 짧은 시간내에 정말 다양한 방법론들이 나오고 있다는 것을 체감하게 되는 것 같습니다. 오늘 보았던 RAFT는 상대적으로 방법론 자체는 단순하지만, 학습을 통해 모델이 관련 Context의 Value를 필터링할 수 있는 힘을 자체적으로 기를 수 있다는 것이 인상 깊은 논문이었던 것 같습니다. 다른 방법론 소개로 다시 찾아뵙겠습니다!
📌 Reference
https://process-mining.tistory.com/217
https://www.youtube.com/watch?v=LrZszoQKm8w