今回の論文
今回取り上げるのは、Weijia Shi、Sewon Min、Michihiro Yasunaga、Minjoon Seo、Richard James、Mike Lewis、Luke Zettlemoyer、Wen-tau Yihによる論文「REPLUG: Retrieval-Augmented Black-Box Language Models」です。2024年6月にNAACL 2024で発表され、ACL Anthologyから公開されています。URLは https://aclanthology.org/2024.naacl-long.463/ 、研究分野は情報検索、RAG、言語モデル推論です。
この論文を選んだ理由は、内部表現に触れずファインチューニングもできないAPI型LLMに対しても、検索を品質改善のレバーにできるからです。モデルを替えられない業務システムで、検索・プロンプト・評価をどう設計すべきかを考える手がかりになります。
どんな技術か
REPLUGは「Retrieve and Plug」の略で、凍結した言語モデルをブラックボックスのまま使い、外部検索器を差し込むRAGです。質問または文脈に近い文書を検索し、各文書を一件ずつ入力の先頭に付けてLLMへ渡します。その複数回の出力確率を、検索スコアに応じて重み付きで平均し、次のトークンを決めます。
ポイントは、検索文書を一つの長いプロンプトへ詰め込むだけではないことです。文書ごとに独立した推論を走らせるため、コンテキスト長の上限に縛られずに上位複数件を使えます。さらにREPLUG LSRでは、どの文書がLLMの正解予測を助けたかをAPIの尤度スコアから測り、検索器だけをそのLLM向けに学習します。
課題
高性能LLMの多くはAPIで提供され、内部の埋め込みや注意機構へアクセスできません。従来のRETROやAtlasのような検索拡張モデルは、検索文書を読むための層を追加したり、LLMと検索器を一緒に学習したりするため、この条件ではそのまま使えません。
単純なRAGにも限界があります。上位文書をすべて一つのプロンプトに入れると、コンテキスト枠を急速に消費し、文書数を増やせません。反対に少数の文書だけを入れると、検索の取りこぼしが直接回答品質へ響きます。無関係な文書は注意をそらし、幻覚を減らすどころか誤答の原因にもなります。
実際のAIシステムでは、社内文書検索、技術サポート、コード支援、最新情報を含む回答などでこの問題が起きます。生成モデルを再学習できなくても、「そのモデルが本当に使えると判断する根拠」を検索器が取れるようにする必要があります。
用語解説
- ブラックボックスLLM
- テキスト入力と生成結果、場合によってはトークン確率だけをAPIで受け取るLLMです。重みや隠れ状態を変更できないため、REPLUGではLLMを固定し、外側の検索器だけを改善対象にします。
- 密検索（dense retrieval）
- 質問と文書をベクトルに変換し、類似度が高い文書を検索する方式です。REPLUGはdual encoderで質問・文書を同じ空間へ埋め込み、FAISSインデックスで上位文書を高速に取り出します。
- 確率アンサンブル
- 複数回のLLM推論で得た次トークン確率を平均する方法です。本論文では多数決ではなく検索スコアをsoftmax化した重みを使うため、より関連する文書の推論を強く反映できます。
- パープレキシティ（perplexity）
- 正解の続きテキストを、モデルがどれほど自然に予測できたかを表す指標で、低いほどよい値です。REPLUG LSRは、文書を加えたときに正解への確率が上がるかを、検索器の教師信号にします。
- KLダイバージェンス
- 二つの確率分布の隔たりを測る損失です。ここでは「検索器が文書へ与える確率」を、「LLMが正解予測に役立つと評価した文書の確率分布」へ近づけるために使います。
技術の仕組み
REPLUGの設計は、検索、文書ごとの推論、確率統合、必要に応じた検索器学習の四段階です。LLMそのものは、推論時も学習時も更新しません。
検索して文書ごとにLLMを呼ぶ
入力文脈をx、文書をdとします。dual encoderは両者の埋め込みを作り、コサイン類似度s(d, x)で上位k件を選びます。文書埋め込みは事前計算してFAISSへ登録するため、オンラインでは質問側を埋め込むだけで検索できます。
次に、上位各文書についてd + xという別々の入力を作ります。LLMはそれぞれから次トークン確率p(y | d, x)を返します。通常の「上位10件を連結した一回の推論」と異なり、10件なら10回の独立推論です。したがって、各入力に使う文書長を保ったまま、より多くの候補を使えます。
検索スコアで出力確率を統合する
最終確率は、各文書のLLM出力を足し合わせたものです。文書dの重みλ(d, x)は、検索類似度をsoftmaxで正規化して作ります。式の意味は、検索器が「この文書は質問に近い」と強く判断するほど、その文書を読んだLLMの予測を採用する、ということです。
p(y | x) = Σ p(y | d, x) × λ(d, x)
この分離は、検索精度と生成精度を別々に観察できる利点もあります。ただし、k件ならAPI呼び出しも概ねk倍になるため、実装では候補数、並列数、キャッシュをレイテンシ予算と一緒に設計します。
LLMを検索器の採点者にするREPLUG LSR
既製の検索器が、LLMにとって役立つ文書を必ず上位に出すとは限りません。そこでLSRは、教師データの入力xと正解の続きyに対し、各候補文書を付けたときのLLMの正解確率を計算します。正解を高確率で予測させる文書ほど、LLM側の高い教師スコアを得ます。
検索器側も上位候補の類似度を温度付きsoftmaxへ変換し、LLM側の教師分布とのKLダイバージェンスを最小化します。LLMは固定の採点APIであり、更新されるのは検索器だけです。論文ではContrieverを初期値に、The Pileから切り出した80万の256トークン列を使い、前半128トークンを入力、後半128トークンを正解続きとして学習しました。
インデックスを非同期に作り直す
検索器を更新すると、既に保存した文書ベクトルが古くなります。毎ステップで全コーパスを再埋め込みすると高コストなので、論文では36 million件の文書ベクトルとFAISSインデックスを3,000ステップごとに再構築します。この「検索器の重みとインデックスの鮮度をずらして運用する」工夫が、大規模な検索器学習を現実的にします。
実験と結果
論文は、言語モデリング、MMLU、オープンドメイン質問応答で評価し、凍結LLMに後付け検索を加える効果と、LSRによる追加効果を分けて検証しました。
言語モデリングではモデル規模を問わず改善
The Pileを検索コーパスにし、bits per byte（BPB、低いほどよい）でGPT-2とGPT-3群を評価しました。推論では128トークンの文脈から検索し、上位10文書を確率アンサンブルしています。GPT-3 Davinci 175Bでは、元の0.80 BPBに対しREPLUGは0.77、LSR付きは0.75となり、LSRは相対6.3%改善でした。
8モデル平均では、検索だけのREPLUGが4.7%改善、検索器をLLM向けに調整したLSRが7.7%改善でした。LLMが大きくても外部情報が無意味になるわけではなく、さらに「意味的に近い文書」より「そのLLMの予測を改善する文書」を探す方が有効だと示しています。
MMLUと質問応答での結果
MMLUではWikipediaから上位10文書を検索し、5-shotの直接プロンプトで評価しました。Codex 175Bは68.3から、REPLUGで71.4、LSRで71.8へ上昇しました。論文の表記では相対4.5%および5.1%の改善です。LLaMA 13Bも55.6から58.8、59.3へ改善しています。
Natural QuestionsとTriviaQAでは、16-shotのCodexにLSRを追加すると、Exact Matchがそれぞれ40.6から45.5、73.6から77.3になりました。前者は相対12.0%、後者は相対5.0%の改善です。一方、フルデータでファインチューニングしたRAGモデルには届かない条件もあり、後付け検索が学習済み専用モデルを常に置き換えるとは結論付けられません。
効いた理由の検証
ランダム文書を同じ数だけアンサンブルした比較では性能が悪化しました。改善は単に「複数回呼んで平均した」ことではなく、関連文書を選んだことによるものです。またWikitext-103では、OPT-125Mでパープレキシティが6.9%、OPT-66Bでも5.6%改善し、GPT-2、OPT、BLOOMの複数サイズで効果が確認されました。
何に使える？
REPLUGは、生成モデルの内部を触れないことを制約ではなく、検索基盤を差し替え可能な改善面として扱う設計です。
API型LLMを使う社内RAG
ベンダー提供LLMで社内規程や障害対応を検索する場合に使えます。まず検索上位数件を別々にLLMへ渡して回答候補の確率を統合し、特に重要な分類や選択式の回答で試せます。検索器だけをドメイン文書・評価質問で調整できるため、LLMのモデル更新を待たずに改善を繰り返せます。
技術サポートとコード補完
エラー文やコード断片を検索クエリにして、類似する過去の障害票、仕様、修正例を取り出す用途です。型番や関数名が同じでも原因が異なる場合、単一の長いコンテキストより文書ごとに候補を評価する方が、どの根拠が有効かを測りやすくなります。生成APIがトークン対数確率を返さない場合は、候補回答の自己評価や下流タスクの正解率を弱い教師として使う実装が考えられますが、これは論文からの発展的な応用です。
モデル乗り換えに強い検索レイヤー
同じ検索器を複数LLMへ接続して評価すると、モデルごとに有効な文書の違いを観察できます。実運用では検索器を一律に共有するより、用途またはLLMごとに重み・再ランキング・評価セットを持つ方がよい場合があります。モデルを乗り換えても、文書整備と検索評価の資産を残しやすい構成です。
開発や事業へのヒント
この論文から得られる実務上の重要な視点は、「検索の正解」を一般的な関連性だけで定義しないことです。最終的に採用するLLMが答えを良く出せるかで検索器を測ると、RAG全体の目的に近づきます。
検索評価を回答品質までつなげる
検索Recall@kだけでは、取得文書がLLMの回答を助けたかは分かりません。代表的な質問について、文書ごとの回答正確性、引用可能性、コストを記録し、検索ランキングと下流品質のずれを見ます。LSRの発想は、検索器を下流モデル固有の目的関数へ合わせる設計です。
まずは候補数を小さくして検証する
論文は上位10件を使いますが、APIコストは推論回数に比例します。小規模プロダクトでは上位3件で始め、単純連結プロンプト、文書別並列推論、検索なしを同じ評価セットで比べるのが現実的です。差分が出る質問タイプを確認してから、候補数や検索器学習へ投資できます。
インデックス更新をプロダクトの機能として扱う
LSRの非同期再インデックスは、文書追加・埋め込みモデル更新・権限変更があるRAGでも必要になる考え方です。埋め込みのバージョン、インデックス作成時刻、評価結果をセットで保持すれば、検索品質の悪化をロールバックしやすくなります。
限界
REPLUGは上位k文書ごとにLLMを呼ぶため、単一プロンプトのRAGより推論費用と遅延が増えます。ストリーミング生成ではトークンごとの確率統合も実装上の課題です。常に10件を使うのではなく、問い合わせの難しさに応じて検索とアンサンブルを発動する制御が必要です。
LSRの学習には、正解の続きテキストとLLMの尤度を取得できるAPIが必要です。チャットAPIのように対数確率を返さない製品では論文どおりには実装できません。また、36 million文書の再埋め込みを含む実験規模は小規模チームには重く、更新時のインデックス不整合も監視しなければなりません。
論文自身も、LLMが回答時に検索文書へ依存したのか、パラメータ知識を使ったのかを明確に追跡できない点を限界に挙げています。取得文書は根拠表示・アクセス制御・鮮度確認と別途結び付ける必要があります。無関係な検索文書がモデルを邪魔する可能性もあるため、空検索時のフォールバックを評価してください。
よくある質問
Q. REPLUGは一般的な「文書を全部プロンプトへ入れるRAG」と何が違いますか？
A. 上位文書を連結して一回だけ生成するのではなく、各文書を別々に付けて複数回推論し、トークン確率を検索スコアで統合します。これによりコンテキスト長と文書数のトレードオフを緩められますが、API呼び出し数は増えます。
Q. LLMをファインチューニングできるならREPLUGは不要ですか？
A. 不要とは限りません。論文でもフルデータで学習したRAGモデルが強い条件はありますが、LLMの再学習は高コストです。モデル更新が難しい、最新文書へ頻繁に追随したい、複数LLMを比較したい場合には、検索レイヤーだけを改善するREPLUGの利点があります。
Q. REPLUG LSRには正解ラベルが必要ですか？
A. 必要です。論文は入力の後半部分を正解の続きとして使い、各候補文書がその正解の尤度をどれだけ上げるかで検索器を学習します。社内RAGなら、確認済みの質問回答、解決済みチケット、編集履歴などを候補にできます。
Q. チャットAPIがlog probabilityを返さない場合でも使えますか？
A. 推論時のREPLUGは、出力確率を得られるほど忠実に実装できます。得られないAPIでは、文書を一件ずつ渡して生成・判定する再ランキング型へ寄せる必要があります。LSRも論文そのままの尤度教師は使えないため、正解率や評価モデルを代替教師にする追加設計が必要です。
Q. 取得文書を根拠として表示できますか？
A. 表示できますが、REPLUGだけで「この回答はこの文書に基づく」と証明できるわけではありません。論文も生成が検索知識とパラメータ知識のどちらに依存したかは解釈しにくいと述べています。回答文の根拠区間を別途検証し、文書ID・版・アクセス権を記録してください。
今日の学び
この論文は、内部を変更できないブラックボックスLLMへ、検索をどう有効に接続するかを扱いました。REPLUGは検索文書ごとにLLM推論を行い、検索スコアで出力確率を統合します。さらにLLMが正解を予測しやすくなる文書を教師にして、検索器だけを改善します。
得られるヒントは、RAGを「モデルに文書を渡す配管」で終わらせず、最終LLMの回答品質まで見て検索を最適化することです。まずは少数候補、並列実行、検索なしとの比較から始め、品質・費用・遅延のバランスを測定するのが実装への近道になります。