Speeding up vector recall 5x with HNSW

Silas Marvin

October 2, 2023

PostgresML makes it easy to use machine learning with your database and to scale workloads horizontally in our cloud. Our SDK makes it even easier.

Introducing HNSW

Underneath the hood our SDK utilizes pgvector to store, index, and recall vectors. Up until this point our SDK used IVFFlat indexing to divide vectors into lists, search a subset of those lists, and return the closest vector matches.

While the IVFFlat indexing method is fast, it is not as fast as HNSW. Thanks to the latest update of pgvector our SDK now utilizes HNSW indexing, creating multi-layer graphs instead of lists and removing the required training step IVFFlat imposed.

The results are not disappointing.

Comparing HNSW and IVFFlat

In one of our previous posts: Tuning vector recall while generating query embeddings in the database we were working on a dataset with over 5 million Amazon Movie Reviews, and after embedding the reviews, performed semantic similarity search to get the closest 5 reviews.

Let's run that query again:

                content_copy
            
WITH request AS (
  SELECT pgml.embed(
    'Alibaba-NLP/gte-base-en-v1.5',
    'query: Best 1980''s scifi movie'
  )::vector(1024) AS embedding
)

SELECT
  id,
  1 - (
    review_embedding_e5_large <=> (
      SELECT embedding FROM request
    )
  ) AS cosine_similarity
FROM pgml.amazon_us_reviews
ORDER BY review_embedding_e5_large <=> (SELECT embedding FROM request)
LIMIT 5;

This query utilized IVFFlat indexing and queried through over 5 million rows in 89.118ms. Pretty fast!

Let's drop our IVFFlat index and create an HNSW index.

                content_copy
            
DROP INDEX index_amazon_us_reviews_on_review_embedding_e5_large;
CREATE INDEX CONCURRENTLY ON pgml.amazon_us_reviews USING hnsw (review_embedding_e5_large vector_cosine_ops);

Now let's try the query again utilizing the new HNSW index we created.

                content_copy
            
WITH request AS (
  SELECT pgml.embed(
    'Alibaba-NLP/gte-base-en-v1.5',
    'query: Best 1980''s scifi movie'
  )::vector(1024) AS embedding
)

SELECT
  id,
  1 - (
    review_embedding_e5_large <=> (
      SELECT embedding FROM request
    )
  ) AS cosine_similarity
FROM pgml.amazon_us_reviews
ORDER BY review_embedding_e5_large <=> (SELECT embedding FROM request)
LIMIT 5;

Not only are the results better (the cosine_similarity is higher overall), but HNSW is over 5x faster, reducing our search and embedding time to 17.465ms.

This is a massive upgrade to the recall speed utilized by our SDK and greatly improves overall performance.

For a deeper dive into HNSW checkout Jonathan Katz's excellent article on HNSW in pgvector.