[论文解读] TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s

Felix Chern, Blake A. Hechtman+4 more|arXiv (Cornell University)|Jun 28, 2022

Algorithms and Data Compression被引用 5

一句话总结

本文提出 TPU-KNN，一种 K-近邻搜索算法，通过利用改进的屋顶线模型优化内存和指令级瓶颈，在 Google 的 TPU 上实现了峰值 FLOP/s 性能。该方法在无需复杂索引结构或超参数调优的情况下，实现了最先进的速度-召回权衡，在 Sift1M 和 Glove1.2M 等标准基准测试中，于相似召回率下优于基于 GPU 的基线方法。

ABSTRACT

This paper presents a novel nearest neighbor search algorithm achieving TPU (Google Tensor Processing Unit) peak performance, outperforming state-of-the-art GPU algorithms with similar level of recall. The design of the proposed algorithm is motivated by an accurate accelerator performance model that takes into account both the memory and instruction bottlenecks. Our algorithm comes with an analytical guarantee of recall in expectation and does not require maintaining sophisticated index data structure or tuning, making it suitable for applications with frequent updates. Our work is available in the open-source package of Jax and Tensorflow on TPU.

研究动机与目标

• 通过建模内存和指令级瓶颈，设计一种在 TPU 硬件上实现峰值性能的 K-NN 算法。
• 开发一种具有期望召回保证的算法，避免依赖复杂索引结构或超参数调优。
• 证明基于 TPU 的 K-NN 可在相似召回条件下超越最先进的基于 GPU 的算法。
• 将屋顶线性能模型扩展以考虑指令吞吐量差异，从而实现对加速器上算法性能的准确预测。
• 为 TPU 部署提供可投入生产的开源实现，支持 JAX 和 TensorFlow。

提出的方法

• 作者将屋顶线模型扩展以考虑不同的指令吞吐量，从而能够准确预测加速器性能瓶颈。
• 他们设计了一种 K-NN 算法，通过最小化数据移动并平衡计算与内存带宽，最大化 FLOP/s 利用率。
• 该算法对所有数据点执行暴力距离计算，确保召回基于顺序统计而非压缩伪影。
• 它利用 TPU 特定的优化技术，如高效向量化、批量处理以及基于性能模型的内存访问模式。
• 该方法使用逐系数指令模型，识别并缓解低吞吐量操作带来的带宽瓶颈。
• 实现代码以 JAX 和 TensorFlow 开源发布，新增一个操作 'approx_max_k'，用于在 TPU 上高效执行 top-K 选择。

实验结果

研究问题

• RQ1 能否通过同时建模内存和指令级瓶颈，设计一种在 TPU 上实现峰值 FLOP/s 利用率的 K-NN 算法？
• RQ2 将屋顶线性能模型扩展以包含指令吞吐量差异，是否能实现对 K-NN 工作负载在加速器上性能的准确预测？
• RQ3 在 TPU 上，是否能够使一种简单、暴力的 K-NN 方法在相似召回水平下超越 IVF-Flat 和 IVF-PQ 等复杂的压缩域搜索方法？
• RQ4 是否可能在避免复杂索引结构或超参数调优的前提下，实现高召回率并具备分析性保证？
• RQ5 在相似召回条件下，TPU 优化的 K-NN 算法性能与最先进的基于 GPU 的 K-NN 方法相比如何？

主要发现

• TPU-KNN 在 Sift1M 和 Glove1.2M 基准测试中实现了最先进的速度-召回权衡，在相似召回水平下优于基于 GPU 的 IVF-Flat、IVF-PQ 和原始搜索方法。
• TPU 实现的性能与扩展后的屋顶线模型预测结果高度吻合，验证了性能建模框架的准确性。
• 该算法通过最小化数据移动并高效平衡计算与内存带宽，在 TPU 上实现了峰值 FLOP/s 利用率。
• 由于该方法依赖于顺序统计而非依赖压缩的启发式方法，因此在不同数据集上保持了稳定的召回率，对数据分布变化具有鲁棒性。
• 在 JAX 和 TensorFlow 中的开源发布，使得该算法可投入生产环境使用，并可轻松集成到现有的 TPU 硬件机器学习流水线中。
• 由于算法计算所有成对距离而不进行降维或压缩，因此对维度灾难具有免疫性。