Posted inAI

妥協しないRAGの構築：BGEとCohereによるリランキングの実践ガイド

AI tutorial - IT technology blog
AI tutorial - IT technology blog

ベクトル検索における精度の課題

RAGシステムを構築した際、AIがドキュメントの全く関係ない部分を自信満々に引用してしまった経験はありませんか？これはよくある悩みです。OpenAIのtext-embedding-3-smallのような高品質なモデルを使用しても、ベクトルデータベースは精度に苦労することがあります。上位10件のドキュメントを返したとしても, 正解が7番目に埋もれていたり、数学的な「類似性」が意味的な文脈と一致しなかったために完全に見落とされたりすることがあります。

標準的なベクトル検索は、Bi-Encoderに依存しています。これらのモデルは、ドキュメントとクエリを独立してベクトルに変換します。この手法は非常に高速で、10ミリ秒未満で数百万のドキュメントを検索できますが、ニュアンスに欠けるという欠点があります。本番環境において、ベクトル類似性のみに頼ることは、ノイズや無関係なコンテキストによって引き起こされる「ハルシネーション（幻覚）」の要因となります。

2段階検索戦略

速度を犠牲にすることなく精度を向上させるために、2段階のパイプラインを使用します。これは、「広く網を張る検索」の後に「詳細な精査」を行うようなイメージです：

  1. 検索 (ステージ1): ベクトルデータベースを使用して、上位50〜100件の候補を素早く抽出します。
  2. リランキング (ステージ2): Cross-Encoder（Reranker）を使用して、それらの候補を再スコアリングします。Rerankerはクエリとドキュメントを同時に分析し、単純な距離計算では見逃してしまう深い関連性を理解します。

Rerankerはベクトル検索よりも計算負荷が高くなります。しかし、上位50件などの少数のドキュメントに対してのみ実行するため、全体のレイテンシは許容範囲内（通常、リクエスト全体に100〜200ミリ秒追加される程度）に収まります。

Bi-Encoder vs. Cross-Encoder：何が違うのか？

Bi-Encoderは、クエリとドキュメントを「夜の海ですれ違う2隻の船」のように別々に扱います。類似性は高次元空間における2点間のコサイン距離に過ぎません。モデルは実際にクエリとドキュメントを「一緒に」読み取っているわけではないのです。

BGE-RerankerのようなCross-Encoderは、クエリとドキュメントを1つのペアとして入力します。Transformerのフルアテンション（Attention）メカニズムを使用して、クエリの各単語をドキュメントの各単語と照らし合わせて評価します。これにより、0から1の間の正確な関連性スコアが生成され、LLMに渡す最終的なコンテキストを選択する際に非常に信頼性の高い指標となります。

オプション1：BGE-Rerankerをローカルで実行する

BGE (Beijing Academy of Artificial Intelligence) Rerankerは、オープンソース界で非常に強力なモデルです。独自のGPUリソースがあり、トークンごとのAPIコストを避けたい場合に最適な選択肢です。特にv2-m3バージョンは、速度と性能のバランスが優れているため人気があります。

1. セットアップ

pip install flash-attn
pip install -U FlagEmbedding

2. 実装

以下のスニペットは、データベースから取得したドキュメントを再ソートする方法を示しています。モデルの効率的な推論には、約2〜4GBのVRAMが必要です。

from FlagEmbedding import FlagReranker

# リランカーの初期化
reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=True) 

query = "RAGシステムを本番環境向けに最適化するにはどうすればよいですか？"

# これらの候補は通常、ベクトル検索（PineconeやMilvusなど）から取得されます
retrieved_documents = [
    "RAGを最適化するには、埋め込みモデルの微調整を検討すべきです。",
    "サンフランシスコの天気は、通常, 朝は霧がかかっています。",
    "リランキングは、RAGパイプラインにおけるコンテキストの関連性を向上させるための重要な手法です。",
    "Pythonは、データサイエンスやAI開発で人気のある言語です。"
]

# ペアをスコアリング
pairs = [[query, doc] for doc in retrieved_documents]
scores = reranker.compute_score(pairs)

# スコアでソート（関連性の高い順）
reranked_results = sorted(zip(retrieved_documents, scores), key=lambda x: x[1], reverse=True)

for doc, score in reranked_results:
    print(f"スコア: {score:.4f} | ドキュメント: {doc}")

オプション2：Cohere Rerank APIを使用する

インフラの管理を避けたい場合は、Cohereが提供するマネージドなRerank APIが便利です。スケーラビリティが高くメンテナンスが不要なため、本番環境のRAGにおけるゴールドスタンダードと広く見なされています。LlamaIndexなどのフレームワークともシームレスに統合できます。

1. クライアントのインストール

pip install cohere

2. 実装

Cohereのrerank-english-v3.0モデルを使用すると、1回のAPIコールで大量のドキュメントを処理できます。

import cohere

co = cohere.Client("YOUR_COHERE_API_KEY")

query = "Cross-Encoderを使用する利点は何ですか？"
docs = [
    "Cross-Encoderは、クエリとドキュメントのペアに対してフルアテンションを実行するため、より正確です。",
    "Bi-Encoderは高速ですが、複雑な意味的一致に関しては精度が低くなります。",
    "ディープラーニングモデルは、多大な計算リソースを必要とします。",
    "運動の利点には、心血管の健康改善が含まれます。"
]

response = co.rerank(
    model="rerank-english-v3.0",
    query=query,
    documents=docs,
    top_n=3
)

for result in response.results:
    print(f"順位: {result.index} | スコア: {result.relevance_score:.4f}")
    print(f"内容: {docs[result.index]}\n")

戦略的な実装方法

RAGアーキテクチャを設計する際、私は「100から5」戦略を推奨しています。ベクトル検索で100件のドキュメントを取得し、それらをリランキングして上位5件を絞り込み、その5件のみをLLM（GPT-4oなど）に送信します。このアプローチは、LLMが長いドキュメントリストの中間に隠れたコンテキストを無視してしまう「Lost in the Middle（中だるみ）」問題を解決します。

ツールの選び方

  • BGE-Reranker: 高いプライバシー要件がある場合や、使用していないGPUリソースがある場合に適しています。ライセンス料は「無料」ですが、ホスティングコストがかかります。
  • Cohere Rerank API: 市場投入までのスピードを重視する場合に適しています。トークン化やモデルのアップデートが自動で行われるため、VRAMの管理ではなくアプリケーションロジックに集中できます。

最後に

リランキングステップを追加することは、趣味のRAGプロジェクトを本番環境で通用するツールへと変える最も効果的な方法です。BGEやCohereのようなモデルを使用することで、LLMに最適なデータを提供できるようになります。RAGの精度は、保有するデータの量ではなく、モデルに提示するために選択するデータの「質」によって決まるのです。

Share:
Tags: