# TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s

arXiv: https://arxiv.org/abs/2206.14286 · HTML: https://mlsanity.com/paper/2206.14286/

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 56 references that are in the map.
- First posted on arXiv 28 Jun 2022.

## Builds on (most cited first, 19 of 19)

- [Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervision](https://mlsanity.com/paper/2102.05918/) (2021)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs](https://mlsanity.com/paper/1603.09320/) (2016)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Universal Sentence Encoder](https://mlsanity.com/paper/1803.11175/) (2018)
- [ANN-Benchmarks: A Benchmarking Tool for Approximate Nearest Neighbor Algorithms](https://mlsanity.com/paper/1807.05614/) (2018)
- [Practical and Optimal LSH for Angular Distance](https://mlsanity.com/paper/1509.02897/) (2015)
- [NVIDIA Tensor Core Programmability, Performance & Precision](https://mlsanity.com/paper/1803.04014/) (2018)
- [Learning to Hash for Indexing Big Data - A Survey](https://mlsanity.com/paper/1509.05472/) (2015)
- [Revisiting the Inverted Indices for Billion-Scale Approximate Nearest Neighbors](https://mlsanity.com/paper/1802.02422/) (2018)
- [Hashing for Similarity Search: A Survey](https://mlsanity.com/paper/1408.2927/) (2014)
- [Statistical Analysis of Nearest Neighbor Methods for Anomaly Detection](https://mlsanity.com/paper/1907.03813/) (2019)
- [On Symmetric and Asymmetric LSHs for Inner Product Search](https://mlsanity.com/paper/1410.5518/) (2014)
- [Hardness of Approximate Nearest Neighbor Search](https://mlsanity.com/paper/1803.00904/) (2018)
- [Hierarchical Roofline Performance Analysis for Deep Learning Applications](https://mlsanity.com/paper/2009.05257/) (2020)
- [Time-Based Roofline for Deep Learning Performance Analysis](https://mlsanity.com/paper/2009.04598/) (2020)

## Built on by (most cited first, 18 of 18)

- [The Faiss library](https://mlsanity.com/paper/2401.08281/) (2024)
- [REVEAL: Retrieval-Augmented Visual-Language Pre-Training with Multi-Source Multimodal Knowledge Memory](https://mlsanity.com/paper/2212.05221/) (2022)
- [The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers](https://mlsanity.com/paper/2210.06313/) (2022)
- [Ask in Any Modality: A Comprehensive Survey on Multimodal Retrieval-Augmented Generation](https://mlsanity.com/paper/2502.08826/) (2025)
- [Sparsity-Constrained Optimal Transport](https://mlsanity.com/paper/2209.15466/) (2022)
- [Chameleon: a Heterogeneous and Disaggregated Accelerator System for Retrieval-Augmented Language Models](https://mlsanity.com/paper/2310.09949/) (2023)
- [Revisiting Neural Retrieval on Accelerators](https://mlsanity.com/paper/2306.04039/) (2023)
- [DeDrift: Robust Similarity Search under Content Drift](https://mlsanity.com/paper/2308.02752/) (2023)
- [TeleRAG: Efficient Retrieval-Augmented Generation Inference with Lookahead Retrieval](https://mlsanity.com/paper/2502.20969/) (2025)
- [AVIS: Autonomous Visual Information Seeking with Large Language Model Agent](https://mlsanity.com/paper/2306.08129/) (2023)
- [Co-design Hardware and Algorithm for Vector Search](https://mlsanity.com/paper/2306.11182/) (2023)
- [Spark Transformer: Reactivating Sparsity in FFN and Attention](https://mlsanity.com/paper/2506.06644/) (2025)
- [Hierarchical Retrieval: The Geometry and a Pretrain-Finetune Recipe](https://mlsanity.com/paper/2509.16411/) (2025)
- [Retrieval with Learned Similarities](https://mlsanity.com/paper/2407.15462/) (2024)
- [Augmentation with Projection: Towards an Effective and Efficient Data Augmentation Paradigm for Distillation](https://mlsanity.com/paper/2210.11768/) (2022)
- [Towards the Identifiability in Noisy Label Learning: A Multinomial Mixture Modelling Approach](https://mlsanity.com/paper/2301.01405/) (2023)
- [Machine learning and high dimensional vector search](https://mlsanity.com/paper/2502.16931/) (2025)
- [Prof-K: Probabilistic One-Pass Filtering for Efficient Top-k Selection](https://mlsanity.com/paper/2608.12573/) (2026)

## Abstract (arXiv)

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.