# SOCKET: SOft Collision Kernel EsTimator for Sparse Attention

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 58 references that are in the map.
- First posted on arXiv 6 Feb 2026.

## Builds on (most cited first, 30 of 41)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models](https://mlsanity.com/paper/2301.12597/) (2023)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [StarCoder: may the source be with you!](https://mlsanity.com/paper/2305.06161/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference](https://mlsanity.com/paper/2406.10774/) (2024)
- [Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs](https://mlsanity.com/paper/2310.01801/) (2023)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [An Introduction to Matrix Concentration Inequalities](https://mlsanity.com/paper/1501.01571/) (2015)
- [Sketching as a Tool for Numerical Linear Algebra](https://mlsanity.com/paper/1411.4357/) (2014)
- [InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory](https://mlsanity.com/paper/2402.04617/) (2024)
- [RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval](https://mlsanity.com/paper/2409.10516/) (2024)
- [MagicPIG: LSH Sampling for Efficient LLM Generation](https://mlsanity.com/paper/2410.16179/) (2024)
- [PQCache: Product Quantization-based KVCache for Long Context LLM Inference](https://mlsanity.com/paper/2407.12820/) (2024)

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

- [ParisKV: Fast and Drift-Robust KV-Cache Retrieval for Long-Context LLMs](https://mlsanity.com/paper/2602.07721/) (2026)
- [Forget Without Compromise: Nexus Sampling for Streaming KV-Cache Eviction Under Fixed Budgets](https://mlsanity.com/paper/2606.23961/) (2026)
- [Inference Time Context Sparsity: Illusion or Opportunity?](https://mlsanity.com/paper/2605.24168/) (2026)

## Abstract (arXiv)

Exploiting sparsity is key to efficient long-context inference, as attention dominates the cost of autoregressive decoding. Sparse attention reduces this cost by restricting computation to a subset of tokens, but its effectiveness hinges on fast and accurate token scoring and selection at inference time. Data-agnostic approaches offer an attractive way to perform this selection, but often incur substantial memory overhead to maintain high recall. We revisit Locality-Sensitive Hashing (LSH) and introduce SOCKET, a SOft Collision Kernel EsTimator that replaces hard bucket matches with probabilistic, similarity-aware aggregation. Traditional LSH relies on binary collision signals, providing limited information for ranking tokens and necessitating many hash tables for accurate retrieval. In contrast, soft LSH accumulates graded collision evidence across hash tables, closely preserving the true top-k ordering with significantly less memory. This reframes LSH from a candidate-generation mechanism into a principled scoring kernel for sparse attention. Building on this insight, SOCKET enables efficient token selection without ad hoc voting and matches or outperforms existing sparse attention methods across multiple long-context benchmarks and diverse language models. With a custom set of CUDA/Triton kernels for scoring, selection, and attention, SOCKET achieves up to approximately 1.5× higher throughput than FlashAttention. Code is open-sourced at https://github.com/amarka8/SOCKET.