Practical Near Neighbor Search via Group Testing

1 Introduction . Nearest neighbor search is a fundamental problem with many applications in machine learning systems . Informally , the task is as follows . Given a dataset D = { x1 , x2 , ... xN } , we wish to build a data structure that can be queried with any point q to obtain a small set of points xi ∈ D that have high similarity ( low distance ) to the query . This structure is called an index . Near neighbor indices form the backbone of production models in recommendation systems , social networks , genomics , computer vision and many other application domains . Applications : In this paper , we focus on algorithms for approximate near neighbor search over highdimensional large scale datasets . Such tasks frequently arise in genomics , web-scale data mining , machine learning , and other large-scale applications . Consider the Yahoo Flickr Creative Commons dataset ( YFCC100M ) which consists of 100 million media embeddings that are derived from the neuron activations for a convolutional neural network [ 21 ] . Each embedding is a 4096-dimensional vector . The dataset is about 1TB in size and presents a substantial challenge for even the most popular algorithms , which struggle with memory , index construction , and query time . Similar issues occur ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . in genomics , where the task is to identify genome sequences with a high Jaccard similarity to the query . Modern genomic datasets can contain millions of reads with billions of possible n-gram sequences [ 17 ] . Many algorithms work well when there are a few hundred dimensions but are inappropriate for such applications . Our experiments demonstrate that for the datasets of interest in this paper , popular indices like HNSW and FAISS can take days to build , require gigabytes of RAM and have a suboptimal precision-recall-query time tradeoff . Since our goal is to perform approximate search , dimensionality reduction is a reasonable strategy . However , dimensionality reduction is costly for ultra-high dimensional data . In genomics applications , n-gram sizes are typically very large ( n > 18 ) . Thus , the one-hot encoding of each sequence can require billions of dimensions ( 418 ≈ 68B ) , making it intractable to learn an embedding model . For embedding applications such as YFCC100M or product search , a large embedding dimension can lead to performance improvements [ 15 ] . Dimensionality reduction can incur a performance penalty , so we may wish to perform the near neighbor search over the original metric space . Ideally , we would choose an algorithm that did not store data points in RAM , evaluate the distance function many times , employ iterative and non-streaming processes such as k-means , or construct complicated structures such as graphs , which are hard to parallelize and distribute . Recent algorithms such as FLASH [ 22 ] provide the ability to search based on aggregate LSH count statistics without computing distances , but these methods are heuristics that do not have theoretical guarantees . On the other hand , algorithms such as LSH , which have a well-established theoretical grounding , tend to perform poorly in practice because of their prohibitive hash table size and post filtering stage , which needs many ( Nρ in theory ) distance computations . In this paper , we present an algorithm having all the practical advantages of a system like FLASH while also being more accurate , theoretically sound , and provably sub-linear . 1.1 Our Contribution . In this paper , we address the computational challenges of high-dimensional similarity search by presenting an index with fast construction time , low memory requirement , and zero query-time distance computations . Our approach is to transform a near neighbor search problem into a group testing problem by designing a test that outputs “ positive ” when a group of points contains a near neighbor to a query . That is , each test answers an approximate membership query over its group . Given a query , our algorithm produces a B × R array of group test results that can be efficiently decoded to identify the nearest neighbors . This is more efficient than statistical aggregation algorithms like FLASH because each test filters out entire groups of non-neighbors with a single test operation . We develop a concrete example of such an algorithm by using Filters to Identify Near Neighbor Groups ( FLINNG ) . We use a standard non-adaptive group testing design with distance-sensitive Bloom filters as tests . We prove that FLINNG solves the randomized nearest neighbor problem in O ( log2 ( 1δ ) log 3 ( N ) N 1 2 +γ ) time , where γ is a query-dependent parameter that characterizes query stability . We also implement FLINNG in C++ and conduct experiments on real-world highdimensional datasets from genomics , embedding search , and URL analysis , where FLINNG achieves up to a 10x query speedup over existing indices with faster construction time and lower memory . 2 Related Work . The near neighbor problem has been the focus of more than four decades of intense research activity . The low dimensional problem is particularly well-understood , with space partitioning trees that can efficiently find the exact k-nearest neighbors . However , exact search in high dimensions is intractable due to the curse of dimensionality - the computational resources needed to solve the exact problem scale exponentially with dimensions . This has led to a diverse set of algorithms to solve the approximate near neighbor problem , which we now describe . Locality Sensitive Hashing : LSH was the first approximate near neighbour algorithm to break the curse of dimensionality . At their core , LSH algorithms use an LSH function to partition the dataset into buckets . The hash function is selected so that the distance between points in the same bucket is likely to be small . To find the near neighbors of a query , we hash the query and compute the distance to every point in the corresponding bucket . Query performance can be improved with replication , which queries multiple independent hash tables , multi-probe methods [ 13 ] , which examine multiple buckets in the hash table , and data-dependent LSH [ 3 ] , which tunes the hash function to the dataset . Recent work shows that machine learning algorithms can also construct effective LSH partitions [ 6 ] . Count-Based LSH : There are several recent algorithms which identify neighbors by counting the number of LSH collisions rather than explicitly computing distances . For example , the algorithm from [ 22 ] uses the count values to quickly identify potential neighbors . The algorithm from [ 5 ] applies compressed sensing techniques to the counts to compress the dataset , and a popular technique in genomics is to simply replace each data point with its hash values [ 17 ] . Graphs : Graph-based methods are another successful family of algorithms . Graph algorithms locate near neighbors by walking the edges of a graph where each point is ( approximately ) connected to its k nearest neighbors . The focus in this area has been to improve graph properties using diversification , pruning , hierarchical structures , and other heuristics [ 14 ] . Graph indices perform well on industrystandard benchmarks but are not theoretically well-understood , despite recent progress [ 18 ] . Graph indices also suffer from long construction times and bloated memory consumption . Sample Compression : A large number of practical methods are based on quantization . Such methods replace points in the dataset with compressed versions of the points . Methods such as scalar quantization , vector quantization and product quantization alias each point to a collection of k-means centroids . One can also use machine learning to obtain learned Hamming codes for the dataset and perform efficient distance computations using bit operations . Advances in quantization are applicable to most other algorithms , but have been particularly effective when combined with brute force search on GPU hardware and partition-based search over billion-scale datasets [ 11 ] . Group Testing : We are not the first algorithm to apply group testing to near neighbor search . However , existing algorithms have key limitations that prevent effective practical implementations and rigorous theoretical analysis . The authors of [ 10 ] propose a group-based filtering algorithm based on group representative vectors , or the vector average of group entries . To query the index , [ 10 ] explicitly compute the distances to all points where the distance between the representative and query exceeds a threshold . The algorithm of [ 20 ] uses the same group representatives , but applies an online backpropagation algorithm to estimate the individual similarities . This work has two shortcomings . First , the methods require many distance calculations against the group representatives ( [ 10 ] requires N/10 distances ) , resulting in poor query time . Second , the average vector can be similar to the query even when all points are far from the query , precluding a theoretical analysis except under restrictive distribution assumptions . In this work , we analyze methods where only the group tests are used to identify the neighbors , as our goal is to avoid performing expensive distance computations . Unsurprisingly , our method is theoretically and practically superior . Theoretical Comparison : In high dimensions , most methods require conditions that depend on the data and query to prove efficient search time . Table 1 shows some representative theoretical results for the classes of algorithms described previously . 3 Background . Formal Problem Statement : In this paper , we solve the randomized nearest neighbor problem . Definition 1 is a stronger version of the well-studied ( R , c ) -approximate near neighbor problem . In particular , any algorithm which solves the randomized nearest neighbor problem also solves the approximate near neighbor problem with c = 1 and any R ≥ the distance to the nearest neighbor . Definition 1 . Randomized Nearest neighbor : Given a dataset D , a distance metric d ( · , · ) and a failure probability δ ∈ [ 0 , 1 ] , construct a data structure which , given a query point y , reports the point x ∈ D with the smallest distance d ( x , y ) with probability greater than 1− δ . 3.1 Group Testing . Suppose we are given a set D of N items , k of which are positive ( “ hits ” ) and N − k of which are negative ( “ misses ” ) . The group testing problem is to identify the hits by grouping items and using a small collection of group tests . A group test is positive if and only if any item from the group is positive . The objective of group testing is to reliably identify the positive items using fewer than N group tests . The problem is noisy if the tests make i.i.d . mistakes with some false positive and false negative rate . The group testing problem may also be adaptive , where we are allowed to design test n based on the results of tests { 1 , 2 , ... n− 1 } , or non-adaptive , where we must perform all tests at once . Since the problem ’ s introduction in 1943 , there has been considerable work toward the construction of test designs under various constraints . For a recent review , see [ 1 ] . In this paper we develop near neighbor search algorithms using the noisy group testing framework . For simplicity , we mainly consider the doubly regular design , where we evenly distribute items among B tests , and we independently repeat this process R times to obtain a B×R grid of group tests ( Figure 1 ) . This block testing design is similar to that of RAMBO [ 8 ] , but with different filtering and inference algorithms . However , our algorithmic framework is compatible with any non-adaptive testing design .

Summary: The paper tackles the problem of near-neighbor (and nearest-neighbor) search in high-dimensional spaces. The authors propose to use group-testing on top of distance-sensitive Bloom filters: to simplify, the returned set is the intersection of the union of small sets, each of these small sets being a hash bucket that the query falls into. The authors derive the theoretical properties of their method FLINNG (sublinear query time), and show the superiority of their method compared to the state of the art (FAISS, HNSW, FLINNG) through extensive experiments on genome, URL and embedding data.