Multi-vector embeddingmodellen, zoals ColBERT, bieden een grote nauwkeurigheid bij zoekopdrachten, maar het aantal opgeslagen vectoren per document doet de indexkosten snel stijgen. Een groeperingstechniek, het token pooling, verkleint deze omvang, en een aangepaste training maakt het volgens de bron mogelijk om bij een gelijk opslagbudget de kwaliteit van de resultaten beter te behouden. Met pool_factor=4 behoudt een normaal getraind model 94,6% van zijn niet-gecomprimeerde nDCG@10 op NanoBEIR, tegenover 97,8% voor het model dat getraind is met inachtneming van pooling.
Waarom multi-vector modellen zoveel vectoren opslaan
Een klassiek embeddingmodel vat een heel document samen in één enkele vector. Multi-vector modellen houden daarentegen een set contextuele vectoren bij, één per token (een tekstuele eenheid, vaak een woord of deel van een woord), waardoor verschillende delen van een zoekvraag kunnen overeenkomen met verschillende delen van een document. De bron geeft het voorbeeld van een zoekvraag over een laptop met een lange batterijduur, die een passage over het apparaat in het ene deel van het document kan vinden en een passage over de batterijduur in een ander deel.
Deze werking heeft een prijs. Volgens de gepresenteerde berekening vertegenwoordigen, bij kwantisatie van 4 bits per residuele coördinaat, 200 vectoren van 128 dimensies 12,8 kB per document, oftewel 12,8 GB voor een miljoen documenten. Door elk document terug te brengen tot 50 vectoren daalt dit totaal naar 3,2 GB, bij gelijke dimensie en precisie.
De methode berust op een observatie: veel vectoren van eenzelfde document kunnen dicht bij elkaar liggen, omdat meerdere tokens vaak vergelijkbare begrippen beschrijven. HierarchicalTokenPooling, de Sentence Transformers-bibliotheek, groepeert deze vectoren via Ward-clustering, die bij elke stap de twee groepen samenvoegt die de kleinste toename van interne spreiding veroorzaken. Elke groep wordt vervolgens vervangen door zijn gemiddelde en genormaliseerd. Met pool_factor=4 behoudt een document ongeveer een kwart van zijn initiële vectoren, zodat een langer document proportioneel meer vectoren houdt.
Het model trainen voor compressie
Hoe sterker de compressie, hoe gemakkelijker nuttige informatie verloren kan gaan. De training die rekening houdt met pooling is er precies op gericht de belangrijke informatie te behouden wanneer vectoren worden samengevoegd. Elke trainingsbatch combineert twee verliezen: zoeken met de originele document-embeddings, en zoeken na compressie tot ongeveer een kwart van de vectoren. Deze twee verliezen worden gemiddeld, en beide doelen werken dezelfde encoder bij.
Omdat clustering discrete beslissingen neemt, geeft de bron aan dat de training een straight-through estimator gebruikt, die de gradients van de gegroepeerde vectoren doorgeeft aan de tokens waaruit ze bestaan, zonder onderscheid te maken in de toewijzing aan clusters. De klasse MultiVectorTokenPoolingLoss verpakt een bestaand zoekverlies, hier CachedMultiVectorMultipleNegativesRankingLoss, met gewichten van 0,5 voor elke variant. Het startmodel is answerdotai/ModernBERT-base, getraind op 120.000 MS MARCO-triplets gedurende één epoch.
Resultaten op NanoBEIR
Bij pool_factor=4 stijgt de nDCG@10 van 49,03 naar 51,23 voor het model getraind met pooling, terwijl het aantal documentvectoren met 74,8% wordt verminderd ten opzichte van een niet-gecomprimeerde zoekopdracht. Het niet-gecomprimeerde model verbetert ook, van 51,84 naar 52,39. Bij een gemiddelde van 29,58 vectoren per document behoudt adaptieve pooling 97,3% van de niet-gecomprimeerde score van het getrainde model, tegenover 93,3% voor het normaal getrainde model.
Om de grootte van een index te verkleinen, suggereert de bron eerst kwantisatie van de embeddings, vervolgens token pooling, eventueel met fine-tuning die rekening houdt met pooling, en tenslotte Matryoshka-embeddings om de dimensie van elke vector te verkleinen. Deze benaderingen kunnen worden gecombineerd. Het werk waarop het artikel is gebaseerd, wordt beschreven in de studie Learn to Pool van Stefan Josef.
