[COLM 2024] RAFT: Adapting Language Model to Domain Specific RAG

최근에 제안서 작업을 하면서 RAG를 팔로업 할 필요성을 느꼈는데요. 해당 논문은 이미 많이 인용된 논문일 뿐만 아니라 코드 또한 공개가 되어 있어 이 논문을 시작으로 계속해서 팔로업 해보고자 합니다. 그럼 리뷰 시작하겠습니다.

1. Introduction

저희가 일반적으로 알고 있는 LLM은 대량의 공개 데이터를 기반으로 학습되기 때문에 일반적인 지식 질문에 뛰어난 성과를 보여주지만, 의학, 법률과 같은 특수한 분야에서는 어려워하는 모습을 보이기도 합니다. 특히나, 특수 도메인인에서의 task는 특정 집합의 문서에 의존하기 때문에 LLM이 맞춤형으로 학습되고 얼마나 전문지식을 갖추는지가 중요해집니다.

최근에 LLM에 전문지식을 갖추기 위해서 사용되는 여러 방법 중에 떠오르는 방법이 있다면 RAG(Retrieval-Augmented Generation)이 있겠습니다. RAG는 간략히 설명하면, LLM이 질문에 답변하기 위해 관련 문서를 참조하는 방법을 말합니다. 그러나 본 논문에서는 이러한 RAG 기반의 기존 방법론(in-context RAG)은 고정된 도메인 설정과 특정 과제를 수행하기 위해 관련 문서나 데이터를 사전에 검토하고 이를 학습하는 과정을 가질 기회를 활용하지 못합니다. (Figure 1에 (b)를 참고하시면 되겠습니다.) 쉽게 말하면, open-book 시험에서 준비 없이 시험을 치르는 것이라 말할 수 있습니다. 반면에 supervised fine-tuning은 모델이 일반적인 패턴을 학습하는데 효과적이지만, 시험의 문서를 참조하는 능력을 잃어버릴 가능성도 있습니다.

본 논문에서는 두 가지 방식의 장점을 모두 활영하여, instruction fine-tuning (IFT)와 RAG를 결합하여 RAFT (Retrieval-Augmented Fine Tuning)을 제안합니다. RAFT는 LLM을 fine-tuning하여 도메인 지식을 통합하는 것과 동시에 도매인 내 RAG 성능을 향상시키는 것을 목표로 합니다.

RAFT는 LLM이 도메인 특정 지식을 학습하면서 방해 문서(도메인과 관련 없는 내용의 문서)에 대해서는 저항하도록 학습합니다. 마치 open-book 시험을 준비하면서 관련 없는 문서를 무시하는 방법을 익히는 것과 유사합니다. 이를 통해서 RAFT는 LLM이 특정 도메인과 관련된 정보를 활용하면서 다양한 방해 문서에 대한 대응 능력을 갖출 수 있습니다.

구체적으로, 모델이 질문($Q$)에 대해서 문서($D^$)에서 답변($A^*$)를 생성하도록 학습합니다. 여기서 $A^*$에는 chain-of-thought reasoning가 포함되어 있고, 문서 사이에 방해문 ($D_k$)가 섞여있습니다. 구체적인 방법은 뒤에서 다뤄보도록 하겠습니다.

2. LLMs for Open-Book Exam

위에서 계속해서 LLM을 이용한 open-book 시험을 언급하고 있는데, 해당 파트에서 closed-book 시험, open-book 시험 그리고 특정 도메인에서의 open-book 시험의 정의와 차이점을 짚고 넘어가고자 합니다.

전통적인 시험 방식인 closed-book 시험은 학생이 추가적인 자료 없이 자신의 기억에 의존하여 질문에 답해야 하는 상황을 의미합니다. LLM도 마찬가지로, 사전 학습 및 지도 학습으로 습득한 지식만으로 질문에 답변합니다.

open-book 시험의 경우, 학생이 시험 중에 외부 자료를 참고할 수 있습니다. LLM도 이와 유사하게, 검색 엔진을 통해 관련 문서(k개)를 가져와 질문에 대해서 답변합니다.

특정 도메인에서의 open-book 시험에서는 LLM이 이미 해당 분야에 대해 충분한 사전 지식을 가지고 시험을 치르는 상황을 의미합니다 예를 들어서 기업의 내부 문서나 코드 repository와 같은 특정 데이터를 기반으로 질문에 답변하는 것이죠. 본 논문에서는 특정 도메인에서의 open-book 시험을 잘 보기 위해서 목표합니다.

3. RAFT

먼저, supervised fine-tuning(SFT) 방식에 대해서 설명하고자 합니다. 전통적인 SFT는 주어진 질문($Q$)에 대해 모델이 답변($A$)를 생성하도록 학습합니다. 이렇게 학습된 모델은 test-time에 RAG와 함께 사용할 수 있고, 질문에 대한 모델의 답변 생성을 돕기 위해 prompt에 문서를 추가할 수 있습니다. 이는 다음과 같이 표현가능합니다.

{Train: $Q \rightarrow A$}, {0-shot Inference: $Q \rightarrow A$}, {RAG Inference: $Q+D \rightarrow A$}

RAFT(Retrieval Augmented Fine-Tuning)에서는 특정 도메인에서 질문에 대한 답변의 정확성을 향상시키기 위해 설계되었습니다. RAFT에서는 각 데이터 포인트가 질문($Q$), 문서 집합($D_k$), 해당 문서($D^*$) 중 하나에서 생성된 chain-of-thought 스타일의 답변($A^*$)를 포함하도록 학습 데이터를 준비합니다.

RAFT는 LLM이 질문에 답변하기 위해 필요한 정보를 더 잘 이해하고, 불필요한 정보를 무시할 수 있도록 만드는데, 이 방법의 핵심은 두 가지 문서 유형을 활용하는 것입니다.

• Golden Document($D*$): 질문에 대한 올바른 답변을 포함하는 문서
• Distractor Document($D_i$): 질문의 답변과 직접적으로 관련이 없는 문서

먼저, 데이터셋에서 $P$ 비율에 해당하는 질문($q_i$)에 대해 golden 문서($d^*_i$)와 districtor 문서($d_{k-1}$)를 유지합니다. ($1-P$) 비율의 질문($q_i$)에는 golden 문서를 포함하지 않고 오직 distroctor 문서($d_k$)만 포함합니다. 이러한 방식으로 모델을 SFT으로 미세 조정하여 주어진 문서와 질문으로부터 답변을 생성하도록 학습합니다. Figure 2를 통해 해당 방식이 어떻게 돌아가는지 파악할 수 있습니다.

위에서 설명한 방식을 통해 본 논문에서 제안한 RAG 방법론은 다른 RAG 방식에 비해서 더 높은 성능을 보이도록 학습됩니다. 일부 경우에서는 golden 문서를 제거함으로써 모델이 문서에서 답변을 도출하는 대신 이를 기억하도록 강요하기도 합니다. RAFT의 학습 데이터는 다음과 같으며, 학습 데이터의 예시는 Figure 3으로 확인할 수 있습니다.

$P$% of data: $Q+D^*+D_1+D_2+ … + D_k \rightarrow A*$

($1-P$)% of data: $Q+D_1+D_2+…+D_k \rightarrow A*$

RAFT의 또 다른 중요 포인트는 Chain-of-Thought(CoT)를 이용하여 답변을 생성하는 것입니다. 이를 통해서 답변의 정확성도 높이며 전체적인 reasoning chain을 만들고 명확하게 인용하여 모델의 질문 응답 정확성을 높입니다.

4. Evalution

실험 테이블에 등장할 베이스라인에 대해서 간단히 설명하면 아래와 같습니다.

• LLaMA2-7B : LlaMA2-7B-chat model with 0-shot prompting을 의미합니다. QA task에 대해서 instruction-finetuning이 진행된 모델이라 말할 수 있습니다.
• LLaMA2-7B + RAG : 위에 셋팅에서 reference document를 포함한 버전입니다.
• DSF : Domain-Specific Finetuning with 0-shot prompting으로, 특정 도메인에 대해서 supervised-finetuning을 수행한 ㅁ델입니다.
• DSF + RAG : Domain-Specific Finetuning with RAG로, 위에 셋팅에서 RAG가 추가된 버전입니다.

Table 1에 있는 데이터셋은 모두 QA Task에서 사용하는 데이터셋이며, 본 논문에서 제안한 RAFT가 일관되게 우수한 성능을 보인다는 것을 확인할 수 있습니다. 특히나 HuggingFace, TensorFlow에서 월등히 뛰어난 성능을 보여주는데, RAFT가 문서에서 정보를 효과적으로 추출하고 노이즈를 처리하는데 강하다는 것을 알 수 있습니다.

Table 2에서는 CoT에 대한 ablation study를 보여줍니다. Table 2에서 볼 수 있듯이, 질문에 대해서 답변을 단순히 제공하는 것만으로는 충분하지 않는 것을 볼 수 있습니다. 또한, 질문에 대해서 답변을 잔순히 제공하는 방법은 loss가 급격하게 감소되어 모델이 과적합되지 쉽다고 합니다. CoT을 적용하게 되면 정확도가 향상될 뿐만 아니라 간결한 답변에 대한 과적합을 방지할 수 있습니다.

Figure 4를 통해서 RAFT와 DSF 간의 답변을 비교할 수 있습니다. DSF 모델의 경우 각본가에 대해 묻는 질문에 대해서 제대로 답변하지 못하는 것을 볼 수 있는데, RAFT의 경우 질문에 정확하게 답변하는 것을 볼 수 있습니다.

논문의 저자는 LLM이 항상 RAG를 위해 golden 문서로만 학습되어야 하는지에 대해서 의문을 가지며 학습 데이터의 어떤 비율(p%)이 golden 문서를 포함해야 하는지?에 대해서 질문합니다. 직관적으로, 문서에서 정보를 읽고 추출하는 학습을 위해서 학습 중에 항상 golden 문서가 포함되어야 한다고 가정할 수 있습니다(P=100%). 하지만 저자들이 실험해 본 결과, golden 문서가 포함되지 않은 학습 데이터의 비율이 80%에 이르렀을 때 모델의 RAG 성능이 오히려 향상되었다고 합니다.

Figure 5에서 해당 실험을 구체적으로 확인할 수 있습니다. P%는 학습 데이터에서 golden 문서가 포함되는 비율을 나타냅니다. 실험 결과에 따르면 데이터셋 마다 최적의 비율이 다르며, 40%, 60%, 100%의 비율에서 효과가 다르다는 것을 확인할 수 있습니다. 이를 통해 데이터셋에 따라 일정 비율의 golden 문서가 포함되지 않은 학습 데이터가 RAG 성능 향상에 도움이 될 수 있음을 확인할 수 있습니다.

추가로 논문의 저자는 RAFT에서 distractor 문서의 수가 모델의 성능에 어떤 영향을 미치는지, 특히 평가 시 top-k RAG의 성능에 대해서 답변을 얻고자 실험을 수행하였습니다. 이전 연구에서 LLM이 관련 없는 텍스트에 취약하다는 점과 top-k RAG가 테스트 시에 높은 재현율을 보장하기 때문입니다. 해당 수치가 높기 위해서는 모델이 관련 없는 내용을 식별하고 무시할 수 있는 능력을 가져야하며, 오직 관련 정보에만 집중해야 합니다.

Figure 6을 보면, golden 문서만으로 fine-tuning 할 경우 다양한 distractor 문서를 포한하는 경우보다 성능이 낮아지는 경우가 많은 것을 볼 수 있습니다. Natural Question데이터셋, Hotpot QA 데이터셋 모두 $$D*$$만 사용했을 때보다 $$D$$를 포함했을 때 성능이 높습니다. 물론 데이터셋 별로 distractor 문서를 몇개 사용했을 때 최적의 성능을 보이는 지는 다릅니다.