FAISS: Exploring Approximate Nearest Neighbours Cell Probe Methods
Date: 2023-09-14 00:44:37
In this post, we'll learn how to do approximate nearest neighbours with FaceBook's FAISS vector search library.

:icons: font

I've been learning about vector search in recent weeks and I came across FaceBook's https://faiss.ai/[FAISS library^].
I wanted to learn the simplest way to do approximate nearest neighbours, and that's what we'll be exploring in this blog post.

[NOTE]
====
I've created a video showing how to do this on https://www.youtube.com/@learndatawithmark[my YouTube channel, Learn Data with Mark^], so if you prefer to consume content through that medium, I've embedded it below:

++++
<iframe width="560" height="315" src="https://www.youtube.com/embed/iY7HuG1r5YM?si=rdN5fF22veQVpvk_" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" allowfullscreen></iframe>
++++

You can also find all the code at the https://github.com/mneedham/LearnDataWithMark/blob/main/faiss-ann/notebooks/ANN-Tutorial.ipynb[ANN-Tutorial.ipynb notebook^].
====

First things first, let's install some libraries:

[source, bash]
----
pip install faiss-cpu pandas numpy
----

We'll be using the following imports:

[source, python]
----
import faiss
import copy
import numpy as np
import pandas as pd
import plotly.graph_objects as go
from plotly_functions import generate_distinct_colors, zoom_in, create_plot, plot_points
----

https://github.com/mneedham/LearnDataWithMark/blob/main/faiss-ann/notebooks/plotly_functions.py[`plotly_functions`^] contains a bunch of helper functions for making it easier to create charts with plot.ly.

== Create vectors

We're going to create 10,000 2D vectors to keep things simple.

[source, python]
----
dimensions = 2
number_of_vectors = 10_000
vectors = np.random.random((number_of_vectors, dimensions)).astype(np.float32)
----

Next, let's create a search vector, whose neighbours we're going to find:

[source, python]
----
search_vector = np.array([[0.5, 0.5]])
----

== Creating a cell probe index

The simplest version of approximate nearest neighbours in FAISS is to use one of the cell probe methods.
These methods partition the vector space into a configurable number of cells using the K-means algorithm.
When we look for our search vector's neighbours, it's going to find the centroid closest to the search vector and then search for all the other vectors that belong to the same cell as that centroid.

We can create a cell probe index like this:

[source, python]
----
cells = 10
quantizer = faiss.IndexFlatL2(dimensions)
index = faiss.IndexIVFFlat(quantizer, dimensions, cells)
----

To create the centroids, we need to call the `train` function:

[source, python]
----
index.train(vectors)
----

We can then find the centroids by querying the quantizer::

[source, python]
----
centroids = index.quantizer.reconstruct_n(0, index.nlist)
centroids
----

.Output
[source, text]
----
array([[0.8503718 , 0.46587527],
       [0.14201212, 0.80757564],
       [0.831061  , 0.82165515],
       [0.5756452 , 0.54481953],
       [0.5543639 , 0.1812697 ],
       [0.84584594, 0.16083847],
       [0.259557  , 0.5097532 ],
       [0.23731372, 0.12491277],
       [0.47171366, 0.8513159 ],
       [0.08305518, 0.30214617]], dtype=float32)
----

== Visualising cells and centroids

Next, let's look at how we can visualise how the vector space has been split.
We can work out which cell each vector has been assigned to by calling the `search` function on the quantizer:

[source, python]
----
_, cell_ids = index.quantizer.search(vectors, k=1)
cell_ids = cell_ids.flatten()
cell_ids[:10]
----

.Output
[source, python]
----
array([0, 4, 3, 1, 1, 8, 9, 4, 0, 9])
----

So far so good.
Now let's create a plot visualising that:

[source, python]
----
color_map = generate_distinct_colors(index.nlist) # <.>

fig_cells = create_plot()

unique_ids = np.unique(cell_ids)
for uid in unique_ids: # <.>
  mask = (cell_ids == uid)
  masked_vectors = vectors[mask]
  plot_points(fig_cells, masked_vectors, color_map[uid], "Cell {}".format(uid), size=6) # <.>

plot_points(fig_cells, centroids, symbol="diamond-tall", color="black", size=15, showlegend=False) # <.>
plot_points(fig_cells, search_vector, symbol="x", color="black", size=15, label="Search Vector") # <.>

fig_cells
----
<.> Get a list of unique colours for each cell
<.> Iterate over the cells
<.> Plot each vector with the colour assigned to its cell id
<.> Plot the centroid of each cell
<.> Plot the search vector

The resulting visualisation is shown below:

.Vectors and their cell assignments
image::{{<siteurl>}}/uploads/2023/09/ann-plot.png[]

When creating the index, we need to specify how many partitions (or cells) we want to divide the vector space into.


== Searching for our vector

It's time to search for our vector.
We'll start by adding the vectors to the index:

[source, python]
----
index.add(vectors)
----

And now let's call the `search` function:

[source, python]
----
distances, indices = index.search(search_vector, k=10)

df_ann = pd.DataFrame({
  "id": indices[0],
  "vector": [vectors[id] for id in indices[0]],
  "distance": distances[0],
})
df_ann
----

.df_ann
[format="csv", options="header"]
|===
include::content/2023/09/14/df_ann.csv[]
|===

We've got a bunch of vectors that are very close to the search vector.
When we ran the `search` function, FAISS first looked for the cell in which it needed to search.
We can figure out which cell it used by asking the quantizer:

[source, python]
----
_, search_vectors_cell_ids = index.quantizer.search(search_vector, k=1)
unique_searched_ids = search_vectors_cell_ids[0]
unique_searched_ids
----

.Output
[source, text]
----
array([3])
----

So the nearest cell to `0.5, 0.5` is the one with index 3.
If we wanted to find the nearest two cells, we could pass in a different `k` value.

We can visualise the nearest neighbours that it's found by running the following code:


[source, python]
----
fig_search = create_plot()

for uid in unique_searched_ids: # <.>
  mask = (cell_ids == uid)
  masked_vectors = vectors[mask]
  plot_points(fig_search, masked_vectors, color_map[uid], label="Cell {}".format(uid)) # <.>
  plot_points(fig_search, centroids[uid].reshape(1, -1), symbol="diamond-tall", color="black", size=10, label="Centroid for Cell {}".format(uid), showlegend=False) # <.>

plot_points(fig_search, points=search_vector, color='black', label="Search Vector", symbol="x", size=10)

ann_vectors = np.array(df_ann["vector"].tolist())
plot_points(fig_search, points=ann_vectors, color='black', label="Approx Nearest Neighbors") # <.>

fig_search
----
<.> Iterate over the cells used in the search (i.e. only cell with index=3)
<.> Plot the vectors in this cell
<.> Plot the centroid for the cell
<.> Plot the nearest neighbours

The resulting visualisation is shown below:

.Approximate nearest neighbours
image::{{<siteurl>}}/uploads/2023/09/ann-search-plot.png[]


== Brute Force vs ANN

It looks like ANN has done pretty well, but let's compare it to the brute force approach where we compare the search vector with every other vector to find its neighbours.
We can create a brute force index like this:

[source, python]
----
brute_force_index = faiss.IndexFlatL2(dimensions)
brute_force_index.add(vectors)
----

And then search like this:

[source, python]
----
distances, indices = brute_force_index.search(search_vector, k=10)

pd.DataFrame({
  "id": indices[0],
  "vector": [vectors[id] for id in indices[0]],
  "distance": distances[0],
  "cell": [cell_ids[id] for id in indices[0]]
})
----

.Brute Force
[format="csv", options="header"]
|===
include::content/2023/09/14/brute_force.csv[]
|===

The results are the same as we got with ANN and we can see that all the neighbours belong to cell 3, which was the one used by ANN.

We can actually tweak ANN to search across more than 1 cell by setting the `nprobe` attribute.
For example, if we wanted to search the two closest cells, we would do this:

[source,python]
----
index.nprobe = 2
----

And then re-run the search code above.
The result for this dataset wouldn't change since it's relatively small and has low dimensionality, but with bigger datasets this is a useful thing to play around with.