Muvera-Py: FDEによる高速・多ベクトル検索

検索に革命を:効率的なマルチベクトル検索のためのMuvera-Pyを導入

情報検索が絶えず進化する中で、ColBERTモデルを基盤とする現代の検索システムは、単一のドキュメントを表現するために何百ものベクトルを利用しています。このアプローチは検索精度を劇的に向上させる一方で、膨大な計算オーバーヘッドと検索時間の遅延という代償を伴うことがよくあります。この課題に取り組むためにGoogleが開発した画期的なアルゴリズムが、MUVERA (Multi-Vector Retrieval via Fixed Dimensional Encodings) です。

本日、MUVERAの新しいオープンソースPython実装であるMuvera-Pyをご紹介できることを大変嬉しく思います。Muvera-Pyは、アクセシビリティと、オリジナルの高度に最適化されたC++実装に対する完全な忠実性を重視して開発されました。これにより、大規模なベクトル検索に取り組む開発者や研究者にとって、状況を一変させる存在となるでしょう。

固定次元エンコーディング(FDE)とは?

FDEの核心は、各ドキュメントが1つではなく、何百ものベクトルで表現される場合に、数十億のドキュメントの中から効率的に検索を行うという根本的な問題を解決することにあります。従来の検索システムは、1つのドキュメントにつき1つのベクトルを使用するため高速ですが、マルチベクトルモデルのような微妙なニュアンスの精度に欠けることがよくあります。一方、マルチベクトル検索は精度が高いものの、非常に遅いことで知られています。

FDEは、これら複数のベクトルを単一の固定サイズベクトルに変換することで、洗練された解決策を提供します。特筆すべきは、その際に重要な類似性関係が驚くほど維持される点です。この魔法は、元のマルチベクトルセット間のChamfer類似性を、FDEベクトル間の単純な内積によって近似する能力にあります。

Muvera-Py: ギャップを埋める

Muvera-Pyは、この強力なアルゴリズムをPythonエコシステムで利用可能にします。Muvera-Pyのすべての関数とパラメータは、C++の対応物と綿密にマッピングされており、同一の動作と信頼性の高いパフォーマンスを保証します。これにより、精度や効率を損なうことなく、Googleの堅牢な研究成果をPythonプロジェクトで活用できます。

主要な機能と実装の詳細:

  • 完全な忠実性: オリジナルのC++実装の動作を再現し、本番環境で信頼できるツールとして機能します。
  • Pythonらしい設計: 設定にはPythonのdataclassesを、効率的な行列操作にはNumPyを使用し、クリーンで直感的なAPIを提供します。
  • 構成可能なエンコーディング: 柔軟なユースケースのために、異なるエンコーディングタイプ(例:クエリ用DEFAULT_SUM、ドキュメント用AVERAGE)とプロジェクションタイプ(例:DEFAULT_IDENTITYAMS_SKETCH)をサポートします。
  • 内部ヘルパー: Grayコード変換やランダム行列生成器(_simhash_matrix_from_seed_ams_projection_matrix_from_seed)などの重要な内部関数を、Pythonのベストプラクティスを用いて実装しています。
  • コアアルゴリズム: _generate_fde_internal()関数は、空間分割、ベクトル集約、最終的なプロジェクションのための洗練されたロジックをカプセル化しています。
  • パブリックAPI: generate_query_fde()generate_document_fde()などの分かりやすい関数を提供し、簡単に統合できます。

FDEの動作原理:ステップバイステップの概要

このアルゴリズムは、マルチベクトル表現を凝縮するために、一連のインテリジェントなステップで動作します:

  1. 空間分割: 各「繰り返し」(異なるランダムシードによる独立した実行)ごとに、SimHashを適用してベクトル空間を分割します。これは、ランダムガウス行列を使用してベクトルを特定のビンに投影します。
  2. ベクトル集約: 各パーティション内でベクトルが集約されます。クエリの場合、ベクトルは合計されます。ドキュメントの場合、平均化されます。
  3. 繰り返しと連結: ステップ1と2は、異なるランダムシードを用いて複数回繰り返されます。各繰り返しからの結果は連結され、最終的な単一のFDEベクトルが形成されます。

このプロセスにより、元の高次元空間で近いベクトルは、同じパーティションに属し、FDEベクトルの同じ部分に貢献する可能性が高くなります。この「局所性感度」こそが、FDEの単純な内積が複雑なChamfer類似性を近似することを可能にしているのです。

パフォーマンスの利点

FDEの最も魅力的な特徴の1つは、そのパフォーマンスプロファイルです。

  • 効率的な生成: FDEの生成時間は、ベクトルの数、次元、繰り返しの数、およびSimHash投影(O(n × d × r × k))に線形にスケールします。
  • 超高速検索: FDEが生成されると、検索時間は標準的な最大内積探索(MIPS)ライブラリを使用して実質的にO(1)となり、リアルタイム検索がスケーラブルになります。
  • メモリ効率: 何百ものベクトルを単一のベクトルに削減する能力により、インデックス作成に必要なメモリ要件が大幅に削減されます。

Muvera-Pyの始め方

Muvera-Pyをプロジェクトに統合するのは簡単です。基本的な例を次に示します。

import numpy as np
from fde_generator import FixedDimensionalEncodingConfig, generate_query_fde, generate_document_fde

# 1. 設定の作成
config = FixedDimensionalEncodingConfig(
    dimension=128, # 元のベクトル次元
    num_repetitions=10, # 独立した分割の数
    num_simhash_projections=6, # 2^6 = 64個のパーティションを作成
    seed=42
)

# 2. モックデータ(例:ColBERT形式の埋め込み)の準備
query_vectors = np.random.randn(32, 128).astype(np.float32) # 32個のベクトルを持つクエリ
doc_vectors = np.random.randn(80, 128).astype(np.float32)   # 80個のベクトルを持つドキュメント

# 3. FDEの生成
query_fde = generate_query_fde(query_vectors, config)
doc_fde = generate_document_fde(doc_vectors, config)

# 4. 類似性の計算(Chamfer類似性を近似)
similarity_score = np.dot(query_fde, doc_fde)
print(f"Similarity: {similarity_score}")

Muvera-Pyは、高度でスケーラブルな情報検索技術をより利用しやすくするための重要な一歩です。開発者が複雑なマルチベクトル埋め込みを効率的に処理できるようにすることで、さまざまな分野でより高速で正確な検索アプリケーションの道を開きます。

詳細、貢献、コードベースの探索については、Muvera-Py GitHubリポジトリをご覧ください。

この記事を共有