Indexation et recherche de similarité avec Faiss

14 min readFeb 6, 2023

Cet article a pour but d’expliquer quelques-uns des algorithmes de recherche de similarité implémentés dans faiss, une bibliothèque développée par facebook, que j’ai découverte à l’occasion d’un projet de NLP. Je passe brièvement sur le contexte du problème et l’intérêt d’utiliser faiss, et je poursuis avec une explication technique des trois principaux indexes proposés par faiss : inverted file index, product quantization, et hierarchical navigable small worlds.

Introduction

Contexte : le projet Grand Débat

Le projet Grand Débat est un projet de R&D mené par La Javaness en automne 2022. Le but était d’appliquer nos compétences et connaissances en NLP et visualisation de données pour créer une application permettant d’explorer les contributions citoyennes récoltées lors du Grand Débat National.

L’une des fonctionnalités que nous souhaitions implémenter était la recherche de contributions similaires à une proposition entrée par l’utilisateur. Plutôt que d’utiliser un moteur de recherche par mots-clés, comme c’est possible par exemple avec Elasticsearch, on souhaite ici permettre à l’utilisateur de visualiser des contributions sémantiquement similaires à la sienne. Nous avons donc décidé de faire usage du moteur d’embedding développé par un de nos data scientists, et d’effectuer une recherche de similarité. (Note : un moteur d’embedding est optimisé pour que des phrases ayant un contenu sémantique similaire soient encodées par des vecteurs proches).

À quoi sert un index ?

Nous avons donc une base de données de plusieurs millions de contributions, découpées en phrases (ou en groupes de phrases) à l’aide d’un tokenizer, puis transformées en vecteurs de taille 1024 par le moteur d’embedding. Lorsque l’utilisateur entre une proposition (aussi appelée requête), cette proposition est elle aussi transformée en vecteur ; on peut ainsi calculer la distance entre notre proposition et la contribution d’un utilisateur (ce qu’on ne pourrait pas faire directement avec du texte).

Pour effectuer une recherche, on pourrait naïvement calculer individuellement cette distance avec toutes les contributions dans notre base de données, et renvoyer la contribution plus proche de la requête. Cela devient vite un casse-tête lorsque la taille de la base de données augmente (non seulement il faut la charger en mémoire, mais en plus on doit faire le calcul de distance pour toutes les contributions).

C’est là qu’intervient le concept d’index : plutôt que de charger toute la base de données en mémoire et de faire plusieurs millions de calculs pour trouver le vecteur le plus proche de notre requête, on va effectuer cette recherche sur une représentation de la base de données : par exemple, on va la découper en “zones” et ne chercher que dans la zone autour de notre requête, ou on va compresser les vecteurs de la base pour en réduire la taille, etc. C’est cette représentation qu’on appelle l’index.

Faiss

Faiss, c’est Fast AI Similarity Search, un module développé par facebook AI pour l’indexation et la recherche dans un ensemble de vecteurs.

Plusieurs types d’indexes sont implémentés dans faiss, à des niveaux de sophistication différents allant de la simple recherche exhaustive, à des indexes compressés et très rapides (permettant la recherche dans des ensembles de l’ordre de plusieurs milliards de vecteurs). Différents indexes présentent différents avantages et inconvénients : en général, il y aura toujours une forme d’équilibre à trouver entre la précision (à quel point la contribution retournée est proche de la requête), et la performance (vitesse et/ou taille en mémoire).

Cet article est accompagné d’un notebook qui présente des exemples/tests sur un dataset généré aléatoirement.

Installation et création d’un index plat (recherche exhaustive).

Il existe deux versions (et donc deux paquets), faiss-cpu et faiss-gpu. On installe avec pip :

pip install faiss-cpu

On commence par créer un index plat(IndexFlatL2) : nos vecteurs sont stockés tels quels, et la recherche est exhaustive. En gros, on ne fait rien de spécial. Ce type d’index renvoie des résultats exacts : comme nos vecteurs ne sont pas compressés, et qu’on cherche dans toute la base, on est sûr que le résultat de notre requête sera LE vecteur le plus proche (et pas un autre). En revanche, cela consomme pas mal de mémoire (on stocke tous les vecteurs) et c’est assez long (on compare la requête avec tout l’ensemble).

Pour initialiser la classe, faiss aura besoin de la dimension d de nos vecteurs. On ajoute ensuite les vecteurs à l’index avec index.add(). Par défaut, faiss s’attend à un array de dimension (n_vecteurs x d).

On effectue une recherche avec index.search(xq, k), où

xq est notre requête. Il peut s’agir d’un seul vecteur ou d’un array de plusieurs vecteurs (si on veut chercher plusieurs vecteurs d’un coup).
k est le nombre de résultats par vecteur à renvoyer.

Donc si on veut les 5 résultats les plus proches de 4 vecteurs donnés, xq sera de dimension (4, d) et k sera égal à 5.

La méthode .search() renvoie deux objets :

une matrice D des distances, de taille (n, k) également, où la i-ème ligne contient les distances (au carré) des k vecteurs les plus proches du i-ème vecteur de la requête.
le résultat de notre recherche : une matrice I de taille (n, k) où n est le nombre de vecteurs dans notre requête. La i-ème ligne contient les indices des k vecteurs les plus proches du i-ème vecteur de la requête.

Un exemple

Avec la commande suivante :

on demande à faiss de chercher les vecteurs les plus proches des trois premiers vecteurs de notre dataset.

Faiss renvoie la matrice I (indices) suivante :

[[    0 18096 23340 69719 76760]
 [    1 21442 14241 16511 88648]
 [    2 31767 77193 90699 42060]]

La première ligne contient les indices des 5 vecteurs les plus proches du vecteur 0 dans la base ; sans surprise, le vecteur le plus proche du vecteur 0 est lui-même. Le deuxième vecteur le plus proches est le vecteur n°18096, etc.

La deuxième ligne contient les indices des 5 vecteurs les plus proches du vecteur 1, et pareil pour la troisième ligne.

La matrice D (distance) renvoyée se présente quant à elle sous la forme suivante :

[[ 0.       14.362421 14.602067 14.767819 15.029004]
 [ 0.       12.729409 13.443033 13.679935 14.520304]
 [ 0.       13.01179  13.469017 13.666645 13.668671]]

La première ligne contient les distances des 5 vecteurs les plus proches du premier vecteur de la base. On peut donc voir que le vecteur 0 est à une distance 0 de lui-même, que le vecteur 18096 est à une distance de 14.36, que le vecteur 88648 est à une distance 14.52 du vecteur 1, etc.

Plus rapide : l’index inversé.

Faiss permet la création d’un index inversé (IVF, inverted file index), qui permet de réduire le nombre de vecteurs qu’on compare à notre requête.

Lors de la création de l’index, on regroupe nos vecteurs en clusters (cellules, dans le langage faiss) avec k-means.
Au moment de la recherche, on compare notre requête aux centroïdes des cellules et on sélectionne les plus proches de la requête.
On compare ensuite notre requête aux vecteurs qui sont situés dans les cellules sélectionnées.

Le nombre de cellules créées au moment de l’indexing, ainsi que le nombre de cellules visitées lors de la recherche, sont des paramètres à déterminer.

Ci-dessous, un exemple avec une dimension d’embedding de 1024, et dix cellules :

Le paramètre quantizer spécifie quelle représentation des vecteurs on va utiliser pour calculer les clusters. Ici on utilise un index plat, c’est-à-dire qu’on va clusteriser les vecteurs sans en réduire la dimension ou les compresser.

On initialise la classe, et on entraîne l’index pour créer nos cellules (note : ça peut être fait sur un extrait représentatif des données). Enfin, on ajoute nos données à l’index, qui va les répartir dans les cellules.

Le paramètre nprobe spécifie combien de cellules on souhaite visiter lors de la recherche (on prendra alors les nprobe cellules dont les centroïdes sont les plus proches de notre requête).

Terminologie, cellules, et équilibre vitesse-précision
La terminologie index inversé vient du fait qu’on assigne d’abord un numéro de cluster à chaque vecteur, puis qu’on écrit une table indexée par numéro de cluster, contenant l’ensemble des vecteurs dans chaque cluster. Ce serait un peu comme un annuaire organisé par adresse plutôt que par nom.
Les clusters sont appelés cellules en référence aux cellules de Voronoï, un pavage de l’espace tel que tout point est assigné à une cellule dont il est plus proche du centroïde que de tous les autres centroïdes. Le clustering est effectué via K-means.
Plus on a de cellules, et plus les cellules sont petites, donc plus la recherche sera rapide. En revanche, on perd en précision : si le vecteur le plus proche de notre requête n’est pas dans la même cellule (par exemple si on est un peu au bord d’une cellule), alors il ne sera pas retourné par la recherche.

Plus compact : la quantization (compression)

Si on ne souhaite pas stocker en mémoire l’intégralité de notre ensemble de vecteurs, on peut se tourner vers une solution de compression proposée par faiss : la product quantization.

On commence par découper nos vecteurs en segments, c’est-à-dire en sous-vecteurs d’une longueur donnée. Par exemple, dans un vecteur de taille 1000, on peut faire 10 segments de taille 100.
On effectue un clustering sur l’ensemble des premiers segments, puis sur l’ensemble des deuxièmes segments, etc… dans notre exemple, cela correspond à effectuer 10 clusterings différents, chacun sur un espace de dimension 100.
Pour stocker un vecteur dans notre index, on le découpe en segments, puis on calcule à quelle cellule appartient chaque segment (disons qu’on a choisi de faire 20 cellules). On peut ensuite remplacer chaque segment par l’ID de la cellule correspondante.
Dans notre exemple, on aura donc remplacé un vecteur de taille 1000 par un vecteur de taille 10, dont chaque coordonnée sera un nombre entre 1 et 20 (correspondant à une cellule).
On ne peut plus retrouver le vecteur d’origine, mais seulement une approximation, en recollant les 10 centroïdes (de taille 100) correspondant aux cellules.

Le processus de quantization illustré avec des vecteurs de dimension 16, et 4 segments. Les segments orange sont répartis en clusters, puis leurs coordonnées sont remplacées par les numéros de cluster dans le vecteur original. L’opération est ensuite répétée pour les autres segments.

La taille de l’index en mémoire est grandement réduite, mais la compression n’est pas sans perte d’information, et la précision de la recherche peut en souffrir : en théorie deux vecteurs peuvent se retrouver avec les mêmes combinaisons de centroïdes, et donc les mêmes coordonnées.

Le code :

Encore une fois, on initialise l’index puis on l’entraîne et on ajoute nos vecteurs. Cette opération peut prendre du temps sur de gros volumes de données.

Note : l’index choisi ici est un IVFPQ , c’est-à-dire un index inversé avec product quantization. Dans ce cas, le paramètre N_CELLS correspondra au nombre de cellules de Voronoï dans l’index inversé, et également au nombre de cellules utilisées dans la quantization. On pourra aussi jouer avec le paramètre nprobe pour la recherche.

Le paramètre quantizer
Le paramètre quantizer détermine comment on va effectuer la recherche parmi nos centroïdes. Ici, on choisit un index plat, c’est-à-dire qu’on va effectuer une recherche exhaustive.
Pour une centaine de centroïdes, c’est assez logique, mais lorsque les jeux de données deviennent conséquents, on va chercher à optimiser cette recherche. On aura donc deux indexes : un index pour les vecteurs, et un index pour les centroïdes.

Un autre paradigme : Hierarchical Navigable Small Worlds (HNSW)

Le fait de découper le jeu de données en cellules permet de se rapporter à un problème plus petit, mais fondamentalement, on effectue la même opération : une recherche exhaustive dans chaque cellule que l’on visite. Notre seul levier d’action en termes de vitesse est alors de diminuer le nombre de vecteurs auxquels on va comparer notre requête (soit en augmentant le nombre de cellules, qui seront alors plus petites, soit en diminuant le nombre de cellules visitées par requête) ; cette méthode a ses limites (notamment quand la dimension devient importante).

Un autre paradigme de recherche implémenté dans faiss permet de court-circuiter certains vecteurs en ne comparant notre requête qu’aux plus pertinents : le Hierarchical Navigable Small Worlds (HNSW).

Le HNSW est construit à partir du NSW (Navigable Small Worlds), qui fonctionne comme suit :

on construit un graphe des vecteurs (note : la construction est détaillée plus bas) ; on s’arrange pour que les paires de vecteurs assez proches soient reliées par une arête, mais aussi pour qu’on ait des arêtes un peu plus longues
on démarre notre recherche en comparant notre requête à un vecteur déterminé à l’avance : le point d’entrée
ensuite, on compare la requête aux vecteurs qui sont reliés à ce premier vecteur, et on choisit le plus proche
on continue jusqu’à ce qu’on ne puisse plus améliorer la distance, ou qu’on ait atteint un nombre pré-déterminé d’itérations

Ci-dessous une illustration de NSW : on entre dans le graphe par le point d’entrée (en forme d’étoile). À chaque étape, parmi les points reliés au nôtre, on choisit le point le plus proche de la requête (en rouge, en haut). On s’arrête quand on a atteint un minimum. Le chemin suivi est représenté par les lignes solides (en jaune), et les comparaisons effectuées en lignes pointillées. Le reste du graphe est représenté en filigrane pour ne pas surcharger l’image.

C’est un algorithme glouton, qui va donc renvoyer un minimum local. Pour qu’il soit aussi proche possible du vrai minimum (LE vecteur le plus proche de notre requête), il faut s’assurer que notre maillage est assez dense. Mais si le maillage est trop dense, on multiplie les comparaisons, et la complexité de la recherche. C’est là qu’entre en jeu la notion de graphes hiérarchiques.

Un graphe HNSW simplifié. La requête est symbolisée par le point rouge en haut à gauche. Le point d’entrée est le point le plus en haut à droite. On voyage le long des arêtes du graphe jusqu’à ce qu’on ne puisse plus améliorer la distance à la requête. Puis on descend d’un étage.

Imaginons que l’on découpe notre graphe en trois étages :

troisième étage : un graphe avec, disons, seulement un vecteur sur 100, pris au hasard, et reliés entre eux selon leur proximité, comme décrit précédemment.
deuxième étage : un graphe comprenant un vecteur sur 10 (les mêmes vecteurs qu’au troisième étage + d’autres vecteurs) ; encore une fois, on les relie entre eux selon leur proximité.
premier étage : un graphe comprenant tous les vecteurs.

Ainsi, quand on veut effectuer une recherche, on commence par naviguer au troisième étage. Une fois qu’on ne peut plus améliorer la distance, on “descend” dans le deuxième étage, et on affine notre résultat. Puis dans le premier étage.

Cette approche a deux avantages :

en moyenne, on prend moins de temps à trouver notre minimum local
on peut se permettre de relier des vecteurs très lointains dans les couches supérieures, sans risquer de trop densifier le maillage

En termes de code, pour générer un HNSW plat :

La taille en mémoire et le temps d’entraînement sont considérables : c’est un index qui est utile pour de gros volumes de données, mais pas vraiment pour un petit dataset (j’ai testé avec ~400k vecteurs, ce n’est pas concluant du tout).

La construction du graphe HSNW en détails
Pour construire le graphe HNSW, on procède par étapes.
1. On détermine (en fonction du dataset) combien de voisins M chaque point doit avoir dans le graphe.
2. Pour chaque niveau, on calcule une probabilité (qui dépend de M) : celle qu'un vecteur appartienne à ce niveau. Le niveau 0, le plus bas, contiendra tous les points (donc la probabilité est de 1). La probabilité décroit exponentiellement avec le niveau (donc plus le niveau est élevé et moins il contient de points).
3. Lorsqu’on insère un nouveau point p dans le graphe, on génère un nombre uniforme entre 0 et 1 pour déterminer le niveau maximum l_maxdans lequel ce point sera présent. Il sera présent dans ce niveau, et dans tous ceux en-dessous.
4. On insère le point p dans chaque niveau jusqu'à l_max, et on le connecte à chaque fois aux M points les plus proches dans ce niveau. Au niveau 0, on le connecte aux 2*M points les plus proches.
5. Petite subtilité : si un voisin de p (disons, q) a déjà un autre voisin r, et que r est plus proche de q que de p, on choisit de ne pas connecter p et r. Cela permet de ne pas multiplier les connexions entre des points très proches au sein d’un même cluster, et de privilégier des connexions avec des distances un peu moins courtes (illustration ci-dessous).

Illustration de la petite subtilité : le point p est plus proche de r que de s. Mais comme il est relié au point q, qui est déjà lui-même très proche de r, on le relie plutôt à s. Cela permet de diversifier la longueur des arêtes.

Combiner les indexes avec le quantizer

Il est possible de combiner HSNW et IVFPQ en se servant du premier comme d’un quantizer ; au moment de chercher les cellules les plus proches de notre requête, l’algorithme fera donc appel à HNSW plutôt que de faire une recherche exhaustive parmi les centroïdes.

Cela permet d’avoir une recherche efficace parmi nos centroïdes, et de bénéficier de la petite taille en mémoire de IVFPQ.

Mention spéciale : Locality Sensitive Hashing

Locality Sensitive Hashing est un autre algorithme de réduction de dimension dont la première version de date de 1999. L’idée est de construire une table de hachage (hash table) de notre jeu de données, en maximisant les collisions (c’est-à-dire, la probabilité que deux vecteurs se retrouvent dans la même alvéole) pour les vecteurs suffisamment similaires.

Cet algorithme souffre de deux défauts principaux :

pour obtenir des résultats suffisamment précis, il faut multiplier le nombre d’alvéoles, ce qui entraîne un surcoût en mémoire
les fonctions de hachage ne sont pas adaptables aux données et à leur distribution : par exemple, on aura autant d’alvéoles dans une région dense de l’espace que dans une région plus vide

La seule version de LSH supportée par faiss est le LSH binaire : un index plat avec des codes binaires. La documentation de faiss n’est pas très claire au sujet de l’algorithme de hachage utilisé ; ci-dessous, l’un des plus populaires, random projection, illustré ci-dessous.

Un LSH en deux dimensions avec quatre hyperplans (droites). Les chiffres du code sont déterminés par le côté de la droite duquel se situent les points. Tous les vecteurs qui sont dans la même “zone” ont le même code.

On produit aléatoirement k vecteurs de l’espace, qu’on va appeler v1, v2, … vk.
Pour chaque vecteur x de notre jeu de données, on construit un vecteur h, de taille k, de la manière suivante : chaque coordonnée est le signe du produit scalaire de x avec l’un des vecteurs aléatoires (1 si le produit est positif, 0 s’il est négatif). Donc, la première coordonnée est le signe de x.v1, la deuxième est le signe de x.v2, etc.

Le signe du produit scalaire indique “de quel côté” du vecteur aléatoire se situe x. Cela revient à découper notre espace en alvéoles avec des hyperplans aléatoires. Note : on ne stocke pas les vecteurs originaux dans l’index, et seulement leur projection.

Lorsque que l’on effectue une recherche, on commence par projeter notre requête sur les vecteurs aléatoires, pour construire le vecteur binaire correspondant, puis on retourne tous les vecteurs correspondant à l’alvéole atteinte. Notre recherche va donc renvoyer plusieurs candidats, que l’on n’a aucune manière d’ordonner (puisque l’on ne stocke pas les vecteurs originaux dans l’index). LSH est très rapide, mais très peu précis.Ci-dessous, le code pour LSH :

Conclusion

L’utilisation d’un index permet d’optimiser les performances d’une recherche de similarité lorsque la recherche exhaustive n’est pas envisageable car trop coûteuse.

Les trois principaux index proposés par faiss sont :

IVF (index inversé), qui répartit les vecteurs en cellules (clusters), et compare la requête aux vecteurs présents dans un ou plusieurs clusters ;
PQ (product quantization), qui compresse les vecteurs en les divisant en segments qui sont ensuite répartis en clusters ;
HNSW (hierarchical navigable small worlds), qui construit un graphe à plusieurs étages permettant de naviguer facilement d’un vecteur de la base à un autre.

Ci-dessous, un tableau récapitulatif des avantages et inconvénients de chacun : un + signifie un avantage, un — : un inconvénient et un — — : un gros inconvénient.

Dans un premier temps, il faudra entraîner l’index sur les données afin de les répartir en clusters, les compresser, ou de construire le graphe (c’est le “AI” de “faiss”). On devra également choisir des hyperparamètres, qui permettront de trouver le juste équilibre entre la précision des résultats renvoyés et les performances.

Enfin, on peut combiner les indexes en utilisant un quantizer afin, par exemple, d’optimiser la recherche parmi les centroïdes de nos clusters.

Remerciement

Merci à nos collègues Agnes GIRARDOT, Edouard LAVAUD, Lu WANG et Nhut DOAN NGUYEN pour la revue de l’article.

A propos de l’auteur

Après un doctorat en mathématiques fondamentales, Béatrice CHETARD a rejoint La Javaness en tant que data scientist en mars 2021. Elle s’intéresse à tous les aspects de la data science (données structurées, NLP, computer vision) ainsi qu’aux question d’éthique et de responsabilité des algorithmes.