Les modèles d’embeddings multi-vecteurs, comme ColBERT, offrent une grande finesse dans la recherche, mais le nombre de vecteurs stockés par document fait vite grimper le coût de l’index. Une technique de regroupement, le token pooling, réduit cette taille, et un entraînement adapté permet, selon la source, de mieux conserver la qualité des résultats à budget de stockage égal. Avec pool_factor=4, un modèle entraîné normalement conserve 94,6 % de son nDCG@10 non compressé sur NanoBEIR, contre 97,8 % pour le modèle entraîné en tenant compte du pooling.
- Les modèles multi-vecteurs gagnent en finesse, mais leur index grossit vite.
- Le token pooling réduit le nombre de vecteurs par document, avec pool_factor=4.
- Un entraînement tenant compte du pooling conserve mieux la qualité, à 97,8 % contre 94,6 %.
- Ces techniques se combinent avec la quantification et les embeddings Matryoshka pour réduire encore l’index.
Pourquoi les modèles multi-vecteurs stockent autant de vecteurs
Un modèle d’embeddings classique résume un document entier en un seul vecteur. Les modèles multi-vecteurs conservent au contraire un ensemble de vecteurs contextuels, un par token (une unité de texte, souvent un mot ou une partie de mot), ce qui permet à différentes parties d’une requête de correspondre à différentes parties d’un document. La source prend l’exemple d’une requête sur un ordinateur portable à longue autonomie, qui peut trouver un passage sur l’appareil dans une partie du document et un passage sur l’autonomie de la batterie dans une autre.
Ce fonctionnement a un coût. Selon le calcul présenté, avec une quantification à 4 bits par coordonnée résiduelle, 200 vecteurs de 128 dimensions représentent 12,8 Ko par document, soit 12,8 Go pour un million de documents. Ramener chaque document à 50 vecteurs ramène ce total à 3,2 Go, à dimension et précision identiques.
La méthode repose sur une observation : beaucoup de vecteurs d’un même document peuvent être proches, car plusieurs tokens décrivent souvent des notions voisines. HierarchicalTokenPooling, la bibliothèque Sentence Transformers, regroupe ces vecteurs par clustering de Ward, qui fusionne à chaque étape les deux groupes provoquant la plus faible augmentation de la dispersion interne. Chaque groupe est ensuite remplacé par sa moyenne, puis normalisé. Avec pool_factor=4, un document conserve environ un quart de ses vecteurs initiaux, de sorte qu’un document plus long garde proportionnellement davantage de vecteurs.
Entraîner le modèle à la compression
Plus la compression est forte, plus il est facile de perdre une information utile. L’entraînement tenant compte du pooling vise précisément à conserver l’information importante lorsque les vecteurs sont fusionnés. Chaque lot d’entraînement combine deux pertes : la recherche avec les embeddings de documents d’origine, et la recherche après compression à environ un quart des vecteurs. Ces deux pertes sont moyennées, et les deux objectifs mettent à jour le même encodeur.
Comme le clustering prend des décisions discrètes, la source indique que l’entraînement utilise un estimateur straight-through, qui transmet les gradients des vecteurs regroupés vers les tokens qui les composent, sans différencier l’affectation aux clusters. La classe MultiVectorTokenPoolingLoss encapsule une perte de recherche existante, ici CachedMultiVectorMultipleNegativesRankingLoss, avec des poids de 0,5 pour chaque variante. Le modèle de départ est answerdotai/ModernBERT-base, entraîné sur 120 000 triplets MS MARCO pendant une époque.
Résultats sur NanoBEIR
À pool_factor=4, le nDCG@10 passe de 49,03 à 51,23 pour le modèle entraîné avec pooling, tout en réduisant de 74,8 % le nombre de vecteurs de documents par rapport à une recherche non compressée. Le modèle non compressé s’améliore également, de 51,84 à 52,39. Pour une moyenne de 29,58 vecteurs par document, le pooling adaptatif conserve 97,3 % du score non compressé du modèle entraîné, contre 93,3 % pour le modèle entraîné normalement.
Pour réduire la taille d’un index, la source suggère d’abord la quantification des embeddings, puis le token pooling, éventuellement avec un réglage fin tenant compte du pooling, et enfin les embeddings Matryoshka pour réduire la dimension de chaque vecteur. Ces approches peuvent se combiner. Les travaux sur lesquels s’appuie l’article sont décrits dans l’étude Learn to Pool de Stefan Josef.
