ColBERT型のモデルで自分の文書を検索しているなら、インデックスは望むより大きくなっているでしょう。こうしたマルチベクトルモデルは文書ごとに一つではなくトークンごとにベクトルを持つため、精度が高い反面、保存コストがかかります。Hugging Faceのブログの新しい記事は、Sentence Transformersで、ベクトルを圧縮した後も精度を保つようにモデルを訓練できるようになったことを示しています。圧縮で普通は失われる品質の大半を取り戻せるといいます。
保存の問題は具体的です。座標あたり4ビットの場合、128次元のベクトル200個で文書1件あたり12.8 KB、100万件では12.8 GBになります。各文書を50ベクトルにまとめれば3.2 GBに減ります。ライブラリのHierarchicalTokenPoolingは、Ward法のクラスタリングで似たトークンベクトルをまとめます。固定の倍率でまとめることも、目標の平均値に合わせて調整できる距離のしきい値でまとめることもでき、クエリ側のベクトルには手を付けません。プーリングはどのマルチベクトルモデルにも使えますが、それを前提に訓練されていないモデルでは重要な区別が失われることがあります。
新しいのは損失関数のラッパーMultiVectorTokenPoolingLossで、各バッチを二つの形で学習させます。文書の全ベクトルと、約4分の1にまとめた版です。考え方は、短くした先頭部分でも役立つように訓練するMatryoshka埋め込みに似ています。著者の実験では、二つのModernBERT-baseモデルを同じ120,000件のMS MARCOトリプレットで1エポック訓練し、違いは損失関数だけにしました。プーリング倍率4では、通常のモデルはプーリングなしのNanoBEIRスコアの94.6%を保ったのに対し、プーリングを意識して訓練したモデルは97.8%を保ち、nDCG@10は49.03から51.23に上がり、保存する文書ベクトルは74.8%少なくなりました。プーリングなしの品質もわずかに向上しました。クラスタリング用のGPUカーネルはSciPyより30-150x速いとされています。
Sentence Transformersを土台とするLightOnのColBERT型検索ライブラリPyLateを使うチームには直接関係があります。この研究はLightOn自身のプーリング正則化の研究を下敷きにしており、訓練手順はHugging Face Jobsを使えばA10G 1枚で4時間以内に収まります。実際的な助言は、量子化、プーリング、Matryoshkaによる切り詰めを組み合わせ、インデックスの大きさと品質を自分のコーパスで測ることです。
注意点は著者自身が挙げています。これは一つのベースモデルによる一回の実験で、ハイパーパラメータの探索はしておらず、評価は小規模なベンチマーク群NanoBEIRで行っています。効果が最もはっきりするのは訓練時と同じプーリング倍率で、他の設定にもある程度引き継がれます。長い文書や特定分野の文書では結果が異なる可能性があります。