Multi-Vektor-Embedding-Modelle wie ColBERT bieten eine hohe Präzision bei der Suche, doch die Anzahl der pro Dokument gespeicherten Vektoren treibt die Indexkosten schnell in die Höhe. Eine Gruppierungstechnik, das Token-Pooling, reduziert diese Größe, und ein angepasstes Training ermöglicht laut Quelle, die Ergebnisqualität bei gleichem Speicherbudget besser zu erhalten. Mit pool_factor=4 behält ein normal trainiertes Modell auf NanoBEIR 94,6 % seines unkomprimierten nDCG@10, gegenüber 97,8 % für das Modell, das unter Berücksichtigung des Poolings trainiert wurde.
Warum Multi-Vektor-Modelle so viele Vektoren speichern
Ein klassisches Embedding-Modell fasst ein gesamtes Dokument in einem einzigen Vektor zusammen. Multi-Vektor-Modelle bewahren dagegen einen Satz kontextueller Vektoren, einen pro Token (eine Texteinheit, oft ein Wort oder Wortteil), sodass verschiedene Teile einer Anfrage mit verschiedenen Teilen eines Dokuments übereinstimmen können. Die Quelle nennt das Beispiel einer Anfrage nach einem Laptop mit langer Akkulaufzeit, die eine Passage über das Gerät in einem Teil des Dokuments und eine Passage über die Akkulaufzeit in einem anderen finden kann.
Diese Funktionsweise hat ihren Preis. Laut der vorgestellten Berechnung entsprechen bei einer Quantisierung mit 4 Bit pro Residualkoordinate 200 Vektoren mit 128 Dimensionen 12,8 KB pro Dokument, also 12,8 GB für eine Million Dokumente. Reduziert man jedes Dokument auf 50 Vektoren, sinkt dieser Gesamtwert bei gleicher Dimension und Präzision auf 3,2 GB.
Die Methode basiert auf einer Beobachtung: Viele Vektoren eines Dokuments können nahe beieinander liegen, da mehrere Token oft ähnliche Konzepte beschreiben. HierarchicalTokenPooling, die Sentence Transformers-Bibliothek, gruppiert diese Vektoren mittels Ward-Clustering, das bei jedem Schritt die beiden Gruppen fusioniert, die den geringsten Anstieg der internen Streuung verursachen. Jede Gruppe wird dann durch ihren Mittelwert ersetzt und normalisiert. Mit pool_factor=4 behält ein Dokument etwa ein Viertel seiner ursprünglichen Vektoren, sodass ein längeres Dokument proportional mehr Vektoren behält.
Das Modell auf Kompression trainieren
Je stärker die Kompression, desto leichter kann nützliche Information verloren gehen. Das Training unter Berücksichtigung des Poolings zielt genau darauf ab, die wichtige Information zu bewahren, wenn Vektoren fusioniert werden. Jeder Trainingsbatch kombiniert zwei Verlustfunktionen: die Suche mit den ursprünglichen Dokument-Embeddings und die Suche nach Kompression auf etwa ein Viertel der Vektoren. Diese beiden Verluste werden gemittelt, und beide Ziele aktualisieren denselben Encoder.
Da das Clustering diskrete Entscheidungen trifft, gibt die Quelle an, dass das Training einen Straight-Through-Schätzer verwendet, der die Gradienten der gruppierten Vektoren an die sie zusammensetzenden Token weitergibt, ohne die Clusterzuweisung zu differenzieren. Die Klasse MultiVectorTokenPoolingLoss kapselt einen bestehenden Suchverlust, hier CachedMultiVectorMultipleNegativesRankingLoss, mit Gewichten von 0,5 für jede Variante. Das Ausgangsmodell ist answerdotai/ModernBERT-base, trainiert auf 120.000 MS-MARCO-Tripeln für eine Epoche.
Ergebnisse auf NanoBEIR
Bei pool_factor=4 steigt der nDCG@10 von 49,03 auf 51,23 für das mit Pooling trainierte Modell, während gleichzeitig die Anzahl der Dokumentvektoren im Vergleich zu einer unkomprimierten Suche um 74,8 % reduziert wird. Das unkomprimierte Modell verbessert sich ebenfalls, von 51,84 auf 52,39. Bei einer durchschnittlichen Anzahl von 29,58 Vektoren pro Dokument bewahrt das adaptive Pooling 97,3 % des unkomprimierten Scores des trainierten Modells, gegenüber 93,3 % für das normal trainierte Modell.
Um die Größe eines Index zu reduzieren, schlägt die Quelle zunächst die Quantisierung der Embeddings vor, dann das Token-Pooling, gegebenenfalls mit einem auf Pooling abgestimmten Fine-Tuning, und schließlich Matryoshka-Embeddings, um die Dimension jedes Vektors zu reduzieren. Diese Ansätze können kombiniert werden. Die Arbeiten, auf denen der Artikel basiert, werden in der Studie Learn to Pool von Stefan Josef beschrieben.
