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

埋め込みモデルのファインチューニングでより良い検索とRAGを実現

要約：ドメイン内のデータで埋め込みモデルをファインチューニングすると、ベクトル検索と検索補強生成（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）を合成データでファインチューニングし、それらを私たちのドメインインテリジェンスベンチマークスイート（DIBS）（FinanceBench、ManufactQA、Databricks DocsQA）からの3つのデータセットで評価しました。その後、それらをOpenAIのtext-embedding-3-largeと比較しました。

主要なポイント

• ファインチューニングにより、データセット全体での取得精度が向上しました、これはベースラインモデルを大幅に上回ります。
• ファインチューニングされた埋め込みは、再ランキングと同等またはそれ以上のパフォーマンスを示しました、これはそれが強力なスタンドアロンの解決策であることを示しています。
• より良い検索がFinanceBenchでのより良いRAGパフォーマンスにつながったことが、エンドツーエンドの利点を示しています。

検索パフォーマンス

3つのデータセットを比較した結果、エンベディングのファインチューニングがこれらのデータセットの2つで精度を向上させることがわかりました。図1は、FinanceBenchとManufactQAについて、ファインチューニングされた埋め込みがそのベースバージョンを上回り、時にはOpenAIのAPIモデル（ライトグレー）をも上回ることを示しています。しかし、Databricks DocsQAについては、OpenAIのテキスト埋め込み-3-ラージの精度がすべてのファインチューニングモデルを上回っています。これは、モデルが公開されているDatabricksのドキュメンテーションで訓練されているためかもしれません。これは、ファインチューニングが効果的である一方で、訓練データセットと評価タスクに強く依存していることを示しています。

ファインチューニング対再ランキング

その後、上記の結果をAPIベースの再ランキングを使用してvoyageai/rerank-1と比較しました（図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（434Mパラメータ、1.75 GB）に基づく人気のある埋め込みモデルです。このモデルはサイズが適度で、オープンライセンスがあるため、私たちはこのモデルで実験を行うことにしました。この埋め込みモデルは現在、Databricks Foundation Model APIでもサポートされています。
• e5-mistral-7b-instructは、強力なLLM（この場合はMistral-7b-instruct-v0.1）の上に構築された新しいクラスの埋め込みモデルに属しています。e5-mistral-7b-instructは、MTEBなどの標準的な埋め込みベンチマークでは優れており、より長く、より微妙なプロンプトを処理することができますが、gte-large-en-v1.5よりも大きく（なぜなら7億のパラメータを持っているからです）、少し遅く、提供するのに少し高価です。

次に、OpenAIのtext-embedding-3-largeと比較しました。

評価データセット

我々は、すべてのモデルを我々のドメインインテリジェンスベンチマークスイート(DIBS)からの次のデータセットで評価しました：FinanceBench、ManufactQA、およびDatabricks DocsQA。

私たちはrecall@10を主要な検索指標として報告します。これは、正しいドキュメントが上位10件の検索結果に含まれているかどうかを測定します。

埋め込みモデルの品質のゴールデンスタンダードは、BEIRのような検索タスクや他の多くの非検索タスクを組み込んだMTEBベンチマークです。gte-large-en-v1.5やe5-mistral-7b-instructのようなモデルはMTEBで良好な結果を出しますが、私たちの内部エンタープライズタスクでどのようにパフォーマンスを発揮するかを見てみました。

学習用データ

上記のベンチマークごとに合成データで個別のモデルを訓練しました：

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

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

私たちは以下の項目でスイープを行いました：

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

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

Databricksでベクトル検索とRAGを改善する方法

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

より良い検索のために：

• より良い埋め込みモデルを使用する：多くのユーザーは、古い埋め込みを知らず知らずのうちに使用しています。単純にパフォーマンスの高いモデルを交換するだけで、即座に利益を得ることができます。トップモデルについては、MTEBリーダーボードをご覧ください。
• ハイブリッド検索を試す：密な埋め込みとキーワードベースの検索を組み合わせて精度を向上させます。Databricks Vector Searchは、ワンクリックソリューションでこれを簡単にします。
• 再ランカーを使用する：再ランカーは、関連性に基づいて結果を並べ替えることで結果を精緻化することができます。Databricksはこれを組み込み機能として提供しています（現在はプライベートプレビュー中）。アカウントエグゼクティブに連絡して試してみてください。

より良いRAGのために：

• プロンプトを最適化する：LLMのプロンプトを少し調整するだけで、レスポンスが大幅に改善します。DSPyはこのプロセスを自動化するのに役立ちます（DatabricksのDSPyを使用してgenAIアプリを構築するを参照）。
• あなたのLLMをアップグレードする：検索が強力であるが回答が弱い場合は、より良い生成モデルを使用することを検討してください。
• LLMのファインチューニング：あなたのドメインがユニークで、十分なデータがある場合、Llama 3のようなモデルをファインチューニングすることで、RAGの品質をさらに向上させることができます。詳細はMosaic AIモデルトレーニング：特殊なタスクと知識のためのDatabricksでのLLMのファインチューニングをご覧ください。

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システムは、リトリーバーと生成モデルの両方から構成されています。リトリーバーは特定のクエリに関連する一連のドキュメントを選択し、それらを生成モデルにフィードします。私たちは最良のファインチューニングされたgte-large-en-v1.5モデルを選択し、それらをシンプルなRAGシステムの最初の検索ステージに使用しました（LLMsのLong Context RAG PerformanceおよびOpenAI o1とGoogle GeminiのLong Context RAG Capabilitiesで説明されている一般的なアプローチに従って）。特に、私たちは最大長512トークンの各k=10ドキュメントを取得し、GPT4oを生成的なLLMとして使用しました。最終的な精度は、LLM-as-a-judge（GPT4o）を使用して評価されました。

FinanceBenchでは、Figure 3は微調整された埋め込みモデルを使用すると、下流のRAGの精度が向上することを示しています。また、それはtext-embedding-3-largeと競争力があります。これは予想されることで、gteの微調整により、ベースラインのgteに対するRecall@10が大幅に向上しました（図1）。この例は、特定のドメインとデータセットでの埋め込みモデルの微調整の有効性を強調しています。

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

このブログ投稿に対するフィードバックをくれたQuinn LengとMatei Zahariaに感謝します。