mlsanityData to 14 Oct 2026

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

arXiv 2206.14286 ↗First posted 28 Jun 2022cs.PF cs.LGcited by 18

Not on the map. Why?

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.

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

19 of 56 references.

References · 19

  1. 2021Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervisioncited by 3,864
  2. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
  3. 2017Billion-scale similarity search with GPUscited by 2,497
  4. 2020REALM: Retrieval-Augmented Language Model Pre-Trainingcited by 2,141
  5. 2021Improving language models by retrieving from trillions of tokenscited by 1,293
  6. 2016Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphscited by 861
  7. 2017In-Datacenter Performance Analysis of a Tensor Processing Unitcited by 618
  8. 2018Universal Sentence Encodercited by 489
  9. 2018ANN-Benchmarks: A Benchmarking Tool for Approximate Nearest Neighbor Algorithmscited by 130
  10. 2015Practical and Optimal LSH for Angular Distancecited by 106
  11. 2018NVIDIA Tensor Core Programmability, Performance & Precisioncited by 79
  12. 2015Learning to Hash for Indexing Big Data - A Surveycited by 39
7 more
  1. 2018Revisiting the Inverted Indices for Billion-Scale Approximate Nearest Neighborscited by 37
  2. 2014Hashing for Similarity Search: A Surveycited by 34
  3. 2019Statistical Analysis of Nearest Neighbor Methods for Anomaly Detectioncited by 29
  4. 2014On Symmetric and Asymmetric LSHs for Inner Product Searchcited by 23
  5. 2018Hardness of Approximate Nearest Neighbor Searchcited by 19
  6. 2020Hierarchical Roofline Performance Analysis for Deep Learning Applicationscited by 11
  7. 2020Time-Based Roofline for Deep Learning Performance Analysiscited by 7

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2022: 4; 2023: 6; 2024: 2; 2025: 5; 2026: 1.

18 papers.

Sort:

By area:

  1. 2024The Faiss librarycited by 376
  2. 2022REVEAL: Retrieval-Augmented Visual-Language Pre-Training with Multi-Source Multimodal Knowledge Memorycited by 123
  3. 2022The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformerscited by 122
  4. 2025Ask in Any Modality: A Comprehensive Survey on Multimodal Retrieval-Augmented Generationcited by 42
  5. 2022Sparsity-Constrained Optimal Transportcited by 32
  6. 2023Chameleon: a Heterogeneous and Disaggregated Accelerator System for Retrieval-Augmented Language Modelscited by 23
  7. 2023Revisiting Neural Retrieval on Acceleratorscited by 14
  8. 2023DeDrift: Robust Similarity Search under Content Driftcited by 11
  9. 2025TeleRAG: Efficient Retrieval-Augmented Generation Inference with Lookahead Retrievalcited by 9
  10. 2023AVIS: Autonomous Visual Information Seeking with Large Language Model Agentcited by 8
  11. 2023Co-design Hardware and Algorithm for Vector Searchcited by 6
  12. 2025Spark Transformer: Reactivating Sparsity in FFN and Attentioncited by 5
6 more
  1. 2025Hierarchical Retrieval: The Geometry and a Pretrain-Finetune Recipecited by 4
  2. 2024Retrieval with Learned Similaritiescited by 3
  3. 2022Augmentation with Projection: Towards an Effective and Efficient Data Augmentation Paradigm for Distillationcited by 2
  4. 2023Towards the Identifiability in Noisy Label Learning: A Multinomial Mixture Modelling Approachcited by 1
  5. 2025Machine learning and high dimensional vector searchcited by 1
  6. 2026Prof-K: Probabilistic One-Pass Filtering for Efficient Top-k Selectioncited by 0