Billion-Scale Similarity Search with GPUs

By Jeff Johnson, Matthijs Douze, Herv'e J'egou. 2017

5,541 citations recorded by the source corpus.

Similarity search finds application in database systems handling complex data such as images or videos, which are typically represented by high-dimensional features and require specific indexing structures. This paper tackles the problem of better utilizing GPUs for this task. While GPUs excel at data parallel tasks such as distance computation, prior approaches in this domain are bottlenecked by algorithms that expose less parallelism, such as <inline-formula><tex-math notation="LaTeX">$k$</tex-math><alternatives><mml:math><mml:mi>k</mml:mi></mml:math><inline-graphic xlink:href="johnson-ieq1-2921572.gif"/></alternatives></inline-formula>-min selection, or make poor use of the memory hierarchy. We propose a novel design for <inline-formula><tex-math notation="LaTeX">$k$</tex-math><alternatives><mml:math><mml:mi>k</mml:mi></mml:math><inline-graphic xlink:href="johnson-ieq2-2921572.gif"/></alternatives></inline-formula>-selection. We apply it in different similarity search scenarios, by optimizing brute-force, approximate and compressed-domain search based on product quantization. In all these setups, we outperform the state of the art by large margins. Our implementation operates at up to 55 percent of theoretical peak performance, enabling a nearest neighbor implementation that is 8.5 × faster than prior GPU state of the art. It enables the construction of a high accuracy <inline-formula><tex-math notation="LaTeX">$k$</tex-math><alternatives><mml:math><mml:mi>k</mml:mi></mml:math><inline-graphic xlink:href="johnson-ieq3-2921572.gif"/></alternatives></inline-formula>-NN graph on 95 million images from the <sc>Yfcc100M</sc> dataset in 35 minutes, and of a graph connecting 1 billion vectors in less than 12 hours on 4 Maxwell Titan X GPUs. We have open-sourced our approach for the sake of comparison and reproducibility.

Subfields: Computer Vision and Pattern Recognition, Representation & Self-Supervised Learning, Evaluation & Benchmarks

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Vision and Pattern Recognition

In Representation & Self-Supervised Learning

In Evaluation & Benchmarks

Citations in the corpus

Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks (2020, methodology): “1 We use the DPR document encoder to compute document embeddings for each document, and we build a single MIPS index using FAISS [19] using Hierarchical Navigable Small World approximation for efﬁcient retrieval [33], which is then used for all experiments.”
Dense Passage Retrieval for Open-Domain Question Answering (2020, background): “TriviaQA (Trivia) (Joshi et al., 2017) contains a set of trivia questions with answers that were originally scraped from the Web.”
Distilling Knowledge from Reader to Retriever for Question Answering (2020, methodology): “These representations are then queried using approximate nearest neighbors (Johnson et al., 2019).”
Beyond English-Centric Multilingual Machine Translation (2020, methodology): “CCMatrix works on the large monolingual corpora in the 91 languages of CCNet (Wenzek et al., 2019), but at this scale, the global search is computationally demanding even with fast indexing from FAISS (Johnson et al., 2019).”
Augmented Language Models: a Survey (2023, methodology): “These algorithms trade off accuracy for speed, allowing significantly faster retrieval while maintaining reasonably accurate results, see Johnson et al. (2021) for further details.”
Galactica: A Large Language Model for Science (2022, methodology): “Retrieval is performed using a FAISS index (Johnson et al., 2019).”
Landmark Attention: Random-Access Infinite Context Length for Transformers (2023, methodology): “When computing the attention weights for the i -th token, landmark tokens for other blocks are placed in the same group as the i -th token.”
Entities as Experts: Sparse Memory Access with Entity Supervision (2020, background): “, 2019) and, while our current implementation does not yet include an efficient implementation of top-k routing (Section 2), we note that it is possible with fast maximum inner product search (Ram and Gray, 2012; Shen et al., 2015; Shrivastava and Li, 2014; Johnson et al., 2017).”
Generalization through Memorization: Nearest Neighbor Language Models (2019, methodology): “For W IKITEXT -103 experiments, we compute squared L 2 distances with full precision keys, but for the other datasets we use the FAISS L 2 distances (not squared) between quantized keys directly, for faster evaluation.”
Large Memory Layers with Product Keys (2019, methodology): “To retrieve key indices efﬁciently, we perform the search over sub-keys with a fast nearest neighbors implementation by Johnson et al. [21].”
ELI5: Long Form Question Answering (2019, methodology): “Speciﬁcally, we perform a nearest neighbor search (Johnson et al., 2017) over the average word embeddings of the question using FAST - TEXT (Bojanowski et al., 2017).”
Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval (2020, methodology): “Our ANN search uses the Faiss IndexFlatIP Index [16].”
ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT (2020, methodology): “Additionally, it enables ColBERT to leverage vector-similarity search indexes (e.g., [1, 16]) to retrieve the top-k results directly from a large document collection, substantially improving recall over models that only re-rank the output of term-based retrieval.”
Progressively Pretrained Dense Corpus Index for Open-Domain Question Answering (2020, methodology): “The MIPS-based retrieval is implemented with FAISS (Johnson et al., 2017).”
Real-Time Open-Domain Question Answering with Dense-Sparse Phrase Index (2019, methodology): “We use faiss (Johnson et al., 2017), open-sourced and large-scale-friendly similarity search package for dense vectors.”
Poly-encoders: Architectures and Pre-training Strategies for Fast and Accurate Multi-sentence Scoring (2019, methodology): “After the context embedding y ctxt is computed, the only operation remaining is a dot product between y ctxt and every candidate embedding, which can scale to millions of candidates on a modern GPU, and potentially billions using nearest-neighbor libraries such as FAISS (Johnson et al., 2019).”
Re2G: Retrieve, Rerank, Generate (2022, methodology): “After DPR Stage 1 training the passages from the corpus are indexed with a Hierarchical Navigable Small World (HNSW) [Malkov and Yashunin, 2018] using FAISS [Johnson et al., 2017].”
The Web Is Your Oyster - Knowledge-Intensive NLP against a Very Large Web Corpus (2021, methodology): “We use a ﬂat FAISS index for Wikipedia.”
Unsupervised Dense Information Retrieval with Contrastive Learning (2021, methodology): “Once a collection of documents, such as Wikipedia articles, is encoded, retrieval is performed with a fast k-nearest neighbors library such as FAISS Johnson et al. (2019).”
You Only Need One Model for Open-domain Question Answering (2021, methodology): “In practice, we retrieve top-N passages by indexing the pre-computed passage representations P 0 using Maximum Inner Product Search tools (MIPS) such as FAISS (Johnson et al., 2021).”