Principe du token pooling
Les modèles d’embedding multi‑vecteurs, comme ColBERT, attribuent à chaque token d’une requête un vecteur distinct, ce qui augmente la capacité de correspondance fine‑grained. Cependant, chaque document peut contenir des centaines de vecteurs : 200 vecteurs de 128 dimensions stockés avec une quantification résiduelle de 4 bits occupent 12,8 KB par document, soit 12,8 GB pour un million de documents. Le token pooling regroupe les vecteurs similaires afin de réduire ce nombre. Avec un pool_factor=4, le nombre de vecteurs est divisé par quatre, passant de 200 à 50 vecteurs et réduisant l’encombrement à 3,2 GB pour le même corpus.
Mécanisme de HierarchicalTokenPooling
Le module HierarchicalTokenPooling de Sentence Transformers applique un clustering hiérarchique de Ward. Chaque token débute dans son propre groupe ; à chaque itération, les deux groupes dont la fusion entraîne la plus petite augmentation de la somme des distances au carré (within‑cluster sum of squares) sont fusionnés. Cette règle favorise des groupes compacts de vecteurs très proches. Le résultat est une hiérarchie où chaque niveau représente un facteur de compression différent. Le vecteur de chaque groupe est la moyenne normalisée de ses membres, prêt à être utilisé avec le scoring MaxSim. Le premier token est protégé par défaut, garantissant que la requête reste inchangée.
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pooling = HierarchicalTokenPooling(pool_factor=4)
Entraînement aware du pooling
Le problème majeur est que le modèle original n’est jamais exposé à des représentations compressées pendant l’apprentissage. Le pooling‑aware training introduit un second objectif : après chaque batch, on calcule la perte de récupération à la fois sur les embeddings complets et sur les embeddings compressés (≈ 25 % des vecteurs). La moyenne des deux pertes alimente le même encodeur, ce qui force les token embeddings à rester discriminants même après agrégation. Le gradient traverse les opérations de clustering grâce à un straight‑through estimator, qui ignore la non‑différentiabilité des affectations de clusters tout en transmettant les signaux d’erreur aux tokens d’origine.
from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
model = MultiVectorEncoder("lightonai/LateOn")
pooling = HierarchicalTokenPooling(pool_factor=4)
# loss wrapper
loss = MultiVectorTokenPoolingLoss(base_loss, pooling)
Évaluation, bénéfices et limites
Sur le benchmark NanoBEIR, le modèle entraîné sans pooling conserve 94,6 % du nDCG@10 du modèle non compressé avec pool_factor=4. En appliquant le pooling‑aware training, ce taux monte à 97,8 % pour le même budget de stockage, soit une amélioration de 3,2 points de pourcentage. La réduction du nombre de vecteurs implique toutefois une perte d’information potentielle : si un document contient plusieurs concepts distincts, le regroupement peut masquer des différences essentielles, surtout lorsque le seuil de distance est trop agressif. Le calibrage du distance_threshold à l’aide de la méthode calibrate_distance_threshold permet de viser un nombre moyen de vecteurs (par ex. 16), mais le résultat dépend fortement de la distribution du corpus et peut varier d’un jeu de données à l’autre.
En résumé, le token pooling offre un compromis stockage‑qualité mesurable, tout en nécessitant un entraînement dédié pour éviter la dégradation des performances de recherche.
