Source-linked AI summary
RAFT: Adapting Language Model to Domain Specific RAG
Tianjun Zhang, Shishir G. Patil, Naman Jain, Sheng Shen, Matei Zaharia, Ion Stoica, Joseph E. Gonzalez
TL;DR
Specialized-domain RAG raises the question of how pretrained LLMs should learn from fixed document collections while handling imperfect retrieval. RAFT combines fine-tuning with golden and distractor documents, chain-of-thought reasoning, and source citations; it consistently improves performance across several in-domain datasets.
Problem
Pretrained LLMs need adaptation for specialized domains, but existing RAG and fine-tuning methods do not fully combine domain learning with document-grounded answering under imperfect retrieval.
Method
RAFT fine-tunes models on questions paired with golden and distractor documents, generating chain-of-thought answers that cite relevant source text.
Results
RAFT consistently outperforms supervised fine-tuning with and without RAG across PubMed, HotpotQA, and Gorilla datasets.
Takeaways & Limitations
RAFT is a post-training recipe for improving pretrained LLMs in domain-specific RAG and open-book question answering.
Abstract
from arXiv · showhide
Pretraining Large Language Models (LLMs) on large corpora of textual data is now a standard paradigm. When using these LLMs for many downstream applications, it is common to additionally bake in new knowledge (e.g., time-critical news, or private domain knowledge) into the pretrained model either through RAG-based-prompting, or fine-tuning. However, the optimal methodology for the model to gain such new knowledge remains an open question. In this paper, we present Retrieval Augmented FineTuning (RAFT), a training recipe that improves the model's ability to answer questions in a "open-book" in-domain settings. In RAFT, given a question, and a set of retrieved documents, we train the model to ignore those documents that don't help in answering the question, which we call, distractor documents. RAFT accomplishes this by citing verbatim the right sequence from the relevant document that would help answer the question. This coupled with RAFT's chain-of-thought-style response helps improve the model's ability to reason. In domain-specific RAG, RAFT consistently improves the model's performance across PubMed, HotpotQA, and Gorilla datasets, presenting a post-training recipe to improve pre-trained LLMs to in-domain RAG. RAFT's code and demo are open-sourced at github.com/ShishirPatil/gorilla.
1 Introduction
Specialized-domain applications require LLMs to answer accurately from fixed document collections, but existing RAG and fine-tuning approaches do not fully combine domain learning with imperfect retrieval. RAFT combines instruction fine-tuning and retrieval-augmented generation to train models to use relevant documents and ignore distractors.
- Specialized-domain applications prioritize accuracy grounded in document collections, including legal, medical, news, enterprise, and software resources.
- The paper asks how to adapt pretrained LLMs for Retrieval Augmented Generation in specialized domains.
- RAG can reference documents but does not leverage fixed-domain learning, while fine-tuning may omit test-time documents or retrieval imperfections.
- Existing approaches resemble either taking an open-book exam without studying or memorizing documents without practicing document-referenced answering.
- RAFT combines instruction fine-tuning with RAG and trains models to learn domain knowledge while recognizing relevant and distracting retrieved information.
- RAFT consistently outperforms supervised fine-tuning with and without RAG across PubMed, HotpotQA, and Gorilla datasets.
2 LLMs for Open-Book Exam
The paper distinguishes closed-book, general open-book, and domain-specific open-book settings. Its focus is the domain-specific case, where the model is fine-tuned for a known document domain and must answer with retrieved information despite varying retrieval inputs.
- Closed-Book Exam: A closed-book setting gives the model no additional documents, so answers draw on knowledge acquired during pretraining and supervised fine-tuning.
- Open Book Exam: An open-book setting appends k retrieved documents or document segments to the user prompt, providing access to domain-specific information.
- Open Book Exam: General open-book performance depends largely on retriever quality and its ability to identify the most relevant information.
- Domain-Specific Open-Book Exam: The domain-specific open-book setting assumes the testing domain is known in advance and answers can be found within its document collection.
- Domain-Specific Open-Book Exam: This setting motivates adapting pretrained LLMs to a specific domain while improving robustness to varying numbers of retrieved documents and distractors.
3 RAFT
RAFT prepares fine-tuning data for domain-specific open-book answering by pairing questions with golden and distractor documents, then training chain-of-thought answers that cite relevant context. The procedure combines supervised training with document reading and imperfect-retrieval conditions.
- RAFT: RAFT is a modified instruction-tuning recipe for training LLMs in domain-specific open-book settings.
- RAFT: Each RAFT example contains a question, documents, and a chain-of-thought answer generated from one or more golden documents.
- RAFT: Distractor documents contain no answer-relevant information, while golden documents contain information from which the answer can be deduced.
- RAFT: For a fraction P of examples, training includes the golden document with distractors; for the remaining examples, it includes only distractors.
- RAFT: The model is fine-tuned with standard supervised training to generate answers from the question and provided documents.
- RAFT: At test time, RAFT receives the question and top-k documents retrieved by a RAG pipeline, independently of the retriever used.
- RAFT: RAFT training data uses reasoning chains that cite source passages and explain how those citations support the answer.
4 Evaluation
RAFT is evaluated across diverse question-answering and document domains against general-purpose, domain-specific fine-tuning, and RAG baselines. It consistently improves performance, while ablations and qualitative analyses examine chain-of-thought reasoning, distractor handling, and golden-document proportions.
- Evaluation setup: RAFT is evaluated on Wikipedia question answering, coding/API documents, and medical question answering against LLaMA, RAG, and domain-specific fine-tuning baselines.The datasets include Natural Questions, TriviaQA, HotpotQA, HuggingFace, Torch Hub, TensorFlow Hub, and PubMed.
- 4.1 Results: RAFT consistently outperforms the evaluated baselines across specialized domains, including gains of 35.25% on HotpotQA and 76.35% on Torch Hub over base Llama-2 instruction tuning.Against DSF, reported gains include 30.87% on HotpotQA and 31.41% on HuggingFace; PubMed QA shows no significant gain over DSF + RAG.
- 4.1 Results: RAFT improves both answer-style alignment and document processing, whereas adding RAG to domain-specific fine-tuning does not invariably improve outcomes.The reported interpretation is that RAFT trains models to use provided context while matching the required answering style.
- 4.2 Effect of CoT: Adding chain-of-thought improves RAFT performance, with gains of 9.66% on HotpotQA and 14.93% on HuggingFace compared with RAFT without CoT.The analysis reports that reasoning chains improve training robustness and help prevent overfitting to concise answers.
- 4.3 Qualitative Analysis: In a HotpotQA example, DSF answers with a film title when asked for a screenwriter, while RAFT extracts the correct answer from the context.The example contrasts DSF’s context-extraction error with RAFT’s accurate response.
- 4.4 Should we train the LLM always with the golden context for RAG?: The optimal golden-document proportion varies across datasets, with reported values of 40%, 60%, and 100%; some training instances omit the golden document.The training setup uses four distractor documents alongside the golden document, and testing maintains this format.
5 RAFT Generalizes to Top-K RAG
RAFT improves robustness to retrieved-document variation by training with distractors rather than only golden documents. Optimal training document counts vary by domain, while distractor training improves resilience to changing test-time document numbers.
- Training solely with golden documents can diminish the model’s ability to disregard irrelevant retrieved information.
- RAFT combines one golden document with distractor documents to prepare models for imperfect retrieval.The reported experimental setup uses one golden document alongside four distractor documents.
- Optimal training document counts differ across domains: four for NQ, three for Trivia QA, and two for Hotpot QA.
- Training with only golden documents leads to poor performance when the number of test-time documents varies.
- Training with distractor documents makes RAFT more resilient to fluctuations in the number of documents presented at test time.The paper presents this variable-document performance as evidence of improved robustness.
6 Related Works
Related work covers retrieval-augmented language models, memorization in neural models, and fine-tuning for RAG. RAFT differs from prior fine-tuning settings by focusing on testing over the same documents as training.
- Retrieval-augmented language models integrate retrieval modules that source relevant information from external knowledge bases for NLP tasks.
- Prior memorization research examines whether neural language models understand text or rely on surface-pattern memorization.Related studies also quantify memorization, its effects on generalization, and associated privacy concerns.
- Recent fine-tuning-for-RAG work constructs training data to improve pretrained LLMs on retrieval-augmented tasks.
- Unlike settings where test domains or documents differ from training, RAFT studies testing on the same set of documents.
7 Conclusion
RAFT is designed for domain-specific open-book question answering. Its key design choices combine distractor documents, examples without golden context, and chain-of-thought answers with direct quotations.
- RAFT targets domain-specific open-book question answering by adapting models to use retrieved documents during evaluation.
- RAFT trains with distractor documents, sometimes without golden documents, and uses chain-of-thought answers containing direct quotations.
- Evaluations on PubMed, HotpotQA, and Gorilla API Bench demonstrate RAFT’s potential.