TPU-KNN：峰值 FLOP/s 下的 K 近邻搜索
性能
2022-07-01 v2 机器学习
摘要
本文提出一种新颖的近邻搜索算法，实现了 TPU（Google Tensor Processing Unit，谷歌张量处理单元）峰值性能，在以相似召回率水平优于当前最先进的 GPU 算法。所提算法的设计动机源于一个精确的加速器性能模型，该模型同时考虑了内存与指令瓶颈。我们的算法带有期望召回率的解析保证，且不需要维护复杂的索引数据结构或进行调参，使其适用于频繁更新的应用。我们的工作已在 TPU 上的 Jax 与 Tensorflow 开源包中提供。
引用
@article{arxiv.2206.14286,
title = {TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s},
author = {Felix Chern and Blake Hechtman and Andy Davis and Ruiqi Guo and David Majnemer and Sanjiv Kumar},
journal= {arXiv preprint arXiv:2206.14286},
year = {2022}
}