A Faster Generalized Two-Stage Approximate Top-K

2025/06/04 by Yashas Samaga, Varun Yerram, Samaga, Yashas +7
Computer Science · #Complexity and Algorithms in Graphs #Data Structures and Algorithms (cs.DS) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning and Algorithms #Stochastic Gradient Optimization Techniques

paper · pdf · doi:10.48550/arxiv.2506.04165

openalex publication_date 2025/06/04 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

We consider the Top-K selection problem, which aims to identify the largest K elements in an array. Top-K selection arises in many machine learning algorithms and often becomes a bottleneck on accelerators, which are optimized for dense matrix multiplications. To address this problem, Chern et al. (2022) proposed a fast two-stage approximate Top-K algorithm that: (i) partitions the input array into equal-sized chunks and selects the top-1 element from each partition; and (ii) sorts the resulting smaller subset and returns the top K elements. In this paper, we generalize the first stage so that each partition selects the top K' elements (for 1 ≤ K' ≤ K). Our contributions include: (i) an expression for the expected recall of this generalized algorithm under random partitioning, and a demonstration that choosing K' > 1 with fewer partitions in the first stage more effectively reduces the input size to the second stage while maintaining the same expected recall as the original algorithm; (ii) a bound on the expected recall of the original algorithm as a function of the algorithm parameters that is provably tighter by a factor of 2 than the bound reported by Chern et al. (2022); and (iii) an implementation of our algorithm on Cloud TPUv5e that achieves approximately an order of magnitude speedup over the original algorithm without sacrificing recall.

Citations

• Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference
• RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
• Exploiting Sparsity for Long Context Inference: Million Token Contexts on Commodity GPUs
• Approximate Top-k for Increased Parallelism
• TidalDecode: Fast and Accurate LLM Decoding with Position Persistent Sparse Attention
• RTop-K: Ultra-Fast Row-Wise Top-K Selection for Neural Network Acceleration on GPUs
• Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference
• Superposed Decoding: Multiple Generations from a Single Autoregressive Inference Pass
• HiRE: High Recall Approximate Top-k Estimation for Efficient LLM Inference
• ReLU² Wins: Discovering Efficient Activation Functions for Sparse LLMs
• DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
• LLM in a flash: Efficient Large Language Model Inference with Limited Memory
• ReLU Strikes Back: Exploiting Activation Sparsity in Large Language Models
• Adaptive Top-K in SGD for Communication-Efficient Distributed Learning
• Treeformer: Dense Gradient Trees for Efficient Attention Computation
• Improving language models by retrieving from trillions of tokens
• Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks
• Efficient Content-Based Sparse Attention with Routing Transformers
• Understanding Top-k Sparsification in Distributed Deep Learning
• Hashing for Similarity Search: A Survey

Cited by

• omni-macos: On-Device Omni-Modal Search on Apple Silicon

Related

• Approximate Top-k for Increased Parallelism
• Communication Efficient Algorithms for Top-k Selection Problems
• How to select the largest k elements from evolving data?
• Parallel Algorithms for Select and Partition with Noisy Comparisons
• Optimally selecting the top k values from X+Y with layer-ordered heaps