Posted in人工知能

BGE-RerankerとFlagEmbeddingによるRAGの検索精度最適化

Artificial Intelligence tutorial - IT technology blog
Artificial Intelligence tutorial - IT technology blog

ベクトル検索が「一見マッチしているようで本質的に違う」結果を返すとき

基本的なRAGパイプラインの多くは、ユーザーからの質問受付 → クエリのエンベディング → ベクトルDBでのコサイン類似度検索 → 上位5件のチャンクをLLMのプロンプトに注入、というお決まりのフローを採用しています。この手順は、PDFが数十ファイル程度のデモ段階であれば非常にスムーズに動作します。

しかし、システムを本番環境（プロダクション）へ投入した途端、状況は一変します。数万ページに及ぶ社内ドキュメントを抱える場合、検索（Retrieval）フェーズのHit Rate（適合率）は壊滅的に低下することが少なくありません。

実際の具体例を見てみましょう：

  • ユーザーの質問: "試用期間中の社員は、会社から健康保険に加入してもらえますか？"
  • ベクトル検索のTop 3結果: 「社員」「健康保険」「支給」といったキーワードの出現頻度が高いため、正社員向けの保険支給規定ばかりが3件返されてしまいました。一方で、試用期間中の社員に関する適用除外条項は、9位まで沈んでしまっています。

その結果どうなるでしょうか？ LLMは誤ったコンテキストを受け取り、「試用期間中の社員も全額保険に加入できる」と自信満々にハルシネーション（嘘の回答）を出力してしまいます。

top_kを15や20に増やしても根本的な解決にはなりません。コンテキストの中央にある重要な情報を見落とすLost in the Middle現象に直面するだけでなく、トークンコストや応答レイテンシも跳ね上がってしまうためです。

なぜBi-Encoderはコンテキストを頻繁に見誤るのか？

この問題を解消するには、Dense Retrievalモデル（text-embedding-3-smallやbge-base-en-v1.5など）のアルゴリズムとハードウェア的な限界を正しく理解する必要があります。

これらのモデルはBi-Encoderアーキテクチャを採用しています：

  1. クエリとドキュメントを独立した2つのブランチに分け、それぞれ固定次元のベクトル（768次元や1536次元など）に圧縮します。
  2. それら2つのベクトル間の内積（Dot Product）やコサイン距離を用いて関連度を測定します。

最大の弱点は、トークンレベルでクエリとドキュメントが直接相互作用しない（Cross-Attentionが存在しない）点にあります。400〜500単語の文章を単一のベクトルに圧縮することで、条件分岐や否定表現といった繊細なニュアンスが失われてしまうのです。

検索精度（Retrieval Accuracy）を向上させる3つのアプローチ

現在、AIエンジニアが検討する主な解決策は以下の3つです：

  • 1. ハイブリッド検索（BM25 + Dense Retrieval）: 完全一致のキーワード検索とベクトルによる意味検索を組み合わせる手法です。エラーコードや型番、固有名詞を含むクエリには効果的ですが、複雑な条件付き推論を要する質問には対応しきれません。
  • 2. LLMベースのフィルタリング: 軽量モデル（Llama-3.1-8BやGPT-4o-miniなど）に上位20件のチャンクを走査させてフィルタリングします。精度は極めて高いものの、レイテンシが大幅に増加し（通常1〜2秒追加）、API利用料も急激に膨らむ点がデメリットです。
  • 3. Cross-Encoder Rerankerを用いた2段階検索（推奨）: 検索パイプラインを明確に2つのステージに分割します：
    • 第1段階（Fast Retrieval）: Bi-Encoderを用いてQdrantやMilvus上の数百万件のレコードを高速スキャンし、約25〜40件の候補を絞り込みます（レイテンシ約10〜20ms）。
    • 第2段階（Re-ranking）: Cross-Encoderを用いてクエリと各候補チャンク間のCross-Attention行列を計算し、最も精度の高いTop 3〜5件を厳選してLLMに渡します。

BGE-RerankerとFlagEmbeddingの実践活用

北京智源人工智能研究院（BAAI）が開発したBGE-Rerankerは、オープンソースのCross-Encoderモデル群の中で最も優れた性能を誇るものの一つです。FlagEmbeddingライブラリと組み合わせることで、10分もかからずにセットアップできます。

1. 環境構築

pip install FlagEmbedding torch

2. FlagRerankerによるコンテキスト識別のクイックテスト

以下のスクリプトは、モデルがいかにコンテキストを細かく識別できるかを実証するコード例です：

from FlagEmbedding import FlagReranker

# bge-reranker-v2-m3は多言語対応に優れており、ベトナム語や日本語にも高精度に対応
# use_fp16=TrueによりVRAM使用量を半減させ、推論速度を向上
reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=True)

query = "試用期間中の社員は、会社負担の健康保険の適用対象となりますか？"
passages = [
    "契約期間が1年以上の全正社員に対し、会社はハイグレードな医療保険パッケージを提供し、法定の健康保険への加入を義務付けています。",
    "2ヶ月の試用期間中、従業員は会社の健康保険の強制加入対象外となり、試用期間中の給与に直接の手当が含まれています。",
    "医療費精算手続き：従業員は退院日から7営業日以内に領収書を経理部に提出してください。"
]

pairs = [[query, p] for p in passages]
scores = reranker.compute_score(pairs)

results = sorted(zip(passages, scores), key=lambda x: x[1], reverse=True)

print("=== リランキング結果 ===")
for doc, score in results:
    print(f"Score: {score:.4f} | 内容: {doc[:80]}...")

実際の実行結果では、キーワードの一致数が多かった第1段落（スコア -2.14）に対し、試用期間社員の正確な規定が書かれた第2段落が圧倒的に高いスコア（約5.82）を獲得していることが確認できます。

3. RetrieverクラスへのRerankerの統合

以下は、RAGバックエンド層に組み込むボイラープレートコードです：

from typing import List, Dict, Any
from FlagEmbedding import FlagReranker

class TwoStageRetriever:
    def __init__(self, vector_store: Any, model_name: str = 'BAAI/bge-reranker-v2-m3'):
        self.vector_store = vector_store
        self.reranker = FlagReranker(model_name, use_fp16=True)

    def retrieve_and_rerank(
        self, 
        query: str, 
        initial_top_k: int = 30, 
        final_top_k: int = 4
    ) -> List[Dict[str, Any]]:
        # ステップ1: ベクトルDBから候補を高速検索
        initial_docs = self.vector_store.similarity_search(query, k=initial_top_k)
        if not initial_docs:
            return []

        # ステップ2: クエリとドキュメントのペアを作成
        doc_texts = [doc.page_content for doc in initial_docs]
        pairs = [[query, text] for text in doc_texts]
        
        # ステップ3: 相互関連度スコアを計算
        scores = self.reranker.compute_score(pairs)
        
        # ステップ4: スコアを付与してTop-Kをソート・取得
        scored_docs = []
        for doc, score in zip(initial_docs, scores):
            doc.metadata["rerank_score"] = float(score)
            scored_docs.append(doc)
            
        scored_docs.sort(key=lambda x: x.metadata["rerank_score"], reverse=True)
        return scored_docs[:final_top_k]

4. 本番環境への導入ノウハウ

  • 最適なモデルサイズの選定:
    • bge-reranker-base（パラメータ数約280M、VRAM使用量約600MB）: 25ms未満の極めて低いレイテンシが求められる場合や、CPU環境で動作させる場合に最適です。
    • bge-reranker-v2-m3（パラメータ数560M、FP16でVRAM約1.2GB）: 多言語対応や複雑な文脈理解において最も優れた選択肢となります。
  • レイテンシの測定: 一般的なT4 GPU環境において、bge-reranker-v2-m3（FP16）を用いて30チャンクをリランクする処理時間はわずか45〜60ms程度です。Hit Rate@3が約68%から91%超へと大幅に向上することを考慮すれば、十分に許容できるトレードオフと言えます。
  • 適切なバッチサイズの閾値: initial_top_kは25〜40の範囲に設定するのが理想的です。80チャンク以上に増やすと、精度の目立った改善が見られないままレイテンシのみが線形に増加してしまいます。
Share:
Tags: