论文标题
部分可观测时空混沌系统的无模型预测
TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s
论文作者
论文摘要
本文提出了一种新颖的邻居搜索算法，可实现TPU（Google Tensor处理单元）的峰值性能，超过了最先进的GPU算法，其召回水平相似。所提出的算法的设计是由准确的加速器性能模型的动机，该模型同时考虑了内存和指令瓶颈。我们的算法具有期望中召回的分析保证，并且不需要维护复杂的索引数据结构或调整，因此它适用于频繁更新的应用程序。我们的工作可在TPU上的Jax和Tensorflow的开源软件包中获得。
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.
