面向大规模向量的k-最大内积搜索高效算法与实现问询

现有查询空间Q与键空间K，均包含N≈10⁶个d≈50维实向量，向量服从正态分布。需为每个查询q找出k≈10个与q内积最高的键k_i，期望在高性能GPU上1秒内完成N=10⁶规模的近似求解。核心诉求：

• 总复杂度为O(N log N)的高效算法或实现
• 获取GAIPS论文的对应实现参考

已尝试方案及存在的问题 #

• naive计算：全量计算内积后取top-k，复杂度过高无法满足性能要求
• pykeops：基于符号矩阵实现，性能接近FAISS，但未达到1秒内的目标
• FAISS：在RTX2060上耗时约20s，复杂度为亚二次而非O(N log N)，不符合复杂度要求
• Annoy：未测试，但据文档需串行处理查询，并行性能无法支撑大规模检索需求
• TorchPQ：因需逐个添加键，性能极慢
• 球树分支定界算法：基于《Maximum Inner-Product Search using Tree Data-structures》自研PyTorch C++扩展，但运行在CPU上，性能比GPU方案慢数个数量级
• ALSH：仅找到无GPU支持的实现，预计性能不及FAISS

性能 profiling 结果（测试环境：Intel i7(9代)+RTX2060） #

• (1) 符号矩阵全量内积计算
• (2) FAISS
• (3) 自研球树C++扩展

内容的提问来源于stack exchange，提问作者Jonas De Schouwer