このブログのコンテンツの概要は、2024年12月18日開催予定の Oracle JAM Session にてご紹介いたします。
はじめに
ベクトル検索結果に対するリランキングの意義を理解しようとすると、ベクトル検索で既に類似度順に並べているのに、なぜ再度並べ替えるのかという疑問が生じます。
ベクトル検索とテキスト検索のハイブリッド検索で、両方の検索結果をまとめて順位付けする場合にリランカーを使うことがあります。これは、それぞれ異なる観点でスコア付けされているため改めて統一された観点で並べ替える必要性があることは理解しやすいですね。しかし、ベクトル検索だけであってもリランキングが有効であることはわかりにくいかもしれません。
そんな疑問にお答えしようというのがこのブログです。
このブログの元ネタは以下の私のX投稿です。
Xの投稿は、リランキングの説明用データを自分へのメモとして書いたものでした。しかし予想外に多くの方に見ていただき、プチバズったため、X投稿で省略していた部分を補足してブログ記事としてまとめました。
リランカーによるスコアの順位は必ずしも人間の感覚と一致するわけではありません。また、リランキングしたチャンク（ドキュメントの断片）をLLMのコンテキストとすることが、ベクトルデータベースの生の出力を使うよりも必ずRAGの回答精度を上げるというわけでもありません。
ところで、そもそもRAGにおいてデータベースから検索した結果をクエリ（質問文）への類似度順に並べてLLMへコンテキストとして与えるのは、コンテキストの先頭、もしくは、末尾にあるデータ程LLMが見つけやすいという性質（Lost in the Middle: How Language Models Use Long Contexts）があることと、類似度が高い検索結果であるほど、クエリの質問に対する正解を含んでいる可能性が高いと期待するからです。しかし、この期待は必ずしも正しいとは限りません。最近、ある論文でLLMが回答を生成する前に、最も回答精度が高くなるであろう検索結果の並び順を予測できることが指摘されました。その論文を紹介した X投稿も論文の著者からメッセージをもらうくらいにはプチバズってしまいましたのでこの分野の関心の高さがうかがえます。
このブログでリランカーについて理解を深めた後は、ぜひ私のX投稿と、そこで紹介している株式会社ナレッジセンスの須藤さんのブログ RAGのハルシネーションを尤度で防ぐと原論文 Likelihood as a Performance Gauge for Retrieval-Augmented Generationをご覧ください。
類似した文の見つけ方
ここからは本題です。さて、そもそも、ある文に類似した文を見つけるにはどうしたらいいでしょうか？
前提
タスク設定
1つのクエリとなる文（長さは50文字）と多数（例えば、5,000万件）の検索対象となる文群（それぞれの長さの500文字）があるとします。この多数のデータの中で、クエリとなる文に最も類似した検索対象文を k 個見つけるというタスクを想定します。
手法
このタスクをLLMなどのTransformerアーキテクチャの機械学習モデルを使って遂行するとします（他の手法はここでは扱いません）。この場合、大きく2つの方法が考えられます。cross-encoder による方法と bi-encoder による方法です。
cross-encoder による方法
cross-encoder とは
2つの文を"[CLS]文1[SEP]文2[SEP]"のように並べて Transformer へ入力して、2つの文の類似性スコアを計算するものです（[CLS]は文頭を表す特殊トークンで[SEP]は文の区切りを表す特殊トークンです）。2つの文を一度に Transformer へ入力することから cross（交差）-encoder と呼ばれます。なお、cross-encoder は類似性スコアだけを出力し、それぞれの文の埋め込みは出力しません。ただし、内部的には、文を埋め込み表現にエンコードするため、encoder と呼ばれます。
cross-encoder による方法の手順
cross-encoder を使って、前述の最も類似性の高い検索対象文を見つけるためには、1つのクエリ文とすべての検索対象文群の類似度を計算して、類似度が最も高い文を見つけることになります。
cross-encoder による方法の問題点
cross-encoderは、前述のとおり2つの文（クエリ文と検索対象文）を同時に入力する必要があります。そのため、クエリの度にクエリと検索対象文群の組み合わせをすべて cross-encoder で評価する必要があります（クエリは実際にユーザーが入力するまで不明であるため）。
この計算コストは膨大過ぎて実用的ではありませんでした。
「Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks」という論文によると10,000センテンスの中で最も類似したペアを見つけるタスクには65時間かかりました。ペアの数は、n(n-1)/2 = 10000 * 9999 / 2 = 49,995,000 個となりますので、約5,000万回の類似性スコアの計算が必要となります。言い換えると、5,000万件のデータからクエリに最も似たデータを見つけ出すには、65時間必要ということになります。
※測定条件の詳細が気になる方もいらっしゃると思いますが、オーダー（桁数）として、非実用的であることには同意いただけるかと思います。
bi-encoder による方法
bi-encoder とは
bi-encoder では、名前の"bi"が示すように2つの文の類似度を計算する際に、それぞれの文の埋め込みを別々に生成します。そのため、検索対象の文群はあらかじめ埋め込み表現にエンコードして保存しておくことができます。
検索対象となる文が5,000万件あれば、5,000万件の埋め込みを生成し、通常、この大量の埋め込みはベクトルデータベースへ格納されます。
次に、検索の段階で、クエリとなる文をTransformerを使って埋め込み表現にエンコードします。これは、通常、1つの埋め込みです。検索対象の文群は既に埋め込みにエンコードされてベクトルデータベースに格納されていますので改めてエンコードする必要はありません。つまり、検索時のエンコード処理は、cross-encoder が 5,000万回のエンコード処理を必要としたのに対して、bi-encoder は、1回で済む こととなります。
このクエリ文の埋め込みに最も類似した検索対象文の埋め込みをベクトルデータベースを使って見つけ出します。
ここで主に２種類の効率化が行われます。1つは、埋め込みを生成する際の Pooling です。そして、もう一つはベクトル検索の際の近似最近傍探索（ANN:Approximate Nearest Neighbor）です。
Pooling とは
bi-encoder（埋め込みモデル）が文を埋め込みに変換するためには、まず最初に文をトークナイザーによりトークンへ分割します。トークナイザーは文字列の出現頻度に応じてトークンを割り当てることで効率的にトークン化を行うよう設計されています。トークナイザによりますが多くの場合、日本語の１文字は平均すると１トークン程度になります。つまり、クエリ文が50文字だったとすると50トークン程度、検索対象文が500文字だったとすると500トークン程度となります。
次に、各トークンは１つの埋め込みに変換されます。このとき、辞書を引くように単純に１トークンを１つの埋め込みへ変換するだけだと、文中の文脈を伝えることができません。そこで、self-attentionなどの機構を使用して文脈を考慮した埋め込みを生成します。つまり、「吾輩は猫である」の「猫」と「猫でもわかるリランキング」の「猫」はトークンは同じですが生成された埋め込みは異なります。この段階では、クエリ文の埋め込みは50個、検索対象文の埋め込みは500個です。
下図は、同じ単語（"bank"）が異なる文脈において異なる埋め込みに変換される例です。
Does BERT Make Any Sense? Interpretable Word Sense Disambiguation with Contextualized Embeddings より異なる文脈における"bank"のコンテキスト埋め込みの例
bi-encoder（埋め込みモデル）は、最終的な埋め込みを生成するために Pooling という処理を行います。これは典型的には、文中のすべての埋め込みの平均操作です。つまり、クエリ文の50個と検索対象文の500個の埋め込みは、それぞれ平均化されて1つの最終的な埋め込みとなります。
下図は、Pooling のイメージです。
こうして Pooling された最終的な埋め込みは、元の50個、もしくは、500個の埋め込みが持っていた特徴をある程度反映しつつ容量を劇的に小さくすることができます。これによりストレージ容量を抑えることができますし、後の検索の際の演算量を少なくすることができます。
近似最近傍探索（ANN:Approximate Nearest Neighbor）とは
クエリの埋め込みと最も類似した検索対象の埋め込みを k 個見つけ出す際、単純にクエリ文の埋め込みと、データベース内のすべての検索対象文の埋め込みとの類似度を計算し、類似度が最も高い k 個の埋め込みを選び出す k-最近傍（kNN:k-Nearest Neighbors）では5,000万回の類似度計算が必要となります。これは計算コストが非常に高く、実用的ではありません。そこで、近似最近傍探索（ANN）という手法を使用します。
ANNは、完全な精度を少し犠牲にする代わりに、劇的に検索速度を向上させます。代表的なアルゴリズムには、HNSW（Hierarchical Navigable Small World）があります。HNSWは多層のグラフ構造を構築し、上位層から徐々に下位層へと探索範囲を絞り込むことで、効率的に類似する埋め込みを見つけ出します。
ANNは、以下のような特徴を持っています。
- インデックス構築：検索前に埋め込みデータに対して特殊なインデックス構造を構築します。これにより、検索時の探索空間を大幅に削減できます
- 近似性：完全な探索と比べて、わずかな精度の低下は許容する代わりに、数桁のオーダーで検索速度を向上させます
- メモリ効率：インデックス構造により、メモリ使用量と検索速度のトレードオフを調整できます
この近傍検索がkNN（Exact Nearest Neighborによる k 個の探索）で見つかる k 個のうちどの程度をカバーしているかという search recall@k というメトリックがあります。ベクトルデータベースの検索精度といった場合にはこれを表しています。この search recall@k が気になる方も多いかと思いますが、最近のToward Optimal Search and Retrieval for RAG という論文では、RAG においては、この search recall@k を 70% 程度まで落としても RAG の性能にはあまり影響しないことが報告されています。また、検索精度を下げることでむしろ検索速度とメモリ効率を向上できるメリットがあることが示唆されています。
なお、類似度のメトリックには多くの場合コサイン類似度が使われますが、埋め込みが正規化されている場合には内積がコサインと同等となりますので計算量の少ない内積が用いられます（kNN であっても ANN であっても同じです）。
bi-encoder による方法の問題点
文を埋め込み空間にマッピングする過程（特に pooling 処理）で情報が圧縮されるため、微妙な文脈や関係性が失われてしまいます。そのため、検索結果を bi-encoder でエンコードされた埋め込みの類似度順に並べると必ずしも最適ではない場合があります。
Bi-Encoder と Cross-Encoder の特徴のまとめ
| 特徴 | Bi-Encoder | Cross-Encoder |
|---|---|---|
| 処理速度 | 高速（文章を個別にエンコード。検索時にはクエリのみエンコード） | 低速（文章ペアごとに処理が必要で検索時にはすべてのペアの類似度計算が必要） |
| スケーラビリティ | 高い（大規模データセットに適している | 低い（文章ペアの組み合わせ数だけ処理が必要） |
| 精度 | 比較的低い | 高い（文章間の関係性をより深く理解） |
| メモリ効率 | 良い（各文章を1回だけエンコード） | 悪い（ペアごとに処理が必要） |
| 用途 | 情報検索、セマンティック検索、クラスタリング | 文章ペアの分類、高精度なランキング |
| 文章の表現方法 | 独立した文章エンベディングを生成 | 内部的に文章ペアを同時にエンコード。埋め込みは出力しない |
| 実用的な制限 | 数十万の文章を効率的に比較可能 | 数十のペア比較に最適 |
リランカーによる解決
そこで、bi-encoder は効率的に候補を絞り込む役割に特化し、絞り込まれた候補文を評価して並べ替える役割を担うのが リランカー です。
- bi-encoderとベクトルデータベースで大規模なデータから候補を絞り込む（例：上位100件）
- 次にリランカーで候補のリランキングを行い、より高い精度で並べ替えた検索結果をLLMのコンテキストとする（例：上位10件）
ベクトルデータベースから取得する件数とLLMのコンテキストに入力する件数は、使用するドキュメントと典型的な想定質問により実際に検証して決定する必要があります。これは、コンテキストには"正解"が含まれている必要があるという観点ではより多くのデータをデータベースから取得した方が良いということになりますが、一方でコンテキスト内にノイズが多かったり、コンテキストの長さが長すぎると精度が低下するという性質もあるためです。
また、実運用に入った後もモニタリングを継続して調整することが望ましいでしょう。
リランカーの仕組み
リランカーは、クエリと候補文をペアとして再評価するために、通常、上記の cross-encoder を使用します。つまり、 リランカーの入力は、情報が圧縮された埋め込みではなく、元の文（クエリ文と絞り込まれた検索対象文）そのもの であるという点が重要です。この仕組みにより、より高い精度で文書をランキングできると期待されます。
リランカーがすべて cross-encoder であるわけではありません。ColBERT のような Multi-vector アーキテクチャを採用するものや LLM でリランキングするものもあります。
今後の展望
RAGの回答精度が高くなる検索結果の配置を質問文の尤度（検索で与えられた文書群を見た後の質問文の条件付き尤度）を計算することで予測することが可能であることが Likelihood as a Performance Gauge for Retrieval-Augmented Generationで報告されています。
また、この研究ではこの最適な配置を効率的に割り出す2つの方法（Naïve likelihood-based selectionとGold Document Reordering）も提案されており、今後の発展が期待されます。
誤解を恐れずに単純化するとLLMにとって検索されたドキュメントの並びと質問の配置が「自然」に感じられる時、回答の精度が高くなる傾向がある。そして、この「自然」な配置を推測する方法が存在する。したがって、データベースを検索したら、検索結果のこの「自然」な配置を割り出して、その順番でLLMへ入力することで高い精度を実現できる可能性があるということですね（ただし、ここで主役となっている尤度は API アクセスのクローズドモデル（商用モデル）では知ることができないので、セルフホストできるオープンウェイトモデルを使う必要があるようです）。
まとめ
類似文検索において、bi-encoderとcross-encoderはそれぞれ長所と短所があります。cross-encoderは高精度な類似度評価が可能ですが、計算コストが膨大で大規模データでの実用が困難です。一方、bi-encoderは効率的な検索が可能ですが、pooling処理による情報の圧縮で精度が低下する可能性があります。
リランカーは、これら両者の利点を組み合わせた実用的なソリューションとして機能します。
具体的には
- まずbi-encoderで効率的に候補を絞り込み
- 次に絞り込まれた少数の候補に対してcross-encoderで精密な再評価を行う
という2段階のアプローチにより、大規模データに対する現実的な処理時間を維持しながら、より正確な類似度評価が可能となります。ただし、このような工夫を施してもリランカーによる順位付けは必ずしも人間の感覚と完全に一致するわけではなく、RAGシステムの回答精度向上も保証されないことには注意が必要です。
補足１：クエリとドキュメントは似ていない件
RAG におけるベクトル検索では、クエリ（質問文）と検索対象文（何らかのドキュメントを分割したチャンク）の類似度が高いものを見つけて LLM への入力とします。
しかし、ここでクエリの文とドキュメントのチャンクはそもそも文の構造が似ていないという問題があります。クエリは、質問文の形をしていることが多いですし、チャンクは平叙文です。この点については大まかに２つのアプローチが試みられています。
ひとつは、HyDE（Hypothetical Document Embeddings）と呼ばれる手法で。ベクトル検索をする際のクエリーに対する仮想的な理想的な回答文書をLLMで生成して、この仮想的な文書に類似したチャンクを検索するものです。Precise Zero-Shot Dense Retrieval without Relevance Labelsという論文で提案されました。
もうひとつは、埋め込みモデルによる対処です。デュアルエンコーダーモデル、もしくは、Two-Tower モデルと呼ばれるものです。
デュアルエンコーダーモデルでは、埋め込みモデルを質問文用のエンコーダーと回答となるドキュメント用エンコーダーの2つのエンコーダーで構成し、ペアとなる質問文と回答（候補）となるドキュメントが埋め込み空間内で近くにマップされるように学習します。
埋め込み生成時には、文が質問文であるのか、その回答（候補）となるドキュメントであるのかを指定して、それぞれに適したエンコーダーを使用して質問文用の埋め込み、もしくは、回答（候補）となるドキュメント用の埋め込みを生成します。
検索時にはこれらの類似性を調べて回答（候補）となるドキュメントを見つけ出します。
Google のVertex AI embeddingsは、このアーキテクチャを採用しているものと思われます。
Vertex AI embeddings 解説記事：Enhancing your gen AI use case with Vertex AI embeddings and task types
Cohere Embed もクエリかドキュメントか等を指定して埋め込みを生成しますが、アーキテクチャについては私が知る限り公開されていないようです。
Cohere ドキュメント：Cohere Embed API Reference(input_type)
補足２：類似度順に並べるとRAGの精度が悪くなる可能性がある件
In Defense of RAG in the Era of Long-Context Language Modelsという論文で、検索されたチャンクを元の文書での出現順序で配置することで、従来のRAGよりも優れた性能を実現できるという報告がなされています。
これは、今後の展望で紹介した質問文の尤度（検索で与えられた文書群を見た後の質問文の条件付き尤度）と RAG の精度に関連がある Likelihood as a Performance Gauge for Retrieval-Augmented Generationことと整合しているように思えます。というのは、直観的には、LLM が Auto-regressive モデルであることの帰結であるように思えるからです。Auto-regressive モデルは、現在位置より前の全てのトークン（入力コンテキストと生成済み出力の両方）を使って、次のトークンを予測するモデルですので、入力コンテキストに配置するチャンクが、クエリ文への類似度で並べ替えられていると実際に入力されたコンテキストは通常存在しえないような文章となるため、生成された文章（回答）の品質が下がるということはあり得そうに思えます。つまり、LLMにとって自然な文章である程、理解し易いが、順番がシャッフルされていると理解度が落ちてしまうと例えることができそうです。昔、同僚のスーパーエンジニアが「石の気持ちになって考える」と言っていましたが、「LLMの気持ちになって考えること」が必要かもしれませんね。