Interroger 3 milliards de vecteurs

• À partir du problème de requête sur 3 milliards de vecteurs de Jeff Dean, compte rendu d’une expérimentation technique qui consiste à implémenter soi-même une solution map-reduce optimale
• En partant d’une implémentation naive calculant le dot product entre 3 milliards de vecteurs float32 de dimension 768 et 1�0 requêtes, des gains de performance progressifs ont été obtenus grâce à la vectorisation avec numpy et à la conversion en float32
• Sur une base de 3�0 vecteurs, le temps est passé d’environ 2 secondes avec l’approche naive à 0,01 seconde après vectorisation, puis jusqu’à 0,0045 seconde avec float32
• En passant à l’échelle de 3 milliards, la matrice de résultats nécessite environ 8,6 To de mémoire, ce qui provoque un problème d’OOM et impose des optimisations supplémentaires comme le traitement par lots, le memory mapping, une réécriture en Rust/C ou l’usage de SimSIMD
• Le texte souligne que, plus encore que la solution technique, la vraie difficulté est de définir les exigences (forme du retour, caractéristiques de la machine, tolérance à la compression, etc.)

---

Point de départ du problème

• Le point de départ est un échange autour du problème de requête sur 3 milliards de vecteurs de Jeff Dean, avec la volonté d’implémenter soi-même une solution map-reduce optimale
• Un vecteur est un tableau de nombres à virgule flottante de dimension n, et la recherche vectorielle sert à trouver des mots ou des éléments sémantiquement similaires
• C’est un schéma courant d’utiliser des embeddings dans les applications de recherche générative, comme la recherche, la recommandation ou des outils comme Cursor

Implémentation naive

• Hypothèse initiale : 3 milliards de vecteurs de documents à interroger et environ 1�0 vecteurs de requête, tous stockés sur disque au format .npy
• La dimension des vecteurs est 768, une taille courante dans de nombreux modèles de requêtes d’embeddings fondés sur la similarité
• L’approche consiste en une double boucle qui compare chaque vecteur de requête à l’ensemble des vecteurs de documents pour calculer le dot product et renvoyer le résultat
• Un premier test sur 3�0 vecteurs a donné environ 2 secondes sur un MacBook M2 (soit une échelle un million de fois plus petite que 3 milliards)

Optimisation par vectorisation

• Application d’une approche vectorisée remplaçant la double boucle par l’opération de multiplication matricielle de numpy (vectors_file @ query_vectors.T)
• Sur 3�0 vecteurs, le temps tombe à 0,0107 seconde
• En passant à 3 millions de vecteurs, l’exécution prend 12,85 secondes

Conversion en float32

• Une optimisation supplémentaire est réalisée en convertissant les données en np.float32
• Sur 3�0 vecteurs, le temps descend jusqu’à 0,0045 seconde
• Comme il faut environ 13 secondes pour 3 millions de vecteurs, on estime qu’il faudrait environ 316 minutes pour 3 milliards, soit 1�0 fois plus

Résumé de la comparaison des performances

• Méthode naive (3�0 vecteurs) : 1,9877 seconde
• Méthode vectorisée (3�0 vecteurs) : 0,0107 seconde
• Méthode vectorisée (3 millions de vecteurs) : 12,8491 secondes
• Méthode float32 (3�0 vecteurs) : 0,0045 seconde

Problème d’OOM et pistes d’optimisation supplémentaires

• Mémoire requise pour traiter 3 milliards d’objets en float32 (4 octets) : environ 8,6 To, ce qui provoque un OOM à l’exécution
• Des optimisations supplémentaires dans l’esprit des pistes proposées par Jeff Dean sont nécessaires :
  • modifier le code pour utiliser des générateurs et comparaisons par lots
  • écrire sur disque à intervalles réguliers ou utiliser le memory mapping
  • réécrire le code d’optimisation bas niveau en Rust ou en C
  • utiliser une bibliothèque spécialisée dans les comparaisons massives de similarité vectorielle comme SimSIMD

Importance de la définition des exigences

• Avant même d’optimiser davantage, plusieurs zones floues dans le problème lui-même apparaissent :
  • s’agit-il d’interroger l’ensemble des 3 milliards avec un seul vecteur de requête et de renvoyer tous les résultats, ou d’une recherche ANN top-k ?
  • faut-il aussi renvoyer les vecteurs d’origine, et un reranking fondé sur la similarité est-il nécessaire ?
  • faut-il interroger l’ensemble avec 1�0 vecteurs de requête en une seule fois ?
  • les vecteurs sont-ils déjà en mémoire, lus un par un depuis le disque, ou dans un mode streaming ?
  • l’exécution est-elle locale, quelles sont les caractéristiques de la machine, et l’usage d’un GPU est-il possible ?
  • quel est l’impact de la taille des embeddings sur la représentation des résultats et sur la taille des vecteurs en entrée/sortie ?
  • la compression de float64 vers float32 est-elle acceptable du point de vue de la précision ?
  • s’agit-il d’un projet ponctuel, et combien de temps est alloué à sa mise en place ?
• Plus que la solution technique elle-même, la tâche la plus difficile est de définir précisément les exigences