ColBERTのようなマルチベクトル埋め込みモデルは、検索の精度が高い一方で、文書ごとに保存されるベクトルの数が増えるため、インデックスのコストが急速に上昇します。トークンプーリングと呼ばれるグループ化技術はこのサイズを削減し、適切な学習を行うことで、同じストレージ予算で結果の品質をより良く維持できるとされています。pool_factor=4の場合、通常学習されたモデルはNanoBEIRにおいて非圧縮時のnDCG@10の94.6%を維持するのに対し、プーリングを考慮して学習されたモデルでは97.8%を維持します。
なぜマルチベクトルモデルは多くのベクトルを保存するのか
従来の埋め込みモデルは、文書全体を単一のベクトルに要約します。一方、マルチベクトルモデルは、トークン（テキストの単位、多くの場合単語またはその一部）ごとに、一連の文脈的ベクトルを保持します。これにより、クエリの異なる部分が文書の異なる部分にマッチすることが可能になります。例として、長いバッテリー寿命のノートパソコンに関するクエリが、文書のある部分ではデバイスについての記述に、別の部分ではバッテリー寿命についての記述にマッチする可能性が挙げられています。
この仕組みにはコストが伴います。提示された計算によると、残差座標ごとに4ビットの量子化を用いた場合、128次元の200個のベクトルは文書あたり12.8KB、つまり100万文書で12.8GBを占めます。各文書を50ベクトルに削減すると、次元と精度を同じに保った場合、この総量は3.2GBまで減少します。
この手法は、一つの文書内の多くのベクトルが互いに近い値を持つ可能性があるという観察に基づいています。これは、複数のトークンがしばしば類似した概念を表すためです。Sentence TransformersライブラリのHierarchicalTokenPoolingは、ウォード法クラスタリングによってこれらのベクトルをグループ化します。この方法は、各ステップで内部分散の増加が最も少ない2つのグループを統合します。各グループはその後、その平均値で置き換えられ、正規化されます。pool_factor=4の場合、文書は初期ベクトルの約4分の1を保持するため、より長い文書は比例してより多くのベクトルを保持します。
圧縮を考慮したモデルの学習
圧縮が強力であればあるほど、有用な情報を失いやすくなります。プーリングを考慮した学習は、ベクトルが統合される際に重要な情報を保持することを特に目的としています。各学習バッチは2つの損失を組み合わせます：元の文書埋め込みを用いた検索と、ベクトル数を約4分の1に圧縮した後の検索です。これら2つの損失は平均化され、両方の目的が同じエンコーダーを更新します。
クラスタリングは離散的な決定を行うため、学習ではストレートスルー推定器が使用されるとされています。これは、クラスタへの割り当てを区別することなく、グループ化されたベクトルの勾配を構成するトークンに伝播させます。MultiVectorTokenPoolingLossクラスは、既存の検索損失（ここではCachedMultiVectorMultipleNegativesRankingLoss）をカプセル化し、各バリアントに0.5の重みを付けます。ベースモデルはanswerdotai/ModernBERT-baseで、MS MARCOの12万トリプレットを用いて1エポック学習されています。
NanoBEIRでの結果
pool_factor=4において、プーリングを考慮して学習されたモデルのnDCG@10は49.03から51.23に向上し、非圧縮検索と比較して文書あたりのベクトル数を74.8%削減しました。非圧縮モデルも51.84から52.39へと改善しています。文書あたり平均29.58ベクトルの場合、適応型プーリングは、通常学習されたモデルの93.3%に対し、プーリング考慮学習モデルの非圧縮スコアの97.3%を維持しました。
インデックスのサイズを削減するために、まず埋め込みの量子化、次にトークンプーリング（必要に応じてプーリングを考慮したファインチューニングを伴う）、そして最後に各ベクトルの次元を削減するMatryoshka埋め込みが提案されています。これらのアプローチは組み合わせることが可能です。本記事が基づく研究は、Stefan Josefによる研究論文Learn to Poolで詳述されています。
