埋め込みモデルのファインチューニングによる検索とRAGの改善

RetrievalとRAGのための埋め込みモデルのファインチューニング

要約: ドメイン固有のデータで埋め込みモデルをファインチューニングすると、ベクトル検索とRetrieval-Augmented Generation (RAG) の精度が大幅に向上します。Databricksを使用すると、埋め込みモデルのファインチューニング、デプロイ、評価が簡単に行え、手動ラベリングなしで合成データを利用して検索を最適化できます。

重要性: ベクトル検索またはRAGシステムが最適な結果を返さない場合、埋め込みモデルのファインチューニングはパフォーマンスを向上させるためのシンプルかつ強力な方法です。財務文書、ナレッジベース、社内コードドキュメントなど、どのようなデータでも、ファインチューニングによってより関連性の高い検索結果と、より優れた下流のLLM応答が得られます。

調査結果: 3つのエンタープライズデータセットで2つの埋め込みモデルをファインチューニングしてテストしたところ、検索メトリクス (Recall@10) と下流RAGパフォーマンスに大幅な改善が見られました。これは、ファインチューニングが、既存のデータのみを利用して、手動ラベリングを必要とせずに精度の向上に大きく貢献できることを意味します。

埋め込みのファインチューニングを試してみませんか？開始に役立つリファレンスソリューションを提供しています。Databricksは、ベクトル検索、RAG、リランキング、埋め込みのファインチューニングを容易にします。詳細については、Databricksアカウントエグゼクティブまたはソリューションアーキテクトにお問い合わせください。

なぜ埋め込みをファインチューニングするのか？

埋め込みモデルは、最新のベクトル検索およびRAGシステムを支えています。埋め込みモデルはテキストをベクトルに変換し、キーワードだけでなく意味に基づいて関連コンテンツを見つけられるようにします。しかし、市販のモデルは特定のドメインに最適化されていない場合があるため、ファインチューニングが役立ちます。

ドメイン固有のデータで埋め込みモデルをファインチューニングすると、いくつかの点で役立ちます。

• 検索精度の向上: カスタム埋め込みにより、データに合わせた検索結果が改善されます。
• RAGパフォーマンスの強化: より良い検索結果は、ハルシネーションを減らし、より根拠のある生成AI応答を可能にします。
• コストとレイテンシの改善: 小規模なファインチューニング済みモデルが、大規模で高価な代替モデルよりも優れたパフォーマンスを発揮する場合があります。

このブログ記事では、埋め込みモデルのファインチューニングが、タスク固有のエンタープライズユースケースにおける検索およびRAGパフォーマンスを向上させる効果的な方法であることを示します。

結果: ファインチューニングは効果的

2つの埋め込みモデル (gte-large-en-v1.5 および e5-mistral-7b-instruct) を合成データでファインチューニングし、当社のDomain Intelligence Benchmark Suite (DIBS) (FinanceBench、ManufactQA、Databricks DocsQA) の3つのデータセットで評価しました。その後、OpenAIのtext-embedding-3-largeと比較しました。

主なポイント:

• ファインチューニングにより、データセット全体で検索精度が向上し、多くの場合、ベースラインモデルを大幅に上回りました。
• ファインチューニング済み埋め込みは、多くの場合、リランキングと同等またはそれ以上のパフォーマンスを示し、スタンドアロンソリューションとしても強力であることを示しました。
• 検索精度の向上は、FinanceBenchでのRAGパフォーマンスの向上につながり、エンドツーエンドのメリットを示しました。

検索パフォーマンス

3つのデータセットを比較した結果、埋め込みのファインチューニングが2つのデータセットで精度を向上させることがわかりました。図1に示すように、FinanceBenchとManufactQAでは、ファインチューニング済み埋め込みはベースバージョンよりも優れたパフォーマンスを発揮し、場合によってはOpenAIのAPIモデル (薄いグレー) を上回りました。しかし、Databricks DocsQAでは、OpenAI text-embedding-3-largeの精度がすべてのファインチューニング済みモデルを上回りました。これは、モデルが公開されているDatabricksのドキュメントでトレーニングされている可能性があるためと考えられます。これは、ファインチューニングが効果的である一方で、トレーニングデータセットと評価タスクに強く依存することを示しています。

ファインチューニング vs. リランキング

voyageai/rerank-1を使用したAPIベースのリランキングの結果と比較しました (図2)。リランカーは通常、埋め込みモデルによって取得された上位k件の結果を取得し、これらの結果を検索クエリとの関連性に基づいて並べ替え、並べ替えられた上位k件 (この場合はk=30、次にk=10) を返します。これは、リランカーが通常、埋め込みモデルよりも大きく強力なモデルであり、クエリとドキュメント間の相互作用をより表現力豊かにモデル化するため機能します。

調査結果は以下の通りです。

• gte-large-en-v1.5のファインチューニングは、FinanceBenchとManufactQAでリランキングを上回りました。
• OpenAIのtext-embedding-3-largeはリランキングの恩恵を受けましたが、一部のデータセットでは改善はわずかでした。
• Databricks DocsQAでは、リランキングの影響は小さかったですが、ファインチューニングでも改善が見られ、これらの手法のデータセット依存性を示しています。

リランカーは通常、埋め込みモデルと比較して、クエリごとの推論レイテンシとコストが追加されます。ただし、既存のベクトルデータベースで使用でき、場合によっては新しい埋め込みモデルでデータを再埋め込みするよりもコスト効率が高いことがあります。リランカーを使用するかどうかの選択は、ドメインとレイテンシ/コスト要件によって異なります。

ファインチューニングはRAGパフォーマンスを向上させる

FinanceBenchでは、検索精度の向上は、GPT-4oと組み合わせた場合のRAG精度の向上に直接つながりました (付録参照)。しかし、Databricks DocsQAのように検索がすでに強力なドメインでは、ファインチューニングによる大きな改善は見られませんでした。これは、ファインチューニングが検索が明確なボトルネックである場合に最も効果的であることを示しています。

埋め込みモデルのファインチューニングと評価方法

ここでは、合成データ生成、ファインチューニング、評価に関する技術的な詳細をいくつか紹介します。

埋め込みモデル

2つのオープンソース埋め込みモデルをファインチューニングしました。

• gte-large-en-v1.5 は、BERT Large (4億3400万パラメータ、1.75 GB) に基づく人気の埋め込みモデルです。 modestなサイズとオープンライセンスのため、このモデルで実験を行うことにしました。この埋め込みモデルは、Databricks Foundation Model APIでも現在サポートされています。
• e5-mistral-7b-instruct は、強力なLLM (この場合はMistral-7b-instruct-v0.1) を基盤とする新しいクラスの埋め込みモデルに属します。e5-mistral-7b-instructは、MTEBなどの標準的な埋め込みベンチマークで優れており、より長くニュアンスのあるプロンプトを処理できますが、gte-large-en-v1.5よりもはるかに大きく (70億パラメータを持つため)、サービングにはわずかに遅く、コストも高くなります。

その後、OpenAIのtext-embedding-3-largeと比較しました。

評価データセット

すべてのモデルを、Domain Intelligence Benchmark Suite (DIBS) の以下のデータセットで評価しました: FinanceBench、ManufactQA、Databricks DocsQA。

主な検索メトリックとしてrecall@10を報告します。これは、正しい文書が検索された上位10件に含まれているかどうかを測定します。

埋め込みモデルの品質のゴールドスタンダードは、MTEBベンチマークです。これには、BEIRなどの検索タスクや、その他の多くの非検索タスクが含まれています。gte-large-en-v1.5やe5-mistral-7b-instructのようなモデルはMTEBで良好な結果を示しますが、社内エンタープライズタスクでどのようにパフォーマンスを発揮するかを確認したいと思いました。

トレーニングデータ

上記の各ベンチマークに合わせて調整された合成データで、個別のモデルをトレーニングしました。

各ドメインのトレーニングセットを生成するために、既存の文書を使用し、Llama 3 405BなどのLLMを使用して、各文書の内容に基づいたサンプルクエリを生成しました。その後、LLM-as-a-judge (GPT4o) によって合成クエリの品質がフィルタリングされました。フィルタリングされたクエリとその関連文書は、ファインチューニングのための対照ペアとして使用されました。対照トレーニングではインバッチネガティブを使用しましたが、ハードネガティブを追加することでパフォーマンスをさらに向上させることができます（付録参照）。

ハイパーパラメータチューニング

以下についてスイープを実行しました。

• 学習率、バッチサイズ、ソフトマックス温度
• エポック数（1〜3エポックをテスト）
• クエリプロンプトのバリエーション（例：「Query:」対指示ベースのプロンプト）
• プーリング戦略（平均プーリング対ラストトークンプーリング）

すべてのファインチューニングは、オープンソースのmosaicml/composer、mosaicml/llm-foundry、およびmosaicml/streamingライブラリを使用して、Databricksプラットフォーム上で行われました。

Databricksでのベクトル検索とRAGの改善方法

ファインチューニングは、ベクトル検索とRAGのパフォーマンスを向上させるためのアプローチの1つにすぎません。以下に、その他のアプローチをいくつか挙げます。

検索精度の向上：

• より優れた埋め込みモデルを使用する：多くのユーザーは、気づかないうちに古い埋め込みを使用しています。パフォーマンスの高いモデルに単純に置き換えるだけで、すぐに改善が得られます。上位モデルについては、MTEBリーダーボードを確認してください。
• ハイブリッド検索を試す：密な埋め込みとキーワードベースの検索を組み合わせて、精度を向上させます。Databricks AI Searchは、ワンクリックソリューションでこれを容易にします。
• リランカーを使用する：リランカーは、関連性に基づいて結果を並べ替えることで、結果を絞り込むことができます。Databricksでは、これを組み込み機能として提供しています（現在プライベートプレビュー中）。試すには、アカウントエグゼクティブにお問い合わせください。

RAGの改善：

• プロンプトを最適化する：LLMプロンプティングのわずかな変更で、応答を劇的に改善できます。DSPyは、このプロセスを自動化するのに役立ちます（DatabricksでDSPyを使用してGenAIアプリを構築するを参照）。
• LLMをアップグレードする：検索は強力だが回答が弱い場合は、より優れた生成モデルの使用を検討してください。
• LLMをファインチューニングする：ドメインがユニークで十分なデータがある場合は、Llama 3のようなモデルをファインチューニングすることで、RAGの品質をさらに向上させることができます。詳細については、Databricksモデルトレーニング：特定のタスクと知識のために独自のLLMをDatabricksでファインチューニングするを参照してください。

Databricksでのファインチューニングを開始する

埋め込みのファインチューニングは、AIシステムでの検索とRAGを改善するための簡単な勝利となる可能性があります。Databricksでは、以下のことが可能です。

• スケーラブルなインフラストラクチャ上で埋め込みモデルをファインチューニングおよびサービングする。
• ベクトル検索、リランキング、RAGのための組み込みツールを使用する。
• ユースケースに最適なモデルを見つけるために、さまざまなモデルを迅速にテストする。

試してみませんか？ファインチューニングを容易にするためのリファレンスソリューションを構築しました。アクセスについては、Databricksのアカウントエグゼクティブまたはソリューションアーキテクトにお問い合わせください。

付録

表1：gte-large-en-v1.5、e5-mistral-7b-instruct、text-embedding-3-largeの比較。図1と同じデータ。

合成トレーニングデータの生成

すべてのデータセットで、トレーニングセットのクエリはテストセットのクエリと同じではありませんでした。ただし、Databricks DocsQA（FinanceBenchやManufactQAは除く）の場合、合成クエリを生成するために使用されたドキュメントは、評価セットで使用されたドキュメントと同じでした。私たちの研究の焦点は、特定のタスクとドメインでの検索精度を向上させること（ゼロショットの汎用的な埋め込みモデルとは対照的に）であるため、これは特定のプロダクションユースケースにおいて有効なアプローチと見なしています。FinanceBenchとManufactQAでは、合成データを生成するために使用されたドキュメントは、評価に使用されたコーパスと重複しませんでした。

対照学習のためのネガティブパッセージの選択には、さまざまな方法があります。ランダムに選択することも、事前に定義することもできます。最初のケースでは、ネガティブパッセージはトレーニングバッチ内から選択されます。これらはしばしば「インバッチネガティブ」または「ソフトネガティブ」と呼ばれます。2番目のケースでは、ユーザーは意味的に難しいテキスト例を事前に選択します。つまり、クエリに関連している可能性があるが、わずかに間違っているか無関係です。この2番目のケースは、時々「ハードネガティブ」と呼ばれます。この作業では、インバッチネガティブのみを使用しました。文献によると、ハードネガティブを使用すると、さらに良い結果が得られる可能性が高いとのことです。

ファインチューニングの詳細

すべてのファインチューニング実験で、最大シーケンス長は2048に設定されました。その後、すべてのチェックポイントを評価しました。すべてのベンチマーキングで、コーパスドキュメントは2048トークンに切り詰められ（チャンク化されず）、これは特定のデータセットにとって妥当な制約でした。クエリプロンプトとプーリング戦略全体をスキャンした後、各ベンチマークで最も強力なベースラインを選択しました。

RAGパフォーマンスの向上

RAGシステムは、リトリーバーと生成モデルの両方で構成されます。リトリーバーは、特定のクエリに関連するドキュメントのセットを選択し、それを生成モデルに渡します。私たちは、単純なRAGシステムの最初のリトリーバルステージ（Long Context RAG Performance of LLMsおよびThe Long Context RAG Capabilities of OpenAI o1 and Google Geminiで説明されている一般的なアプローチに従って）に使用するために、最高のファインチューニング済みgte-large-en-v1.5モデルを選択しました。具体的には、最大長512トークンのドキュメントをk=10個取得し、GPT4oを生成LLMとして使用しました。最終的な精度は、LLM-as-a-judge（GPT4o）を使用して評価されました。

FinanceBenchでは、図3に示すように、ファインチューニングされた埋め込みモデルを使用すると、下流のRAG精度が向上します。さらに、text-embedding-3-largeとも競争力があります。これは、gteのファインチューニングがベースラインgte（図1）と比較してRecall@10の大幅な改善につながったため、予想されることです。この例は、特定のドメインとデータセットにおける埋め込みモデルのファインチューニングの有効性を強調しています。

Databricks DocsQAデータセットでは、ベースラインgteを上回るファインチューニング済みgteモデルを使用しても、改善は見られませんでした。これは、図1および図2におけるベースラインモデルとファインチューニング済みモデルの差が小さいことを考えると、ある程度予想されることです。興味深いことに、text-embedding-3-largeは、gteモデルのいずれよりも（わずかに）高いRecall@10を持っていますが、それにより下流のRAG精度が高くなるわけではありません。図1に示すように、すべての埋め込みモデルはDatabricks DocsQAデータセットで比較的高いRecall@10を持っていました。これは、検索がRAGのボトルネックではない可能性が高く、このデータセットで埋め込みモデルをファインチューニングすることが必ずしも最も有益なアプローチではないことを示唆しています。

このブログ記事に関するフィードバックをいただいたQuinn Leng氏とMatei Zaharia氏に感謝いたします。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事