はじめに

RAGを使用する際、Contextにはどうしても無関係なドキュメントが入ってしまいます。そのため、RAGで使用するLLMは、このようなノイズに対して頑健である必要があります。また、LLMは無関係なドキュメントを無視し、ゴールデンドキュメント(回答の根拠が含まれるドキュメント)を正確に抽出する必要があります。
では、どうやって実現すれば良いのでしょうか？このような精度や頑健性を向上させるための手法として RAFT があります。

今回は、LLMをトレーニングする手法であるRAFT (Retrieval Augmented Fine-Tuning) について解説します。RAFTの原論文である RAFT: Adapting Language Model to Domain Specific RAG をベースに解説していきます。

本記事はRAGに関する知識を前提としています。RAGに関して不安のある方は、【RAG入門】RAGって何？なぜRAGが注目を浴びるのか？を読んでみてください！

目次

RAFTのモチベーション
RAFTの手法
実験詳細
結果
CoT (Chain-of-Thought) に効果はあるのか
ゴールデンドキュメントを常に含めるべきか？
RAGへの応用
所感
参考文献

RAFTのモチベーション

なぜ、RAFTという手法を開発しようと思ったのか、その背景について説明します。

近年、LLMは医療や法律などの専門領域で活用される機会が増えてきています。これらの領域では一般知識への推論能力はそれほど重要ではなく、事実性を担保した状態での正確な回答が重要となります。実際、LLMを専門領域に適応させることは多くの新しいアプリケーションにとって不可欠です。
では、実際に開発する際に、RAGを専門的な領域に適応させるには、事前学習済みのLLMをどのようにチューニングすればよいのでしょうか？

LLMを専門領域に適応させる一般的な方法は以下の2つです。

RAGを用いたインコンテキスト学習
教師ありファインチューニング

RAGを用いたインコンテキスト学習では、固定されたドメイン環境やテストドキュメントへの事前アクセスによる学習機会を活かすことができません。 一方で、教師ありファインチューニングでは、ドキュメント内のより一般的なパターンを学習し、最終的なタスクやユーザーの好みに適合させることができます。しかし、「テスト時にドキュメントを活用しない（RAGを統合していない）」「学習時に検索プロセスを考慮していない」といった課題があります。

これらの課題を克服するための手法がRAFTです。

RAFTの手法

RAFTは、LLMをファインチューニングする際に、専門領域の知識を獲得しながら、RAGの性能を向上させることを目的としています。
ここにおけるRAGの性能向上とは、

誤った検索結果(ディストラクタードキュメント)に対しても頑健な推論ができるようになること
LLMの推論精度の向上

を指します。つまり、RAFTは

LLMが適切なドキュメントと不要なドキュメントを識別して推論できるようになるためのトレーニング
LLM本体の推論精度を向上させるようなトレーニング

を行います。

ここから、具体的な手法の話をしていきますが、小難しい話はしません。なぜなら、RAFTは特別なトレーニング手法を行っていないからです。実際に、標準的な教師ありファインチューニングの手法を用いています。 では、どうやって「専門領域の知識を獲得しながら、RAGの性能を向上させる」のでしょうか？
その答えは、トレーニングデータの工夫にあります。
まず、RAFTのトレーニングデータは以下の構成になっています。

質問（$Q$）
ドキュメントセット（$D_k$）
- ゴールデンドキュメント ($D*$)
- ディストラクタードキュメント($D_i$)
Chain-of-Thought形式の回答（$A*$）

ただし、ゴールデンドキュメントは必ずしも1つではなく、データセットによっては複数のドキュメントから成ることもあります。

適切なドキュメントと不要なドキュメントを識別するためのトレーニング

適切なドキュメントと不要なドキュメントを識別するためのトレーニングは、以下のようなトレーニングセットで行われます。

\displaylines{
P\% \: of \: data: Q+D*+D_1+D_2+...+D_k → A*\\
(1-P)\% \: of \: data: Q+D_1+D_2+...+D_k → A*
}

つまり、$P$%のトレーニングセットにはゴールデンドキュメントが存在し、残りの($1-P$)%のトレーニングセットはディストラクタードキュメントのみで構成されるように設計しています。実際には、$P$%のトレーニングセットにおいては、ゴールデンドキュメント+４つのディストラクタードキュメントで構成されています。
また、テスト時には、質問（$Q$）と、RAGパイプラインで検索されたトップ$K$のドキュメントを使用します。ポイントは、RAFTが特定のRetriever（検索システム）に依存しないという点です。 というのも、あくまで学習しているのは、「適切なドキュメントと不要なドキュメントを識別すること」なので、Retrieverが何で、どんなドキュメントを持ってこようが関係ないといった具合です。

Retrieverの品質は、質問に対して「わからない」という回答ではなく、意味ある回答をする上では極めて重要です。しかし、RAFTの本質はそこではありません。

LLM本体の推論精度を向上させるためのトレーニング

LLM本体の推論精度を向上についてはトレーニングデータの回答形式をChain-of-Thought形式の回答（$A*$）にすることで実現します 。$A*$にCoTを組み込むことで、

完全な推論チェーンの生成
##begin_quote## タグと ##end_quote## タグを用いた回答の根拠となるドキュメントの明示的な引用

をトレーニングするよう促します。そうすることで、元のContextを適切に参照しながら推論チェーンを形成するようになり、回答の説明性と正確性を向上させているのです。

トレーニングデータのイメージとしては、以下の図のようになります。

Tianjun Zhang et al.(2024) RAFT: Adapting Language Model to Domain Specific RAG

実験詳細

ここでは、実験で使用したデータセットとRAFTモデルおよび各ベースラインモデルについて説明します。

データセット

Wikipedia・コーディング/APIドキュメント・医療文書に関する質問応答のデータセットを用いて評価しています。Wikipediaベースのデータセットは、主に映画やスポーツなどの一般知識が含まれます。一方で、コーディング/APIドキュメントと医療文書のデータセットに関しては、専門領域のドキュメントであり、専門領域における応答性能を測るために使用しています。

Wikipediaベースのオープンドメイン質問応答データセット
- Natural Questions
- Trivia QA
- HotpotQA
コーディング/APIドキュメントのデータセット
- APIBenchに含まれるHuggingFace・Torch Hub・TensorFlow Hub
医療文書に関する質問応答のデータセット
- PubMed QA

RAFTモデルおよびベースラインモデル

RAFTには、LLaMA2-7B-chatモデルを使用します。 このモデルは、LLaMA2-7Bに対して質問応答タスク向けにInstruction Tuningが施されたモデルです。
ベースラインモデルとしては4種類用意されています。

LLaMA2-7B-chatモデルを0-shot promptingで使用
LLaMA2-7B-chatモデルにRAGを組み合わせたLLama2 + RAGモデル
ドメイン固有のファインチューニングを施したモデル（DSF: Domain-Specific Fine-Tuning）
ドメイン固有のファインチューニングを施したモデルにRAGを組み合わせたDSF + RAGモデル

DSFモデルには、おそらくLLaMA2-7B-chatモデルが使用されていると思いますが、明記されていませんでした。また、DSFを使用する理由としては、特定のドメインのContextに慣れさせることが目的です。一般的には、DSFモデルの方が精度は良くなります。
これらのベースラインモデルに加えて、本実験では参考としてGPT-3.5+RAGでの評価も行っています。やはり評価を行ううえで、GPTシリーズの結果は皆さん気になりますし、無視できないモデルになってきている印象です。

結果

Tianjun Zhang et al.(2024) RAFT: Adapting Language Model to Domain Specific RAG

ここからは、ベースラインとRAFTを比較し、RAFTの有効性について説明します。
まず、上記の表からRAFTが一貫してベースラインを上回っていることがわかります。

※全てのデータセットにおける結果が載っているわけではありません。以降、上図をもとに議論していきます。

次に、個々のベースラインと比較します。全体として、LLaMA-7B-chatは、RAGを用いるかどうかに関係なく、回答スタイルのズレにより、単体では十分な性能を発揮しなかったようです。(※表からは読み取れません。論文から引用しています)
しかし、ドメイン固有のファインチューニングを適用することで、性能が大幅に向上することが確認できます。 一方で、RAGをDSFモデルに適用しても、必ずしも性能向上にはつながらないことが確認できます。これは、モデルがContextからゴールデンドキュメントを抽出できていない可能性を示唆しています。
したがって、RAFTでは

適切な回答スタイルの学習
ドキュメントの抽出能力の向上とディストラクタードキュメントに対する耐性の強化

を実現できるということがわかります。
最後にGPT-3.5+RAGとの比較では、HotPot QAについては差が開いていますが、他のデータセットでは上回っています。モデルサイズが異なることを踏まえて考えると、RAFTに大きな優位性があるように思えます。

CoT (Chain-of-Thought) に効果はあるのか

Tianjun Zhang et al.(2024) RAFT: Adapting Language Model to Domain Specific RAG

結果からRAFTがモデルの性能向上に寄与していることがわかりました。ここでは、トレーニングデータの回答($A*$)をChain-of-Thought形式にしたことがどの程度有効だったのかを評価します。
上図から、単純に質問に対する回答を提供するだけでは、十分な性能を発揮できない場合があることがわかります(w.oはwithoutの略です)。
少し技術的な話になりますが、CoTなしの回答形式ではLossの急激な低下が生じ、モデルが短い回答に過度に適応してしまうそうです。つまり、ファインチューニングの本質である知識の学習ではなく、回答形式に過度に適応してしまう状況に陥ります。例えば、Yes/Noであれば、2~3文字を出力することを学習し、本当にYes/Noであったかどうかは学習していないといった感じです。
繰り返しになりますが、上図から回答にCoTを組み込むことで全体的な精度が向上していることがわかります。つまり、CoT形式の回答はモデルが適切な学習を行う上で、極めて重要な要素であることがわかります。

ゴールデンドキュメントを常に含めるべきか？

本研究のトレーニングセットは、

\displaylines{
P\% \: of \: data: Q+D*+D_1+D_2+...+D_k → A*\\
(1-P)\% \: of \: data: Q+D_1+D_2+...+D_k → A*
}

でした。この$P$は何％にするべきなのでしょうか？

一見すると、RAGタスクにおいてゴールデンドキュメントを抽出する能力を向上させるためには、トレーニング時に常にゴールデンドキュメントを含めるべき（$P = 100%$）だと考えられるでしょう。
しかし、結果は私たちの予想に反して、必ずしもそうとは限らないそうです。具体的には、最適な$P$%はデータセットごとに異なり、40%、60%、100%といった範囲で変動します。したがって、モデルは完全な正解情報がない状況でも推論を行う能力を養い、より頑健な回答生成が可能になると考えられます。

Tianjun Zhang et al.(2024) RAFT: Adapting Language Model to Domain Specific RAG

RAGへの応用

ここでは、RAGのRetrieveにおける次の問いについて考えます。

RAFTにおけるディストラクタードキュメントの数がRAGの$Top$-$K$検索にどのような影響を与えるのでしょうか？

前提として、これまでの研究でLLMが無関係なテキストを入れると効率性と性能に問題が生じると指摘されています。(Shi et al.(2023) Large Language Models Can Be Easily Distracted by Irrelevant Context; Weston & Sukhbaatar (2023) System 2 Attention (is something you might need too); Liu et al.(2023) Lost in the middle: How language models use long contexts.)

結果は下図の通りです。この図では、トレーニング時に異なる数のディストラクタードキュメントを含めたモデルが、テスト時のドキュメント数の変化に対してどのような反応をするのか示しています。結果としては、トレーニング時にディストラクタードキュメントを含めた場合、テスト時のドキュメント数の変動に対してモデルの性能がより安定することがわかりました。
この結果は、RAFTの頑健性を裏付けるものであり、テスト時に遭遇するドキュメント数が変化しても、モデルの性能が大きく低下しないことを示しています。

ただし、データセットによって必要となるディストラクタードキュメントの数は異なるそうで、場合によってはゴールデンドキュメントのみの場合より性能が低下しています。ディストラクタードキュメント数の見極めが大切ですね。

Tianjun Zhang et al.(2024) RAFT: Adapting Language Model to Domain Specific RAG

所感

LLMにRAFTを適用することで、無関係なドキュメントに対してより頑健なモデルを手に入れることができるようになります。これはRAGにおいて非常に重要です。なぜなら、Retrieveにおいてほぼ確実に無関係なドキュメントが含まれてしまうからです。
また、私個人としての大きな発見として、

ゴールデンドキュメントを含めないデータセットを作った方が精度が上がる
CoT形式の回答が過学習に対して効果的であった
実際にRAGに適応した際にも安定して精度向上が見られた

というものがありました。PrePrintということもあり、少し荒さも見えましたが、非常に有益な結果を示していると思います。(※記事執筆時の2025/3/1現在)
私自身、LLMのファインチューニングにおいて、過学習にかなり悩まされたためCoT形式の回答は打開策になりそうです！

参考文献