はじめに

RAGのデータストアにどのような構造でデータを格納するかは極めて重要です。その理由は、RAGの仕組みに起因します。RAGにおいて、LLMは検索結果の情報をもとに回答を生成します。関連していない情報をLLMに渡してしまうと期待通りの回答を得られません。そのため、文字情報同士(クエリとドキュメント)の関連度を正確に測ることが重要です。また、仮に関連している情報を検索できたとしても、回答する上で十分な情報が載っている(LLMが理解しやすい形で)必要があります。前者は、Embedding ModelのチューニングやBM25のようなキーワード検索手法の改善が主な方法です。

今回は後者の 「回答する上で十分な情報が載っている必要がある」 という課題に対してアプローチした新しい方法である、RAPTOR (Recursive Abstractive Processing for Tree-Organized Retrieval) について解説します。RAPTORの原論文である RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval をベースに色々補足しながら解説していきます。

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

目次

RAGの課題
RAPTORとは？
実験詳細
結果
所感
参考文献

RAGの課題

RAGの課題は色々ありますが、今回はデータストアに紐づいた課題について考えていこうと思います。

まず RAG の課題を説明する前に、クエリの種類について説明したいと思います。大きく分けると、

シングルホップ型
1つの文や事実のみに基づいて答えが導ける質問のことを指します。また、このような質問は情報のつながりをたどる必要がないため、RAGにおいては比較的扱いやすい質問になります。例えば、「リンカーン大統領の生年月日は？」や「地球は何番目の惑星ですか？」などが当てはまります。
マルチホップ型
シングルホップ型の逆です。複数の情報源（文や段落）を結びつけて推論しないと答えが導けない質問のことを指します。RAGにおいては扱いにくい質問になります。例えば、「シンデレラはどのようにしてハッピーエンドに辿り着いたのか？」や「オバマ元大統領の妻の出身地は？」などが当てはまります。後者は一見、マルチホップ型のように見えませんが、「Step1: オバマ元大統領の妻は誰？→Step2: ミシェル・オバマの出身地は？」となるため、マルチホップ型に分類されます。

の２つに大別されます。上記の説明で、マルチホップ型の質問の方が難しいと書きました。直感的にもそのように感じると思いますが、具体的になぜ難しいのでしょうか？ その理由を理解するためには「どのようにデータストアを構築するのか？」を理解する必要があります。

では、データストアはどのよう構築するのでしょうか？
RAGシステムの種類によって様々なのですが、大体こんな感じという話をしようと思います。
データストアは以下の3つのステップで構築されます。

ドキュメントを任意のチャンクサイズで分割する
Embedding Model を用いて各チャンクをベクトル化する
各チャンクをデータストアに格納する

チャンクサイズをどの程度にするかは扱うドキュメントの種類に依存しますが、大概は512トークン以下にします。これは transformer のトークン制限に起因します。しかし、近年では gemini_embedding-001 のように512トークンを超えて扱うことのできるモデルも出てきています。ちなみに gemini_embedding-001 は2048トークンまで扱うことができます。

また、検索は以下の3つのステップで行われます。

クエリを Embedding Model を用いてベクトル化する
ベクトル化したクエリと各チャンクのベクトルを$cos$類似度や$k$近傍法を用いて関連度評価する
top-k を検索結果として取得する

つまり、このように構築されたデータストアに対しては、短く連続したテキストチャンクしか検索できません。 したがって、マルチホップ型の質問に対しては、適切ではありません。なぜなら、上位 $k$ 件の短く連続したテキストを取得するだけでは十分な文脈が得られないからです。

短く連続したテキストチャンクとは？

「短く」とは、ドキュメントをそのまま扱わずチャンキングしていることを指しています。チャンクの方法は様々ですが、センテンスや固定長で分割します。固定長の場合は100~500トークン or ワードが多いのではないでしょうか？

「連続した」とは、スライドしながらドキュメントを分割していることを指しています。そのため、200トークン分あるドキュメント ($doc$)を100トークンずつ分割すると、$doc = chank_1 + chank_2$となります。これが「連続した」という意味です

では、チャンクサイズを大きくすれば良いのでしょうか？
これはこれで問題があります。理由は、大きく2つあります。
1つ目は、ベクトル化していることに起因します。 ベクトル化するということは、チャンクを固定次元のベクトル空間に埋め込むことを意味します。例えば、1024次元の空間に埋め込む場合、100トークン分のテキストでも10,000トークン分のテキストでも強制的に1024次元のベクトルに変換されます。したがって、チャンクサイズが大きければ大きいだけ、情報が圧縮され失われます。
2つ目は生成時に使用する LLM に起因します。 RAG では検索結果を Context として扱いますが、一般的にLLMは「離れた位置にある情報を十分に活用できない」「関連情報が長大なコンテキスト内に埋もれている場合、性能が低下する」という傾向があります。つまり、長大な Context をLLMに与えたとしても、冒頭の情報と末尾の情報を結びつけて推論するタスクや、生年月日のような特定の情報を抽出するタスク (シングルホップ型) においては、うまく回答生成することができません。

つまり、マルチホップ型の質問は短く連続したテキストチャンクでもうまくいかないし、チャンクサイズを大きくしても、Embedding ModelやLLMの特性によりうまくいかない可能性があります。

RAPTORとは？

このような課題に対する研究は RAPTOR 以外にも色々あります。では、RAPTOR の特徴は何でしょうか？
それは、マルチホップ型の質問だけではなく、シングルホップ型の質問にも同時に答えることができるようなデータストアを構築することにあります。

ここからは、実際にどのようにして実現するのかという具体的な話をしていこうと思います。
まず、RAPTOR は2つのフェーズに大別されます。

データストア構築フェーズ
検索フェーズ

です。
まずは、データストア構築フェーズについて説明します。ざっくり説明しているバージョンと詳しく説明しているバージョンの両方を書いたので、目的に合わせてどちらを読むかを選んでください。

ざっくり説明

RAPTORのデータストアはツリー構造をしています。

Parth Sarthi et al.(2024) RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval

このツリーは以下の3つのステップで構築されます。

テキストのチャンキング
クラスタリング
LLMによる要約

1 → 2 → 3 → 2 → 3 → ・・・といったように繰り返していくことで下位層から上位層へとツリー構造を構築していきます。つまり、何らかの類似性によってテキストをクラスタリングして、そのクラスタ内のテキスト集合を要約し、さらにその要約同士を何らかの類似性によってクラスタリングし、要約する・・・といったことを繰り返します。したがって、上位層ほどより広範囲な要約を保持し、下位層ほど詳細情報 (最下層ノードは原文) を保持するようになります。
ちなみに、RAPTORの論文では要約のLLMには、GPT-3.5 turbo を使っています。

詳しく説明

ここでは、RAPTORについて詳しくしていきます。

Parth Sarthi et al.(2024) RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval

まず、RAPTORはツリー構造をしたデータストアですが、このツリーは以下の3つのステップで構築されます。(上図中央)

テキストのチャンキング
ガウス混合モデルによるクラスタリング
LLMによる要約

1 → 2 → 3 → 2 → 3 → ・・・といったように繰り返していくことで下位層から上位層へとツリー構造を構築していきます。

テキストのチャンキング

固定長のチャンキングを行います。具体的には、長さ100の短い連続テキストに分割します。 固定長の欠点として、文章の途中で切れてしまう可能性があり、そうなると文脈の保持が難しいというものがあります。そのため、RAPTORでは、各チャンク内の文脈的・意味的な一貫性を保つために、文章が100トークンの制限を超える場合、その文全体を次のチャンクに移動させます。

100トークンを超える文章のチャンキングの例。100トークンを超える場合、手前にあるピリオドで文章を区切り、途切れてしまった文章を切り離します。

ちなみに、Embedding ModelにはBERTベースのエンコーダーであるSentence-BERTを使用しています。

ガウス混合モデルによるクラスタリング

クラスタリングでは、各チャンク(葉ノード)や要約(その他ノード)を何らかの類似性によりグルーピングします。RAPOTRでは、ソフトクラスタリングを採用し、アルゴリズムにはガウス混合モデル (GMM) を使用しています。 ソフトクラスタリングとガウス混合モデルの説明は後述します。
下図が、クラスタリングの概要です。

RAPTORのクラスタリングの詳細

まず、Embedding Modelにより取得したテキストチャンクのベクトルは高次元です。次元が低いモデルでは384次元程度、高いモデルでは、4096次元程度です。このような高次元ベクトルにGMMを使用すると距離指標(マハラノビス距離)がうまく機能しない可能性があります。マハラノビス距離というのは、クラスタの分布がある方向に広がっており、そこに沿って遠くても“普通”とみなし、狭い方向に外れていたら“遠い”とみなすという距離関数です。しかし、高次元空間においては全ての点が似たような距離になってしまう結果、「近い」「遠い」の差が曖昧になるため、機能しないことがあります。 そのため、RAPTORでは Uniform Manifold Approximation and Projection (UMAP) を用いて次元削減を行い、ベクトルを低次元にします。 そして、この低次元ベクトルに対してGMMによるクラスタリングを行います (大域的クラスタリング)。
しかし、このままでは後続処理がうまくいかない場合があります。それは、クラスタが大き過ぎる場合です。つまり、子ノードを多く含む結果、それらの合計トークン数が要約モデルのトークン制限を超えてしまうということが起こり得ます。そのような場合には、該当クラスタにおいて更なるクラスタリングを行います (局所的クラスタリング)。gpt-3.5-turbo の context window は16,385なので、大体164ノードが1つのクラスタに含まれる子ノード数の限界となります。

ここからは、ソフトクラスタリングとガウス混合モデルについて説明します。読み飛ばしていただいても構いません。
ソフトクラスタリングとは、データが複数のクラスタに所属することができるというクラスタリング手法です。物語や書籍など、トピックが複数分類にまたがるようなものにおいて、柔軟に対応できるというメリットがあります。

ソフトクラスタリングの例。5番のテキストチャンクが複数クラスタに属しています。これがソフトクラスタリングの特徴です。

次に、ガウス混合モデルについて説明します。
データが複数のグループ（クラスタ）に分かれていて、それぞれのグループはガウス分布（正規分布）に従っていると仮定してクラスタリングするモデルです。また、ガウス混合モデルの目標はデータ$x_i$がどのガウス分布 (クラスタ) から来た可能性が高いのか？を推定することです。
ガウス混合モデルでは、以下の式に従ってクラスタリングが行われます。以下の式は、データ$x$がクラスタ$k$に属する確率です。

\displaylines{
P(k\ |\ x)= \frac{\pi_k・P(x\ |\ k)}{P(x)}
}

ここで、$\pi_k$はクラスタ$k$を選ぶ確率であり、$\sum \pi_k=1$となります。また、$P(x\ |\ k)$と$P(x)$はそれぞれ以下のように求まります。

\displaylines{
P(x)=\sum_{k=1}^{K}\pi_k・P(x\ |\ k)\\
P(x\ |\ k)=\mathcal{N}(x;\  \mu_k, \Sigma_k)
}

言葉で説明すると、$P(x\ |\ k)$はクラスタ$k$におけるデータ$x$の生成確率です。これがガウス分布 (正規分布) に従っていると仮定しています。また、$P(x)$は全てのクラスタを考慮したデータ$x$の確率密度です。つまり、それぞれのガウス分布におけるデータ$x$の生成確率を合算したものという意味です。

ここで、実際にクラスタリングすることを考えてみましょう。2つの疑問が生じると思います。
クラスタ数$K$とモデルのパラメータ ($\pi_k,\ \mu_k,\ \Sigma_k$)はどのように決めるのでしょうか？
この2つを決めるために、ベイズ情報量規準 (BIC)とEMアルゴリズムを用います。
BICとは、モデルの当てはまりの良さと、モデルの複雑さをバランスよく評価する手法です。EMアルゴリズムとは、観測できない（隠れた）変数を含む確率モデルのパラメータ推定に使われる手法です。GMMでクラスタリングをする場合においては、パラメータ$\pi_k,\ \mu_k,\ \Sigma_k$を求めます。
クラスタ数$K$とモデルのパラメータは、BICとEMアルゴリズムを用いて以下の手順で決定していきます。

クラスタ数$k_0$でEMアルゴリズムを用いて、パラメータ$\pi_{k_0},\ \mu_{k_0},\ \Sigma_{k_0}$を求める
BICを用いてクラスタ数$k_0$のときのモデルを評価する
この手順を繰り返し、$k_0, k_1, ・・・, k_n$まで繰り返してBICのスコアが最も良かったクラスタ数$k_i$を選択する。このとき、パラメータは$\pi_{k_i},\ \mu_{k_i},\ \Sigma_{k_i}$です

上記の手順により、クラスタ数$K$とモデルのパラメータ ($\pi_k,\ \mu_k,\ \Sigma_k$) が決まったため、GMMによりクラスタリングを実行します。

LLMによる要約

GMMでクラスタリングしたノードの集合に対して、gpt-3.5-turboで要約を生成します。

要約生成時に使用したプロンプト。
Parth Sarthi et al.(2024) RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval

このようにすることで、クラスタ内の情報を統合することができます。例えば、リンカーン大統領に関する情報というラベルでクラスタリングが行われた場合、リンカーン大統領に対する概要が要約の内容になるといったイメージです。

LLMに精通した人であれば、1つの疑問が生じると思います。
それは、ハルシネーションは大丈夫なのだろうか？ だと思います。
LLMはたまに情報をでっちあげて根拠のないことを生成します (これを「ハルシネーション：日本語では幻覚」と言います) 。つまり、今回の場合においては要約がうまくできていない、または不自然なものになっている 可能性があるということです。RAPTOR においては実際に検証したところ、約4%の要約に軽微なハルシネーションが含まれていました。 具体的には、

モデルがトレーニングデータに含まれていた可能性のある細かな情報を要約に加えたケース
要約の際に誤って情報を拡張してしまうケース

です。しかし、全件を確認したところ、これらの誤りは親ノードには伝播せず、QAタスクへの影響も認められませんでした。つまり、LLMを用いることで稀に要約にハルシネーションが起きますが、ツリー構築やQAタスクを行う上では全く問題はないと言えます。 ("全く"は言い過ぎかもしれませんが...)
ハルシネーション分析における詳しい実験詳細や結果を知りたい方は原論文の付録 Eを読んでみてください。

次に、検索フェーズから説明していこうと思います。検索においては、2つの方法が提案されています。

Tree Traversal Retrieval
Collapsed Tree Retrieval

です。それぞれ、詳しく説明していきます。

Parth Sarthi et al.(2024) RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval

Tree Traversal Retrieval

この方法は、ツリー構造を活かした検索方法です。
まず、最上位ノードにおいてクエリとノードを$cos$類似度に基づいて比較し、$Top_k$のノードを選択します。さらに、選択されたノードの子ノードに対して$cos$類似度を計算し、$Top_k$のノードを選択します。これを最下層ノード (葉ノード) まで繰り返し、最終的に選択された全てのノードを検索結果とします。こうすることで、上位層ノードによる大まかな文脈把握と、下位層ノードによる詳細情報の把握を両立することが可能になります。 しかし、ツリーの深さ$d$と選択されたノード数$k$が固定されているため、常に同じレベル比でノードが選択されます。その結果、クエリの性質にかかわらず要約情報と詳細情報の比率が一定に保たれるという欠点があります。

$Top_k=2$の例。黄色が選択されたノードであり、赤枠が検索対象となるノードです。簡単のためクラスタ内のノード数は2つで固定しています。

Collapsed Tree Retrieval

これは、非常にシンプルな方法です。全てのツリーノードを単一レイヤーに平坦化し、一括で比較・選択するという方法です。

$Top_k=6$の例。黄色が選択されたノードです。簡単のためクラスタ内のノード数は2つで固定しています。

実験詳細

ここでは、データセットと予備実験について説明します。

データセット

本研究では、長文テキストの理解能力を評価することを目的としたデータセットを使用しています。

Narrative QA
- 書籍の全文や映画のスクリプトに基づいた質問応答ペアを含む
QASPER
- 1,585本のNLP論文にまたがる5,049の質問で構成され、各質問は論文本文に埋め込まれた情報を問う
QuALITY
- Project Gutenbergのフィクション作品やSlate誌の記事などで構成される、長文読解を目的とした多肢選択式の質問応答データセット

Tree Traversal Retrieval vs Collapsed Tree Retrieval

ここからは、本格的な実験を行う前の予備実験となります。RAPTORでは、検索手法としてTree Traversal Retrieval と Collapsed Tree Retrieval の2つを紹介しました。では、どちらの手法が優れているのでしょうか？

この実験では、QASPERの20個の論文を用いてContextの長さを変えながら、F1 Scoreで評価しています。F1 ScoreとはRecallとPrecisionのバランスを考慮した評価指標です。

Recall
適合率とも呼びます。関連ドキュメントのうち、実際に関連していると予測されたドキュメントの割合です

Precision
再現率とも呼びます。取得ドキュメントのうち、関連しているドキュメントの割合です

Parth Sarthi et al.(2024) RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval

結果は、上図の通りです。Collapsed Tree Retrieval の方が一貫して優れた性能を示しました。 Collapsed Tree Retrievalは全ノードを一度に比較できるため、質問に応じた適切な抽象度の情報を取得しやすいという柔軟性があります。一方で、Tree Traversal Retrieval では深さ$d$と検索数$k$が固定されているため、質問の性質にかかわらず要約情報と詳細情報の比率が一定に保たれるという制約があります。この質問の粒度に対する柔軟性の差が性能に現れたのだと考えられます。
上記の実験を踏まえて、以下の実験では全て Collapsed Tree Retrieval を使用し、最大トークン数は2000に設定します (上位20ノード相当) 。

ここでまた疑問が生まれます。結局、Collapsed Tree Retrieval はRAGの課題で説明した一般的なRetrievalの方法と大して変わらないということです。Collapsed Tree Retrieval では、せっかく作ったツリー構造を壊して検索するため、ツリー構造の恩恵を享受できない可能性があります。

Collapsed Tree Retrieval vs Dense Passage Retrieval

上記の疑問を明らかにするために、一般的なDense Passage Retrievalと比較します。この実験ではあらかじめ1500語に要約された「シンデレラ」の物語を使用します。これに対して、テーマを問うようなマルチホップ型の質問を２つします。

Q1: シンデレラはどのようにしてハッピーエンドを迎えるのか？
Q2: この物語の中心的なテーマは何か？

実際には、Q1・Q2ともに英語で記述されています

Q1: How does Cinderella find a happy ending?

Q2: What is the central theme of the story?

結果は、Q1・Q2ともにCollapsed Tree Retrievalの方が優れた性能を示しました。 具体的にQ1においては、Dense Passage Retrievalは「与えられたコンテキストでは、回答できない」と出力され、正答できませんでした。Q2においては、Collapsed Tree Retrievalは物語全体を反映した回答をしているのに対し、Dense Passage Retrievalでは、物語の前半部分だけに基づいて回答しました。また、下図からはDense Passage Retrievalが取得する情報は、Collapsed Tree Retrievalで取得する情報の中に直接あるいは間接的に含間れていることが読み取れます。 以上から、Collapsed Tree Retrievalはより関連性が高く網羅的な情報の取得ができるため、下流タスクにおいても優れた性能を発揮することが示されました。 したがって、ツリー構造の恩恵は平坦化しても、抽象的な情報を扱うことができるという大きなメリットがあるということがわかりました。

Parth Sarthi et al.(2024) RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval

結果

ここでは、3つの実験結果を紹介します。

ツリー構造の有用性評価

まずは、RAPTORの有用性を示すために、ツリー構造の有無でどの程度性能に差が出るかを評価しています。 結果は下図の通りです。

上段がNarrativeQAの結果であり、下段がQuALITYとQASPERの結果である。
Parth Sarthi et al.(2024) RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval

回答生成モデルには、UnifiedQA 3Bを使用していおり、Retrieverの種類を変えて実験を行っています。その結果、RAPTORを導入した方が全ての場合において性能が向上しました。 この結果から、RAPTORの手法の有効性が示唆されます。

回答生成モデルの違いによる性能評価

次に、回答生成モデルを変えたときに、RAPTORを導入したモデルはどの程度性能に差が出るのかを評価しています。 結果は下図の通りです。

上段がQASPERの結果であり、下段がQuALITYの結果である。
Parth Sarthi et al.(2024) RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval

QASPERとQuALITYにおいて、全てのLLMでRAPTORを使用した手法が一貫して他の手法の性能を上回りました。これはLLMに依存せずにRAPTORの手法が効果を発揮するということがわかります。 QASPERにおいては、論文内の情報統合が必要なため、要約情報を抽出する手法が優れた結果を示すのは当然だと言えるでしょう。しかし、想像通りにうまく機能し、実証できたことはRAPTORの有効性をより強めます。

SOTA(State-of-the-art)モデルとの比較

RAPTORの有効性は十分理解できました。では、各データセットにおけるSOTAモデルと比較したときに競争力のある手法となっているのでしょうか？ 結果は下図の通りです。

上段がQASPERの結果であり、中段がNarrativeQAの結果であり、下段がQuALITYとQASPERの結果である。
Parth Sarthi et al.(2024) RAPTOR : Recursive Abstractive Processing for Tree-Organized Retrieval

QASPERとQuALITYについては、最先端の性能を達成しました。 特に、QuALITYのTest SetとHard Subsetにおいては大きく性能が向上しました。ここで、Hard Subsetとは、人間が簡単には答えられない難しい質問を集めたデータセットです。また、NarrativeQAにおいては、METEORスコアにおいて最先端の性能を達成しました。 METEORスコアとは、BLEUやROUGE同様に生成文と元の文の一致度のスコアを算出するだけでなく、同義語や語順の違いも考慮して評価するという評価指標です。そのため、ROUGEやBLUEではRetriever+Readerに劣っていますが、METEORで最先端の性能を達成したことは、大きな意味を持ちます。以上から、RAPTORはSOTAモデルに大きく上回る、もしくは匹敵する性能を発揮できる手法であるということが言えます。

所感

RAPOTRという手法について紹介しました。いかがでしたでしょうか？
RAPTORは長文テキストを理解するというタスクにおいて大きな性能向上をもたらします。そして、ただ性能向上をもたらすだけでなく、最先端の性能も達成しています。これは、ツリー構造にすることで、さまざまな抽象レベルの要約情報から原文レベルの情報まで質問の粒度に合わせて取得できるためだと考えられます。
しかし、RAPTORは万能ではありません。私が、PuBMedQAというデータセットでRAPTORを使った際に、クラスタリングがうまく機能しませんでした。というのも、意味のあるまとまりでグルーピングしてくれず、そのため、上層の要約があまり意味あるものになりませんでした。やっていることが、長文テキストの文脈理解というタスクと少し違うので、当たり前と言われればそうかもしれませんが...少し期待していただけに残念でした。ただし、これはクラスタリングの問題であり、RAPTORの概念自体は素晴らしいもので、依然として活用できると考えています。クラスタリングの部分を少し改良して挑戦してみたいと考えています！

参考文献