Nearest neighbor search is a fundamental problem in computational geometry with applications in computer vision, recommendation systems, and natural language processing. However, as data dimensions increase, the curse of dimensionality causes distances between points to lose their discriminative power, making exact search expensive and infeasible. To address this, approximate nearest neighbor (ANN) methods have been proposed, offering faster query times with a trade-off in accuracy. These methods, including tree-based, hashing, and graph-based approaches, perform well on dense low-dimensional datasets. However, on sparse high-dimensional data, they suboptimally utilize compute power and memory bandwidth, resulting in low query throughput. Therefore, we work toward a low latency ray-tracing-inspired algorithm specialized for high-dimensional sparse ANN search.
We begin with a dataset of search points $\mathcal{X}=\{x_1, x_2, \dots, x_n\}$ and queries $Q=\{q_1, q_2, \dots, q_{n'}\}$ with $x_i, q_j \in \mathbb{R}^D$. We would like to find for each query $q_j$: $$ { \arg \min_{x_i} \text{d}(x_i, q_j) } $$ for some distance metric $\text{d}(\cdot, \cdot)$ over $\mathbb{R}^D$ which is usually the Euclidean metric. For maximum inner product similarity (MIPS) where higher is better, we instead find for each $q_j$: $$ { \arg \max_{x_i} \text{MIPS}(x_i, q_j) } $$
The search accuracy is measured with recall metric. $1$-recall@$R$ means the algorithm returns $R$ candidates for the true (i.e. $1$) nearest neighbor, and if the 1-nearest-neighbor is contained in this set, then an indicator score is incremented. Indicator scores are then accumulated and averaged across all queries. $$ { 1\text{-recall@}R = \frac{1}{n'}\sum_{j=1}^{n'}\mathbb{I}_j\left[ \text{1-nearest-neighbor is in $q_j$'s candidate set of size $R$ } \right] } $$
These definitions are similar for $k$NN search where accuracy is typically measured by $k$-recall@$R$. High accuracy ANN search algorithms have high $R$-recall@$R$, but of course this is difficult to achieve and balance with performance.
In ray tracing, rays are cast from the camera through each pixel of the projected image into a virtual scene of objects. Ray tracing engines simulate object interactions in the scene to render pixels with accurate colors. Here, we will construct a scene of spheres during the offline step and cast rays toward them during the online step.
During offline scene construction, we fix the subspace dimensionality to $d=2$ and hence the number of subspaces is $m=D/d$. In each subspace, we train $c$ cluster centroids using $k$-means and create corresponding spheres ($m \times c$ in total) which encapsulate $r$-neighborhoods of candidate search points. We fix the subspace dimensionality to create an interpretable scene: dimensions $1, 2$ correspond to $x, y$ coordinates of the cluster centroid and sphere. The $z$-dimension is chosen such that spheres do not overlap between subspaces. Finally, to accelerate scene traversal time, we construct a Bounding Volume Hierarchy on the sphere scene, allowing us to rapidly prune the candidate search space at runtime.
During online search, A query $q_j \in \mathbb{R}^D$ is decomposed into $q_{j0}, q_{j1}, \dots, q_{jm} \in \mathbb{R}^d$. Since $d=2$, the components of the query subvector in a given subspace are the $x, y$ coordinates for the corresponding ray. We generate a ray for each subspace and cast them into their respective subspace scenes. A first-stage candidate set is generated by taking the union of points contained in the intersected spheres across all subspaces. Intuitively, we only consider in each subspace points no more than a distance of $r$ from the sphere centroid as if a point is encapsulated in an intersected $r$-neighborhood, then it is likely a good nearest-neighbor canadidate. The final candidate set is generated by selecting from the first-stage set the closest $R$ points to $q_j$, through selective distance accumulation.
The core application is accelerating approximate nearest neighbor search for text embeddings, some of which exhibit these sparse, high-dimensional characteristics. As a result, our algorithm is widely applicable to semantic search which finds documents that are conceptually similar to a query beyond keyword matching, recommendation systems which suggest related products, articles, or videos based on user preferences, and question answering systems which identify the most relevant passage in a document to answer a given question. It is useful in any scenario in which we want to quickly query the embedding space!