Pooling conscient des tokens dans Sentence Transformers

Un nouveau dispositif d’entraînement combine des variantes avec et sans regroupement de tokens pour apprendre à des modèles multi‑vecteurs à résister à la compression. Associé à un pooling hiérarchique calibrable, il vise une moyenne de vecteurs par document tout en conservant la dimension et la précision annoncées. Les gains mémoire sont concrets à 4 bits quand des centaines de vecteurs par document doivent être indexés.

Indexer moins lourd sans repasser par l’encodeur

Les variantes de pooling réutilisent les embeddings de tokens déjà calculés ; ajouter un pooler supplémentaire n’exige pas de nouveau passage d’encodeur et ajoute seulement du travail de clustering et de scoring. Un script d’entraînement peut fonctionner avec des mini‑lots de 64 et une taille de lot de 256, et la perte de base s’instancie avec un mini_batch_size de 64. Côté indexation, réduire chaque document à 50 vecteurs ramène l’empreinte à 3,2 Go pour un million de documents, contre 12,8 Go quand 200 vecteurs de 128 dimensions sont stockés à 4 bits par coordonnée. L’objectif de compression s’additionne à l’objectif de récupération, et les poids choisis moyennent leurs pertes.

MultiVectorTokenPoolingLoss associe regroupement et apprentissage

MultiVectorTokenPoolingLoss regroupe une fonction de perte de récupération multi‑vecteurs et l’utilise avec les poolers spécifiés. La fonction de perte de récupération apprend à chaque requête à attribuer un score supérieur à son document pertinent par rapport aux négatifs du lot, en utilisant par défaut la mesure de similarité définie dans le modèle. Un modèle dont le paramètre est réglé sur meanmaxsim emploierait également MeanMaxSim lors de l’entraînement. Un exemple de configuration utilise token_poolings=[None, pooling] et des poids [0,5, 0,5] avec un encodeur tel que "answerdotai/ModernBERT-base" et un pooler HierarchicalTokenPooling avec pool_factor=4. Le caching des gradients permet de traiter de petits mini‑lots tout en conservant davantage de négatifs efficaces.

Un entraînement pensé pour des représentations regroupées

Chaque lot peut contribuer à deux pertes pondérées 0,5 : l’une sans pooling et l’autre avec pooling. Le premier terme apprend la récupération sur les représentations complètes, le second pousse le modèle à structurer ses embeddings de tokens pour que les correspondances utiles survivent à la fusion. Les deux objectifs mettent à jour le même encodeur, avec un clustering de Ward qui prend des décisions discrètes. Pendant l’entraînement, les gradients des vecteurs regroupés sont transmis directement aux tokens membres sans différencier les attributions, permettant d’apprendre des erreurs observées sur les versions compressées. Cette approche transpose à la réduction de tokens l’idée des embeddings Matryoshka, en appliquant l’objectif à des représentations moins détaillées, afin d’exploiter le regroupement déjà utile sur des modèles multi‑vecteurs classiques tout en limitant ses effets indésirables.

Le pooling hiérarchique par Ward en pratique

HierarchicalTokenPooling s’appuie sur le clustering hiérarchique de Ward pour regrouper des embeddings de tokens similaires. Il part d’un cluster par token, fusionne itérativement les deux groupes dont l’union accroît le moins l’étalement interne, mesuré par la somme des distances au carré à la moyenne. Les fusions construisent une hiérarchie qui autorise différents budgets de compression. Chaque cluster est représenté par sa moyenne, puis les vecteurs sont renormalisés et utilisables directement avec MaxSim. Cette mécanique s’inscrit dans le paradigme multi‑vecteurs où les différentes parties d’une requête peuvent correspondre à des segments distincts d’un document, à l’opposé d’un unique embedding pour tout le texte.

Fixer un budget : plafond, fraction ou seuil de distance

Trois contrôles sont proposés pour maîtriser le nombre de vecteurs par document : un maximum fixé à 16, la conservation d’environ un quart des vecteurs initiaux ou un seuil de distance partagé pour les fusions de Ward. Ce paramètre joue un rôle clé lorsque la taille des documents fluctue : avec un pool_factor=4, un document plus long garde près d’un quart de ses vecteurs, tandis qu’en utilisant un seuil de distance, le nombre de vecteurs peut croître si certains embeddings distincts doivent être maintenus. Le regroupement adaptatif qui repose sur un seuil commun procède à une fusion des similitudes de façon plus poussée et retient davantage de vecteurs dans les zones où le contenu présente plus de diversité. Cette approche fonctionne avec tout modèle multi‑vecteurs, y compris ceux non entraînés spécifiquement pour le pooling.

Calibrer le seuil pour une moyenne de vecteurs visée

Plutôt que de fixer manuellement un seuil, il est possible de le calibrer pour atteindre une moyenne cible de vecteurs. À partir d’un échantillon représentatif, on encode les documents, on calcule un seuil avec HierarchicalTokenPooling.calibrate_distance_threshold (target_tokens=16) puis on réutilise ce seuil lors de l’indexation. La calibration ne requiert que des embeddings de documents et vise une moyenne : chaque texte peut recevoir plus ou moins de 16 vecteurs, et la moyenne obtenue peut varier sur un autre corpus. Cette démarche se combine avec des index compressés via une quantification résiduelle de 2 ou 4 bits, contexte dans lequel le coût mémoire des centaines de vecteurs par document s’accumule rapidement, comme l’illustrent les ordres de grandeur à 4 bits.