SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Filtering: We find that for large  $N$  it is useful to have an initial lossy filtering step that can cheaply reduce the total number of sets we consider with a low false-negative rate. We use an inverted index on the documents for this filtering step.
To build the inverted index, we first perform  $k$ -means clustering on a representative sample of individual item vectors at the start of indexing. The inverted index we will build is a map from centroid ids to document ids. As we add each set  $S_{i}$  to  $\mathcal{D}$  in Algorithm 1, we also add it into the inverted index: we find the closest centroid to each vector  $x \in S_{i}$ , and then we add the document id  $i$  to all of the buckets in the inverted index corresponding to those centroids.
This method is similar to PLAID, the recent optimized ColBERT implementation [37], but our query process is much simpler. During querying, we query the inverted index buckets corresponding to the closest filter Probe centroids to each query vector. We aggregate the buckets to get a count for each document id, and then only rank the filter  $k$  documents with DESSERT that have the highest count.
Space Optimized Sketches: DESSERT has two features that constrain the underlying hash table implementation: (1) every document is represented by a hash table, so the tables must be low memory, and (2) each query performs many table lookups, so the lookup operation must be fast. If (1) is not met, then we cannot fit the index into memory. If (2) is not met, then the similarity approximation for the inner aggregation step will be far too slow. Initially, we tried a naive implementation of the table, backed by a std::vector, std::map, or std::unordered_map. In each case, the resulting structure did not meet our criteria, so we developed TinyTable, a compact hash table that optimizes memory usage while preserving fast access times. TinyTables sacrifice  $O(1)$  update-access (which DESSERT does not require) for a considerable improvement to (1) and (2).
A TinyTable replaces the universal hash table in Algorithm 1, so it must provide a way to map pairs of (hash value, hash table id) to lists of vector ids. At a high level, a TinyTable is composed of  $L$  inverted indices from LSH values to vector ids. Bucket  $b$  of table  $i$  consists of vectors  $x_{j}$  such that  $h_{i}(x_{j}) = b$ . During a query, we simply need to go to the  $L$  buckets that correspond to the query vector's  $L$  lsh values to find the ids of  $S_{i}$ 's colliding vectors. This design solves (1), the fast lookup requirement, because we can immediately go to the relevant bucket once we have a query's hash value. However, there is a large overhead in storing a resizable vector in every bucket. Even an empty bucket will use  $3 * 8 = 24$  bytes. This adds up: let  $r$  be the hash range of the LSH functions (the number of buckets in the inverted index for each of the  $L$  tables). If  $N = 1M$ ,  $L = 64$ , and  $r = 128$ , we will use  $N \cdot L \cdot r \cdot 24 = 196$  gigabytes even when all of the buckets are empty.
Thus, a TinyTable has more optimizations than make it space efficient. Each of the  $L$  hash table repetitions in a TinyTable are conceptually split into two parts: a list of offsets and a list of vector ids. The vector ids are the concatenated contents of the buckets of the table with no space in between (thus, they are always some permutation of 0 through  $m - 1$ ). The offset list describes where one bucket ends and the next begins: the  $i$ th entry in the offset list is the (inclusive) index of the start of the  $i$ th hash bucket within the vector id list, and the  $i + 1$ th entry is the (exclusive) end of the  $i$ th hash bucket (if a bucket is empty, indices $[i] =$  indices $[i + 1]$ ). To save more bytes, we can further concatenate the  $L$  offset lists together and the  $L$  vector id lists together, since their lengths are always  $r$  and  $m_i$  respectively. Finally, we note that if  $m \leq 256$ , we can store all of the the offsets and ids can be safely be stored as single byte integers. Using the same hypothetical numbers as before, a filled TinyTable with  $m = 100$  will take up just  $N(24 + L(m + r + 1)) = 14.7\mathrm{GB}$ .
The Concatenation Trick: In our theory, we assumed LSH functions such that  $p(h(x) = h(y)) = \mathrm{sim}(x,y)$ . However, for practical problems such functions lead to overfull buckets; for example, GLOVE has an average vector cosine similarity of around 0.3, which would mean each bucket in the LSH table would contain a third of the set. The standard trick to get around this problem is to concatenate  $C$  hashes for each of the  $L$  tables together such that  $p(h(x) = h(y)) = \mathrm{sim}(x,y)^C$ . Rewriting, we have that
$$
\operatorname {s i m} (x, y) = \exp \left(\frac {\ln [ p (h (x) = h (y)) ]}{C}\right) \tag {1}
$$
During a query, we count the number of collisions across the  $L$  tables and divide by  $L$  to get  $\hat{p}(h(x) = h(y))$  on line 11 of Algorithm 2. We now additionally pass count/L into Equation 1 to get an accurate similarity estimate to pass into  $\sigma$  on line 12. Furthermore, evaluating Equation 1 for every collision probability estimate is slow in practice. There are only  $L + 1$  possible values for the count/L, so we precompute the mapping in a lookup table.
Datasets: We tested DESSERT on both synthetic data and real-world problems. We first examined a series of synthetic datasets to measure DESSERT's speedup over a reasonable CPU brute force algorithm (using the PyTorch library [35] for matrix multiplications). For this experiment, we leave out the prefiltering optimization described in Section 5 to better show how DESSERT performs on its own. Following the authors of [25], our synthetic dataset consists of random groups of Glove [36] vectors; we vary the set size  $m$  and keep the total number of sets  $N = 1000$ .
We next experimented with the MS MARCO passage ranking dataset (Creative Commons License) [32],  $N \approx 8.8M$ . The task for MS MARCO is to retrieve passages from the corpus relevant to a query. We used ColBERT to map the words from each passage and query to sets of embedding vectors suitable for DESSERT [37]. Following [37], we use the development set for our experiments, which contains 6980 queries.
Finally, we computed the full resource-accuracy tradeoff for ten of the LoTTE out-of-domain benchmark datasets, introduced by ColBERTv2 [38]. We excluded the pooled dataset, which is simply the individual datasets merged.
Experiment Setup: We ran our experiments on an Intel(R) Xeon(R) CPU E5-2680 v3 machine with 252 GB of RAM. We restricted all experiments to 4 cores (8 threads). We ran each experiment with the chosen hyperparameters and reported overall average recall and average query latency. For all experiments we used the average of max similarities scoring function.
6.1 Synthetic Data
The goal of our synthetic data experiment was to examine DESSERT's speedup over brute force vector set scoring. Thus, we generated synthetic data where both DESSERT and the brute force implementation achieved perfect recall so we could compare the two methods solely on query time.
The two optimized brute force implementations we tried both used PyTorch, and differed only in whether they computed the score between the query set and each document set individually ("Individual") or between the query set and all document sets at once using PyTorch's highly performant reduce and reshape operations ("Combined").
In each synthetic experiment, we inserted 1000 documents of size  $m$  for  $m \in [2,4,8,16,\ldots,1024]$  into DESSERT and the
Figure 2: Query time for DESSERT vs. brute force on 1000 random sets of  $m$  glove vectors with the  $y$ -axis as a log scale. Lower is better.
brute force index. The queries in each experiment were simply the documents with added noise. The DESSERT hyperparameters we chose were  $L = 8$  and  $C = \log_2(m) + 1$ . The results of our experiment, which show the relative speedup of using DESSERT at different values of  $m$ , are in Figure 2. We observe that DESSERT achieves a consistent 10-50x speedup over the optimized Pytorch brute force method and that the speedup increases with larger  $m$  (we could not run experiments with even larger  $m$  because the PyTorch runs did not finish within the time allotted).
6.2 Passage Retrieval
Passage retrieval refers to the task of identifying and returning the most relevant passages from a large corpus of documents in response to a search query. In these experiments, we compared DESSERT to PLAID, ColBERT's heavily-optimized state-of-the-art late interaction search algorithm, on the MS MARCO and LoTTE passage retrieval tasks.
We found that the best ColBERT hyperparameters were the same as reported in the PLAID paper, and we successfully replicated their results. Although PLAID offers a way to trade off time for accuracy, this tradeoff only increases accuracy at the sake of time, and even then only by a fraction of a percent. Thus, our results represent points on the recall vs time Pareto frontier that PLAID cannot reach.
MS MARCO Results For MS MARCO, we performed a grid search over DESSERT parameters  $C = \{4,5,6,7\}$ ,  $L = \{16,32,64\}$ , filterProbe = {1,2,4,8}, and filter_k = {1000,2048,4096,8192,16384}. We reran the best configurations to obtain the results in Table 2. We report two types of results: methods tuned to return  $k = 10$  results and methods tuned to return  $k = 1000$  results. For each, we report DESSERT results from a low latency and a high latency part of the Pareto frontier. For  $k = 1000$  we use the standard R@1000 metric, the average recall of the top 1 passage in the first 1000 returned passages. This metric is meaningful because retrieval pipelines frequently rerank candidates after an initial retrieval stage. For  $k = 10$  we use the standard MRR@10 metric, the average mean reciprocal rank of the top 1 passage in the first 10 returned passages. Overall, DESSERT is 2-5x faster than PLAID with only a few percent loss in recall.
Method	Latency (ms)	MRR@10	Method	Latency (ms)	R@1000
DESSERT	9.5	35.7 ± 1.14	DESSERT	22.7	95.1 ± 0.49
DESSERT	15.5	37.2 ± 1.14	DESSERT	32.3	96.0 ± 0.45
PLAID	45.1	39.2 ± 1.15	PLAID	100	97.5 ± 0.36
Table 2: MS MARCO passage retrieval, with methods optimized for  $\mathrm{k} = {10}$  (left) and  $\mathrm{k} = {1000}$  (right). Intervals denote 95% confidence intervals for average latency and recall.
LoTTE Results For LoTTE, we performed a grid search over  $C = \{4,6,8\}$ ,  $L = \{32,64,128\}$ , filterProbe = {1,2,4}, and filter_k = {1000,2048,4096,8192}. In Figure 3, we plot the full Pareto tradeoff for DESSERT on the 10 LoTTE datasets (each of the 5 categories has a "forum" and "search" split) over these hyperparameters, as well as the single lowest-latency point achievable by PLAID. For all test datasets, DESSERT provides a Pareto frontier that allows a tradeoff between recall and latency. For both Lifestyle test splits, both Technology test splits, and the Recreation and Science test-search splits, DESSERT achieves a 2-5x speedup with minimal loss in accuracy. On Technology, DESSERT even exceeds the accuracy of PLAID at half of PLAID's latency.
We observe a substantial speedup when we integrate DESSERT into ColBERT, even when compared against the highly-optimized PLAID implementation. While the use of our algorithm incurs a slight recall penalty – as is the case with most algorithms that use randomization to achieve acceleration – Table 2 and Figure 3 shows that we are Pareto-optimal when compared with baseline approaches.
We are not aware of any algorithm other than DESSERT that is capable of latencies in this range for set-to-set similarity search. While systems such as PLAID are tunable, we were unable to get them to operate in this range. For this reason, DESSERT is likely the only set-to-set similarity search algorithm that can be run in real-time production environments with strict latency constraints.
Figure 3: Full Pareto frontier of DESSERT on the LoTTE datasets. The PLAID baseline shows the lowest-latency result attainable by PLAID (with a FAISS-IVF base index and centroid pre-filtering).
We also ran a single-vector search baseline using ScaNN, the leading approximate kNN index [14]. ScaNN yielded 0.77 Recall@1000, substantially below the state of the art. This result reinforces our discussion in Section 1.2 on why single-vector search is insufficient.
Broader Impacts and Limitations: Ranking and retrieval are important steps in language modeling applications, some of which have recently come under increased scrutiny. However, our algorithm is unlikely to have negative broader effects, as it mainly enables faster, more cost-effective search over larger vector collections and does not contribute to the problematic capabilities of the aforementioned language models. Due to computational limitations, we conduct our experiments on a relatively small set of benchmarks; a larger-scale evaluation would strengthen our argument. Finally, we assume sufficiently high relevance scores and large gaps in our theoretical analysis to identify the correct results. These hardness assumptions are standard for LSH.
In this paper, we consider the problem of vector set search with vector set queries, a task understudied in the existing literature. We present a formal definition of the problem and provide a motivating application in semantic search, where a more efficient algorithm would provide considerable immediate impact in accelerating late interaction search methods. To address the large latencies inherent in existing vector search methods, we propose a novel randomized algorithm called DESSERT that achieves significant speedups over baseline techniques. We also analyze DESSERT theoretically and, under natural assumptions, prove rigorous guarantees on the algorithm's failure probability and runtime. Finally, we provide an open-source and highly performant  $\mathrm{C + + }$  implementation of our proposed DESSERT algorithm that achieves 2-5x speedup over ColBERT-PLAID on the MS MARCO and LoTTE retrieval benchmarks. We also note that a general-purpose algorithmic framework for vector set search with vector set queries could have impact in a number of other applications, such as image similarity search [42], market basket analysis [22], and graph neural networks [43], where it might be more natural to model entities via sets of vectors as opposed to restricting representations to a single embedding. We believe that DESSERT could provide a viable algorithmic engine for enabling such applications and we hope to study these potential use cases in the future.
This work was completed while the authors were working at ThirdAI. We do not have any external funding sources to acknowledge.
[33] Priyanka Nigam, Yiwei Song, Vijai Mohan, Vihan Lakshman, Weitian Ding, Ankit Shingavi, Choon Hui Teo, Hao Gu, and Bing Yin. Semantic product search. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 2876-2885, 2019.
[34] Steve Olenski. Why brands are fighting over milliseconds, Nov 2016.
[35] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[36] Jeffrey Pennington, Richard Socher, and Christopher D. Manning. Glove: Global vectors for word representation. In Empirical Methods in Natural Language Processing (EMNLP), pages 1532-1543, 2014.
[37] Keshav Santhanam, Omar Khattab, Christopher Potts, and Matei Zaharia. Plaid: An efficient engine for late interaction retrieval. arXiv preprint arXiv:2205.09707, 2022.
[38] Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. Colbertv2: Effective and efficient retrieval via lightweight late interaction. arXiv preprint arXiv:2112.01488, 2021.
[39] Aneesh Sharma, C Seshadhri, and Ashish Goel. When hashes met wedges: A distributed algorithm for finding high similarity vectors. In Proceedings of the 26th International Conference on World Wide Web, pages 431-440, 2017.
[40] Anshumali Shrivastava and Ping Li. Improved asymmetric locality sensitive hashing (alsh) for maximum inner product search (mips). In Proceedings of the Thirty-First Conference on Uncertainty in Artificial Intelligence, pages 812–821, 2015.
[41] Yiqiu Wang, Anshumali Shrivastava, Jonathan Wang, and Junghee Ryu. Flash: Randomized algorithms accelerated over cpu-gpu for ultra-high dimensional similarity search. arXiv preprint arXiv:1709.01190, 2017.
[42] Jun Yang, Yu-Gang Jiang, Alexander G Hauptmann, and Chong-Wah Ngo. Evaluating bag-of-visual-words representations in scene classification. In Proceedings of the international workshop on Workshop on multimedia information retrieval, pages 197-206, 2007.
[43] Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William L. Hamilton, and Jure Leskovec. Graph convolutional neural networks for web-scale recommender systems. CoRR, abs/1806.01973, 2018.
Proof. Take some  $\mathbf{x} \in D$  (so each  $x_i \in I$ ). Since in  $\mathbb{R}$ ,  $\max(x) = x$ , we have from the definition of  $(\alpha, \beta)$ -maximal that
$$
\beta x \leq x \leq \alpha x
$$
For the upper bound, we have
$$
\sigma (\mathbf {x}) = \frac {1}{m} \sum_ {i = 1} ^ {m} \varphi (x _ {i}) \leq \frac {1}{m} \sum_ {i = 1} ^ {m} \alpha x _ {i} = \frac {\alpha}{m} \sum_ {i = 1} ^ {m} x _ {i} \leq \frac {\alpha}{m} (m \max  (\mathbf {x})) = \alpha \max  (\mathbf {x})
$$
where the second inequality follows by the properties of the max function.
For the lower bound, we have that
$$
\sigma (\mathbf {x}) = \frac {1}{m} \sum_ {i = 1} ^ {m} \varphi (x _ {i}) \geq \frac {1}{m} \sum_ {i = 1} ^ {m} \beta x _ {i} = \frac {\beta}{m} \sum_ {i = 1} ^ {m} x _ {i} \geq \frac {\beta}{m} \max \mathbf {x}
$$
where the second inequality again follows by the properties of the max function.
Proof. We first apply a generic Chernoff bound to  $\sigma(\hat{\mathbf{s}})$ , which gives us the following bounds for any  $t > 0$ :
$$
\operatorname * {P r} [ \sigma (\hat {\mathbf {s}}) \geq \tau ] = \operatorname * {P r} [ e ^ {t \sigma (\hat {\mathbf {s}})} \geq e ^ {t \tau} ] \leq \frac {\mathbb {E} [ e ^ {t \sigma (\hat {\mathbf {s}})} ]}{e ^ {t \tau}}
$$
We now proceed by continuing to bound the numerator. Because  $\sigma$  is  $(\alpha, \beta)$ -maximal, we can bound  $\sigma(\hat{\mathbf{s}})$  with  $\alpha \max \hat{\mathbf{s}}$ . We can further bound  $\max \hat{\mathbf{s}}$  by bounding the maximum with the sum and the sum with  $m$  times the maximal element. We are now left with the formula for the moment generating function for  $\hat{s}_{\max}$ .  $\hat{s}_{\max} \sim$  scaled binomial  $L^{-1}\mathcal{B}(s_{\max}, L)$ , so we can directly substitute the binomial moment generating function into the expression:
$$
\begin{array}{l} \mathbb {E} \left[ e ^ {t \sigma (\hat {\mathbf {s}})} \right] \leq \mathbb {E} \left[ e ^ {t \alpha \max  _ {j} \hat {s} _ {j}} \right] \leq \sum_ {j = 1} ^ {m _ {i}} \mathbb {E} \left[ e ^ {t \alpha \hat {s} _ {j}} \right] \leq m \mathbb {E} \left[ e ^ {t \alpha \hat {s} _ {\max }} \right] \\ = m \big (1 - s _ {\mathrm {m a x}} + s _ {\mathrm {m a x}} e ^ {\frac {\alpha t}{L}} \big) ^ {L} \\ \end{array}
$$
Combining these two equations yields the following bound:
$$
\operatorname * {P r} [ \sigma (\hat {\mathbf {s}}) \geq \tau ] \leq m e ^ {- t \tau} (1 - s _ {\max} + s _ {\max} e ^ {\frac {\alpha t}{L}}) ^ {L}
$$
We wish to select a value of  $t$  to minimize the upper bound. By setting the derivative of the upper bound to zero, and imposing  $0 < \tau < \alpha$ ,  $\alpha \geq 1$ , and  $0 < s_{\max} < 1$ , we find that
$$
t ^ {\star} = \frac {L}{\alpha} \ln \left(\frac {\tau (1 - s _ {\mathrm {m a x}})}{s _ {\mathrm {m a x}} (\alpha - \tau)}\right)
$$
This is greater than zero when the numerator inside the  $\ln$  is greater than the denominator, or equivalently when  $\tau > s_{\max} \alpha$ . Thus the valid range for  $\tau$  is  $(s_{\max} \alpha, \alpha)$  (and similarly the valid range for  $s_{\max}$  is  $(0, \tau / \alpha)$ ). These bounds have a natural interpretation: to be meaningful, the threshold must be between the expected value and the maximum value for  $\alpha$  times a  $p = s_{\max}$  binomial. Substituting  $t = t^{\star}$  into our upper bound, we obtain:
$$
\operatorname * {P r} [ \sigma (\hat {\mathbf {s}}) \geq \tau ] \leq m \left(\left(\frac {\tau (1 - s _ {\max})}{s _ {\max} (\alpha - \tau)}\right) ^ {- \frac {\tau}{\alpha}} \left(\frac {\alpha (1 - s _ {\max})}{\alpha - \tau}\right)\right) ^ {L}
$$
Thus we have that
$$
\gamma = \left(\frac {s _ {\mathrm {m a x}} (\alpha - \tau)}{\tau (1 - s _ {\mathrm {m a x}})}\right) ^ {\frac {\tau}{\alpha}} \left(\frac {\alpha (1 - s _ {\mathrm {m a x}})}{\alpha - \tau}\right)
$$
We will now prove our claims about  $\gamma$  viewed as a function of  $s_{\mathrm{max}} \in (0, \tau / \alpha)$  and  $\tau \in (s_{\mathrm{max}}\alpha, \alpha)$ . We will first examine the limits of  $\gamma$  with respect to  $\tau$  at the ends of its range. Since  $\gamma$  is continuous, we can find one of the limits by direct substitution:
$$
\lim  _ {\tau \searrow s _ {\max } \alpha} \gamma = \lim  _ {s _ {\max }, \nearrow \tau / \alpha} \gamma = \left(\frac {\tau / \alpha (\alpha - \tau)}{\tau (1 - \tau / \alpha)}\right) ^ {\frac {\tau}{\alpha}} \left(\frac {\alpha (1 - \tau / \alpha)}{\alpha - \tau}\right) = 1 ^ {\frac {\tau}{\alpha}} * 1 = 1
$$
The second limit is harder; we merge  $\gamma$  into one exponent and then simplify:
$$
\begin{array}{l} \lim  _ {\tau \nearrow \alpha} \gamma = \lim  _ {\tau \nearrow \alpha} \left(\frac {s _ {\operatorname* {m a x}} (\alpha - \tau) ^ {1 - \alpha} \tau \alpha^ {\alpha / \tau}}{\tau (1 - s _ {\operatorname* {m a x}}) ^ {1 - \alpha / \tau}}\right) ^ {\frac {\tau}{\alpha}} \\ = \lim  _ {\tau \nearrow \alpha} \frac {s _ {\operatorname* {m a x}} (\alpha - \tau) ^ {1 - \alpha \tau} \alpha^ {\alpha / \tau}}{\tau (1 - s _ {\operatorname* {m a x}}) ^ {1 - \alpha / \tau}} \\ = \lim  _ {\tau \nearrow \alpha} s _ {\max } (\alpha - \tau) ^ {1 - \alpha / \tau} = \lim  _ {\tau \nearrow \alpha} \left(s _ {\max } (\alpha - \tau) ^ {\alpha - \tau}\right) ^ {- 1 / \tau} \\ = s _ {\max } \left(\lim  _ {\alpha - \tau \searrow 0} \left((\alpha - \tau) ^ {\alpha - \tau}\right)\right) ^ {- 1 / \alpha} \\ = s _ {\max } (1) ^ {- 1 / \alpha} = s _ {\max } \\ \end{array}
$$
where we use the fact that  $\lim_{x\searrow 0}x^{x} = e^{\lim_{x\searrow 0}x\ln (x)} = 1$  (we can see that  $\lim_{x\to 0}x\ln (x) = \lim_{x\to 0}\ln (x) / (1 / x) = 0$  with L'Hopital's rule). We next find the partial derivatives of  $\gamma$ :
$$
\frac {\delta \gamma}{\delta s _ {\mathrm {m a x}}} = \frac {\left(\tau - \alpha s _ {\mathrm {m a x}}\right) \left(\frac {\alpha s _ {\mathrm {m a x}} - s _ {\mathrm {m a x}} \tau}{\tau - s _ {\mathrm {m a x}} \tau}\right) ^ {\frac {\tau}{\alpha}}}{s _ {\mathrm {m a x}} (\alpha - \tau)} \qquad \frac {\delta \gamma}{\delta \tau} = \frac {\left(s _ {\mathrm {m a x}} - 1\right) \left(\frac {\alpha s _ {\mathrm {m a x}} - s _ {\mathrm {m a x}} \tau}{\tau - s _ {\mathrm {m a x}} \tau}\right) ^ {\frac {\tau}{\alpha}} \ln \left(\frac {\tau - s _ {\mathrm {m a x}} \tau}{\alpha s _ {\mathrm {m a x}} - s _ {\mathrm {m a x}} \tau}\right)}{\alpha - \tau}
$$