SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Joshua Engels
ThirdAI
josh.adam.engels@gmail.com
Benjamin Coleman
benjamin-ray.coleman@gmail.com
Vihan Lakshman
vihan@thirdai.com
Anshumali Shrivastava
ThirdAI, Rice University
anshu@thirdai.com, anshumali@rice.edu
We study the problem of vector set search with vector set queries. This task is analogous to traditional near-neighbor search, with the exception that both the query and each element in the collection are sets of vectors. We identify this problem as a core subroutine for semantic search applications and find that existing solutions are unacceptably slow. Towards this end, we present a new approximate search algorithm, DESSERT (DESSERT Efficiently Searches Sets of Embeddings via Retrieval Tables). DESSERT is a general tool with strong theoretical guarantees and excellent empirical performance. When we integrate DESSERT into ColBERT, a state-of-the-art semantic search model, we find a 2-5x speedup on the MS MARCO and LoTTE retrieval benchmarks with minimal loss in recall, underscoring the effectiveness and practical applicability of our proposal.
Similarity search is a fundamental driver of performance for many high-profile machine learning applications. Examples include web search [16], product recommendation [33], image search [21], de-duplication of web indexes [29] and friend recommendation for social media networks [39]. In this paper, we study a variation on the traditional vector search problem where the dataset  $D$  consists of a collection of vector sets  $D = \{S_1, \ldots, S_N\}$  and the query  $Q$  is also a vector set. We call this problem vector set search with vector set queries because both the collection elements and the query are sets of vectors. Unlike traditional vector search, this problem currently lacks a satisfactory solution.
Furthermore, efficiently solving the vector set search problem has immediate practical implications. Most notably, the popular ColBERT model, a state-of-the-art neural architecture for semantic search over documents [23], achieves breakthrough performance on retrieval tasks by representing each query and document as a set of BERT token embeddings. ColBERT's current implementation of vector set search over these document sets, while superior to brute force, is prohibitively slow for real-time inference applications like e-commerce that enforce strict search latencies under 20-30 milliseconds [34, 5]. Thus, a more efficient algorithm for searching over sets of vectors would have significant implications in making state-of-the-art semantic search methods feasible to deploy in large-scale production settings, particularly on cost-effective CPU hardware.
Given ColBERT's success in using vector sets to represent documents more accurately, and the prevailing focus on traditional single-vector near-neighbor search in the literature [1, 41, 19, 12, 14, 28, 18], we believe that the potential for searching over sets of representations remains largely untapped. An efficient algorithmic solution to this problem could enable new applications in domains
where multi-vector representations are more suitable. To that end, we propose DESSERT, a novel randomized algorithm for efficient set vector search with vector set queries. We also provide a general theoretical framework for analyzing DESSERT and evaluate its performance on standard passage ranking benchmarks, achieving a 2-5x speedup over an optimized ColBERT implementation on several passage retrieval tasks.
1.1 Problem Statement
More formally, we consider the following problem statement.
Definition 1.1. Given a collection of  $N$  vector sets  $D = \{S_1, \ldots, S_N\}$ , a query set  $Q$ , a failure probability  $\delta \geq 0$ , and a set-to-set relevance score function  $F(Q, S)$ , the Vector Set Search Problem is the task of returning  $S^*$  with probability at least  $1 - \delta$ :
$$
S^{\star} = \operatorname *{argmax}_{i\in \{1,\ldots N\}}F(Q,S_{i})
$$
Here, each set  $S_{i} = \{x_{1},\ldots x_{m_{i}}\}$  contains  $m_{i}$  vectors with each  $x_{j}\in \mathbb{R}^{d}$ , and similarly  $Q = \{q_{1},\dots q_{m_{q}}\}$  contains  $m_{q}$  vectors with each  $q_{j}\in \mathbb{R}^{d}$ .
We further restrict our consideration to structured forms of  $F(Q, S)$ , where the relevance score consists of two "set aggregation" or "variadic" functions. The inner aggregation  $\sigma$  operates on the pairwise similarities between a single vector from the query set and each vector from the target set. Because there are  $|S|$  elements in  $S$  over which to perform the aggregation,  $\sigma$  takes  $|S|$  arguments. The outer aggregation  $A$  operates over the  $|Q|$  scores obtained by applying  $A$  to each query vector  $q \in Q$ . Thus, we have that
$$
F (Q, S) = A (\{I n n e r _ {q, S}: q \in Q \})
$$
$$
I n n e r _ {q, S} = \sigma (\{\operatorname {s i m} (q, x): x \in S \})
$$
Here,  $\operatorname{sim}$  is a vector similarity function. Because the inner aggregation is often a maximum or other non-linearity, we use  $\sigma(\cdot)$  to denote it, and similarly since the outer aggregation is often a linear function we denote it with  $A(\cdot)$ . These structured forms for  $F = A \circ \sigma$  are a good measure of set similarity when they are monotonically non-decreasing with respect to the similarity between any pair of vectors from  $Q$  and  $S$ .
1.2 Why is near-neighbor search insufficient?
It may at first seem that we could solve the Vector Set Search Problem by placing all of the individual vectors into a near-neighbor index, along with metadata indicating the set to which they belonged. One could then then identify high-scoring sets by finding near neighbors to each  $q \in Q$  and returning their corresponding sets.
There are two problems with this approach. The first problem is that a single high-similarity interaction between  $q \in Q$  and  $x \in S$  does not imply that  $F(Q, S)$  will be large. For a concrete example, suppose that we are dealing with sets of word embeddings and that  $Q$  is a phrase where one of the items is "keyword." With a standard near-neighbor index,  $Q$  will match (with 100% similarity) any set  $S$  that also contains "keyword," regardless of whether the other words in  $S$  bear any relevance to the other words in  $Q$ . The second problem is that the search must be conducted over all individual vectors, leading to a search problem that is potentially very large. For example, if our sets are documents consisting of roughly a thousand words and we wish to search over a million documents, we now have to solve a billion-scale similarity search problem.
Contributions: In this work, we formulate and carefully study the set of vector search problems with the goal of developing a more scalable algorithm capable of tackling large-scale semantic retrieval problems involving sets of embeddings. Specifically, our research contributions can be summarized as follows:
1. We develop the first non-trivial algorithm, DESSERT, for the vector set search problem that scales to large collections  $(n > 10^{6})$  of sets with  $m > 3$  items.
2. We formalize the vector set search problem in a rigorous theoretical framework, and we provide strong guarantees for a common (and difficult) instantiation of the problem.
3. We provide an open-source  $\mathrm{C + + }$  implementation of our proposed algorithm that has been deployed in a real-world production setting<sup>1</sup>. Our implementation scales to hundreds of millions of vectors and is 3-5x faster than existing approximate set of vector search techniques. We also describe the implementation details and tricks we discovered to achieve these speedups and provide empirical latency and recall results on passage retrieval tasks.
Near-Neighbor Search: Near-neighbor search has received heightened interest in recent years with the advent of vector-based representation learning. In particular, considerable research has gone into developing more efficient approximate near-neighbor (ANN) search methods that trade off an exact solution for sublinear query times. A number of ANN algorithms have been proposed, including those based on locality-sensitive hashing [1, 41], quantization and space partition methods [19, 12, 14], and graph-based methods [28, 18]. Among these classes of techniques, our proposed DESSERT framework aligns most closely with the locality-sensitive hashing paradigm. However, nearly all of the well-known and effective ANN methods focus on searching over individual vectors; our work studies the search problem for sets of entities. This modification changes the nature of the problem considerably, particularly with regards to the choice of similarity metrics between entities.
Vector Set Search: The general problem of vector set search has been relatively understudied in the literature. A recent work on database lineage tracking [25] addresses this precise problem, but with severe limitations. The proposed approximate algorithm designs a concatenation scheme for the vectors in a given set, and then performs approximate search over these concatenated vectors. The biggest drawback to this method is scalability, as the size of the concatenated vectors scales quadratically with the size of the vector set. This leads to increased query latency as well as substantial memory overhead; in fact, we are unable to apply the method to the datasets in this paper without terabytes of RAM. In this work, we demonstrate that DESSERT can scale to thousands of items per set with a linear increase (and a slight logarithmic overhead) in query time, which, to our knowledge, has not been previously demonstrated in the literature.
Document Retrieval: In the problem of document retrieval, we receive queries and must return the relevant documents from a preindexed corpus. Early document retrieval methods treated each documents as bags of words and had at their core an inverted index [30]. More recent methods embed each document into a single representative vector, embed the query into the same space, and performed ANN search on those vectors. These semantic methods achieve far greater accuracies than their lexical predecessors, but require similarity search instead of inverted index lookups [15, 33, 26].
ColBERT and PLAID: ColBERT [23] is a recent state of the art algorithm for document retrieval that takes a subtly different approach. Instead of generating a single vector per document, ColBERT generates a set of vectors for each document, approximately one vector per word. To rank a query, ColBERT also embeds the query into a set of vectors, filters the indexed sets, and then performs a brute force sum of max similarities operation between the query set and each of the document sets. ColBERT's passage ranking system is an instantiation of our framework, where  $\mathrm{sim}(q,x)$  is the cosine similarity between vectors,  $\sigma$  is the max operation, and  $A$  is the sum operation.
In a similar spirit to our work, PLAID [37] is a recently optimized form of ColBERT that includes more efficient filtering techniques and faster quantization based set similarity kernels. However, we note that these techniques are heuristics that do not come with theoretical guarantees and do not immediately generalize to other notions of vector similarity, which is a key property of the theoretical framework behind DESSERT.
At a high level, a DESSERT index  $\mathcal{D}$  compresses the collection of target sets into a form that makes set to set similarity operations efficient to calculate. This is done by replacing each set  $S_{i}$  with a sketch  $\mathcal{D}[i]$  that contains the LSH values of each  $x_{j} \in S_{i}$ . At query time, we compare the corresponding
$^{1}$ https://github.com/ThirdAIResearch/Dessert
Figure 1: The DESSERT indexing and querying algorithms. During indexing (left), we represent each target set as a set of hash values ( $L$  hashes for each element). To query the index (right), we approximate the similarity between each target and query element by averaging the number of hash collisions. These similarities are used to approximate the set relevance score for each target set.
LSH values of the query set  $Q$  with the hashes in each  $\mathcal{D}[i]$  to approximate the pairwise similarity matrix between  $Q$  and  $S$  (Figure 1). This matrix is used as the input for the aggregation functions  $A$  and  $\sigma$  to rank the target sets and return an estimate of  $S^{*}$ .
We assume the existence of a locality-sensitive hashing (LSH) family  $\mathcal{H} \subset (\mathbb{R}^d \to \mathbb{Z})$  such that for all LSH functions  $h \in \mathcal{H}$ ,  $p(h(x) = h(y)) = \mathrm{sim}(x, y)$ . LSH functions with this property exist for cosine similarity (signed random projections) [8], Euclidean similarity ( $p$ -stable projections) [11], and Jaccard similarity (minhash or simhash) [6]. LSH is a well-developed theoretical framework with a wide variety of results and extensions [4, 3, 20, 40]. See Appendix C for a deeper overview.
Algorithm 1 describes how to construct a DESSERT index  $\mathcal{D}$ . We first take  $L$  LSH functions  $f_{t}$  for  $t\in [1,L]$ ,  $f_{t}\in \mathcal{H}$ . We next loop over each  $S_{i}$  to construct  $\mathcal{D}[i]$ . For a given  $S_{i}$ , we arbitrarily assign an identifier  $j$  to each vector  $x\in S_i$ ,  $j\in [1,|S_i|]$ . We next partition the set  $[1,m_i]$  using each hash function  $h_t$ , such that for a partition  $p_t$ , indices  $j_1$  and  $j_2$  are in the same set in the partition iff  $h(S_{j_1}) = h(S_{j_2})$ . We represent the results of these partitions in a universal hash table indexed by hash function id and hash function value, such that  $\mathcal{D}[i]_{t,h} = \{j|x_j\in S_i\wedge f_t(x_j) = h\}$ .
Algorithm 2 describes how to query a DESSERT index  $\mathcal{D}$ . At a high level, we query each sketch  $\mathcal{D}_i$  to get an estimate of  $F(Q,S_i)$ ,  $score_i$ , and then take the argmax over the estimates to get an estimate of  $\operatorname{argmax}_{i\in \{1,\dots N\}}F(Q,S_i)$ . To get these estimates, we first compute the hashes  $h_{t,q}$  for each query  $q$  and LSH function  $f_t$ . Then, to get an estimate  $score_i$  for a set  $S_i$ , we loop over the hashes  $h_{t,q}$  for each query vector  $q$  and count how often each index  $j$  appears in  $\mathcal{D}[i]_{t,h_{t,q}}$ . After we finish this step, we have a count for each  $j$  that represents how many times  $h_t(q) = h_t(x_j)$ . Equivalently, since  $p(h(x) = h(y)) = \sin (x,y)$ , if we divide by  $L$  we have an estimate for  $\sin (x_j,q)$ . We then apply  $\sigma$  to these estimates and save the result in a variable  $\mathrm{agg}_q$  to build up the inputs to  $A$ , and then apply  $A$  to get our final estimate for  $F(Q,S_i)$ , which we store in  $score_i$ .
In this section, we analyze DESSERT's query runtime and provide probabilistic bounds on the correctness of its search results. We begin by finding the hyperparameter values and conditions that are necessary for DESSERT to return the top-ranked set with high probability. Then, we use these results to prove bounds on the query time. In the interest of space, we defer proofs to the Appendix.
Notation: For the sake of simplicity of presentation, we suppose that all target sets have the same number of elements  $m$ , i.e.  $|S_i| = m$ . If this is not the case, one may replace  $m_i$  with  $m_{\max}$  in our analysis. We will use the boldface vector  $\mathbf{s}(q, S_i) \in \mathbb{R}^{|S_i|}$  to refer to the set of pairwise similarity calculations  $\{\mathrm{sim}(q, x_1), \ldots, \mathrm{sim}(q, x_{m_i})\}$  between a query vector and the elements of  $S_i$ , and we will drop the subscript  $(q, S_i)$  when the context is clear. See Table 1 for a complete notation reference.
Table 1: Notation table with examples from the document search application.
Notation	Definition	Intuition (Document Search)
D	Set of target vector sets	Collection of documents
N	Cardinality |D|	Number of documents
D	DESSERT index of D	Search index data structure
Si	Target vector set i	ith document
Q	Query vector set	Multi-word query (e.g., a question)
S*	See Definition 1.1	The most relevant document to Q
xj∈Si	jth vector in target set Si	Embedding from document i
qj∈Q	jth vector in query set Q	Embedding from a query
d	sj, xj ∈ Rd	Embedding dimension
mi, m	Cardinality |Si|, mi = m	Number of embeddings in ith document
F(Q, Si)	Q and Si relevance score	Measures query-document similarity
scorei	Estimate of F(Q, Si)	Approximation of relevance score
D[i]	Sketch of ith target set	Estimates relevance score for Si and any Q
sim(a, b)	a and b vector similarity	Embedding similarity
A, σ	See Section 1.1	Components of relevance score
L	Number of hashes	Larger L increases accuracy and latency
fi	ith LSH function	Often maps nearby points to the same value
s(q, Si), s	sim(q, xj) for xj ∈ Si	Query embedding similarities with Si
Algorithm 1 Building a DESSERT Index
4: for  $i = 1$  to  $N$  do
1: Input:  $N$  sets  $S_{i},|S_{i}| = m_{i}$    
2: Output:A DESSERT index  $\mathcal{D}$    
3:  $\mathcal{D} =$  an array of  $N$  hash tables, each indexed by  $x\in \mathbb{Z}^2$
7:  $h = f_{t}(x_{j})$
5: for  $x_{j}$  in  $S_{i}$  do  
6: for  $t = 1$  to  $L$  do
9: Return  $\mathcal{D}$
8:  $\mathcal{D}[i]_{t,h} = \mathcal{D}[i]_{t,h}\cup \{j\}$
Algorithm 2 Querying a DESSERT Index
2: Output: Estimate of  $\operatorname{argmax}_{i\in \{1\dots N\}}A\circ \sigma (Q,S_i)$
1: Input: DESSERT index  $\mathcal{D}$ , query set  $Q$ ,  $|Q| = m_q$ .
4:  $h_{1,q},\ldots ,h_{L,q} = f_1(q),\ldots ,f_L(q)$
3: for  $q$  in  $Q$  do
5: for  $i = 1$  to  $N$  do
6: for  $q$  in  $Q$  do
8: for  $t = 1$  to  $L$  do
7:  $\hat{\mathbf{s}} = 0$
10:  $\hat{\mathbf{s}}_j = \hat{\mathbf{s}}_j + 1$
9: for  $j$  in  $\mathcal{D}[i]_{t,h_{t,q}}$  do
11:  $\hat{\mathbf{s}} = \hat{\mathbf{s}} /L$
12:  $agg_{q} = \sigma (\hat{\mathbf{s}})$
14: Return  $\operatorname{argmax}_{i\in \{1\dots N\}}\text{score}_i$
13:  $score_{i} = A(\{agg_{q} \mid q \in Q\})$
4.1 Inner Aggregation
We begin by introducing a condition on the  $\sigma$  component of the relevance score that allows us to prove useful statements about the retrieval process.
Definition 4.1. A function  $\sigma(\mathbf{x}) : \mathbb{R}^m \to \mathbb{R}$  is  $(\alpha, \beta)$ -maximal on  $U \subset \mathbb{R}^m$  if for  $0 < \beta \leq 1 \leq \alpha$ ,  $\forall x \in U$ :
$$
\beta \max {\mathbf {x}} \leq \sigma (\mathbf {x}) \leq \alpha \max {\mathbf {x}}
$$
The function  $\sigma (\mathbf{x}) = \max \mathbf{x}$  is a trivial example of an  $(\alpha ,\beta)$ -maximal function on  $\mathbb{R}^m$ , with  $\beta = \alpha = 1$ . However, we can show that other functions also satisfy this definition:
Lemma 4.1.1. If  $\varphi(x): \mathbb{R} \to \mathbb{R}$  is  $(\alpha, \beta)$ -maximal on an interval  $I$ , then the following function  $\sigma(x): \mathbb{R}^m \to \mathbb{R}$  is  $\left(\alpha, \frac{\beta}{m}\right)$ -maximal on  $U = I^m$ :
$$
\sigma (\mathbf {x}) = \frac {1}{m} \sum_ {i = 1} ^ {m} \varphi (x _ {i})
$$
Note that in  $\mathbb{R}$ , the  $(\alpha, \beta)$ -maximal condition is equivalent to lower and upper bounds by linear functions  $\beta x$  and  $\alpha x$  respectively, so many natural functions satisfy Lemma 4.1.1. We are particularly interested in the case  $I = [0,1]$ , and note that possible such  $\varphi$  include  $\varphi(x) = x$  with  $\beta = \alpha = 1$ , the exponential function  $\varphi(x) = e^x - 1$  with  $\beta = 1$ ,  $\alpha = e - 1$ , and the debiased sigmoid function  $\varphi(x) = \frac{1}{1 + e^{-x}} - \frac{1}{2}$  with  $\beta \approx 0.23$ ,  $\alpha = 0.25$ . Our analysis of DESSERT holds when the  $\sigma$  component of the relevance score is an  $(\alpha, \beta)$  maximal function.
In line 12 of Algorithm 2, we estimate  $\sigma(\mathbf{s})$  by applying  $\sigma$  to a vector of normalized counts  $\hat{\mathbf{s}}$ . In Lemma 4.1.2, we bound the probability that a low-similarity set (one for which  $\sigma(\mathbf{s})$  is low) scores well enough to outrank a high-similarity set. In Lemma 4.1.3, we bound the probability that a high-similarity set scores poorly enough to be outranked by other sets. Note that the failure rate in both lemmas decays exponentially with the number of hash tables  $L$ .
Lemma 4.1.2. Assume  $\sigma$  is  $(\alpha, \beta)$ -maximal. Let  $0 < s_{\max} < 1$  be the maximum similarity between a query vector and the vectors in the target set and let  $\hat{\mathbf{s}}$  be the set of estimated similarity scores. Given a threshold  $\alpha s_{\max} < \tau < \alpha$ , we write  $\Delta = \tau - \alpha s_{\max}$ , and we have
$$
\Pr [ \sigma (\hat {\mathbf {s}}) \geq \alpha s _ {\max} + \Delta ] \leq m \gamma^ {L}
$$
for  $\gamma = \left(\frac{s_{\max}(\alpha - \tau)}{\tau(1 - s_{\max})}\right)^{\frac{\tau}{\alpha}}\left(\frac{\alpha(1 - s_{\max})}{\alpha - \tau}\right) \in (s_{\max}, 1)$ . Furthermore, this expression for  $\gamma$  is increasing in  $s_{\max}$  and decreasing in  $\tau$ , and  $\gamma$  has one-sided limits  $\lim_{\tau \searrow \alpha} s_{\max} \gamma = 1$  and  $\lim_{\tau \nearrow \alpha} \gamma = s_{\max}$ .
Lemma 4.1.3. With the same assumptions as Lemma 4.1.2 and given  $\Delta >0$ , we have:
$$
P r [ \sigma (\hat {\mathbf {s}}) \leq \beta s _ {\max } - \Delta ] \leq 2 e ^ {- 2 L \Delta^ {2} / \beta^ {2}}
$$
4.2 Outer Aggregation
Our goal in this section is to use the bounds established previously to prove that our algorithm correctly ranks sets according to  $F(Q, S)$ . To do this, we must find conditions under which the algorithm successfully identifies  $S^{\star}$  based on the approximate  $F(Q, S)$  scores.
Recall that  $F(Q, S)$  consists of two aggregations: the inner aggregation  $\sigma$  (analyzed in Section 4.1) and the outer aggregation  $A$ . We consider normalized linear functions for  $A$ , where we are given a set of weights  $0 \leq w \leq 1$  and we rank the target set according to a weighted linear combination of  $\sigma$  scores.
$$
F (Q, S) = \frac {1}{m} \sum_ {j = 1} ^ {m} w _ {j} \sigma (\hat {\mathbf {s}} (q _ {j}, S))
$$
With this instantiation of the vector set search problem, we will proceed in Theorem 4.2 to identify a choice of the number of hash tables  $L$  that allows us to provide a probabilistic guarantee that the algorithm's query operation succeeds. We will then use this parameter selection to bound the runtime of the query operation in Theorem 4.3.
Theorem 4.2. Let  $S^{\star}$  be the set with the maximum  $F(Q, S)$  and let  $S_{i}$  be any other set. Let  $B^{\star}$  and  $B_{i}$  be the following sums (which are lower and upper bounds for  $F(Q, S^{\star})$  and  $F(Q, S_{i})$ , respectively)
$$
B ^ {\star} = \frac {\beta}{m _ {q}} \sum_ {j = 1} ^ {m _ {q}} w _ {j} s _ {\max } (q _ {j}, S ^ {\star}) \quad B _ {i} = \frac {\alpha}{m _ {q}} \sum_ {j = 1} ^ {m _ {q}} w _ {j} s _ {\max } (q _ {j}, S _ {i})
$$
Here,  $s_{\max}(q, S)$  is the maximum similarity between a query vector  $q$  and any element of the target set  $S$ . Let  $B'$  be the maximum value of  $B_i$  over any set  $S_i \neq S$ . Let  $\Delta$  be the following value (proportional to the difference between the lower and upper bounds)
$$
\Delta = (B ^ {\star} - B ^ {\prime}) / 3
$$
If  $\Delta > 0$ , a DESSERT structure with the following value of  $L$  solves the search problem from Definition 1.1 with probability  $1 - \delta$ .
$$
L = O \left(\log \left(\frac {N m _ {q} m}{\delta}\right)\right)
$$
${}^{2}L$  additionally depends on the data-dependent parameter  $\Delta$  ,which we elide in the asymptotic bound; see the proof in the appendix for the full expression for  $L$  .
4.3 Runtime Analysis
Theorem 4.3. Suppose that each hash function call runs in time  $O(d)$  and that  $|\mathcal{D}[i]_{t,h}| < T\forall i,t,h$  for some positive threshold  $T$ , which we treat as a data-dependent constant in our analysis. Then, using the assumptions and value of  $L$  from Theorem 4.2, Algorithm 2 solves the Vector Set Search Problem in query time
$$
O \left(m _ {q} \log (N m _ {q} m / \delta) d + m _ {q} N \log (N m _ {q} m / \delta)\right)
$$
This bound is an improvement over a brute force search of  $O(m_q mNd)$  when  $m$  or  $d$  is large. The above theorem relies upon the choice of  $L$  that we derived in Theorem 4.2.
5 Implementation Details