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能高概率返回正确结果,同时将$d$的复杂度从$O(\sqrt{d})$改进为$O(1)$。在四个合成数据集和真实数据集上的实验表明,BanditMIPS优于现有最优算法。例如,在Movie Lens数据集($n$=4,000,$d$=6,000)上,BanditMIPS比次优算法快20倍,同时返回相同结果。BanditMIPS无需数据预处理,并包含一个超参数,供实践者在准确性与运行时间之间进行权衡。我们还提出算法变体BanditMIPS-$\alpha$,通过采用跨坐标非均匀采样进一步加速。最后,我们展示了如何利用已知预处理技术进一步加速BanditMIPS,并讨论了其在匹配追踪和傅里叶分析中的应用。