メインコンテンツまでスキップ

RAG (Retrieval-Augmented Generation) とは?

概要 (画像は、Geminiで作成されたものです)

RAGの概要

Retrieval-Augmented Generation(検索拡張生成、以下RAG)は、2020年にFacebook AI Research(現在のMeta AI)やロンドン大学(UCL)、ニューヨーク大学(NYU)の研究者らによって発表された論文『Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks』において提唱された概念です。

大規模言語モデル(LLM)の能力を外部データベースからの知識と組み合わせることで拡張する、強力かつ有望なソリューションとして広く普及しています。

LLMは驚異的なテキスト生成能力を示しますが、知識集約型のタスクにおいては、 外部のノンパラメトリックメモリ(検索可能なドキュメントのインデックスなど)とパラメトリックメモリ(事前学習済みのニューラルネットワークの重み)を組み合わせることで、 パフォーマンスを劇的に向上させることができます。 RAGは、モデルが学習時に見たことのない最新の情報や、ドメイン特化の専門知識を動的に統合し、正確で信頼性の高いテキスト生成を可能にします。

これまでのパラダイムの課題とRAGの位置づけ

言語モデルの活用において、RAGの登場以前(あるいは純粋なLLM単体での運用)には主に以下のような課題が存在していました。

  • ハルシネーション(幻覚)の生成: LLMは、自身のトレーニングデータに含まれないクエリや最新の情報を求められた際、事実とは異なる内容を生成してしまう傾向があります。
  • 知識の陳腐化と更新の困難さ: モデルのパラメータに知識を保存する(パラメトリックな)アプローチでは、世界の変化に合わせて知識を簡単に拡張または修正することが困難です。
  • 推論プロセスのブラックボックス化: LLM単体では、その回答を導き出した根拠や決定の出処(プロベナンス)を提示することができません。

RAG は、これらの課題を解決するためのハイブリッドなアプローチです。 事前学習済みの生成モデルに対し、微分可能なアクセス機構を備えた外部メモリ(Wikipediaの密ベクトルインデックスなど)を付与します。 これにより、知識の直接的な修正や拡張が可能になり、モデルが参照した知識をユーザーが検証・解釈できるようになります。

RAGの処理概要(Naive RAG)

処理概要 (画像は、Geminiで作成されたものです)

RAGの最も基本的な研究パラダイムである「Naive RAG」は、「Retrieve-Read(検索・読解)」フレームワークとして特徴づけられ、主に以下の3つのステップで構成されます。

  1. Indexing(インデックス作成):
  • PDF、HTML、Wordなどの多様なフォーマットの生データをクリーニングし、統一されたプレーンテキストに変換します。
  • 言語モデルのコンテキスト制限に合わせるため、テキストを小さなチャンク(塊)に分割します。
  • 埋め込みモデル(Embedding model)を用いてチャンクをベクトル表現にエンコードし、ベクトルデータベースに保存します。
  1. Retrieval(検索):
  • ユーザーからクエリを受け取ると、インデックス作成時と同じエンコードモデルを使用してクエリをベクトル表現に変換します。
  • クエリのベクトルとコーパス内のチャンクのベクトル間で類似度スコアを計算し、最も類似性の高い上位 KK 個のチャンクを抽出します。
  1. Generation(生成):
  • 抽出されたドキュメントチャンクと元のクエリを合成し、一貫したプロンプトを作成します。
  • このプロンプトをLLMに入力し、タスクに応じた最終的な回答を生成させます。

RAGの構成要素(詳細)

初期のRAGモデルは、以下の2つの事前学習済みコンポーネントをエンドツーエンドで組み合わせて構築されました。

  1. Retriever(検索器):
  • クエリ xx が与えられた際に、テキストパッセージの分布を返すコンポーネントです。
  • 具体的には、Dense Passage Retriever (DPR) のアーキテクチャが採用されています。
  • ドキュメントとクエリはそれぞれBERTベースのエンコーダによって表現され、Maximum Inner Product Search (MIPS) を用いて最も関連性の高い上位のドキュメントが高速に検索されます。
  1. Generator(生成器):
  • 過去のトークン、元の入力クエリ、および検索されたパッセージに基づいて、現在のトークンを生成するコンポーネントです。
  • 事前学習済みのseq2seqトランスフォーマーである BART-large(4億パラメータ)などが用いられます。
  • 入力クエリと検索されたコンテンツを単に連結(concatenate)してBARTに入力する形で機能します。

2つの周辺化の仕方 ―― RAG-Sequence と RAG-Token

RAG-SequenceおよびRAG-Tokenは、「Generator」単体や「Retriever」単体のことを指す名称ではありません。これらは 「Retrieverが抽出したドキュメントを、Generatorがどのように組み合わせて最終的な回答を生成するか」という、全体の結合戦略(アーキテクチャ) を指します。

RAGのオリジナル論文では、抽出した潜在的なドキュメント(zz)をどのように周辺化(Marginalize)してテキスト(yy)を生成するかによって、2つの異なるモデル形式(結合戦略)が提案されています。

現代の実用的RAGシステムとの関係

現在のChatGPTやLangChainなどで広く使われているRAGシステムは、これから紹介する数式をそのままプログラムで計算しているわけではありません。 現代のLLMはコンテキストウィンドウ(一度に読める文章量)が非常に大きいため、検索した複数ドキュメントを1つのプロンプトに単純に連結して入力し、情報の統合をLLM自身の内部メカニズム(Self-Attention)に委ねるアプローチが主流です(実際の開発における実装手法については、後述の現代のLLM開発における実際のアプローチセクションをご参照ください)。

しかし、「複数文書の情報を組み合わせて1つの回答を作る」とは数学的にどういうことなのか、なぜ複数文書が必要になるケースと単一文書で済むケースで振る舞いが変わるのかを理解する上で、オリジナル論文で示された以下の2つの定式化(RAG-Sequence / RAG-Token)はRAGの根本原理を知るための不可欠な理論的基盤となります。本章ではあえてこの基本に立ち返り、RAGの本質的なメカニズムを解き明かします。

RAG-Sequence モデル

生成するシーケンス全体を通して、検索された「同一のドキュメント」を使用するアプローチです。
トップ KK のドキュメントを検索し、それぞれのドキュメントに基づく出力シーケンスの確率を計算した上で、それらを周辺化します。

pRAG-Sequence(yx)ztop-k(p(x))pη(zx)iNpθ(yix,z,y1:i1)p_{\text{RAG-Sequence}}(y\vert{}x) \approx \sum_{z \in \text{top-k}(p(\cdot\vert{}x))} p_\eta(z\vert{}x) \prod_i^N p_\theta(y_i\vert{}x, z, y_{1:i-1})

RAG-Token モデル

ターゲットとなるトークンを生成するごとに、それぞれ「異なるドキュメント」からコンテンツを選択できるアプローチです。
これにより、ジェネレータは回答を生成する際に、複数のドキュメントから情報を柔軟に組み合わせて出力することが可能になります。

pRAG-Token(yx)iNztop-k(p(x))pη(zx)pθ(yix,z,y1:i1)p_{\text{RAG-Token}}(y\vert{}x) \approx \prod_i^N \sum_{z \in \text{top-k}(p(\cdot\vert{}x))} p_\eta(z\vert{}x) p_\theta(y_i\vert{}x, z, y_{1:i-1})

数式の変数の意味:

  • xx: 入力クエリ(ユーザーからの質問など)
  • yy: 最終的に生成される出力テキスト(シーケンス)
  • yiy_i: ii番目に生成されるターゲットトークン
  • y1:i1y_{1:i-1}: 既に生成された過去のトークン列(11番目からi1i-1番目まで)
  • zz: Retrieverによって抽出された潜在的なドキュメント(テキストパッセージ)
  • pη(zx)p_\eta(z\vert{}x): パラメータ η\eta を持つRetrieverによる、クエリ xx に対するドキュメント zz の検索確率
  • pθ(yix,z,y1:i1)p_\theta(y_i\vert{}x, z, y_{1:i-1}): パラメータ θ\theta を持つGeneratorによる、クエリ xx、ドキュメント zz、過去のトークン列 y1:i1y_{1:i-1} に基づく、トークン yiy_i の生成確率
  • NN: 生成されるターゲットシーケンスの長さ(トークン数)
  • top-k(p(x))\text{top-k}(p(\cdot\vert{}x)): クエリ xx に対して検索確率が上位の KK 個のドキュメント集合
補足:現代のRAGを同じ記法で書くと

Naive RAGは検索結果を1本のコンテキスト c=[z(1);z(2);;z(K)]c = [z_{(1)}; z_{(2)}; \dots; z_{(K)}] に連結して一度に入力するため、上の2式にあった z\sum_z が消え、

p(yx)=iNpθ(yix,c,y1:i1)p(y\vert{}x) = \prod_i^N p_\theta(y_i\vert{}x, c, y_{1:i-1})

という単一の条件付き確率になります。zz はもはや確率変数ですらありません。pη(zx)p_\eta(z\vert{}x) も消滅したわけではなく、top-KK の足切りと並び順にだけ使われて生成側では 1[ztop-k]\mathbf{1}[z \in \text{top-k}] という0/1に退化します(=検索スコアのソフトな重みが生成確率に反映されない)。リランキングやコンテキスト圧縮が効くのは、この重みが捨てられる前に文書集合そのものを整形しているためだと解釈できます。

なお「Attentionに周辺化を委ねる」というのはあくまで比喩であり、RAG-Tokenと数学的に等価ではありません。RAG-Tokenが文書ごとに独立に計算した出力分布を最後に混合するのに対し、Self-Attentionは層ごとに中間表現を混合し、連結された文書同士も相互に参照し合います。表現力はむしろ広い一方で zpηpθ\sum_z p_\eta p_\theta という分解構造が失われるため、後ほどRAG-Tokenのトレースで観察する事後確率 p(zx,y1:i1)p(z\vert{}x, y_{1:i-1}) にあたる量を取り出せず、どのトークンがどの文書に由来するかを追跡できなくなります

RAG-Sequence / RAG-Token を最小の数値例で確かめる

前節の2つの式は、よく見ると 「総和 \sum と総乗 \prod の順序が入れ替わっているだけ」 です。 しかしこの順序の違いが、実際の振る舞いにどのような差を生むのかは、式を眺めているだけではなかなか掴めません。 ここでは、実際に手元で動かして挙動を観察できるコードを用意しました。

まずはモデルを一切使わず、確率の数値だけで2つの定式化を比較します。 文書は2件(K=2K=2)、生成するトークンは3個とし、 pθ(yix,z,y1:i1)p_\theta(y_i\vert{}x, z, y_{1:i-1}) を手で与えます。

この例の読み方
  • 動物1匹を1トークンとみなしています。 実際には「フラミンゴ」のような語は複数のサブワードに分割されますが(後の言語モデルによる検証では 'Mount' 'Fuji' のような実トークンが登場します)、ここでは行列 P の1列を1つの動物に対応させ、\sum\prod の順序の効果だけを見やすくしています。
  • 検索確率を両ケースとも pη(zx)=0.5p_\eta(z\vert{}x) = 0.5 に固定しています。 本来「暑い場所に住んでいる生き物は?」というクエリなら z2z_2(寒い場所)の検索確率は下がるはずですが、ここでは検索側の条件を揃えることで、スコアの差が周辺化の順序だけから生じるようにしています。
import numpy as np

def scores(p_z, P):
"""p_z: (K,) 文書の検索確率 / P: (K, N) 各文書のもとでの各トークンの生成確率"""
# RAG-Sequence: 文書ごとに全トークンの確率を掛け合わせ(prod)、その後に検索確率を掛けて文書について足す(sum)
seq = float(np.sum(p_z * np.prod(P, axis=1)))
# RAG-Token: トークンごとに文書について期待値を計算し(sum)、その後にトークン全体で掛け合わせる(prod)
tok = float(np.prod(np.sum(p_z[:, None] * P, axis=0)))
return seq, tok

# 文書の検索確率はどちらのケースも0.5(等確率)に固定
p_z = np.array([0.5, 0.5])

# 前提知識:
# z1: 暑い場所に住む生き物 [フラミンゴ, オオハシ, ラクダ]
# z2: 寒い場所に住む生き物 [ホッキョクグマ, トナカイ, エンペラーペンギン]

# ケースA: クエリ「暑い場所に住んでいる生き物は?」 -> ターゲット回答: [フラミンゴ, オオハシ, ラクダ]
# 文書z1だけが回答全体を強く支持している(z2は関係ないため確率が低い)
A = np.array([[0.9, 0.9, 0.9],
[0.1, 0.1, 0.1]])

# ケースB: クエリ「鳥類を教えて?」 -> ターゲット回答: [フラミンゴ, オオハシ, エンペラーペンギン]
# 前半のトークン(フラミンゴ, オオハシ)はz1が支持し、後半(エンペラーペンギン)はz2が支持する(1文書では回答しきれない)
B = np.array([[0.9, 0.9, 0.1],
[0.1, 0.1, 0.9]])

# 2つのケースをループで回してスコアを比較
for name, P in [("A: 単一文書で回答が完結する(暑い場所の生き物)", A),
("B: 複数文書の情報を組み合わせる必要がある(鳥類)", B)]:
s, t = scores(p_z, P)
print(f"{name}")
print(f" RAG-Sequence = {s:.6f}")
print(f" RAG-Token = {t:.6f}")
print(f" -> {'RAG-Sequence' if s > t else 'RAG-Token'} の方が高いスコア\n")

scores 関数では、上で解説した2つの数式をNumPyのみでシミュレートしています。

  • RAG-Sequence: np.prod(P, axis=1) で文書ごとに全トークンの確率 pθ\prod p_\theta を掛け合わせた後、np.sum() で文書の検索確率 pηp_\eta と掛けて総和 \sum を取ります。これは「特定の文書1つを使って回答全体を生成する」プロセスを表現しています。
  • RAG-Token: 逆に、np.sum(p_z[:, None] * P, axis=0) でトークンごとに全文書にわたる確率の期待値 pηpθ\sum p_\eta p_\theta を計算してから、np.prod() でシーケンス全体の積 \prod を取ります。これにより「トークンごとに担当文書を切り替える」ことが可能になります。 同じ確率行列 P を渡していても、この集約の順序が入れ替わるだけでスコアがどう変わるかを見てみます。

実行結果

A: 単一文書で回答が完結する(暑い場所の生き物)
RAG-Sequence = 0.365000
RAG-Token = 0.125000
-> RAG-Sequence の方が高いスコア

B: 複数文書の情報を組み合わせる必要がある(鳥類)
RAG-Sequence = 0.045000
RAG-Token = 0.125000
-> RAG-Token の方が高いスコア
確認したいポイント
  • ケースA(暑い場所の生き物) では、z1という1つの文書がターゲットとなる回答(フラミンゴ、オオハシ、ラクダ)全体を通して高い確率を与えるため、その文書に基づく確率の積が生き残り、RAG-Sequenceの方が高くなります。
    • RAG-Sequenceの計算:
      • z1z_1ベース: 0.5×(0.9×0.9×0.9)=0.36450.5 \times (0.9 \times 0.9 \times 0.9) = 0.3645
      • z2z_2ベース: 0.5×(0.1×0.1×0.1)=0.00050.5 \times (0.1 \times 0.1 \times 0.1) = 0.0005
      • 合計(\sum): 0.3645+0.0005=0.3650.3645 + 0.0005 = 0.365
    • RAG-Tokenの計算:
      • 各トークンの期待値(\sum): (0.5×0.9)+(0.5×0.1)=0.5(0.5 \times 0.9) + (0.5 \times 0.1) = 0.5
      • 全体の積(\prod): 0.5×0.5×0.5=0.1250.5 \times 0.5 \times 0.5 = 0.125
  • ケースB(鳥類) では、どの文書も回答(フラミンゴ、オオハシ、エンペラーペンギン)の一部分しか支持できません。RAG-Sequenceは「1文書で全トークンを通す」ことを要求するため確率の積が急速に小さくなる一方、RAG-Tokenはトークンごとに担当する文書を「z1 → z1 → z2」と乗り換えられるため、スコアが落ちません。
    • RAG-Sequenceの計算:
      • z1z_1ベース: 0.5×(0.9×0.9×0.1)=0.04050.5 \times (0.9 \times 0.9 \times 0.1) = 0.0405
      • z2z_2ベース: 0.5×(0.1×0.1×0.9)=0.00450.5 \times (0.1 \times 0.1 \times 0.9) = 0.0045
      • 合計(\sum): 0.0405+0.0045=0.0450.0405 + 0.0045 = 0.045
    • RAG-Tokenの計算:
      • 1・2番目(z1z_1寄り)の期待値: (0.5×0.9)+(0.5×0.1)=0.5(0.5 \times 0.9) + (0.5 \times 0.1) = 0.5
      • 3番目(z2z_2寄り)の期待値: (0.5×0.1)+(0.5×0.9)=0.5(0.5 \times 0.1) + (0.5 \times 0.9) = 0.5
      • 全体の積(\prod): 0.5×0.5×0.5=0.1250.5 \times 0.5 \times 0.5 = 0.125
  • つまり RAG-Tokenの優位性は「複数文書の情報を組み合わせる必要があるとき」に現れる ことが、具体的なトークンの例からも明確に確認できます。

RAG-Sequence / RAG-Token を実際の言語モデルで確かめる

最小の数値例では確率の数値だけを比べましたが、ここでは実際のseq2seqモデルに回答を生成させてみます。 実際にコードを書いてみると、そもそもデコードの手続き自体が2つで全く別物であることがはっきりと見えてきます。

ポイントを先に述べると、次の非対称性があります。

デコードの実体model.generate()
RAG-Sequence文書ごとに独立に生成 → 候補を集約 → 再スコアリングそのまま使えるKK 回呼ぶ)
RAG-Token1トークンごとに文書方向へ周辺化しながら1本を生成使えない(自前ループが必要)

RAG-Token が generate() で書けないのは、Hugging Faceのデコードループが「単一の分布から次トークンを選ぶ」ことを前提にしているためです。 RAG-Tokenは各ステップで KK 個の分布を pη(zx)p_\eta(z\vert{}x) で混合してから選ぶ必要があるので、その混合を挟み込む場所がありません。 この違いは実装の都合ではなく、\sum\prod の内側にあるか外側にあるかという定式化の差がそのまま現れたものです。

(なお、原論文ではGeneratorにBART-large、RetrieverにDPRが使われていますが、 ここでは手元で軽量に動かすことを優先し、Generatorに google/flan-t5-base、 RetrieverにTF-IDFによる簡易的な類似度検索を用いています。 周辺化の挙動を見ることが目的なので、Retrieverの実装は本質ではありません。)

必要なライブラリのインストール

!pip install torch transformers scikit-learn sentencepiece

このセクションの検証を手元の環境で実行するには、上記のライブラリが必要です。

  • torch: PyTorch。深層学習モデル(今回はT5)を動作させ、内部の対数確率(log_softmax)やテンソル演算、周辺化のための数学的集約(logsumexp)を行う基盤として使用します。
  • transformers: Hugging Faceのライブラリ。学習済みの言語モデル(Generatorとして使用する google/flan-t5-base)本体とそのトークナイザーを簡単にダウンロードし、推論処理を回すために必須となります。
  • scikit-learn: 機械学習ライブラリ。今回はドキュメント群からTF-IDFベクトルを構築し、入力クエリとのコサイン類似度を高速に計算する「簡易Retriever」を作成するために使用します。
  • sentencepiece: テキストをサブワードに分割するためのトークナイザ・バックエンド。T5系のモデル(Flan-T5など)が内部でテキストをトークンIDに変換する際に裏側で要求されるためインストールが必要です。

セットアップとRetriever

まずは、素朴な類似度検索(TF-IDF)によるRetrieverだけでは「ディストラクタ(もっともらしいが不正解の文書)」が上位に来てしまうことを確認します。 この不完全な検索結果をそのままGeneratorに渡すことで、後続のスコアリングでRAG-SequenceとRAG-Tokenが「ノイズや分割された情報」をどう捌くのかを比較するための舞台設定を行います。

検索対象のコーパスには、先ほどのケースBと同じ構造 —— 回答に必要な情報がわざと2つの文書に分割されている ものを用意します。 z1z_1 は「日本で最も高い山は富士山」という山名だけを、z2z_2 は「日本で最も高い山は3,776m」という標高だけを持ち、 どちらも単独では Mount Fuji, 3,776 meters という回答を作れません。 さらに、エベレストに関する文書をディストラクタ(紛らわしい不正解文書) として混ぜています。

注記

Generatorに使う flan-t5-base は英語での性能が安定しているため、コーパスとクエリは英語で記述しています。 また、出力確率や生成トークンは transformerstorch のバージョン、およびモデル本体のアップデートによって(乱数が絡まなくとも)変動し、後述するトレース表の数値が変わる可能性があります。本記事と完全に同じ実行結果を再現できるよう、コード内では検証時に使用した特定のコミットハッシュを指定してモデルをロードしています。

import numpy as np
import torch
import torch.nn.functional as F
from sklearn.feature_extraction.text import TfidfVectorizer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# 今回のGeneratorとして軽量なFlan-T5を使用
# (※実行結果の完全な再現性を担保するため、特定のコミットハッシュを指定)
MODEL_NAME = "google/flan-t5-base"
MODEL_REVISION = "7bcac572ce56db69c1ea7c8af255c5d7c9672fc2"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, revision=MODEL_REVISION)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, revision=MODEL_REVISION).eval()

# 検索対象となるコーパス(知識ベース)を定義
DOCS = [
"Mount Fuji is the highest mountain in Japan.", # z1: 山名のみ
"The highest mountain in Japan rises 3,776 meters above sea level.", # z2: 標高のみ(山名を含まない)
"Mount Fuji straddles the border of Shizuoka and Yamanashi prefectures.",# z3: 富士山に関する別の情報
"Mount Everest is the highest mountain in the world, 8,849 meters tall.",# z4: ディストラクタ(紛らわしい不正解文書)
"Cherry blossoms usually bloom in Tokyo at the end of March.", # z5: ノイズ(無関係な文書)
]
# クエリ: 日本で一番高い山は?そしてそれは何メートル?
QUESTION = "What is the highest mountain in Japan, and how many meters tall is it?"

# 簡易的なRetrieverとしてTF-IDFベクトル化器を初期化・学習
_vec = TfidfVectorizer().fit(DOCS)
_mat = _vec.transform(DOCS)

def retrieve(question, k=4, temperature=0.3):
"""簡易Retriever。上位k件と、その検索確率 p_eta(z|x)、DOCS上の元番号を返す"""
# クエリとコーパスのコサイン類似度を計算
sims = (_vec.transform([question]) @ _mat.T).toarray()[0]
# 類似度が高い上位k件のインデックスを取得
idx = np.argsort(-sims)[:k]
# temperature付きのソフトマックスで検索確率分布 p_eta(z|x) を算出
p_z = torch.softmax(torch.tensor(sims[idx] / temperature, dtype=torch.float), dim=0)
return [DOCS[i] for i in idx], p_z, sims[idx], idx

# 検索の実行
docs, p_z, sims, doc_ids = retrieve(QUESTION, k=4)

# 表示ラベルは「検索順位」ではなく DOCS 上の元番号にそろえる(例: z1, z2)
LABELS = [f"z{i + 1}" for i in doc_ids]
for lab, d, p, s in zip(LABELS, docs, p_z.tolist(), sims):
print(f" {lab} p(z|x)={p:.3f} sim={s:.3f} {d}")

コードの主な処理内容は以下の通りです。

  • GeneratorとRetrieverの初期化 google/flan-t5-base をGeneratorとして初期化し、簡易的なTF-IDFを用いたRetrieverを retrieve 関数で実装しています。
  • 検索対象(知識ベース)の設計 検索対象の DOCS には、質問の回答に必要な情報が「日本で一番高い山は富士山(=山名)」(z1)と「日本で一番高い山は3,776m(=標高)」(z2)という2つの文書に意図的に分割して格納されており、さらにダミーの文書(z4 のエベレストなど)も混ぜています。
  • 検索確率 pη(zx)p_\eta(z\vert{}x) の算出 retrieve 関数内では TfidfVectorizer でクエリと各文書のコサイン類似度を計算し、Softmax関数 torch.softmaxtemperature パラメータ付き)を適用することで、検索確率の分布を算出しています。

実行結果

z1 p(z|x)=0.399 sim=0.639 Mount Fuji is the highest mountain in Japan.
z4 p(z|x)=0.360 sim=0.608 Mount Everest is the highest mountain in the world, 8,849 meters tall.
z2 p(z|x)=0.160 sim=0.364 The highest mountain in Japan rises 3,776 meters above sea level.
z3 p(z|x)=0.081 sim=0.160 Mount Fuji straddles the border of Shizuoka and Yamanashi prefectures.

回答に必要な z1z_1(山名)と z2z_2(標高)が両方とも上位に入っています。 一方で、ディストラクタである z4z_4(エベレスト)が z2z_2 より高い検索確率を得ている点に注目してください。 「日本で最も高い山」という表現を共有しているためで、TF-IDFのような表層的な検索器が引っかかりやすい典型例です。 以降の表では、この z1,z4,z2,z3z_1, z_4, z_2, z_3 という検索順に列が並びます。

先に結論:2つの定式化はこう答えた

ここまでは数式とスコアの話でしたが、ここからは RAG-Sequence と RAG-Token に実際に回答を書かせ、出てくる文章がどう違うかを見ていきます。

先に結論を出します。同じ質問・同じ文書・同じ設定で、2つの定式化はこう答えました。

生成された回答
RAG-Sequence3,776 meters above sea level.
RAG-TokenMount Fuji is the highest mountain in Japan, and its height is 3,776 meters.

RAG-Sequenceは標高しか答えられていません。質問は「日本で一番高い山は何で、標高は何メートルか」ですから、半分しか答えていないことになります。 一方RAG-Tokenは山名と標高の両方を含む完全な回答を生成しました。

この差がどこから来るのかを、順に確認していきます。

共通の下準備

どちらの定式化でも、クエリと各文書を連結して KK 本のプロンプトを作るところは共通です。

import torch
import torch.nn.functional as F

# 「山名と標高の両方を1文で答える」ことを指示するプロンプト
# (英語での推論性能を安定させるため、プロンプトは英語で記述)
PROMPT = ("Answer in one full sentence. Include both the name of the mountain "
"and its height in meters. question: {q} context: {d}")

MIN_NEW_TOKENS = 8 # 短答で打ち切らせないための下限 (短すぎると比較できないため)
MAX_NEW_TOKENS = 32 # 生成する最大トークン数

def _encode(question, docs):
"""
クエリとドキュメントのリストから、T5モデル用の入力トークン(バッチ)を生成する
"""
# 複数ドキュメントそれぞれに対してプロンプト文字列を作成
prompts = [PROMPT.format(q=question, d=d) for d in docs]
# トークナイザでテンソルに変換し、長さが異なる場合はパディングしてバッチ化
return tokenizer(prompts, return_tensors="pt", padding=True)
  • プロンプトと生成トークンの制約: まず、PROMPTでは回答が単語レベルで終わってしまわないように、「1文で、山名と標高の両方を含めて答える」ようにモデルに指示しています。(google/flan-t5-baseは抽出的な短答を好む傾向があるため、これを防ぐ意図があります。)また、MIN_NEW_TOKENSMAX_NEW_TOKENSで生成トークン数の制約を定義しています。
  • 入力データのバッチ化とテンソル変換: _encode関数では、入力されたquestionと複数ドキュメントdocsのリストから、PROMPTの雛形に埋め込んだ上でtokenizerを呼び出し、モデル入力用のテンソルを作成しています。その際、return_tensors="pt"でPyTorchテンソルを指定し、padding=Trueで各プロンプトの長さを最大長に合わせてパディング(バッチ化)しています。
なぜ「1文で答えろ」と指示し、最低長を設けるのか

flan-t5-base は指示チューニングの過程で、QA形式のプロンプトに対して抽出的な短答を返すよう学習しています。 何も指定せずに投げると、RAG-Tokenは 'Mount Fuji' の2トークンで生成を終えてしまいます。 山名を答えた時点で質問に答えたつもりになり、標高に触れないまま終了するのです。 これでは担当する文書が交代する場面がそもそも訪れず、2つの定式化の差を観察できません。 そこで「1文で両方答えること」を指示し、最低8トークンはEOS(終了トークン)を禁止しています。 なお貪欲デコードで長さを強制すると同じ語を繰り返しやすいため、3-gramの反復を禁止するガードも入れています。

重要なのは、「最低長(MIN_NEW_TOKENS)」と「反復の禁止(n-gram ペナルティ)」の制約は両方の定式化で揃えている点です。 なお、以降の実装コードでは、RAG-Sequence側はビームサーチ(num_beams=4)を用いた上で全候補の再スコアリングを行い、RAG-Token側は単なる貪欲デコード(argmax)を用いています。探索戦略としてはむしろRAG-Sequence側により有利な条件を与えているにもかかわらず、それでもRAG-Sequenceからは統合された回答が候補にすら現れないという結果が、定式化の根本的な違いをより強く浮き彫りにします。

docsp_zretrieve() が返したものを使いますが、 ここではまずディストラクタ z4z_4(エベレスト)を外した3文書で試します。 z4z_4 を含めるとどうなるかは、後半で改めて確認します。

# ディストラクタ z4 を除いて検索確率を再正規化する
keep = [i for i, lab in enumerate(LABELS) if lab != "z4"]
docs_c = [docs[i] for i in keep]
p_z_c = p_z[keep] / p_z[keep].sum()
LABELS_C = [LABELS[i] for i in keep]

print("p_z_c =", [f"{l}:{p:.3f}" for l, p in zip(LABELS_C, p_z_c.tolist())])
# p_z_c = ['z1:0.624', 'z2:0.249', 'z3:0.127']

RAG-Sequence:文書ごとに生成してから集約する

pRAG-Sequence(yx)=ztop-kpη(zx)iNpθ(yix,z,y1:i1)p_{\text{RAG-Sequence}}(y\vert{}x) = \sum_{z \in \text{top-}k} p_\eta(z\vert{}x) \prod_{i}^{N} p_\theta(y_i \mid x, z, y_{1:i-1})

これを K=3K=3 の場合について展開して書くと、以下のようになります。それぞれの文書(z1z3z_1 \sim z_3)ごとに計算した「生成確率の総乗」を、後から足し合わせている構造がより明確になります。

pRAG-Sequence(yx)=pη(z1x)iNpθ(yix,z1,y1:i1)+pη(z2x)iNpθ(yix,z2,y1:i1)+pη(z3x)iNpθ(yix,z3,y1:i1)\begin{aligned} p_{\text{RAG-Sequence}}(y\vert{}x) &= p_\eta(z_1\vert{}x) \prod_{i}^{N} p_\theta(y_i \mid x, z_1, y_{1:i-1}) \\ &\quad + p_\eta(z_2\vert{}x) \prod_{i}^{N} p_\theta(y_i \mid x, z_2, y_{1:i-1}) \\ &\quad + p_\eta(z_3\vert{}x) \prod_{i}^{N} p_\theta(y_i \mid x, z_3, y_{1:i-1}) \end{aligned}

\sum が一番外側にあるということは、内側の \prod を文書ごとに独立に計算してよいということです。 したがって生成も文書ごとに独立に行えます。つまり、model.generate() によるテキスト生成処理は、数式における iNpθ(yix,z,y1:i1)\prod_{i}^{N} p_\theta(y_i \mid x, z, y_{1:i-1}) の部分(特定のドキュメント zz に基づくシーケンス生成確率) に直接該当します。これを KK 回(文書の数だけ)呼び出して回答の候補を作成し、最後に出てきた候補を \sum で集約して1つ選ぶ、という2段構えになります。

ただしここで問題が生じます。文書 zkz_k から生成された候補 yy は、他の文書 zkz_{k'} からは生成されていないかもしれません。 そのままでは pθ(yx,zk)p_\theta(y\vert{}x, z_{k'}) が手元になく、\sum を取れないのです。

具体例:なぜ \sum が取れないのか?

例えば、文書 z1z_1(富士山の名前)からは「富士山」、文書 z2z_2(標高)からは「3,776m」、文書 z3z_3(県境)からは「静岡県と山梨県の県境」という候補がそれぞれ生まれたとします。 数式の \sum は「すべての文書からの支持(確率)を足し合わせる」ことを要求するため、候補 y=y=「富士山」の総合スコアを出すには、以下の式を計算する必要があります。

p(y=富士山x)=pη(z1x)iNpθ(富士山ix,z1,y1:i1)計算済みで手元にある+pη(z2x)iNpθ(富士山ix,z2,y1:i1)手元にないので計算できない!+pη(z3x)iNpθ(富士山ix,z3,y1:i1)手元にないので計算できない!\begin{aligned} p(y=\text{富士山}\vert{}x) &= \underbrace{p_\eta(z_1\vert{}x) \prod_{i}^{N} p_\theta(\text{富士山}_i \mid x, z_1, y_{1:i-1})}_{\text{計算済みで手元にある}} \\ &\quad + \underbrace{p_\eta(z_2\vert{}x) \prod_{i}^{N} p_\theta(\text{富士山}_i \mid x, z_2, y_{1:i-1})}_{\text{手元にないので計算できない!}} \\ &\quad + \underbrace{p_\eta(z_3\vert{}x) \prod_{i}^{N} p_\theta(\text{富士山}_i \mid x, z_3, y_{1:i-1})}_{\text{手元にないので計算できない!}} \end{aligned}

1行目の確率は z1z_1 から「富士山」を生成した際に計算済みですが、2行目や3行目の「z2z_2z3z_3 を読んだ時に『富士山』と答える確率」は手元にありません(z2z_2z3z_3 が実際に生成したのは別の回答だからです)。この「手元にない確率」が pθ(yx,zk)p_\theta(y\vert{}x, z_{k'}) であり、これがないと足し算(クロスチェック)が完了できないというジレンマです。

原論文はこの問題に対して、以下の2つの現実的なアプローチ(近似)を用意しています。

  • Thorough Decoding(徹底的なデコード):全候補について pθ(yx,zk)p_\theta(y\vert{}x, z_{k'}) を追加の順伝播で計算し直し、すべての文書で周辺化するアプローチ。「z2z_2 を読んだ状態で、強制的に『富士山』と出力させたら確率はいくらか?」を追加計算するため正確です(下記のPythonコードで行っているのはこちらです)。
  • Fast Decoding(高速なデコード):その文書から生成されなかった候補は「どうせ生成確率はほぼゼロだろう」と見なして、pθ(yx,zk)0p_\theta(y\vert{}x, z_{k'}) \approx 0 と省略するアプローチ。追加計算が不要なので高速です。

以下のコードでは、Hugging Faceの transformers ライブラリとPyTorchを用いて、RAG-Sequenceにおける Thorough Decoding(徹底的なデコード) を実際にシミュレートします。各文書から独立して回答候補を生成した後、集められた全候補をすべての文書に対して再度順伝播させて欠けていた確率を補い、最終的な総合スコア(周辺化)を計算して比較する一連の流れを実装しています。

@torch.no_grad()
def sequence_logprobs(question, docs, answer):
"""
ある回答候補 y (answer) が、すべての文書 z_k からそれぞれ生成される
確率(対数確率) log p_θ(y | x, z_k) を計算して (K,) のテンソルで返す関数。
(Thorough Decodingにおける「手元にない確率」を補うための追加計算)
"""
enc = _encode(question, docs)
# 回答候補(例:「富士山」)をトークナイズし、文書の数(K個)だけ複製してラベルとする
labels = tokenizer(answer, return_tensors="pt").input_ids.repeat(len(docs), 1)

# 順伝播(Forward)を実行し、出力語彙に対するロジットを取得
logits = model(**enc, labels=labels).logits # (K, N, V)
lp = F.log_softmax(logits, dim=-1) # ロジットを対数確率に変換

# 実際に正解ラベル(候補のトークン列)が選ばれる対数確率だけを抽出 (K, N)
tok_lp = lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
mask = (labels != tokenizer.pad_token_id).float() # 無効なPADトークンの寄与を除外

# トークン方向に和を取る(対数空間での足し算 = 通常の確率空間での掛け算 ∏)
# これにより、シーケンス全体の生成確率が文書ごとに求まる
return (tok_lp * mask).sum(dim=1)

@torch.no_grad()
def rag_sequence_generate(question, docs, p_z, num_beams=4, return_all=False):
enc = _encode(question, docs)
log_pz = torch.log(p_z) # Retrieverが計算した各文書の検索確率を対数化

# ---------------------------------------------------------
# 1) 文書ごとに独立にビームサーチし、候補集合(プール)を作る
# ---------------------------------------------------------
per_doc, cands, origin = [], [], {}
for k in range(len(docs)):
# K番目の文書だけを入力として、独立してテキストを生成させる
out = model.generate(input_ids=enc["input_ids"][k:k + 1],
attention_mask=enc["attention_mask"][k:k + 1],
num_beams=num_beams, num_return_sequences=num_beams,
min_new_tokens=MIN_NEW_TOKENS,
max_new_tokens=MAX_NEW_TOKENS,
no_repeat_ngram_size=3)
# 生成されたトークンIDを文字列にデコード
ys = [y.strip() for y in tokenizer.batch_decode(out, skip_special_tokens=True)]
per_doc.append(ys)

# 重複を省きながら、全文書から生まれた回答候補を一つのプール(cands)にまとめる
for y in ys:
if not y:
continue
if y not in origin:
origin[y] = set()
cands.append(y)
origin[y].add(k) # どの文書がこの候補を生み出したかを記録

# ---------------------------------------------------------
# 2) 全候補を全文書で再スコアリングし、文書方向へ周辺化(= Σ)
# ---------------------------------------------------------
ranked = []
for y in cands:
# sequence_logprobs を呼んで、すべての文書に対する生成確率を再計算する
# log_pz (検索確率) + sequence_logprobs (生成確率) = 対数空間での (p_η × p_θ)
# logsumexp で文書方向に和をとる = 対数空間での Σ
score = torch.logsumexp(log_pz + sequence_logprobs(question, docs, y), dim=0)
ranked.append((float(score), y, sorted(origin[y])))

# 総合スコア(周辺化された確率)が高い順にソートしてランキングを確定
ranked.sort(reverse=True)
return (ranked, per_doc) if return_all else ranked[0][1]

# 実際に処理を実行する
# QUESTION = "What is the highest mountain in Japan, and how many meters tall is it?"
# p_z_c = ['z1:0.624', 'z2:0.249', 'z3:0.127']
# docs_c = [
# "Mount Fuji is the highest mountain in Japan.",
# "The highest mountain in Japan rises 3,776 meters above sea level.",
# "Mount Fuji straddles the border of Shizuoka and Yamanashi prefectures."
#]
ranked, per_doc = rag_sequence_generate(QUESTION, docs_c, p_z_c, return_all=True)

print(f"RAG-Sequence : {ranked[0][1]!r}\n")
print("各文書が単独で生成した候補(1位のみ表示)")
for lab, ys in zip(LABELS_C, per_doc):
print(f" {lab}: {ys[0]!r}")
  • Thorough Decodingの追加計算: sequence_logprobs関数は、ある回答候補answerがすべての文書から生成される確率(対数確率)を再計算するための関数です。モデルに正解ラベルlabelsとして候補のトークン列を与え、F.log_softmaxgatherを用いて対象トークンの確率のみを抽出します。最後にsum(dim=1)でトークン方向の和(対数空間での積\prod)をとり、各文書に対する生成確率を計算しています。
  • 候補の独立生成: rag_sequence_generate関数の前半では、各文書に対して個別にmodel.generateを呼び出し、ビームサーチによって回答候補のプール(cands)を構築しています。ここではそれぞれの文書が持つ情報のみに依存した回答が生成されます。
  • 再スコアリングと周辺化: 関数の後半では、集められたすべての候補に対してsequence_logprobsを呼び出し、再スコアリングを行います。torch.logsumexpを用いることで、検索確率log_pzと生成確率の和(対数空間での積)について、全文書のスコアを安定して足し合わせる(数式における\sum)処理を実現し、全候補の最終的なランキングを決定しています。

実行結果

RAG-Sequence : '3,776 meters above sea level.'

各文書が単独で生成した候補(1位のみ表示)
z1: 'Mount Fuji is the highest mountain in Japan, and its height is 100 meters.'
z2: '3,776 meters above sea level.'
z3: 'Mount Fuji straddles the border of Shizuoka and Yamanashi prefectures'

集約前の候補を見ると、何が起きているかがはっきりします。

  • 各文書が持つ情報の限界: 各文書は、自分が持っている情報しか書けません。z1z_1 は山名を知っているので Mount Fuji と書けますが、標高を持っていません。z2z_2 は標高を知っているので 3,776 meters と書けますが、山名を持っていません。z3z_3 は県境の話しかできません。
  • 統合された回答が生成されない根本原因: つまり、Mount Fuji is ... 3,776 meters という統合された回答は、候補集合のどこにも存在しないのです。RAG-Sequenceの集約は、生成された候補の中から1つを選ぶ作業にすぎません。どれだけ丁寧に再スコアリングしても、候補として存在しない回答は選びようがないのです。これは実装の不備ではなく、\sum\prod の外に置いた定式化の必然的な帰結です。
  • ハルシネーション(捏造)の発生: さらに注目すべきは z1z_1 の出力です。標高を持っていないにもかかわらず「1文で標高も答えろ」と指示された結果、100 meters という数値を捏造しています。文書に無い情報を求められたモデルが埋め合わせようとする、典型的なハルシネーションです。
  • 最終的な選択結果とその理由: 最終的に選ばれたのは z2z_2 の候補でした。検索確率は z1z_10.6240.624 で最も高いにもかかわらず、z1z_1 の候補は選ばれていません。\prod はトークンを掛けるたびに縮むため長い候補が不利になること、そして 100 meters という誤った数値を他の文書が支持しないことの、両方が効いていると考えられます。捏造は回避できたものの、代わりに山名が失われたという結果です。

RAG-Token:毎ステップ周辺化しながら生成する

pRAG-Token(yx)=iNztop-kpη(zx)pθ(yix,z,y1:i1)p_{\text{RAG-Token}}(y\vert{}x) = \prod_{i}^{N} \sum_{z \in \text{top-}k} p_\eta(z\vert{}x)\, p_\theta(y_i \mid x, z, y_{1:i-1})

RAG-Sequenceと違い \sum\prod の内側にあるため、文書ごとに独立に生成することができません。 1トークン進めるたびに全文書の分布を混ぜる必要があるので、生成ループそのものを書き下します。

デコーダ側の系列 decKK 本すべてで共有し(同じ回答 yy を書いているので当然です)、 各ステップで得られる (K, V) のロジットを logsumexp(V,) に潰してから次トークンを選びます。 選ばれたトークンは KK 本すべてに追記されるため、次のステップでは「全文書が同じ続きを見ている」状態になります。

def _banned_by_ngram(tokens, n=3):
"""
直近 n-1 トークンと同じ並びが過去にあれば、その次のトークンを禁止する。
生成時に同じフレーズを無限に繰り返すのを防ぐためのペナルティ機構。
"""
if len(tokens) < n:
return []
# 直近の n-1 トークンを取得
prefix = tuple(tokens[-(n - 1):])
# 過去の履歴を走査し、同じ prefix が出現した直後のトークンをリストアップ
return [tokens[i + n - 1] for i in range(len(tokens) - n + 1)
if tuple(tokens[i:i + n - 1]) == prefix]

@torch.no_grad()
def rag_token_generate(question, docs, p_z, trace=False):
"""
RAG-Tokenの生成ロジック。
各ステップごとにすべての文書を評価し、確率を周辺化(合算)してから次のトークンを1つだけ選ぶ。
これを繰り返すことで、複数文書の知識を動的に切り替えながら1つの文を生成する。
"""
K = len(docs)
enc = _encode(question, docs)

# 1. エンコーダのキャッシュ
# プロンプト(質問+各文書)は生成中ずっと変わらないため、エンコーダは最初に1回だけ実行する。
# ここで K 個の異なるコンテキスト表現が作られ、デコーダ側から参照される。
enc_out = model.get_encoder()(input_ids=enc["input_ids"],
attention_mask=enc["attention_mask"])

# 検索確率 p_η(z|x) を対数化しておく (計算時のアンダーフローを防ぐため)
log_pz = torch.log(p_z).unsqueeze(-1) # 形状: (K, 1)

# 2. デコーダの初期化
# デコーダ入力は K 本すべてで共有する(全文書が「全く同じ生成途中の文 y」を書き進めていくため)。
# 最初はBOS (Begin Of Sequence) トークンのみを K 個用意する。
dec = torch.full((K, 1), model.config.decoder_start_token_id, dtype=torch.long)

tokens, rows = [], []

for step in range(MAX_NEW_TOKENS):
# 3. 各文書ごとの次トークン予測
# K個のデコーダを同時に1ステップ進める。出力 logits の形状は (K, 語彙数 V)
logits = model(encoder_outputs=enc_out,
attention_mask=enc["attention_mask"],
decoder_input_ids=dec).logits[:, -1, :]
lp = F.log_softmax(logits, dim=-1) # p_θ(y_i | x, z, y_{1:i-1})

# 4. ★ RAG-Token の核心部分:文書方向への周辺化 ★
# 各文書の検索確率 log_pz と、各文書が予測した次トークンの確率 lp を足し合わせ、
# 全文書 K にわたって logsumexp で周辺化(実数空間での加算)を行う。
# これにより「全文書の意見を総合した、次トークンの確率分布 marg」が得られる。
marg = torch.logsumexp(log_pz + lp, dim=0) # 形状: (V,)

# 5. ペナルティの適用
# 指定された文字数に満たない場合はEOS(生成終了)を禁止する
if step < MIN_NEW_TOKENS:
marg[tokenizer.eos_token_id] = -float("inf")
# 過去の繰り返し(n-gram)ループを検出した場合はそのトークンを禁止する
for b in _banned_by_ngram(tokens):
marg[b] = -float("inf")

# 6. 次のトークンの決定 (Greedy Decoding)
# 統合された確率分布 marg の中から最も確率の高いトークンを1つ選ぶ
nxt = int(marg.argmax())

if nxt == tokenizer.eos_token_id:
break

if trace:
# 7. プロベナンス(情報源)のトラッキング
# 選ばれたトークン nxt に対して、各文書がどれだけ貢献したか(事後確率)を計算する
# p(z | x, y_{1:i}) ∝ p_η(z|x) · p_θ(y_i | x, z, y_{1:i-1})
post = torch.softmax(log_pz.squeeze(-1) + lp[:, nxt], dim=0)
rows.append((tokenizer.decode([nxt]), post))

tokens.append(nxt)

# 8. 選ばれた 1 つのトークンを、K 本すべてのデコーダ入力 (dec) の末尾に一斉に追加する。
# これにより、次のステップでも全文書が「同じ続きのテキスト」を評価することになる。
dec = torch.cat([dec, torch.full((K, 1), nxt, dtype=torch.long)], dim=1)

text = tokenizer.decode(tokens, skip_special_tokens=True)
return (text, rows) if trace else text

tok_text, trace = rag_token_generate(QUESTION, docs_c, p_z_c, trace=True)
print(f"RAG-Token : {tok_text!r}\n")
print("各トークンを支持した文書の割合")
for t, post in trace:
dist = " ".join(f"{l}={p:.2f}" for l, p in zip(LABELS_C, post))
print(f" {t!r:14} {dist}")
  • エンコーダの事前計算とキャッシュ: model.get_encoder()を使用して、プロンプトと文書を結合した入力に対するエンコーダの出力 enc_out を最初に1回だけ計算しています。プロンプトは生成中に変化しないため、事前に計算してキャッシュしておくことで推論時の計算コストを大幅に削減しています。
  • デコーダ入力の共有化: デコーダの入力となる dec テンソルは、すべての文書数(K)分だけ同じ初期トークンを複製して用意されます。RAG-Tokenでは文書ごとに別々の候補を生成するのではなく、「全文書が協力して1つの回答を紡ぎ出す」ため、すべての文書が全く同じ生成履歴(コンテキスト)を参照し続ける必要があります。
  • ステップごとの周辺化(RAG-Tokenの核心): ループ内の torch.logsumexp(log_pz + lp, dim=0) がこのモデルの最重要部分です。各ステップで各文書から出力された語彙の対数確率 lp(次元 (K, V))と、各文書の検索スコア log_pz を対数空間で加算(確率空間での積)し、文書方向(dim=0)に周辺化(和を取る)しています。これにより、全文書の知見を統合した単一の語彙分布 marg(次元 (V,))がトークンごとに生成されます。
  • トークンの選択とコンテキストの更新: 統合された分布 marg から argmax() で最も確率の高いトークン nxt を選び、torch.cat を用いて全 K 本のデコーダ入力 dec に一斉に追記します。これにより、次のステップでも全文書が同じ続きのテキストを予測できるようになります。
  • n-gramペナルティによる反復の回避: _banned_by_ngram関数は、言語モデルが同じフレーズを無限に繰り返すループに陥るのを防ぐためのペナルティ機構です。n=3(3-gram)としており、これまでに生成されたtokensの中から直近2トークンの並び(prefix)を調べ、過去に全く同じ並びが存在した場合、それに続いたトークンを禁止リストとして返します。これを生成ループ内に組み込むことで、貪欲なデコードにおける無駄な反復を強制的に回避しています。
  • 事後確率(プロベナンス)のトラッキング: trace=True の場合、torch.softmax(log_pz.squeeze(-1) + lp[:, nxt], dim=0) を用いて「選ばれたトークンに対して、どの文書がどれだけ生成確率に貢献したか」を表す事後確率 post (数式における p(zx,y1:i)p(z\vert{}x, y_{1:i}) )を記録します。これはRAG-Tokenならではの利点であり、生成された各単語の出処(プロベナンス)を細かくトレースすることを可能にします。

実行結果

コードを実行すると以下の結果が得られます。

RAG-Token : 'Mount Fuji is the highest mountain in Japan, and its height is 3,776 meters.'

各トークンを支持した文書の割合
'Mount' z1=0.83 z2=0.00 z3=0.17
'Fuji' z1=0.72 z2=0.13 z3=0.15
'is' z1=0.99 z2=0.00 z3=0.00
'the' z1=0.71 z2=0.28 z3=0.01
'highest' z1=0.60 z2=0.32 z3=0.07
'mountain' z1=0.63 z2=0.25 z3=0.13
'in' z1=0.62 z2=0.25 z3=0.13
'Japan' z1=0.62 z2=0.25 z3=0.13
',' z1=0.57 z2=0.25 z3=0.18
'and' z1=0.64 z2=0.23 z3=0.13
'its' z1=0.71 z2=0.14 z3=0.15
'height' z1=0.63 z2=0.24 z3=0.13
'is' z1=0.59 z2=0.28 z3=0.14
'3,' z1=0.01 z2=0.98 z3=0.01
'7' z1=0.06 z2=0.92 z3=0.01
'76' z1=0.12 z2=0.85 z3=0.03
'meters' z1=0.61 z2=0.27 z3=0.12
'.' z1=0.69 z2=0.17 z3=0.14
  • 事前分布からの乖離と「担当の交代」: トレース表を注意深く見ると、mountain, in, Japan, height といった単語の事後確率は z10.63,z20.25,z30.13z_1 \approx 0.63, z_2 \approx 0.25, z_3 \approx 0.13 となっており、初期の検索確率(事前分布 p(z)=[0.624,0.249,0.127]p(z) = [0.624, 0.249, 0.127])とほぼ完全に一致しています。これは「z1z_1 が文を主導している」のではなく、「どの文書もこの汎用的なトークンに対して尤度の差をつけていない(=文書に依存せず次に来るべき単語として妥当なため、事前分布がそのまま残っている)」状態を意味します。一方、各文書が実際に強い意見を表明し、事後確率が事前分布から大きく乖離しているのは Mount (0.83), Fuji (0.72) といった固有名詞や、3, 7 76 の数値トークンです。
  • 必要なトークンでのみピンポイントに文書が交代する: 汎用的な単語では事前分布の重みが維持されますが、標高の数値(3, 7 76 の3トークン)に入った瞬間、事後確率が大きく変動し z2z_20.850.980.85 \sim 0.98 を占め、数値を書き終えると再び元の分布に戻っています。1つの回答を書いている最中に、特定の事実が必要なトークンにおいてのみ担当が入れ替わっているわけです。これが「\sum\prod の内側に置いた」ことの実利であり、回答全体で1回しか文書の重みが決まらないRAG-Sequenceには原理的に存在しない挙動です。
  • 周辺化がハルシネーションを上書きする: 先ほど見たとおり、z1z_1 単独では 100 meters と捏造しました。しかしRAG-Tokenでは、数値の場面で z1z_1 の尤度が下がり(支持率 0.010.01)、正しい数値を知っている z2z_2 が高い尤度を出して事後分布を 0.980.98 まで引き寄せたため、生成された標高は正しい 3,776 となりました。文書ごとに独立に生成すると捏造で埋められる部分を、トークンごとに全文書で投票させることで**「z1z_1 の捏造を z2z_2 が上書きして打ち消した」**という構図です。これが、RAG-Tokenが単に「複数文書を統合できる」以上の価値を持つ理由です。

ディストラクタを戻すとどうなるか

ここまでは、ディストラクタである z4z_4(エベレスト)を除外していました。 セットアップとRetriever で確認したとおり、TF-IDFは z4z_4pη(zx)=0.360p_\eta(z\vert{}x)=0.360 という高い検索確率を与えてしまっています。 これを戻すために、z4z_4 を除外しない以下のコードを実行します。

# ディストラクタ z4 を除外せずに検索確率をそのまま使用する
keep = list(range(len(LABELS)))
docs_c = [docs[i] for i in keep]
p_z_c = p_z[keep] / p_z[keep].sum()
LABELS_C = [LABELS[i] for i in keep]

print("p_z_c =", [f"{l}:{p:.3f}" for l, p in zip(LABELS_C, p_z_c.tolist())])
# p_z_c = ['z1:0.399', 'z4:0.360', 'z2:0.160', 'z3:0.081']

実行結果

同じ変数 docs_c, p_z_c を使用して、先ほどのRAG-SequenceとRAG-Tokenの生成コードを再度実行すると、以下の結果が得られます。

RAG-Sequence : 'Mount Everest is 8,849 meters tall.'

各文書が単独で生成した候補(1位のみ表示)
z1: 'Mount Fuji is the highest mountain in Japan, and its height is 100 meters.'
z4: 'Mount Everest is the highest mountain in the world, 8,849 meters tall.'
z2: '3,776 meters above sea level.'
z3: 'Mount Fuji straddles the border of Shizuoka and Yamanashi prefectures'

RAG-Token : 'Mount Fuji. Mount Fuji is the highest mountain in Japan, and its height is 8,849 meters.'

各トークンを支持した文書の割合
'Mount' z1=0.51 z4=0.39 z2=0.00 z3=0.10
'Fuji' z1=0.72 z4=0.00 z2=0.13 z3=0.15
'.' z1=0.29 z4=0.66 z2=0.05 z3=0.00
'Mount' z1=0.68 z4=0.29 z2=0.03 z3=0.00
'Fuji' z1=0.40 z4=0.35 z2=0.16 z3=0.08
'is' z1=0.54 z4=0.42 z2=0.04 z3=0.00
'the' z1=0.55 z4=0.42 z2=0.02 z3=0.00
'highest' z1=0.34 z4=0.46 z2=0.15 z3=0.05
'mountain' z1=0.40 z4=0.36 z2=0.16 z3=0.08
'in' z1=0.40 z4=0.36 z2=0.16 z3=0.08
'Japan' z1=0.41 z4=0.34 z2=0.16 z3=0.08
',' z1=0.25 z4=0.54 z2=0.12 z3=0.09
'and' z1=0.41 z4=0.37 z2=0.14 z3=0.08
'its' z1=0.53 z4=0.25 z2=0.09 z3=0.13
'height' z1=0.41 z4=0.36 z2=0.15 z3=0.08
'is' z1=0.38 z4=0.35 z2=0.18 z3=0.09
'8,' z1=0.00 z4=1.00 z2=0.00 z3=0.00
'8' z1=0.05 z4=0.94 z2=0.01 z3=0.01
'49' z1=0.01 z4=0.99 z2=0.00 z3=0.00
'meters' z1=0.39 z4=0.38 z2=0.16 z3=0.08
'.' z1=0.44 z4=0.33 z2=0.14 z3=0.09
  • RAG-Tokenの機構の正常な動作と誤った交代先: 数値トークンの位置で担当文書が交代する挙動自体は先ほどと同じであり、z4z_4 の支持率は 1.001.00 に達しています。ここでの問題は、機構が機能しなかったわけではなく、乗り換えた先が z2z_2(正解)ではなく z4z_4(ディストラクタ)だったということです。結果として「富士山の名前」と「エベレストの標高」がキメラのように合成されてしまいました。
  • 誤った文書が選ばれた理由: 最大の要因は検索確率の差です。pη(z4x)=0.360p_\eta(z_4\vert{}x)=0.360 に対し pη(z2x)=0.160p_\eta(z_2\vert{}x)=0.160 と2倍以上のスコア差がありました。RAG-Tokenは「山名」を生成する段階では最も確率の高い z1z_1(富士山)に従いましたが、「標高」を生成する段階になると、具体的な数値を持つ文書の中で最も検索確率の高い z4z_4(エベレストの8,849m)へ乗り換えてしまったのです。
  • Retrieverの質に依存する限界: つまり RAG-Tokenは「正しい文書」ではなく「確率の高い文書」のパーツを繋ぎ合わせるため、周辺化処理自体はRetrieverの検索エラーを訂正してくれません(RAG-Sequenceに至っては、確率の高いエベレストの文をそのまま最終回答に選んでしまっています)。TF-IDFが z4z_4z2z_2 より上位に置いた時点で、この誤答はすでに決まっていました。どちらの定式化を選ぼうと、最終的な性能の上限はRetrieverの質で決まるというのが、ここから読み取るべき重要な教訓です。
注記

冒頭の Mount Fuji. Mount Fuji is ... という重複は、MIN_NEW_TOKENS でEOSを禁止したことの副作用です。 モデルは Mount Fuji. で答え終えたつもりでしたが、終了を許されずに続きを書き始めました。 最低長を強制している間はピリオドも併せて禁止すると解消します。

RAG-SequenceとRAG-Tokenの比較

RAG-SequenceRAG-Token
生成の実体文書ごとに独立に生成し、候補を集約1トークンごとに全文書で投票しながら1本を生成
model.generate()そのまま使える(KK 回呼ぶ)使えない(自前のループが必要)
複数文書の統合できない(統合回答が候補に現れない)できる(必要な場面で担当が交代する)
ハルシネーション文書に無い情報は捏造されうる他の文書の投票で打ち消されうる
計算コストKK 回の生成 + 再スコアリング毎ステップ KK 本の順伝播
Retrieverの誤り訂正されない訂正されない
注記

Hugging Faceには RagSequenceForGeneration / RagTokenForGeneration も用意されていますが、 DPR + BART の組み合わせと専用インデックスが前提です。上のように自前でデコードを書けば、 任意のencoder-decoderモデルと任意の検索器で同じ定式化を試せます。 decoder-onlyモデルで再現する場合は、encoder_outputs の代わりに KK 本のプロンプトを それぞれKVキャッシュに詰め、生成したトークンを KK 本すべてにappendする形になります (logsumexp の部分は全く同じです)。

現代のLLM開発における実際のアプローチ

ここまでRAG-SequenceとRAG-Tokenの挙動を詳しく見てきましたが、「それではChatGPTやLangChainなど、現在実用化されている最新のLLM・RAGシステムではどちらのアーキテクチャが使われているのか?」 という疑問が生じるかもしれません。

結論から言うと、現在の多くの実用的なRAGシステムでは、厳密にはどちらの定式化もそのままの形では使われていません。 代わりに、「In-Context Learning(コンテキスト内学習)」を利用して検索結果をプロンプトに単純に連結するアプローチが主流となっています。近年、この最も基礎的で標準的な手法は、後述する発展型と区別するために 「Naive RAG(素朴なRAG)」 と総称されるようになっています。

  • オリジナルのRAG(Sequence/Token)が生まれた背景: 論文発表当時(2020年頃)の主流な生成モデル(BART等)は、一度に入力できるコンテキスト長(トークン数)が1024程度と短く、検索された複数のドキュメントをすべてプロンプトに詰め込むことが困難でした。そのため、ドキュメントを1つずつモデルに入力して出力を計算し、最後にプログラム側で「確率を周辺化(計算して足し合わせる)」する複雑なアーキテクチャが必要だったのです。
  • 現代のアプローチ(巨大なコンテキストウィンドウの活用): GPT-4やGeminiなどの最新LLMは、数万〜数百万トークンという巨大なコンテキストウィンドウを持っています。そのため、Retrieverが取得した上位 KK 個のドキュメントをすべて1つのプロンプトの中に単純に連結して(例: 参考資料: [doc1] [doc2] ...)、一度にLLMへ入力してしまいます。

現代のLLMは強力な自己注意機構(Self-Attention)を備えているため、プロンプト内にある全てのドキュメントを同時に俯瞰し、必要な情報をトークンごとに拾い上げながら(複数文書の情報を統合しながら)回答を生成できます。 つまり、現代の実用システムは 「巨大なプロンプト空間を利用して、LLM自身のAttentionメカニズムにRAG-Tokenのような複数文書の統合を丸投げしている」 手法をとっていると言えます。確率を厳密に計算して周辺化するオリジナルのRAG-Token/Sequenceは、現在では主にRAGの基本概念を学ぶための理論的基盤として、あるいは極めて厳格に出典をコントロールしたい特殊な研究ケースで参照されています。

実践:LangChainとGeminiを用いた現代のRAG (Naive RAG)

実際にLangChainとGemini(gemini-3.7-flash)を用いて、これまでと全く同じ質問および4つのドキュメントで実装すると以下のようになります。

まず、APIキーを設定するためにプロジェクトのルートに .env ファイルを作成します。

.env
GEMINI_API_KEY=your_api_key_here
ヒント

Gemini APIキーは、Google AI Studioから無料で取得することができます。APIキーの取得・管理方法の詳細については、Googleの公式ドキュメントをご参照ください。 なお、無料枠(Free tier)を利用する場合、入力したデータはGoogleの製品改善やモデル学習のために利用される可能性がある点にご留意ください(機密情報の入力は避け、詳細は公式の利用規約をご確認ください)。

続いて、以下のコードを実行します(事前に pip install langchain langchain-google-genai python-dotenv などを実行して必要なパッケージをインストールしておきます。なお、.env から GEMINI_API_KEY を自動で読み込む挙動は langchain-google-genai バージョン4.0.0以降を前提としています)。

import os
from dotenv import load_dotenv
from langchain_core.prompts import PromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# .envファイルから環境変数を読み込む
load_dotenv()

# 今回使用した4つのドキュメント(エベレストのディストラクタを含む)
DOCS = [
"Mount Fuji is the highest mountain in Japan.",
"Mount Everest is the highest mountain in the world, 8,849 meters tall.",
"The highest mountain in Japan rises 3,776 meters above sea level.",
"Mount Fuji straddles the border of Shizuoka and Yamanashi prefectures."
]

QUESTION = "What is the highest mountain in Japan, and how many meters tall is it?"

# ドキュメントをリストから箇条書きテキストに結合する関数
def format_docs(docs):
return "\n".join([f"- {doc}" for doc in docs])

# プロンプトのテンプレートを作成
prompt = PromptTemplate.from_template(
"Answer in one full sentence. Include both the name of the mountain and its height in meters.\n\n"
"question: {question}\n"
"context:\n{context}"
)

# Geminiモデルの初期化
llm = ChatGoogleGenerativeAI(
model="gemini-3.7-flash",
temperature=0.0
)

# LangChainのLCEL (LangChain Expression Language) でチェーンを構築
chain = (
{"context": lambda x: format_docs(DOCS), "question": RunnablePassthrough()}
| prompt
| llm
| StrOutputParser()
)

# 実行
result = chain.invoke(QUESTION)
print("Naive RAG :", repr(result))

実行結果:

Naive RAG : 'The highest mountain in Japan is Mount Fuji, which is 3,776 meters tall.'

実践:Google公式SDK(google-genai)を用いた場合

LangChainのような抽象化フレームワークを使わず、Googleの新しい公式SDKである google-genai を使ってシンプルに実装することも可能です(必要なパッケージ:google-genai python-dotenv)。

import os
from dotenv import load_dotenv
from google import genai
from google.genai import types

# .envファイルから環境変数を読み込む
load_dotenv()

# クライアントの初期化(環境変数 GEMINI_API_KEY が自動で読み込まれます)
client = genai.Client()

# 今回使用した4つのドキュメント
DOCS = [
"Mount Fuji is the highest mountain in Japan.",
"Mount Everest is the highest mountain in the world, 8,849 meters tall.",
"The highest mountain in Japan rises 3,776 meters above sea level.",
"Mount Fuji straddles the border of Shizuoka and Yamanashi prefectures."
]

QUESTION = "What is the highest mountain in Japan, and how many meters tall is it?"

# プロンプトを文字列として構築
context = "\n".join([f"- {doc}" for doc in DOCS])
prompt = f"""Answer in one full sentence. Include both the name of the mountain and its height in meters.

question: {QUESTION}
context:
{context}"""

# 生成を実行
response = client.models.generate_content(
model='gemini-3.7-flash',
contents=prompt,
config=types.GenerateContentConfig(
temperature=0.0
)
)

print("Naive RAG (SDK) :", repr(response.text))

実行結果:

Naive RAG (SDK) : 'Mount Fuji is the highest mountain in Japan, rising 3,776 meters above sea level.'

上記のように、現代の強力なモデルであれば、プロンプト内に「富士山」の断片的な情報と「エベレスト(ディストラクタ)」の情報が混在していても、事前の確率計算なしに自身のAttentionメカニズムのみで正確に情報を取捨選択し、文脈に合った完璧な統合回答を出力することができます。

RAGパラダイムの進化と発展

サーベイ論文: Retrieval-Augmented Generation for Large Language Models: A Survey (Gao et al., 2023)

RAGの技術は、その後の研究により「Advanced RAG」や「Modular RAG」へと急速な進化を遂げています(「Naive / Advanced / Modular」という3つのRAGパラダイムの分類は、上記のサーベイ論文に由来します)。

Advanced RAG(高度なRAG)

Advanced RAG(高度なRAG)は、Naive RAGにおける検索精度や生成品質の限界を克服するため、検索前(Pre-retrieval)および検索後(Post-retrieval)に最適化戦略を組み込んだパラダイムです。

  • 検索前処理: スライディングウィンドウアプローチや細粒度のセグメンテーション、メタデータの追加などにより、インデックスの品質を向上させます。クエリの書き換えや拡張も含まれます。
  • 検索後処理: 検索された情報をLLMに入力する前に、最も関連性の高いコンテンツをプロンプトの端(冒頭と末尾)に再配置するリランキング(Reranking、Lost in the Middle 対策)や、不要な情報を削ぎ落とすコンテキスト圧縮(Context compressing)を行います。

Modular RAG(モジュール式RAG)

さらに柔軟性と適応性を高め、単なる「検索して読む」というシーケンシャルな処理にとらわれないアーキテクチャです。

  • 新モジュールの導入: 検索エンジンや知識グラフ(KG)を直接検索する「Searchモジュール」や、マルチクエリストラテジーを用いて多様な視点から検索を行う「RAG-Fusion」などの新しい機能モジュールが導入されています。
  • 適応的・反復的な検索: 必要に応じて反復的(Iterative)に検索を行ったり、LLM自身に検索が必要かどうかを適応的(Adaptive)に判断させる(FlareやSelf-RAGなど)パターンも発展しています。

まとめ

本記事では、RAG(検索拡張生成)の基礎概念から、オリジナルの論文で提案された2つの定式化(RAG-Sequence / RAG-Token)の数学的なメカニズム、そして現代の実用システムにおけるアプローチまでを詳しく解説しました。

この記事の主なポイント

  • RAGの有用性: 外部の知識ベース(ノンパラメトリックメモリ)とLLM(パラメトリックメモリ)を結合することで、ハルシネーションの抑制と最新情報の動的な参照が可能になります。
  • 定式化の違いによる振る舞い: RAG-Sequenceは「文書ごとに独立して回答を生成し、後から確率を集約する」ため、単一の文書に答えが揃っていないと正答できません。一方、RAG-Tokenは「1トークンごとに全文書の確率を周辺化する」ため、複数文書から柔軟に情報を合成できます。
  • 検索品質の重要性(キメラ効果): 柔軟な合成能力を持つRAG-Tokenであっても、Retrieverが誤った文書(ディストラクタ)を上位にランク付けしてしまうと、正解と不正解のパーツが継ぎ接ぎされた「キメラ」のような回答を生み出してしまいます。最終的なRAGの性能上限は、常に検索品質(Retriever)に依存します。
  • 現代のLLMにおけるアプローチ(Naive RAG): 現在のChatGPTやGeminiなどのシステムでは、巨大なコンテキストウィンドウを活かし、複雑な確率計算(周辺化)を行わずに検索結果をそのままプロンプトに詰め込みます。情報の取捨選択や統合は、強力なSelf-Attentionメカニズムを備えたLLM自身に委ねるのが現代の主流です。

RAGは現在も「Advanced RAG」や「Modular RAG」といった形で急速に進化を続けています。しかし、複数の知識をどのように統合し、ノイズをどう処理するかという根本的な課題は変わりません。本記事で確認した「生成モデルが複数文書の情報をどう捌くのか」という原理的なメカニズムへの理解は、最先端のRAGアーキテクチャを設計・最適化する上でも重要な指針となるはずです。

本記事の文章・構成の一部に生成AIを使用しています。