Learning Space Partitions for Nearest Neighbor Search
Authors: Yihe Dong
Piotr Indyk
Ilya Razenshteyn
Tal Wagner
Venue:  (2020)
Keywords: space partition
lsh
locality sensitive hashing
nearest neighbor search

Abstract: Space partitions of $\mathbb{R}^d$ underlie a vast and important
class of fast nearest neighbor search (NNS) algorithms. Inspired by recent theoretical work on NNS for general metric spaces (Andoni et al. 2018b,c), we develop a new framework for building space partitions reducing the problem to balanced graph partitioning followed by supervised classification.
We instantiate this general approach with the KaHIP graph partitioner (Sanders and Schulz 2013) and neural networks, respectively, to obtain a new partitioning procedure called Neural Locality-Sensitive Hashing (Neural LSH). On several standard benchmarks for NNS (Aumuller et al. 2017), our experiments show that the partitions obtained by Neural LSH consistently outperform partitions found by quantization-based and tree-based methods as well as classic, data-oblivious LSH.

Published as a conference paper at ICLR 2020
LEARNING SPACE PARTITIONS FOR NEAREST NEIGHBOR SEARCH
Yihe Dong' Microsoft
Piotr Indyk MIT
llya Razenshteyn Microsoft Research
Tal Wagner MIT
ABSTRACT
Space partitions of I2 underlie vast and important class of fast nearest neighbor search (NNS) algorithms. Inspired by recent theoretical work on NNS for general metric spaces (Andoni et al., 2018b;c), we develop a new framework for building space partitions reducing the problem to balanced graph partitioning followed by supervised classification. We instantiate this general approach with the KaHIP graph partitioner (Sanders & Schulz, 2013) and neural networks, respectively, to obtain a new partitioning procedure called Neural Locality- Seaiiii---esiivvveeee Hashing (Neural LSH). On several standard benchmarks for NNS (Aumüller et al., 2017), our experiments show that the partitions obtained by Neural LSH consistently outperform partitions found by quantization- and tree- beseb methods as well as classic, data-oblivious LSH.
1 Introduction
The Neurest Neighbor Search (NNS) problem is defined as follows. Given an n-point dataset P in d-dimensional Fuclidean space R2, we would like to preprocess P to answer k-nearest neighbor queries quickly. That is, given query point q E Rdd, we want to find the k data points from P that are closest to q. NNS is cornerstone of the modern data analysis and, at the same time, a fundamental geometric data structure problem that led to many exciting theoretical developments over the past decades. See, e.g., Wang al. (2016); Andoni et al. (201 8a) for an overview.
The main two approaches to constructing efficient NNS data structures are indexing and sketching. The goal of indexing is to construct data structure that, given query point, produces small subset of P (called candidate set) that includes the desired neighbors. Such a data structure can be stored on a single machine, or (if the data set is very large) distributed among multiple machines. In contrast, the goal of sketching is compute compressed representations of points to enable computing approximate distances quickly (c.g., compact binary hash codes with the Hamming distance used as an estimator, see the surveys Wang al. (2014; 2016)). Indexing and sketching can be (and often are) combined to maximize the overall performance (Wu et al., 2017; Johnson et al., 2017).
Both indexing and sketching have been the topic of vast amount of theoretical and empirical literature. In this work, we consider the indexing problem. In particular, we focus on indexing based on space partitions. The overarching idea is to build a partition of the ambient space Rd and split the dataset P accordingly. Given a query point q. we identify the bin containing and form the resulting list of candidates from the data points residing in the same bin (or, to boost the accuracy, nearby bins as well). Some of the popular space partitioning methods include locality-sensitive hashing (LSH) (Lv et al., 2007; Andoni al., 2015; Dasgupta et al., 2017); quantization-base are obtained via k-means clustering of the dataset (Jégou et al., 2011: I; Babenko & Lempitsky, 2012); and tree-based methods such as aandom-projectin trees or PCA trees (Sproull, 1991; Bawa et al., 2005; Dasgupta & Sinha, 2013; Keivani & Sinha, 2018).
approaches, where partitions
Compared to other indexing methods, space partitions have multiple benefits. First, they are naturally applicable in distributed settings, as different bins can be stored on different machines (Bahmani et al., 2012; Ni et al., 2017; Li et al., 2017; Bhaskara & Wijewardena, 2018). Moverover, the computational efficiency of search can be further improved by using any nearest neighbor search algorithm locally on each machine. Second, partition-based indexing is particularly suitable for GPUs due the simple and predictable memory access pattern (Johnson et al., 2017). Finally, partitions can be combined with cryptographic techniques to yield efficient secure similarity search algorithms (Chen et al., 2019). Thus, in this paper we focus on designing space partitions that optimize the trade-off between their key metrics: the number of reported candidates, the fraction of the true nearest neighbors among the candidates, the number of bins, and the computational efficiency of the point location.
Author names are ordered alphabetically.

---

Published as a conference paper at ICLR 2020
Recently, there has been a large body of work that studies how modern machine learning techniques (such as neural networks) can help tackle various classic algorithmic problems (a partial list includes Mousavi et al. (2015); Baldassarre et al. (2016); Bora et al. (2017); Dai et al. (2017); Metzler et al. (2017); Kraska et al. (2018); Balcan et al. (2018); Lykouris & Vassilvitskii (2018); Mitzenmacher (2018); Purohit et al. (2018)). Similar methods under the name "leam to hash"-have been used to improve the sketching approach to NNS (Wang et al., 2016). However, when it comes to indexing, while some unsupervised techniques such as PCA or k-means have been successfully applied, the full power of modern tools like neural networks has not yet been harnessed. This state of affairs naturally leads to the following general question: Can we employ modern (supervised) machine learning techniques to find good space partitions for nearest neighbor search?
1.1 Our contribution
In this paper we address the aforementioned challenge and present a new framework for finding high-quality space partitions of RRd Our approach consists of three major steps:
1. Build the K-NN graph G of the dataset by connecting each data point nearest neighbors;
2. Find a balanced partition P of the graph G' into m parts of nearly-equal size such that the number of edges between different parts is as small as possible;
3. Obtain a partition of TRdd by training a classifier on the data points with labels being the parts of the partition P found in the second step.
See Figure 1 for illustration. The new algorithm directly optimizes the performance of the partition-based nearest neighbor data structure. Indeed, if a query is chosen as a uniformly random data point, then the average k-NN accuracy is exactly equal to the fraction of edges of the k--N graph G' whose endpoints separated by the partition P. This generalizes to out-of-sample queries provided that the query and dataset distributions are close, and the test accuracy of the trained classifier is high.
At the same time, our approach is directly related to and inspired by recent theoretical work (Andoni et al., 2018b;c) on NNS for general metric spaces. In particular, using the framework of (Andoni al., 2018b;e), we prove that, under mild conditions on the dataset P, the k-NN graph of P can be partitioned with a hyperplane into two parts of comparable size such that only few edges get split by the hyperplane. This gives a partial theoretical justification of our method.
The new framework is very flexible and uses partitioning and learning in a black-box way. This allows us to plug various models (linear models, neural networks, etc.) and explore the trade-off between the quality and the algorithmic efficiency of the resulting partitions. We emphasize the importance of balanced partitions for the indexing problem, where all bins contain roughly the same number of data points. This property crucial in the distributed setting, since we naturally would like to assign a similar number of points to each machine. Furthermore, balanced partitions allow tighter control of the number of candidates simply by varying the number of retrieved parts. Note that a priori, it is unclear how to partition Idd so as to induce balanced bins of given dataset. Here the combinatorial portion of our approach is particularly useful, as balanced graph partitioning is a well-studied problem, and our supervised extension to Rd naturally preserves the balance by virtue of attaining high training accuracy.
We speculate that the new method might be potentially useful for solving the NNS problem for non-Eucliin Euclidean metrics, such as the edit distance (Zhang & Zhang, 2017) optimal transport distance (Kusner et al., 2015). Indeed, for any metric space, one can compute the k-NN graph and then partition it. The only step that needs to be adjusted to the specific metric at hand is the learning step.
Let us finally put forward the challenge of scaling our method up to billion-sizei or even larger datasets. For such scale, one needs build an approximate k-NN graph as well as using graph partitioning algorithms that are faster than KaHIP. We leave this exciting direction to future work. For the current experiments (datasets of size 10 points), preprocessing takes several hours. Another important challenge is obtain NNS algorithms based on the above partitioning with provable guarantees in terms of approximation and running time. However, we expect it to be difficult, in particular, since all the current state-of-ther NNS algorithms lack such guarantees (e.g., k-means-bart based (Jégou et al., 2011) or graph methods (Malkov & Yashunin, 2018 8), see also (Aumüller et al., 2017) for a recent SOTA survey).
Evaluation We instantiate our framework with the KaHIP algorithm (Sanders Schulz, 2013) for the graph partitioning step, and either linear models or small-size neural networks for the learning step. We evaluate it on several standard benchmarks for NNS (Aumüller et al., 2017) and conclude that in terms of quality of the resulting partitions, it consistently outperforms quantization-bast based and tree-based partitioning procedures, while maintaining comparable
2

---

Published as a conference paper at ICLR 2020
(a) Datasct
(b) K-NN graph together with balanced partition
(c) Learned partition
Figure 1: Stages of our framework
algorithmic efficiency. the high accuracy regime, our framework yields partitions that lead to processing up to 2.3x fewer candidates than the strongest baseline.
As a baseline method we use k-means clustering (Jégou et al., 2011). It produces a partition of the dataset into bins, in a way that naturally extends to all of Rd, by assigning a query point g to its nearest centroid. (More generally, for multi-probe querying, we can rank the bins by the distance of their centroids to q). This simple scheme yields very high-quality results for indexing. Besides k-means, evaluate LSH (Andoni al., 2015). ITQ (Gong et al., 2013), PCA tree (Sproull, 1991), RP tree (Dasgupta & Sinha, 2013), and Neural Catalyzer (Sablayrolles et al., 2019).
1.2
Related work
On the empirical side, currently the fastest indexing techniques for the NNS problem are graph-based (Malkov & Yashunin, 2018). The high-level idea is to construct a graph the dataset (it can be the k-NN graph, but other constructions are also possible), and then for each query perform a walk, which eventually converges to the nearest neighbor. Although very fast, graph-based approaches have suboptimal "locality of reference", which makes them less suitable for several modern architectures. For instance, this is the case when the algorithm is run on GPU (Johnson et al., 2017), or when the data is stored in external memory (Sun et al., 2014) or distributed manner (Bahmani al., 2012; Ni et al.. 2017). Moreover, graph-baser based indexing requires many rounds of adaptive access to the dataset, whereas partition- based indexing accesses the dataset in one shot. This is crucial, for example, for nearest neighbor search over encrypted data (Chen et al., 2019). These benefits justify further study of partition-base based methods.
Machine learning techniques are particularly useful for the sketching approach, leading to a vast body of research under the label "learning to hash" (Wang et al., 2014; 2016). In particular, several recent works employed neural networks to obtain high-quality sketches (Liong et al., 2015; Sablayrolles et al., 2019). The fundamental difference from our work is that sketching is designed speed up linear scans over the dataset, by reducing the cost of distance evaluation, while indexing is designed for sublinear time searches, by reducing the number of distance evaluations. We note that while sketches not designed for indexing, they can be used for that purpose, since b-bit hashing scheme induces a partition of Rd into 2 parts. Nonetheless, our experiments show that partitions induced by these methods (such as Iterative Quantization (Gong et al., 2013)) not well-suited for indexing, and underperform compared to quantization--bsed indexing, as well as to our methods.
We highlight in particular the recent work of Sablayrolles et al. (2019), which uses neural networks learn mapping f: Idd ndd that improves the geometry of the dataset and the queries facilitate subsequent sketching. is natural to ask whether the same family of maps can be applied to enhance the quality of partitions for indexing. However, as our experiments show, in the high accuracy regime the maps learned using the algorithm of Sablayrolles al. (2019) consistently degrade the quality of partitions.
Finally, we mention that here some prior work on learning space partitions: Cayton & Dasgupta (2007); Ram & Gray (2013); Li et al. (2011). However, all these algorithms learn hyperplane partitions into two parts (then applying them recursively). Our method, on the other hand, much more flexible, since neural networks allow us to learn a much richer class of partitions.

---

Published as a conference paper at ICLR 2020
2 Our method
Given a dataset P C Rd of points, and a number of bins m > 0, our goal is to find a partition R. of Rd into m bins with the following properties:
1. Balanced: The number of data points in each bin is not much larger than n/m..
2. Locality sensitive: For a typical query point q IRd most of its nearest neighbors belong to the same bin of R. We assume that queries and data points come from similar distributions.
3. Simple: The partition should admit a compact description and, moreover, the point location process should be computationally efficient. For example, we might look for space partition induced by hyperplanes.
Formally, we want the partition R that minimizes the loss By (peNi (R()) s.t. VEE [R(p)] - (1 + ()/mm)) where q is sampled from the query distribution, (g) C P is the set of its k nearest neighbors in P, 17 >> 0 is a balance parameter, and R(p) denotes the part of R that contains
First, suppose that the query is chosen as uniformly random data point, q ~ P. Let G be the N-NN graph of P, whose vertices are the data points, and each vertex is connected to its k nearest neighbors. Then the above problem boils down to partitioning vertices of the graph G into m bins such that each bin contains roughly n/m vertices, and the number of edges crossing between different bins is as small as possible (see Figure 1(b)). This balanced graph partitioning problem extremely well-studied, and there are available combinatorial partitioning solvers that produce very high-quality solutions. In our implementation, we use the open-source solver KaHIP (Sanders & Schulz, 2013), which based on a sophisticated local search.
More generally, we need to handle out-of-sample queries, i.e., which are not contained P. Let R. denote the partition of G (equivalently, of the dataset P) found by the graph partitioner. To convert R. into a solution to our problem, we need to extend it to a partition R of the whole space Rd that would work well for query points. In order to accomplish this, we train a model that, given a query point q Rd predicts which of the m bins of R the point q belongs to (see Figure I()). We use the dataset P as a training set, and the partition R as the labels - i.e., each data point is labeled with the ID of the bin of R containing it. The method is summarized in Algorithm 1. The geometric intuition for this learning step is that - even though the partition R is obtained by combinatorial means, and in principle might consist of ill-behaved subsets of Rd - in most practical scenarios, we actually expect be close to being induced by a simple partition of the ambient space. For example, if the dataset is fairly well- distrribtd on the unit sphere, and the number of bins m 2, balanced cut of G should be close to hyperplane.
The choice of model to train depends on the desired properties of the output partition R. For instance, if we are interested in a hyperplane partition, we can train a linear model using SVM or regression. In this paper, we instantiate the learning step with both linear models and small-sized neural networks. Here, there is natural tension between the size of the model we train and the accuracy of the resulting classifier, and hence the quality of the partition we produce. A larger model yields better NNS accuracy, at the expense of computational efficiency. We discuss this in Section 3.
Multi patti-probe querying Given a query point q, the trained model can be used to assign to a bin of a partition R, and search for nearest neighbors within the data points in that part. In order to achieve high search accuracy, we actually train the model to predict several bins for a given query point, which are likely to contain nearest neighbors. For neural networks, this can be done naturally by taking several largest outputs of the last layer. By searching through more bins (in the order of preference predicted by the model) we can achieve better accuracy, allowing for a trade-off between computational resources and accuracy.
Hierarchical partitions When the required number of bins m is large, in order to improve the efficiency of the resulting partition, it pays off to produce it in hierarchical manner. Namely, we first find a partition of Rd into m, bins, then recursively partition each of the bins into m2 bins, and so on, repeating the partitioning for I, levels. The total number of bins in the overall partition is m mi m2- ... m.,. See Figure 2 for illustration. The advantage of such a hierarchical partition is that it is much simpler to navigate than a one-shot partition with m bins.
Neural LSH with soft labels In the primary instantiation of our framework, we set the supervised learning component to a neural network with a small number of layers and constrained hidden dimensions (the exact parameters are specified in the next section). In order to support effective multi multi-probe querying, we need to infer not just the bin that contains the query point, but rather distribution over bins that are likely to contain this point and its neighbors. A T T-probe candidate list is then formed from all data points in the most likely bins. In order to accomplish this, we use soft labels for data points generated as follows. For S > and a data point p. the soft label P = (pl. P2. .......P. Pm)) is

---

Published as a conference paper at ICLR 2020
R3
RS
B
P
Pb
Pu
1
P
P
Figure Hierarchical partition into 9 bins with my = m2 = 3. R2's are partitions, P,'s are the bins the dataset. Multi-probe query procedure, which descends into bins, may visit the bins marked in bold.
Preprocessing
Input: Dataset C Rd, integer parameter > 0, number of bins m > 0
1: Build a k-NN graph G of P.
2: Run a balanced graph partitioning algorithm on G into m parts. Number the parts arbitrarily as 1,..., m. Let
(p) {I,..., m denote the part containing p, for every p P.
3: Train a machine learning model M with training set P and labels {p)))pEP. For every a E IR4, let (() e
mm. denote the prediction of M on ø.
M(-) defines our m-way partition of Rdd. Note that it is possible that xp) * M(p) for some p P. if M attains imperfect training accuracy.
Query
Input: query point q E IRd, number of bins search
1: Run inference on M to compute (q).
2: Search for a near neighbor of in the bin (q). i.e., among the candidates {p E P M(p) M(g))-
If M furthermore predicts distribution over bins, search for neighbor in the b top-ranked bins according to the ranking induced by the distribution (i.e., from the most likely bin to less likely ones).
Algorithm 1: Nearest neighbor search with learned space partition
distribution over the bin containing point chosen uniformly at random among S nearest neighbors of p (including p itself). Now, for predicted distribution = (gt, (2 ......, pm ). seek minimize the KL divergence between P and @: 2Th pi log Pi, Intuitively, soft labels help guide the neural network with information about multiple bin ranking. S' is hyperparameter that needs to be tuned; we study its setting in the appendix (cf. Figure 6b).
3 Sparse hyperplane-induced cuts in k:-N graphs
We state and prove theorem that shows, under certain mild assumptions, that the k-NN graph of dataset P C I&d can be partitioned by hyperplane such that the induced cut is sparse (i.c., has few crossing edges while the sizes of two parts are similar). The theorem based on the framework of (Andoni al., 201 2018b;c and uses spectral techniques.
We start with some notation. Let (2(p) be the set of k nearest neighbors of in P. The degree of p in the k-NN graph is dog(p) N& U (p' P p (p))) Let be the distribution over the dataset P. where point p Pis sampled with probability proportional to its degree deg(p). Let Delose be the distribution over pairs (p.p') where p is uniformly random, and p' is uniformly random element of Ng (p). Denote a (pprr' Detue [|lp pp'ppp'lm [[]] and Dann>/Daag-pllllll - P211)) 113] We will proceed assuming that a (typical distance between a data point and its nearest neighbors) is noticeably smaller than 3 (typical distance between two independent data points).
) E P x P.
The following theorem implies, informally speaking, that a & 3, then there exists a hyperplane which splits the dataset into two parts of not too different size while separating only few pairs of (v, p'), where p' is one of the k nearest neighbors of p. For the proof of the theorem, see Appendix C.
5

---

Published as a conference paper at ICLR 2020
Theorem 3.1. There exists a hyperplane H = fi E R² (a, = by such that the following holds. Les P P= P-PPPP-
be the partition of P induced by H: Pi = {p E PI (a,p) S b). P2 - {v E P (a.p) > b3. Then, one has:
Pr(p.r') |p and pf are separated by H]
2a
min(Prp-Dly Pil, Pip=Dlp E Pb]] - V 3
(1)
4 Experiments
Datasets For the experimental evaluation, we use three standard ANN benchmarks (Aumüller et al., 2017): SIFT (image descriptors, 1M 128- 18-dimensional points), GloVe (word embeddings (Pennington et al., 2014), approximately 1.2M 100-dimensional points, normalized), and MNIST (images of digits, 60K 784-dimensional points), All three datasets come with 10 000 query points, which are used for evaluation. We include the results for SIFT and GloVe in the main text, and MNIST in Appendix A.
Evaluation metries We mainly investigate the trade-off' between the number of candidates generated for query point, and the k-NN accuracy, defined as the fraction of its k nearest neighbors that are among those candidates. The number of candidates determines the processing time of individual query. Over the entire query set, report both the average as well as the 0.95-th quantile of the number of candidates. The former measures the throughput of the data structure, while the latter measures its latency. 2 We focus on parameter regimes that yield k-NN accuracy of at least 0.75, in the setting k- = 10. Additional results with broader regimes of accuracy and of k are included in the appendix.
Our methods We evaluate two variants of our method, with two different choices of the supervised learning component:
Neural LSH: In this variant we use small neural networks. We compare this method with k-means clustering, Iterative Quantization (ITQ) (Gong et al., 2013). Cross-polytope LSH (Andoni et al., 2015), and Neural Cat- alyzer (Sablayrolles al., 2019) composed over k-means clustering. We evaluate partitions into 16 bins and 256 bins. We test both one-level (non-bierarchical) and two-level (hierarchical) partitions. Queries are multi-probe. result,
Regression LSH: This variant uses logistic regression the supervised learning component and, produces very simple partitions induced by hyperplanes. We compare this method with PCA trees (Sproull, 1991; Kumar et al., 2008; Abdullah et al., 2014), random projection trees (Dasgupta & Sinha, 2013), and recursive bisections using 2-means clustering. We build trees of hierarchical bisections of depth up to 10 (thus total number of leaves up to 1024). The query procedure descends single root-to-leaf path and returns the candidates in that leaf.
4.1 Implementation details
Neural LSH uses fixed neural network architecture for the top-level partition, and fixed architecture for all second- level partitions. Both architectures consist of several blocks, where each block is a fully-connected layer + batch normalization (Ioffe & Szegedy, 2015) ReLU activations. The final block is followed by a fully-connected layer and softmax layer. The resulting network predicts distribution over the bins of the partition. The only difference between the top-level network the second-leve network architecture is their number of blocks (b) and the size of their hidden layers (s). In the tep-level network we use b = 3 and 512. In the second-level networks we use b = 2 and s = 390. To reduce overfitting, we use dropout with probability 0.1 during training. The networks are trained using the Adam optimizer (Kingma& & Ba, 2015) for under 20 epochs on both levels. We reduce the learning rate multiplicatively at regular intervals. The weights are initialized with Glorot initialization (Glorot & Bengio, 2010). To tune soft labels, we try different values of S between 1 and 120.
We evaluate two settings for the number of bins in each level, m = 16 and m = 256 (leading to a total number of 162 = 256 and 256 65 536, respectively). In the
bins of the total number of bins in the two-level experiments two-level setting with m = 256 the bottom level of Neural LSH uses k-means instead of a neural network, to avoid overfitting when the number of points per bin is tiny. The other configurations (two-levels with m 16 and one-level with either m = 16 or m 256) we use Neural LSH at all levels.
We slightly modify the KaHIP partitioner make it more efficient on the k-NN graphs. Namely, introduce hard threshold of 2000 on the number of iterations for the local search part of the algorithm, which speeds up the partitioning dramatically, while barely affecting the quality of the resulting partitions.
'Nubeer queries per second. Maximum time per query, modulo small fraction of outliers.
6

---

Published as a conference paper at ICLR 2020
GloVe Ve
SIFT
Averages
0.95-quantiles 2.125
Averages
0.95-quantilss
One level
16 bins
1.745
1,031
1.240
256 bins
1.491
1.752
1.047
1.348
Two levels
16 bins
2.176
2.308
1.113
1.306
256 bins
1.241
1.154
1.182
1.192
Figure 3: Largest ratio between the number of candidates for Neural LSH and k-means over the settings where both attain the same target 10-NN accuracy, over accuracies of at least 0.85. See details in Section 4.2.
4.2 Comparison with multi-bin methods
Figure 4 shows the empirical comparison of Neural LSH with k-means clustering, ITQ, Cross-polytope LSH, and Neural Catalyzer composed over k-means clustering. It turns out that k-means is the strongest among these baselines." ? The points depicted in Figure are those that attain accuracy 0.75. In the appendix (Figure 10) we include the full accuracy range for all methods.
In all settings considered, Neural LSH yields consistently better partitions than k-means.* Depending on the setting, k-means requires significantly more candidates to achieve the same accuracy:
Up to 117% more for the average number of candidates for GloVe;
Up to 130% more for the 0.95-quantiles of candidates for GloVe;
Up to 18% more for the average number of candidates for SIFT;
Up to 34% more for the 0.95- -.95-quantiles of candidates for SIFT;
Figure lists the largest multiplicative advantage in the number of candidates of Neural LSH compared to k-means, for accuracy values of at least 0.85. Specifically, for every configuration of k-means, we compute the ratio between the number of candidates in that configuration and the number of candidates Neural LSH in its optimal configuration, among those that attained at least the same accuracy as that k-means configuration.
We also note that in all settings except two-level partitioning with m = 256,5 Neural LSH produces partitions for which the 0.95-quantiles for the number of candidates are very close to the average number of candidates, which indicates very little variance between query times over different query points. contrast, the respective gap the partitions produced by k-means is much larger, since unlike Neural LSH, it does not directly favor balanced partitions. This implies that Neural LSH might be particularly suitable for latency-criticee NNS applications.
Model sizes. The largest model size learned by Neural LSH is equivalent to storing about 5700 points for SIFT, or se 7100 points for Ve. is considerably larger than k-means with & < 256, which stores at most 256 points. Nonetheless, we believe the larger model size is acceptable for Neural LSH, for the following reasons. First, in most of the NNS applications, especially for the distributed setting, the bottleneck in the high accuracy regime is the memory accesses needed to retrieve candidates and the further processing (such distance computations, exact or approximate). The model size is not hindrance as long as does not exceed certain reasonable limits (e.g.. it should fit into a CPU cache). Neural LSH significantly reduces the memory access cost, while increasing the model size by an acceptable amount. Second, we have observed that the quality of the Neural LSH partitions is not too sensitive to decreasing the sizes the hidden layers. The model sizes we report are, for the sake of concreteness, the largest ones that still lead to improved performance. Larger models do not increase the accuracy, and sometimes decrease it due to overfitting.
'It is important to note that ITQ is not designed produce space partitions; as explained in Section 1, it does side-effect. Simiarly, Neural Catalyzer is not designed to enhance partitions. The comparison intended to show that they do not outperform indexing techniques despite being outside their intended application.
*We note that two-level partitioning with = 256 is the best performing contiguration of k-means, for both SIFT and GloVe, in terms of the minimum number of candidates that attains 0.9 accuracy. Thus evaluate this baseline its optimal performance. 'As mentioned earlier, in this setting Neural LSH uses k-means at the second level, due to the large overall number of bins compared to the size of the datasets. This explains why the gap between the average and the quantile number of candidates of Neural LSH larger for this setting.
7

---

Published as conference paper at ICLR 2020
0.95
0.95
0.9
0.9
0.85
0.8
0.75
0.75
50000
150000
250000
350000
450000
550000
50000
70030
110000 130000
-50000
170000
210000 230000 250000
Ve, one level, 16 bins
(b) SIFT. level, 16 bins
0.95
0.95
0.85
0.75
0.75
20000
120000 140000 160000 180003
40000
55000
(c) GloVe, one level, 256 bins
(d) SIFT, one level, 256 bins
0.95
0.8
0.75
075
(e) Glo Ve, two levels, 16 bins
(f) SIFT, two levels, bins
0.95
Neural average)
Neurai LSH (o.95 3535qunntile
0.89
0.85
m-mans average)
k-meen (0.95 (3.55quantile))
ISH averagel
LSH 10, 10.30-quantlie)
0.8
TTQ average)
TOQ (0.85 quantile)
Catalyzer-kk-men (average) Cutalyzer-i-mesns 10 (0.5--uanttle)
2000
4000
6000
8000
10000
12000
24000 16003 18000 20000
(g) GloVe, two levels, 256 bins, k-means at 2nd level
(h) SIFT, two levels, 256 bins. k-means at 2nd level
Figure 4: Comparison of Neural LSH with baselines; x-axis is the number of candidates, y-axis is the 10-NN accuracy
8

---

Published as conference paper at ICLR 2020
SIFT
GloVe
Regrendion LSI jnveraget
0.0
Regrwadin LSM (3-55- 1885-quantile
2 meanns (average)
2 Zimeans 30 qaaddie)
0.7
PGA tree Eaveragei
PCA (5:55 quantik))
tree (wverage)
0.5
RP quuantaan
Learned (D [average)
Learned (D-tre (0--quunttl)))
03
03
fracted Search
Search (5-quantia))
190000
200009
300000
400000
500000
600030
700000
Figure 5: Comparison of decision trees built from hyperplanes: x-axis number of candidates, y-axis 10-NN accuracy
4.3 Comparison with tree-based methods
Next compare binary decision trees, where in each tree node hyperplane is used to determine which of the two subtrees to descend into. We generate hyperplanes with the following methods: Regression LSH, the Learned KD-tree of Cayton & Dasgupta (2007), the Boosted Search Forest of Li et al. (2011), cutting the dataset into two equal halves along the top PCA direction (Sproull, 1991; Kumar et al., 2008), 2-means clustering, and random projections of the centered dataset (Dasgupta & Sinha, 2013; Keivani & Sinha, 2018). We build trees of depth up 10, which correspond to hierarchical partitions with the up to 20 1024 bins. Results for GloVe and SIFT are summarized in Figure 5 (see appendix). For random projections, we run each configuration 30 times and average the results.
For Ge, Regression LSH significantly outperforms 2-means, while for SIFT, Regression LSH essentially matches 2-means in terms of the average number of candidates, but shows a noticeable advantage terms of the -95-peecss In both instances, Regression LSH significantly outperforms PCA tree, and all of the above methods dramatically improve upon random projections.
Note, however, that random projections have an additional benefit: in order to boost search accuracy, one can simply repeat the sampling process several times and generate an ensemble of decision trees instead of single tree. This allows making each individual tree relatively deep, which decreases the overall number of candidates, trading space for query time. Other considered approaches (Regression LSH, 2-means, PCA tree) are inherently deterministic, and boosting their accuracy requires more care: for instance, one can use partitioning into blocks as in Jégou et al. (201 or alternative approaches like Keivani Sinha (2018). Since we focus on individual partitions and not ensembles, we leave this issue out of the scope.
4.4 Additional experiments
In this section we include several additional experiments.
First, study the effect of setting k. We evaluate the 50-NN accuracy of Neural LSH when the partitioning step is run on either the 10-NN or the 50-NN graph. We compare both algorithms to k-means with 50. Figure 6a compares these three algorithms on GloVe for 16 bins reporting average numbers of candidates. From this plot, can that for k = 50, Neural LSH convincingly outperforms k-means, and whether we use 10-NN or 50-NN graph matters very little.
Second, we study the effect of varying S (the soft labels parameter) for Neural LSH on GloVe for 256 bins. See Figure 6b where we report the average number of candidates. As we can see from the plot, the setting S = 15 yields much better results compared to the vanilla case of S = 1. However, increasing S beyond 15 brings diminishing returns on the overall accuracy.
'Neural LSH can solve N-NNS by partitioning the K'NN graph, for any k,, k/: they do not have to be equal.
9

---

Published as conference paper at ICLR 2020
0.9
0.85
Neural LSH, 10-NN graph
5
0.7
S0-NN graph
k-mmaass
5
0.80
100000
200000
300000
400000
500000
40000
50000
120000
160000
GloVe, one level, 16 bins, 50-NN accuracy using 10-NN and 50-NN graphs
(b) Glo level, 256 bins, varying S
Figure 6: Effect of various hyperparameters
5 Conclusions and future directions
We presented a new technique for finding partitions of Rd which support high-performance indexing for sublinear-time NNS. It proceeds in major steps: (1) We perform combinatorial balanced partitioning of the k-NN graph of the dataset; (2) We extend the resulting partition to the whole ambient space I by using supervised classification (such as logistic regression, neural networks, etc.). Our experiments show that the new approach consistently outperforms quantization-bazt based and tree-based partitions. There is number of exciting open problems we would like to highlight:
Can we use our approach for NNS over non-Euclidean geometries, such the edit distance (Zhang & Zhang. 2017) or the optimal transport distance (Kusner et al., 2015)? The graph partitioning step directly carries through, but the learning step may need to be adjusted.
Can we jointly optimize graph partition and classifier the same time? By making the two components aware of each other, we expect the quality of the resulting partition of Rd to improve. A related approach has been successfully applied in Li al. (2011) for hyperplane tree partitions.
Can our approach be extended to learning several high-quality partitions that complement each other? Such an ensemble might be useful to trade query time for memory usage (Andoni al., 2017).
Can we use machine learning techniques improve graph-bused indexing techniques (Malkov Yashunin, 2018) for NNS? (This is in contrast to peatition-bbsed indexing, as done in this work).
Our framework is example of combinatorial tools aiding "continuous" learning techniques. A more open-ended question is whether other problems can benefit from such symbiosis.
References
Amirali Abdullah, Alexandr Andoni, Ravindran Kannan, and Robert Krauthgamer. Spectral approaches nearest neighbor search. arXiv preprint arXiv:1408.00 07514000007,, 2014.
Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya Razenshteyn, and Ludwig Schmidt. Practical and optimal Ish for angular distance. Advances in Neural Information Processing Systems, pp. 1225-1233, 2015.
Alexandr Andoni, Thijs Laarhoven, Ilya Razenshteyn, and Erik Waingarten. Optimal hashing-based time-space trade-offs for approximate Discrete Algoritiuns, pp. 47-66. Society for Industrial and Applied Mathematics, 2017.
neighbors. In Proceedings of the Eivent-Eiggttl Annual ACM-SIAM Symposium
Alexandr Andoni, Piotr Indyk, and Ilya Razenshteyn. Approximate nearest neighbor search in high dimensions. arXiv preprint 18Xiv:1806.00 09823, 2018a.
Alexandr Andoni, Assaf Naor, Aleksandar Nikolov, Ilya Razenshteyn, and Erik Waingarten. Data-dependent hashing via nonlinear spectral gaps. Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing, pp. 787-800, 2018b.
10

---

Published as a conference paper at ICLR 2020
Alexandr Andoni, Assaf Naor, Aleksandar Nikolov, Ilya Razenshteyn, and Erik Waingarten. Hölder homeomorphisms and approximate nearest neighbors. In 2018 IEEE 59th Annual Symposium on Foundations of Computer Science (FOCS), pp. 159-169. IEEE, 2018c.
Martin Aumüller, Erik Bernhardsson, and Alexander Faithfull. Ann-benchmarks: A benchmarking tool for approximate nearest neighbor algorithms. In International Conference on Similarity Search and Applications, pp. 34 49. Springer, 2017.
Artem Babenko and Victor Lempitsky. The inverted multi-index. In Computer Vision and Pattern Recognition (CVPR), 2012 IEEE Conference on, pp. 3069-3076. IEEE, 2012.
Bahman Bahmani, Ashish Goel, and Rajendra Shinde. Efficient distributed locality sensitive hashing. In Proceedings of the 21st ACM international conference on Information and knowledge management, pp. 2174-2178. ACM, 2012.
Maria-Florina Balcan, Travis Dick. Tuomas Sandholm, and Ellen Vitercik. Learning to branch. In International Conference on Machine Learning, 2018
Luca Baldassarre, Yen-Huan Li, Jonathan Scarlett, Baran Gözeü, Ilija Bogunovic, and Volkan Cevher. Learning-based compressive subsampling. IEEE Journal of Selected Topics in Signal Processing, 10(4):809-822, 2016.
Mayank Bawa, Tyson Condie, and Prasanna Ganesan. Lsh forest: self-tuning indexes for similarity search. Proceedings of the 14th international conference World Wide Web, pp. 651-660. ACM, 2005.
Aditya Bhaskara and Maheshakya Wijewardena. Distributed clustering via Ish based data partitioning. In International Conference on Machine Learning, pp. 569-578, 2018.
Ashish Bora, Ajil Jalal, Eric Price, and Alexandros G Dimakis. Compressed sensing using generative models. In International Conference on Machine Learning, pp. 537-546, 2017.
Lawrence Cayton and Sanjoy Dasgupta. A learning framework for nearest neighbor search. In Advances in Neural Information Processing Systems, pp. 233-240, 2007.
Hao Chen, llaria Chillotti, Yihe Dong, Oxana Poburinnaya, Ilya Razenshteyn, and M Sadegh Riazi. Sanns: Scaling up secure approximate k-nearest neighbors search. arXiv preprint wiXiv:194.02003: 2019.
Fan RK Chung. Laplacians of graphs and cheegers incqualities. Combinatorics, Paul Erdos is Eighty, 23157-122):1--2 1996.
Hanjum Dai, Elias Khalil, Yuyu Zhang, Bistra Dilkina, and Le Song. Learning combinatorial optimization algorithms over graphs. In Advances in Neural Information Processing Systems, pp. 6351-6366 2017.
Sanjoy Dasgupta and Kaushik Sinha. Randomized partition trees for exact nearest neighbor search. In Conference on Learning Theory, pp. 317-337, 2013.
Sanjoy Dasgupta, Charles F Stevens, and Saket Navlakha. A neural algorithm for a fundamental computing problem. Science, 358(636))7933-99 2017.
Xavier Glorot and Yoshua Bengio.
Understanding the difficulty of training deep feedforward neural networks. In International Conference on Artificial Intelligence and Statistics, pp. 249-256, 2010.
Yunchao Gong, Svetlana Lazebnik, Albert Gordo, and Florent Perronnin. Iterative quantization: A procrustean approach
learning binary codes for large-scale image retrieval. JEEE Transactions on Pattern Analysis and Machine Intelligence, 35(12):29 16-2299, 2013.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprini 15Xiv:1502.03167 2015.
Herve Jégou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-128, 2011.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus. arXiv preprint arXiv:/702.08734. 2017.
Omid Keivani and Kaushik Sinha. Improved nearest neighbor search using auxiliary information and priority functions. In International Conference Machine Learning, pp. 2578-2580 2018

---

Published as a conference paper at ICLR 2020
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference for Learning Representations, 2015.
Tim Kraska, Alex Beutel, Ed H Chi, Jeffrey Dean, and Neoklis Polyzotis. The case for learned index structures. In Proceedings of the 2018 International Conference on Management of Data, pp. 489-504. ACM, 2018.
Neeraj Kumar, Li Zhang, and Shree Nayar. What is a good nearest neighbors algorithm for finding similar patches in images? In European conference on computer vision, pp. 378. Springer, 2008
Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger. From word embeddings to document distances. In International Conference on Machine Learning, pp. 957-966, 2015.
Jinfeng Li, James Cheng, Fan Yang, Yuzhen Huang, Yunjian Zhao, Xiao Yan. and Ruihao Zhao. Losha: A general framework for scalable locality sensitive hashing. In Proceedings of the 40th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 635-644--6 644. ACM, 2017.
Zhen Li, Huazhong Ning, Liangliang Cao, Tong Zhang. Yihong Gong, and Thomas S Huang. Learning to search efficiently in high dimensions. In Advances in Neural Information Processing Systems, pp. 1710-1718, 2011.
Venice Erin Liong, Jiwen Lu, Gang Wang, Pierre Moulin, and Jie Zhou. Deep hashing for compact binary codes learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2475-2483, 2015.
Qin Lv, William Josephson, Zhe Wang, Moses Charikar, and Kai Li. Multi-probe Ish: efficient indexing for high- dimensional similarity search. In Proceedings of the 33rd international conference Very large data bases, pp. 950-961. VLDB Endowment, 2007.
Thodoris Lykouris and Sergei Vassilvitskii. Competitive caching with machine learned advice. In International Conference on Machine Learning, 2018
Yury A Malkov and Dmitry A Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence, 2018.
Chris Metzler, Ali Mousavi, and Richard Baraniuk. Learned d-amp: image recovery. In Advances in Neural Information Processing Systems, pp. 1772-1783, 2017.
Principled neural network based compressive
Michael Mitzenmacher. A model for learned bloom filters and optimizing by sandwiching. In Advances in Neural Information Processing Systems, 2018.
Ali Mousavi, Ankit B Patel, and Richard G Baraniuk. A deep learning approach to structured signal recovery. In Communication, Control, and Computing (Allerton). 2015 53rd Annual Allerton Conference on, pp. 1336-1343. IEEE, 2015.
Y Ni, K Chu, and J Bradley. Detecting abuse at scale: Locality sensitive hashing at uber engineering, 2017.
Jeffrey Pennington, Richard Socher, and Christopher Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), pp. 1532-1543, 2014.
Manish Purohit, Zoya Svitkina, and Ravi Kumar. Improving online algorithms via ml predictions. In Advances in Neural Information Processing Systems, pp. 9661-9670, 2018.
Parikshit Ram and Alexander Gray. Which space partitioning to use for search? In Advances in Neural Information Processing Systems, pp. 656-664, 2013.
Alexandre Sablayrolles, Matthijs Douze, Cordelia Schmid, and Herve Jégou. Spreading vectors for similarity search. In International Conference on Learning Representations, 2019. URL
http:/ewwwwwwwweeeeeeeeeeeeeeeee net/forun?
Sd--SGGG2R5tmm.
Peter Sanders and Christian Schulz. Think Locally, Act Globally: Highly Balanced Graph Partitioning. In Proceedingy of the 12th International Symposium on Experimental Algorithms (SEA' 13), volume 7933 of LNCS, pp. 164-175. Springer, 2013.
Robert Sproull. Refinements to neireot-n searching ink-dimenssoonal trees. Algorithmica, 6(1-6):579-589 1991.
12

---

Published as conference paper at ICLR 2020
Yifang Sun, Wei Wang, Jianbin Qin, Ying Zhang, and Xuemin Lin. Sis: solving c-approximate nearest neighbor queries in high dimensional euclidean space with a tiny index. Proceedings of the VLDB Endowment, 8(1):1-12, 2014.
Jingdong Wang. Heng Tao Shen, Jingkuan Song, and Jianqiu Ji. Hashing for similarity search: A survey. arXiv preprint arXiv: 1408.9927 2014.
Jun Wang. Wei Liu, Sanjiv Kumar, and Shih-Fu Chang. Learning to hash for indexing big data a survey. Proceedings of the IEEE, 104(1):34-57, 2016.
Xiang Wu, Ruiqi Guo, Ananda Theertha Suresh, Sanjiv Kumar, Daniel N Holtmann Rice, David Simcha, and Felix Yu. Multiscale quantization for fast similarity search. In Advances in Neural Information Processing Systems, pp. 5745-5755, 2017.
Haoyu Zhang and Qin Zhang. Embedjoin: Efficient edit similarity joins via embeddings. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 585-594, ACM, 2017.
2017.
A Results for MNIST
We include experimental results for the MNIST dataset, where all the experiments are performed exactly the same way for SIFT and GloVe. Consistent with the trend observed for SIFT and GloVe, Neural LSH consistently outperforms k-means (see Figure 7) both in terms of average number of candidates and especially in terms of the 0.95-th quantiles. We also compare Regression LSH with recursive 2-means, well as PCA tree and random projections (see Figure 8), where Regression LSH consistently outperforms the other methods.
1.00
0.90
0.90
(average)
(average)
0.85
Neural LSH (0.95 quantiit)le)
0.85
Neural LSH (0.95 (--utuilellle
k-means (average)
k-means (average)
k-means (0.95-quantile)
k-means (0. (0.95-quantile)
0.80
0.80
(a) One level, 16 bins
(b) Two levels, 16 bins
Figure 7: MNIST, comparison of Neural LSH with k-means; x-axis number of candidates, y-axis 10-NN accuracy
2-maans 43S-qunntlle)
(average)
PCA tree (average, 0.35-quantile)
Random projections (0.95-quse) (0.95-quantile)
Random projections (averago)
Regression LSH (0.95 (0.95-quantile)
Regression (average)
0.2
Figure 8: MNIST, comparison of trees built from hyperplanes; x-axis number of candidates, y-axis 10-NN accuracy
13

---

Published as a conference paper at ICLR 2020
B Effect of Neural Catalyzer on Space Partitions
In this section we compare vanilla k-means with k-means run after applying a Neural Catalyzer map (Sablayrolles et al., 2019). The goal is to check whether the Neural Catalyzer which is designed boost up the performance of sketching methods for NNS by adjusting the input geometry could also improve the quality of space partitions for NNS. See Figure 9 for the comparison on GloVe and SIFT with 16 bins. On both datasets (especially SIFT), Neural Catalyzer in fact degrades the quality of the partitions. We observed similar trend for other numbers of bins than the setting reported here. These findings support our observation that while both indexing and sketching for NNS can benefit from learning-basee based enhancements, they are fundamentally different approaches and require different specialized techniques.
0.95
0.95
0.90
0.90
0.85
k--means (average)
k-means (average)
k-means (0.95-quantile
k-means (0.95-quantile))
0.80
Catalyzer- k-means (average)
Catalyzer k-means (average)
0.80
Catalyzer k-means (0.95 quanniiuat
* Catalyzer k-means (0.95-quantile)
0.75
200000
400000
600000
200000
400000
600000
Glo one level, 16 bins
(b) SIFT, level, 16 bins
Figure 9: Comparison of k-means and Catalyzer + k-means
C Proof of Theorem 3.1
Proof. Consider an undirected graph G = (V, F), where the set of vertices V P, and the (multi-)set of edges contains an edge (p.p') for every p' E (p). The graph contains n vertices and kn edges, and some of the edges might be double (if p' (p) and p E N& (p') at the same time). Let Ag be the symmetric adjacency matrix of G normalized by 2kn (so that the sum of all the entries equals to 1. thus giving a probability distribution over P x P. which can be seen to be equal Delose). The rows and columns of AG can naturally be indexed by the points of P. Denote Pa (p) = Ep (pr (Ag) p.pp- (AG)p.p' is immediate to check that pa yields a distribution over P. which can be seen to be equal to D. Denote DG diag(PG). Denote LG DG Ag the Laplacian of Ag. Due the equivalence of pG and D and Ag and D Delose. we have:
a Epwer(Ac)ps' IIp p'12
3
= Eprep PC --l
(2)
By considering all possible coordinate projections and using additivity of || 11 over coordinates, we conclude that there exists coordinate 1" [d] such that:
Epver(Ac)po' (pi- p/- a E pa(P)po(p') (pi- p/.) 3
(3)
Define a vector RRP by U/ Pr. We now apply the following standard fact from spectral graph theory: If A is the weighted adjacency matrix of a graph, and L is its Laplacian matrix, then Lx = Σij=1 Atj - æz)² for all R", Thus the numerator of (3) becomes y Lay. For the denominator, consider the graph H on P in which every pair p. p' is connected by an edge of weight pc (p)po((p))
Its weighted adjacency matrix An is given by (An) pa(p)po(p') for + p/, and with on the diagonal. Thus A# Pop' D& (recall that DG diag(po)).
Pa (v') = pc(p))(pa(p)) (recall that Epep Po = 1). Therefore the diagonal degree matrix of H is DH DG DG.
The degree of each node p in H equals pa (p)
14

---

Published as conference paper at ICLR 2020
1
0.9
0.8
0.7
0.
Noural LSH Dogg
..lle.r (0.95 quantilc)
(aog)
quantte)
0.6
0.5
(Iva)
quanttleee
lavg)
00 quantile)
0.5
0.4
Catalysarekk-mm
Levereed RCS-SSH (eag)
00 (u-uie
50000
250000
450000
550000
850000
1090000
50000
100000
150000
200000
250000
300000
350000
400000
450003
500000
(a) Ve, one level, 16 bins
(b) SIFT. level, 16 bins
0.9
0.8
0.85
0.7
0.8
0.6
0.75
0.5
0.7
100000
200000
300000
400000
500000
60C000
700003
800000
50000
100000
150000
200000
250000
300000
(c) GloVe, one level, 256
(d) SIFT, one level. 256 bins
0.8
0.6
0.6
0.5
0.4
0.4
100000 150000 200000 250000 300000 350000 400000 450000 500000
20000
40000
60000
80000
100000 120000 140000 160003 180000 200000
(e) GloVe, two levels, 16 bins
(I) SIFT, two levels, 16 bins
0.8
0.7
0.7
Alayya
03
0.6
0.5
0.8
0.
0.4
20000
30000
40000
50000
70000
80000
GloVe, two levels, 256 bins, k- m-maans at 2nd level
(h) SIFT, levels, 256 bins, k-means 2nd level
Figure 10: Results from Figure with broader candidate and accuracy regimes. The "Learned RCS-LSH" baseline is the learned rectilinear cell structure locality sensitive hashing method of Cayton & Dasgupta (2007).
15

---

Published as conference paper at ICLR 2020
Together, the Lapacian of H is LH (D - popp)yy. Overall, we have:
Dμ Ah
DG PaPP'
Therefore the denominator of (3) becomes
y' LGy $ 3
y PGPC)y
Next, we define y = y - c. 1, where is the all-1's vector, and c is the scalar c c= (v'pa)/(1ppc) This scalar is chosen to render y 1. Pa- Furthermore, since is in the kernel of every Laplacian matrix, have LGğ Ley and LHÿ = Luy. Together, we get
Laÿ y' Day
y LGY
D' - paPC)y
3
Now by the Cheeger's inequality (Chung, 1996), conclude that there exists threshold yo R such that:
Epi Suo No (Ac)p1.2 min Epp Suo Pa(P), Ep pG(D) S v22
D Lou
/2a
S
V3
(4)
1 Day
One trace back all the definitions and observe that the {p E P: Dp S yo is induced by (axis-aligned) hyperplane, and the left-hand side of (4) is nothing else but the left- left-hand side of (1).
16


Reviews:
[{"review_id": "r1eWb4FrFS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a technique for partitioning (presumably large number of) points in a d-dimensional, such that approximate nearest neighbor search in this space can be efficiently performed by reducing to search in a smaller partition. To this end, the paper provides the following procedure:\n1. Build an exact k nearest neighbor graph of the provided dataset.\n2. Partition this graph into m balanced sets while minimize inter-partition edges.\n3. Train a neural net or linear model to predict the partition for each point. They also propose creating hierarchical partitions in the same way and also propose training with \"soft\" partition labels.\n\nAt test time, the trained model is used to predict probably partitions which are then searched over for nearest neighbors. The authors compare this technique to, most notably, a k-means based technique where centroids are used to partition the space. They show that on Glove and SIFT datasets, the neural partitioning approach is more accurate i.e. it can correctly identify correct nearest neighbors significantly more than k-means.\n\n######\nPros:\n1. This proposed approach of using neural techniques to build partitions of space for fast nearest neighbors is novel. While k-means-style techniques have been proposed to learn index structures, this is the first application of neural techniques as far as I can tell.\n2. The paper is well-written and is easy to read. The experiments are sufficiently detailed and well documented.\n\nCons:\n1. Experiments: The paper reports the number of correct neighbors obtained while searching on average and at 0.95 quantile. However, the core motivation of the paper is *fast* nearest neighbors search. In that vein, the metric of interest must incorporate time e.g. queries per second at a fixed level of accuracy. This is not new -- e.g. the benchmarks provided by Andoni and co-authors are evaluated on a QPS basis (https://github.com/erikbern/ann-benchmarks). The same holds for Amueller et al. \n2. Baselines: A related point is, what happens to baseline algorithms when they are given more \"capacity\". E.g. comparing with k-means baselines with a significantly large value of k or a multi-level k-means algorithm (e.g. considered by \"FAST APPROXIMATE NEAREST NEIGHBORS WITH AUTOMATIC ALGORITHM CONFIGURATION\" by Muje and Lowe). The comparison in the paper with k-means baseline with k=50 does not seem fair here. E.g. a hierarchical neural LSH approach with 256 partitions at top and 10 k means partitions at the second level performs 2560 distance computations. Alternatively, the authors can directly compare QPS as stated above.\n3. The proof provided in the appendix is not entirely clear. E.g. after step 3, how does the next step (before (4)) follow? It is easy to see the denominator but the numerator is not clear.\n\n\nSuggestions:\n1. Please cite the following paper which is highly related and highly cited: FAST APPROXIMATE NEAREST NEIGHBORS WITH AUTOMATIC ALGORITHM CONFIGURATION.\n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "S1gSN8yaKr", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a learnable space partitioning nearest-neighbour algorithm where the model learns a hierarchical space partition of the data (using graph partitioning) and learns to predict which buckets a query will reside in (using a learnable model).\n\nThe paper is an interesting combination of classical optimization and deep learning.\n\nThe authors make a concerted effort to compare to a range of space-partitioning-based nearest neighbour algorithms, however crucially none of these algorithms are state-of-the-art in comparison to some of the best PQ NNS or hierarchical Navigable Small World Graphs. \n\nIn terms of writing, the model description requires significant overhaul. I had to re-read it several times and sketch out what exactly you are doing. I would recommend an algorithm box to summarize both the hierarchical graph partitioning *and* the learnable query network. \n\n4.3 \"additional experiments\" is some of the most interesting parts of the paper, I would recommend promoting these experimental results to the main paper.\n\nI would recommend accepting this paper because I think it may lead to state-of-the-art NNS algorithms with theoretical guarantees, especially with a more powerful query network.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "HyeSOPl2tr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nThe paper proposes a scheme to learn space partitions for improved\nnearest neighbor search by first converting the search problem into a\nsupervised classification problem by graph-partitioning the\nk-nearest-neighbor graph, and then using some machine learning model\nto learn space partitions corresponding to the supervised problem. The\nscheme is theoretically motivated and the empirical results\ndemonstrates the utility of the proposed scheme against various\nbaselines.  \n\n\nWhile the paper presents a promising new direction for\nnearest-neighbor search with space partitioning schemes, I am somewhat\non the border, leaning towards reject. This is mostly due to my lack\nof understanding why Cayton & Dasgupta (2007) and Li et al. (2011) are\nnot valid baselines and are dismissed as \"hyperplane\npartitions\". It would be very helpful to me as a reader to understand\nwhy this literature does not warrant to be a valid\nbaseline. Especially given the following connection points:\n\n- Firstly, the theoretical results presented by the authors also\n  limit to hyperplane splits. \n- Secondly, the authors present empirical evaluation of recursive\n  hyperplane partitions where the proposed scheme is the only learning\n  based scheme while the techniques presented in Cayton & Dasgupta\n  (2007) and Li et al. (2011) would have been the learning based\n  partitioning baseline to beat.\n- Thirdly, LSH space partitioning is considered as a baseline in\n  Section 4.2 with the Neural LSH and Neural Catalyzer + k-means being\n  the only learning based baseline; Cayton & Dasgupta (2007)'s learned\n  LSH scheme would have been a valid baseline to the best of my\n  understanding. \n- Finally, the authors mention the need for coupling the graph\n  partitioning and the supervised learning phase. Li et al., 2011\n  utilizes a loss function that essentially promotes a sparse graph\n  cut while preserving near-neighborhoods, providing the closely\n  coupled scheme. In the original paper, the authors utilize a\n  recursive hyperplane partitioning scheme. However, a neural network\n  with a softmax layer can be trained to minimize this closely coupled\n  loss directly, obviating the need for the separate graph\n  partitioning and supervised learning phases proposed in this\n  paper. But it is possible that I have misinterpreted the connection\n  and any clarification here would be very helpful.\n\nMinor:\n\n- It would be very helpful as a reader to understand how sensitive the\n  neural network based partioning scheme is to the network\n  architecture and hyperparameters. Some intuition behind the choices\n  would be very useful.\n- It is a little unclear how the plots in Figure 2 are generated --\n  for each query, we would have access to the accuracy vs. number of\n  candidates curve. In that case, it is unclear how the curves are\n  aggregated across queries to generate average number of candidates\n  (and 0.95 quantile) vs. accuracy curves. Given that not all queries\n  in all the methods (including the proposed ones) will have the same\n  sequences of accuracies, the quantile curve is especially unclear.\n- The authors mention that neural LSH is only useful when the bin\n  sizes are large, and otherwise k-means is more useful. Is there a\n  straightforward way to decide when to switch between the two? Is it\n  another hyperparameter we need to tune over while generating these\n  space partitions?\n- Theorem 3.1 considers only hyperplane partitions while neural\n  networks usually generate nonlinear partitions. What are the\n  conditions on the nonlinear splits that allows the guarantees to\n  improve (or at least not degrade)? Is it obviously always improved?\n", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a new framework for improved nearest neighbor search by learning a space partition of the data, allowing for better scalability in distributed settings and overall better performance over existing benchmarks.

The two reviewers who were most confident were both positive about the contributions and the revisions. The one reviewer who recommended reject was concerned about the metric used and whether comparison with baselines was fair. In my opinion, the authors seem to have been very receptive to reviewer comments and answered these issues to my satisfaction. After author and reviewer engagement, both R1 and myself are satisfied with the addition of the new baselines and think the authors have sufficiently addressed the major concerns. For the final version of the paper, I’d urge the authors to take seriously R4’s comment regarding clarity and add algorithmic details as per their suggestion.

Decision: Accept (Poster)