Fast Exact Retrieval for Nearest-neighbor Lookup ( FERN ) (2024)

Abstract: Exact nearest neighbor search is a computationally intensive process, and even its simpler sibling — vector retrieval — can be computationally complex. This is exacerbated when retrieving vectors which have high-dimension d relative to the number of vectors, N , in the database. Exact nearest neighbor retrieval has been generally acknowledged to be a O(Nd) problem with no sub-linear solutions. Attention has instead shifted towards Approximate Nearest-Neighbor (ANN) retrieval techniques, many of which have sub-linear or even logarithmic time complexities. However, if our intuition from binary search problems (e.g. d=1 vector retrieval) carries, there ought to be a way to retrieve an organized representation of vectors without brute-forcing our way to a solution. For low dimension (e.g. d=2 or d=3 cases), kd-trees provide a O(d\log N) algorithm for retrieval. Unfortunately the algorithm deteriorates rapidly to a O(dN) solution at high dimensions (e.g. k=128 ), in practice. We propose a novel algorithm for logarithmic Fast Exact Retrieval for Nearest-neighbor lookup (FERN), inspired by kd-trees. The algorithm achieves O(d\log N) look-up with 100% recall on 10 million d=128 uniformly randomly generated vectors.

## introduction vector retrieval is pervasive, underlying search engines, transformers, and open-book language models. at heart, one of key attributes of computing systems lie in their ability to retrieve knowledge. sometimes, when this knowledge is sufficiently broad -and a powerful enough retrieval architecture is built -these computing systems may even be so good at retrieving relevant knowledge that they appear to reason @xcite . given the power of knowledge retrieval for both commercial and academic pursuits, significant en-ergy has been devoted towards effectively converting various types of data into vectors, from words @xcite , images @xcite , and audio @xcite to documents, sentences, and paragraphs @xcite . in this work, we differentiate between look-up and search. look-up involves the retrieval of vectors guaranteed to be contained in the database, while search involves the retrieval of queries not necessarily contained in the database. retrieving queries without exact matches may involve instead retrieving that vector's nearest neighbors. the definition of nearest can be further disambiguated into euclidean or cosine similarity measures, among others. note that under the hood, a euclidean distance-based nearest neighbor algorithm can be easily adapted to be cosine similarity-based simply by dividing each vector in the database by its magnitude during insertion. during look-up, the query vector is then also divided by magnitude. the resulting nearest neighbors we obtain are also such by cosine similarity. the scalable and effective look-up of large numbers of high dimensional vectors is thus desired. while the vanilla hashmap algorithm provides o(1) time complexity for scalars, extending to o(d) for vectors in d-dimensional space, this holds only when the cardinality of the hash function range is large relative to the number of elements, n . when the number of elements becomes large relative to the number of bins, b, finding the key within each bin becomes a linear search problem. since we expect each bin to have n b collisions, the time complexity for look-up is o( n d b ). while we can also get fast look-up with a heap in low dimensions, fern can be thought of as an extension of the heap to high dimensions while also presenting a novel approach that could lead to sub-linear vector search. ## related work prior work has taken 3 major directions to attain fast nearest neighbor search. the approaches involve bucketing, divide and conquer, or graphbased approaches. these techniques are specifically tuned to work well for vectors in high dimensional space, which should be unsurprising since each of these techniques is really just an extension of familiar 1-d concepts: hashmaps, binary search, and breadth first search. i believe there are key learnings that can be taken from both exact and approximate retrieval and search settings, even though we are interested in the exact variant. some new techniques such as certified cosine certificates @xcite offer structure-agnostic tweaks to speed up performance. ## fern goal we aim to build an algorithm that satisfies two primary goals: we must be able to perform quick look-up on vectors guaranteed to be contained in the database and we must be able to quickly insert vectors. when designing the algorithm, we initially assume our database contains n vectors, each spanning d dimensional space. the i-th vector v i is defined as follows this process effectively generates vectors lying within a d dimensional "ball"-like shape of radius 1 in each direction. the directions in which the vectors point are also evenly distributed direction-wise. we discover later however, that the algorithm we design with this simplifying assumption maintains logarithmic time lookup for v i of any arbitrary direction and length. in terms of time complexity, we define quick as anything taking logarithmic time -this means that lookup in a database of ten billion vectors should only take seven times longer than lookup in a database of a thousand vectors. that is remarkable, because a naive linear search would take ten million times longer, a nearly intractable time difference at scale. since a logarithmic lookup time and linear space complexity is state of the art (sota), we believe a key contribution of our work is an alternative data structure and algorithm that achieves sota while simultaneously providing the capacity to be extended to logarithmic-time nearest neighbor search, given its unique approach to dividing the vectors by hyperplanes defined based on the vectors in the database rather than measuring along a specific direction like the traditional kd-trees process. each node is an object that stores a vector, pointers to the left and right children, and a pointer to that node's parent node. this results in a binary tree with undirected edges. while we ultimately implement retrieval using a queue structure, this bidirectional edge only adds marginal complexity to the algorithm and underlying data structure while enabling a backtracking-based traversal method. the queue-based method emulates a level-order traversal of candidate nodes while a stack-based backtracking-based traversal method (that fully explores a specific path before backtracking; exploring each sibling node that could not be pruned 2 without potentially missing the nearest neighbor) emulates a depth-first search. methodology we design a novel algorithm that is a variant of kd-trees, but has the capacity to perform better at higher dimensions. broadly, the structure is a binary tree. each node that has both left and right children defines a hyperplane using the vectors of its left and right children as support vectors. the tree is constructed so that all vectors in each child's subtree are on the same side of the hyperplane as that child. this allows us to perform vector look-up in logarithmic time, provided that vectors are added to the database in a sufficiently random manner. it is feasible, however, for an adversarial insertion process to result in a heavily imbalanced tree and consequently worst-case linear look-up time. while we don't observe this as an issue in practice, we can resolve the issue by implementing a slightly more complicated variant of fern -using a variant of the red-black tree technique @xcite to guarantee balanced trees, logarithmic depth, and thus logarithmic retrieval time complexity. there are two key components to our algorithm: one method for insertion (algorithm 1) and another for lookup (algorithm 2). the insertion algorithm (algorithm 1) is fairly concise. when inserting a vector into the tree, it is placed at the root if the tree has not been initialized yet. otherwise, if the current node is missing a left or right child, we insert the vector as a child node. if the node is a leaf node -that is, missing both left and right children -then the left child is always inserted first, before the right child. during insertion, if a node has both left and right children, then we set the current node instead to the child node that is closest to the vector we are inserting. that is, if we form a hyperplane from the set of points equidistant to both left and right children, then we set the current node to the left child if the vector to be inserted lies on the same side of the hyperplane as the left child, otherwise we set the current node to the right child. the result of this insertion algorithm is thatfor a balanced binary tree -we get a maximum tree depth of o(log n ) where n is the number of elements in the database. insertion time per vector is thus o(log n ) since we only visit one node per depth level. algorithm 1 fern insertion 1: function insert(vector) 2: if root not initialized then 3: root ← vectornode(null) 4: end if 5: node ← root 6: while true do 7: if no left child then 8: set left node to vector 9: break 10: else if no right child then 11: set right node to vector 12: break 13: else if vector closer to left child then 14: node ← node.left 15: else 16: node ← node.right 17: ostensibly, we would expect a non-negligible proportion of vectors to be sufficiently far from the hyperplane to be pruned. however, we notice that in practice, as the dimensionality of the vectors increase, so too does the proportion of vectors lying close to the boundary. this makes sense intuitively since we are trying to project increasingly higher dimensions of vectors onto a 1-d line (the line normal to the hyperplane and passing through both support vectors). during retrieval, we effectively perform the depth-first or level-order search, as described previously. for a balanced tree with strong boundaries (that is, most queries lie far away from the hyperplane), per-vector time should be logarithmic with respect to the number of elements already present in the database (another word for our proposed data structure) at insertion-time. however, it becomes linear otherwise. we first define mip and mip_vec, the distance to the nearest neighbor found thus far and the vector representing the near-3 44 est vector retrieved thus far. we then create a queue and insert the root node. algorithm 2 fern lookup require: query ensure: mip_vec 1: mip, mip_vec ← ∞, none 2: curr ← none 3: queue ← [self.root] 4: while queue not empty do 5: curr ← oldest element in queue 6: update mip, mip_vec if curr closer to query 7: if curr has both left and right children then 8: if query is closer to left child then 9: queue.append(curr.left) 10: if query close to boundary then 11: queue.append(curr.right) 12: end if 13: else 14: queue.append(curr.right) 15: if query close to boundary then 16: queue.append(curr.left) 17: end if 18: end if 19: else if curr has left child only then 20: queue.append(curr.left) 21: else if curr has right child only then 22: queue.append(curr.right) 23: for mapping applications, we can add an additional variable, data (a byte array), to the node class. in algorithms 1 and 2, "closeness" is quantified by euler distance, and the "boundary" is the midpoint between left and right child nodes. ## experimental results during experiments, we typically utilize @xmath0. while this may not be characteristic of all data distributions, we note that our architecture is actually agnostic to the distribution of the vectors being inserted. what matters (in terms of potential effects on performance) is the order in which vectors are inserted based on their relative positions. to run our experiments, we use the intel xeon platinum 8380 cpu (2.30 ghz), the same processor used for running the popular ann-benchmark @xcite . for values of n we use 10 4 , 5 * 10 4 , 10 5 , 5 * 10 5 , and 10 6 , 5 * 10 6 , 10 7 . the last setting, equivalent to look-up on 10 million vectors, has comparable values of n and d to many of the euclidean distance based benchmarks in ann-benchmark. we notice a nearly perfect logarithmic time complexity, and at @xmath1. ## discussion during our experiments, we noticed that in the search modality, system dynamics can change drastically based on the dimensionality of the vectors. we experiment with different ways of defining the hyperplanes and various algorithms that would balance or repair the tree to try to guarantee logarithmic retrieval for large d and n . unfortunately, while these algorithms almost universally gave logarithmic time complexity for 100% recall, the performance broke down drastically beyond @xmath2. in particular, we note the importance of having well-defined hyperplane boundaries. boundary sharpness we want to be able to maximize pruning since we achieve o(log 2 n ) time complexity when we prune 50% of the nodes in the database each time we measure the distance between a node and the query). indeed, note that since a k-means based nearest neighbors search allows us to prune up to n/k nodes per comparison, we might wonder why 2-means search doesn't achieve o(log n ) complexity. it's because k-means has an o(n/k + k) time complexity, the n/k term means that we would need to do a linear number of searches regardless of the number of clusters. we observe empirically that the proportion of nodes in the database that are visited increases sharply when there are more nodes that are closer to the hyperplane than the support vectors that define the plane. that is to say, during the insertion process, there will be vectors that lie between a support vector and the hyperplane. now when we're retrieving, the query may lie on one side of the hyperplane but its nearest neighbor may be one of these "in-between" nodes on the other side of the hyperplane. this means we now must be much more prudent when pruning which decreases the proportion of vectors that are pruned and thus increases the time complexity. however in the look-up modality, we achieve logarithmic time complexity on both a vector database of dimension @xmath3, using the larger train splits (@xmath4, 000 to @xmath5, 000, 000). we observe logarithmic time complexity over a diverse dimensions and vector distributions. ## conclusion we are able to achieve our goal of creating a novel vector database structure that achieves state of the art look-up time complexity that is logarithmic in the number of vectors. the algorithm presented here, fern, maintains 100% recall while performing lookup on vectors in high-dimensional space (e.g. @xmath6, and presents a potential path towards a data structure and algorithm that will allow for the first sub-linear exact nearest neighbor retrieval process. we believe that the exact process for attempting to perform binary search on a vector database requires carefully defined hyperplanes, which presents an area for further work. we find the "fixing" step of the red-black tree algorithm to be particularly inspirational as a direction of future work. evaluation on additional datasets and further investigations of recall-retrieval-time tradeoffs could help in the pursuit of sub-linear search. we further also believe that an alternative for hyperplanes is to use a graph based approach, similar to the approach taken in many recent works @xcite @xcite , since this could allow us to more easily divide the database in a well-defined and easy to update way. overall we are excited by the potential and hope to further develop this algorithm in pursuit of sub-linear exact search.