Maximum Inner Product Search (MIPS) is a ubiquitous task in machine learning applications such as recommendation systems. Given a query vector and $n$ atom vectors in $d$-dimensional space, the goal of MIPS is to find the atom that has the highest inner product with the query vector. Existing MIPS algorithms scale at least as $O(\sqrt{d})$, which becomes computationally prohibitive in high-dimensional settings. In this work, we present BanditMIPS, a novel randomized MIPS algorithm whose complexity is independent of $d$. BanditMIPS estimates the inner product for each atom by subsampling coordinates and adaptively evaluates more coordinates for more promising atoms. The specific adaptive sampling strategy is motivated by multi-armed bandits. We provide theoretical guarantees that BanditMIPS returns the correct answer with high probability, while improving the complexity in $d$ from $O(\sqrt{d})$ to $O(1)$. We also perform experiments on four synthetic and real-world datasets and demonstrate that BanditMIPS outperforms prior state-of-the-art algorithms. For example, in the Movie Lens dataset ($n$=4,000, $d$=6,000), BanditMIPS is 20$\times$ faster than the next best algorithm while returning the same answer. BanditMIPS requires no preprocessing of the data and includes a hyperparameter that practitioners may use to trade off accuracy and runtime. We also propose a variant of our algorithm, named BanditMIPS-$\alpha$, which achieves further speedups by employing non-uniform sampling across coordinates. Finally, we demonstrate how known preprocessing techniques can be used to further accelerate BanditMIPS, and discuss applications to Matching Pursuit and Fourier analysis.

翻译：最大内积搜索（MIPS）是机器学习应用（如推荐系统）中的一项基础任务。给定一个查询向量和$d$维空间中的$n$个原子向量，MIPS的目标是找到与查询向量内积最大的原子。现有MIPS算法的复杂度至少为$O(\sqrt{d})$，在高维场景下计算成本过高。本文提出BanditMIPS，一种新颖的随机化MIPS算法，其复杂度与$d$无关。BanditMIPS通过子采样坐标估计每个原子的内积，并对更有潜力的原子自适应地评估更多坐标。这种自适应采样策略受多臂老虎机启发。我们提供理论保证：BanditMIPS能以高概率返回正确结果，同时将复杂度从$O(\sqrt{d})$改进为$O(1)$。在四个合成和真实数据集上的实验表明，BanditMIPS优于现有最优算法。例如，在Movie Lens数据集（$n$=4,000, $d$=6,000）中，BanditMIPS在返回相同结果的前提下，速度比次优算法快20倍。该算法无需数据预处理，并包含一个可供实践者权衡精度与运行时间的超参数。我们还提出变体算法BanditMIPS-$\alpha$，通过跨坐标非均匀采样实现进一步加速。最后，我们展示了已知预处理技术如何进一步加速BanditMIPS，并讨论了其在匹配追踪和傅里叶分析中的应用。