Efficient Multivector Retrieval with Token-Aware Clustering and Hierarchical Indexing

1. Énoncé du problème

Les modèles de recherche par multivecteurs (par exemple, ColBERT) ont atteint une efficacité de pointe en Recherche d'Information (RI) en encodant les documents comme des ensembles de vecteurs au niveau des tokens et en utilisant une « interaction tardive » (MaxSim) pour calculer la pertinence. Cependant, leur déploiement est entravé par :

• Coûts de calcul et de mémoire élevés : Le stockage et la recherche de milliards de vecteurs de tokens nécessitent des ressources massives.
• Clustering inefficace : Les solutions actuelles reposent sur le clustering $k$-means standard pour générer des vecteurs centroïdes en vue de la compression et de la collecte de candidats.
  • Goulot d'étranglement de l'évolutivité : Le $k$-means s'adapte mal à la taille des jeux de données et au nombre de clusters ($k$), nécessitant souvent des dizaines d'heures ou des GPU coûteux pour entraîner même des ensembles modérément grands (par exemple, 262k centroïdes).
  • Déséquilibre des tokens : Le $k$-means standard traite tous les vecteurs de manière égale. Par conséquent, les tokens fréquents (mots vides) dominent la fonction de perte et consomment la majorité des centroïdes, tandis que les tokens rares et hautement discriminants sont sous-représentés, dégradant la qualité de la recherche.

2. Méthodologie

Les auteurs proposent Tachiom, un système de recherche comprenant deux innovations principales : le Clustering Sensible aux Tokens (Tac) et un Index Hiérarchique pour une recherche optimisée.

A. Clustering Sensible aux Tokens (Tac)

Tac remédie aux limitations du $k$-means standard en exploitant la propriété structurelle de l'identité des tokens dans les embeddings multivecteurs. Au lieu de regrouper tous les vecteurs globalement, Tac décompose le problème en sous-problèmes indépendants par token.

• Stratégie d'allocation des centroïdes : Tac alloue le budget de centroïdes en fonction de la fréquence des tokens et de la variance sémantique, plutôt que d'une distribution uniforme des vecteurs.
  • Gestion de la queue : Les tokens extrêmement rares ($n < \mu$) garantissent au moins 1 centroïde ; les petits tokens ($\mu \le n < \tau$) en obtiennent 2.
  • Score amorti : Pour les tokens actifs, un poids $w_j$ est calculé comme $\sqrt{n_j} \cdot s_j$, où $n_j$ est la fréquence et $s_j$ la variance sémantique (variance composante par composante). La racine carrée atténue l'influence des tokens à haute fréquence, empêchant qu'ils ne monopolisent le budget.
  • Bornes et réconciliation : Des planchers et plafonds stricts assurent une représentation minimale pour les tokens rares et empêchent la sur-allocation. Le budget final est réconcilié pour correspondre à la cible $k$.
• Parallélisation : En résolvant des problèmes $k_j$-means indépendants pour chaque type de token, Tac permet un parallélisme massif, réduisant drastiquement la complexité computationnelle par rapport au $k$-means global.
• Compression des résidus : Après le clustering, les résidus (vecteur - centroïde) sont normalisés et compressés à l'aide de la Quantification Produit (PQ).

B. Index Hiérarchique pour la recherche (Tachiom)

Tachiom exploite les centroïdes haute résolution générés par Tac pour optimiser le pipeline de recherche en deux phases :

1. Phase de collecte (Score uniquement sur les centroïdes) :
   • Un graphe Hnsw est construit sur les centroïdes.
   • Chaque centroïde maintient une liste inversée d'identifiants de documents.
   • Pour une requête, le système parcourt le graphe Hnsw pour trouver les centroïdes les plus proches, récupère les documents associés et calcule un score MaxSim approximatif en utilisant uniquement les interactions de centroïdes.
   • Cela contourne entièrement les calculs coûteux au niveau des tokens lors de la sélection initiale des candidats.
2. Phase d'affinement (MaxSim optimisé par PQ) :
   • Les meilleurs candidats sont affinés en utilisant à la fois les centroïdes et les résidus compressés par PQ.
   • Disposition optimisée pour le cache : Les auteurs introduisent une disposition mémoire spécialisée pour les tables de distances PQ. Au lieu d'accès mémoire dispersés pour chaque token de requête, les distances sont organisées en Macro-blocs (sous-espaces), Blocs de centroïdes et Micro-blocs de tokens de requête. Cela permet de récupérer les distances pour tous les tokens de requête contre un centroïde spécifique comme un seul bloc mémoire contigu, améliorant significativement la localité du cache et l'utilisation SIMD.

3. Contributions clés

1. Clustering Sensible aux Tokens (Tac) : Un algorithme novateur qui décompose le clustering global en sous-problèmes par token. Il garantit théoriquement une accélération minorée par le rapport entre le poids total du vocabulaire et le poids maximal d'un seul token.
2. Architecture Tachiom : Un système de recherche qui combine la collecte de centroïdes basée sur Hnsw avec une disposition PQ optimisée pour le cache, adaptée à l'interaction tardive, permettant un score complet efficace.
3. Validation théorique et empirique : L'article fournit une borne inférieure théorique pour l'accélération de Tac et démontre que la méthode s'adapte à des millions de centroïdes tout en maintenant ou en améliorant l'efficacité de la recherche.

4. Résultats expérimentaux

Des expériences ont été menées sur Ms Marco-v1 (8,8M de passages) et LoTTE (2,4M de passages) en utilisant des embeddings ColBERTv2.

• Efficacité du clustering :
  • Tac a atteint jusqu'à 247 fois plus rapide en entraînement que le $k$-means standard (Faiss) et 84 fois plus rapide que Faiss avec accélération MKL.
  • Tac a réussi à regrouper 600M de vecteurs en 262k clusters en 8 minutes, tandis que les concurrents atteignaient un délai d'attente de 24 heures avant d'atteindre 131k–524k clusters.
  • Tac s'adapte de manière transparente à plus de 4 millions de centroïdes en ~100 minutes.
• Efficacité de la recherche :
  • Tac a obtenu un MRR@10 égal ou supérieur à celui du $k$-means standard avec des budgets de centroïdes fixes, prouvant que la priorisation des tokens rares améliore la qualité de l'approximation.
• Vitesse de recherche :
  • Tachiom a atteint jusqu'à 9,8 fois plus rapide en recherche de bout en bout par rapport aux systèmes de pointe (Emvb, Warp, Igp).
  • Il a égalé l'efficacité maximale d'Emvb (qui utilise des variantes PQ supervisées coûteuses) tout en utilisant une PQ standard et en obtenant une accélération de 5,5 fois par rapport à Emvb.

5. Importance

Ce travail modifie fondamentalement le paradigme de l'efficacité de la recherche par multivecteurs :

• Évolutivité : Il élimine le goulot d'étranglement computationnel du clustering, permettant aux systèmes d'utiliser des millions de centroïdes (granularité plus fine) qui étaient auparavant impossibles à entraîner.
• Compromis Qualité-Efficacité : En alignant l'allocation des centroïdes sur les objectifs de recherche (favorisant les tokens rares discriminants) plutôt que sur l'erreur de reconstruction pure, il atteint une meilleure efficacité avec moins de ressources.
• Déploiement pratique : La combinaison d'un entraînement rapide et d'une inférence ultra-rapide (via la collecte Hnsw et la PQ optimisée) rend la recherche par multivecteurs haute performance viable pour des environnements de production à grande échelle.

Les auteurs ont publié leur implémentation Rust de Tac et Tachiom sur GitHub pour assurer la reproductibilité.