AI・機械学習
ハニー、埋め込みを縮小しました:マトリョーシュカ対PCA
Honey, I shrunk the embeddings: Matryoshka vs. PCA (dylancastillo.co)
要約
LLMで独自のドキュメントを扱う際、埋め込み(embeddings)の効率的な保存と検索が課題となっています。この課題に対し、AI研究者は次元数を減らしつつ検索精度を維持するMatryoshka Representation Learning (MRL) を開発しました。本記事では、このMRLと、より古くシンプルな手法である主成分分析(PCA)を比較実験し、どちらが埋め込み次元削減において優れているかを検証しています。
全文翻訳
人々がLLMを独自のドキュメントと共に使い始めると、新しい問題が現れました。その情報をすべて効率的に保存し、検索するにはどうすればよいでしょうか? ベクトルデータベースがすぐに標準的な解決策となりました。しかし、ベクトルは数千の次元を持つことがあり、数百万ものベクトルを保存すると、検索が遅く高価になる可能性があります。AIラボは、Matryoshka Representation Learning (MRL) と呼ばれる技術でこれに対応しました。これは、検索精度をあまり犠牲にすることなく、より少ない埋め込み次元を使用できるものです。これは、ベクトルデータベースの請求額が小さくなり、クエリが速くなることを意味します。ハッピーエンドです。ほとんど。私は通常、水曜日にベクトルデータベースの請求額を心配して過ごすことはありません。しかし、ダグ・ターンブル氏がベクトルの次元を削減するために主成分分析(PCA)を使用することについての記事は、私に好奇心を抱かせました。この古く、よりシンプルな技術はMRLと比較してどうでしょうか?それを知るために、私は8つの標準的な検索品質データセットで2つの方法を比較しました。この記事では、実験の過程と私が発見したことを共有します。すべてのコードとデータはGitHubで入手可能です。MRLとPCAとは何ですか?どちらの方法も、フルベクトルとほぼ同様に機能する小さなベクトルを生成します。しかし、それらは異なる方法でそれを実現します。MRLはトレーニング中に機能します。あなたは、最初の64次元、最初の128次元など、複数のプレフィックス長で同時に適用される損失でモデルをトレーニングします。これにより、入れ子になったマトリョーシュカ人形のように、最も重要な情報をベクトルの先頭に詰め込むように教えられます。推論時には、最初のd次元を保持し、結果のベクトルを再正規化するだけです。多くの最新の埋め込みモデルはこの方法でトレーニングされていますが、古いモデルはそうではないため、MRLベースの切り捨ては多くの人気のある埋め込みモデルでは利用できません。PCAはトレーニング後に機能するため、どのモデルでも使用できます。あなたは埋め込みのサンプルを取り、それらが最も変動する方向を見つけ、それらの上位d個を射影行列として保持します1。PCAはより小さなベクトルを提供しますが、追加の運用上の複雑さも伴います。PCA変換を保存およびバージョン管理し、インデックスへの追加とクエリの際に同じバージョンを一貫して適用する必要があります。実験をどのように実行したか実験の背後にある考えは単純でした。各方法で埋め込みを縮小し、ベンチマークを実行し、各サイズで検索品質がどの程度低下するかを確認します。私はOpenRouterを通じてすべての埋め込みを生成し、8つのBEIRデータセットで検索を評価しました。SciFact、NFCorpus、ArguAna、FiQA、SciDocs、Quora、TREC-COVID、およびWebis-Touché 2020です。各データセットについて、埋め込みを512、256、128、64、および32次元に削減しました。注データセットの詳細データセットタスクコーパスクエリSciFact科学的主張検証5.2K 300 NFCorpus医療検索3.6K 323 ArguAna反論検索8.7K 1,406 FiQA-2018金融質問応答57K 648 SciDocs引用推薦25K 1,000 Quora重複質問検索523K 10,000 TREC-COVID生物医学検索(COVID-19)171K 50 Touché 2020ウェブドキュメントからの議論検索382K 49 答えたかった3つの質問がありました。Q1:次元をカットするにつれて、どちらの方法がより多くの検索品質を維持するか?これに答えるために、私は2つのMRLトレーニング済みモデルを使用しました。OpenAIのtext-embedding-3-small(1,536次元)と、オープンウェイトMTEB BEIRリーダーボードのトップに位置するAlibabaのqwen3-embedding-8b(4,096次元)です。各モデルの埋め込みを2つの方法で削減しました。切り捨てでは、最初のd次元を保持し、再正規化しました。これがMRLで次元を削減する方法です。PCAでは、同じデータセットのフル次元ドキュメント埋め込みで射影を適合させ、上位d個のコンポーネントを保持しました。それ以降、すべてのドキュメントとクエリの埋め込みは、検索前にその射影行列で乗算され、再正規化されます。次に、8つのデータセットでベンチマークを実行し、埋め込みが小さくなるにつれてどちらの方法がより多くの検索品質を維持するかを確認するために、両方のモデルで2つの方法を比較しました。Q2:MRLトレーニングだけなのか?PCAがうまくいく場合、それはMRLトレーニングがすでに埋め込み空間を便利な方法で整理しているからかもしれません。それをテストするために、MRLを使用してトレーニングされていない古い1,536次元モデルであるtext-embedding-ada-002をコントロールとして追加しました。両方のモデルにPCAを適用しました。PCAがtext-embedding-3-smallとtext-embedding-ada-002で同様の検索品質のシェアを維持する場合、そのパフォーマンスはMRLトレーニングに依存しないことを示唆しています。Q3:適合データは重要か?PCAは何かの上に適合させる必要があるため、2つの実践的な質問が生じます。適合データはどのくらい必要か、そして検索するコーパスから来る必要があるか?最初の質問については、FiQAの57Kドキュメントのランダムサンプル(1,000、5,000、20,000、およびフルコーパス)でPCAを適合させ、適合サンプルのサイズが検索品質をどの程度変更するかを確認しました。これは、データの一部でPCAを適合させ、インデックスが増加しても更新しないシナリオをカバーします。2番目の質問については、通常のPCAとドメイン外PCAを比較しました。100,000のMS MARCOパスで一度適合され、変更なしで他のすべてのデータセットで使用される射影です。これは最も極端なバージョンです。あなたは実際にインデックスを作成するものとはかなり異なるデータでPCAを適合させます。しかし、まず、健全性チェック実験を実行する前に、評価パイプラインが妥当な結果を生成していることを確認したかったのです。そのため、最初にモデルの公式MTEBスコアを再現しようとしました。以下の表は、フル次元でのNDCG@10を比較しています。公式MTEBスコア対私の評価パイプラインによって生成されたスコアです。データセットada-002 3-small qwen3-8b MTEB mine MTEB mine MTEB mine SciFact 0.7275 0.7277 0.7337 0.7296 0.7846 0.7863 NFCorpus 0.3697 0.3705 0.3833 0.3847 0.4145 0.4150 ArguAna 0.5744 0.5757 0.5549 0.5573 0.7685 0.7689 FiQA 0.4441 0.4440 0.4491 0.4484 0.6457 0.6492 SciDocs 0.1836 0.1837 0.2080 0.2077 0.3274 0.3268 Quora 0.8760 0.8759 0.8883 0.8880 0.8890 0.8901 TREC-COVID 0.6847 0.6884 0.7790 0.7775 0.9499 0.9492 Touché 2020 0.2161 0.2143 0.2428 0.2433 0.3593 0.3596 公式数値はMTEB結果リポジトリ(text-embedding-3-small、text-embedding-ada-002、qwen3-embedding-8b)から来ています。私の結果は公式数値に十分に近いため、パイプラインが期待どおりに機能していると確信しています。Q1:どちらの方法がより多くの検索品質を維持するか?PCAは、両方のMRLトレーニング済みモデルで、ほぼすべての次元でMRL切り捨てに匹敵するか、それを上回りました。以下の各セルは、削減によってどのくらいの検索品質が維持されるかを示しています。フル次元(text-embedding-3-smallの場合は1,536、qwen3-embedding-8bの場合は4,096)でのNDCG@10を、8つのデータセットで平均したものです。次元 3-small (MRL) qwen3-8b (MRL) 切り捨て PCA 切り捨て PCA 512 98% 97% 99% 98% 256 94% 95% 96% 96% 128 86% 90% 91% 91% 64 71% 82% 83% 84% 32 46% 65% 68% 71% 3-smallqwen3-8b 512次元では、MRL切り捨てがわずかに優位です。しかし、ベクトルが小さくなるにつれて、PCAは追いつき、MRLを上回るか同等になります。text-embedding-3-smallでは、PCAは256次元未満で明らかに優れています(32次元で65%対46%を維持)。qwen3-embedding-8bでは、レースははるかに接戦です(32次元で71%対68%)。興味深い発見の1つは、qwen3-embedding-8bのMRLトレーニングが十分に強力であるため、4,096次元から切り取られた32次元ベクトルの品質が、1,536次元から切り取られた3-smallよりも多く維持されることです。データセットごとに見ると、いくつかの違いがあります。3-smallqwen3-8b 最小次元では、PCAはtext-embedding-3-smallの8つのデータセットのうち7つで、qwen3-embedding-8bの8つのうち6つで勝利します。次元全体で見ると、PCAは両方のモデルでSciFact、FiQA、およびNFCorpusでより良いパフォーマンスを発揮します。ArguAnaとSciDocsもPCAをtext-embedding-3-smallで支持していますが、qwen3-embedding-8bでは結果はそれほど決定的ではありません。切り捨てはQuoraでは両方のモデルでより良いパフォーマンスを発揮します。Touchéはqwen3-embedding-8bでは切り捨てを支持しますが、text-embedding-3-smallではほぼ同等であり、TREC-COVIDの結果は次元に依存します。Q2:MRLトレーニングだけなのか?PCAはtext-embedding-3-smallとqwen3-embでしか良く見えないのではないかと心配していました。