SOCKET: SOft Collision Kernel EsTimator for Sparse Attention
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , , , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2026
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866914541669974016 |
|---|---|
| author | Joshi, Sahil Chowdhury, Agniva Bellinger, Wyatt Kanakamedala, Amar Singh, Ekam Le, Hoang Anh Duy Desai, Aditya Shrivastava, Anshumali |
| author_facet | Joshi, Sahil Chowdhury, Agniva Bellinger, Wyatt Kanakamedala, Amar Singh, Ekam Le, Hoang Anh Duy Desai, Aditya Shrivastava, Anshumali |
| contents | Exploiting sparsity during long-context inference is key to scaling large language models, as attention dominates the cost of autoregressive decoding. Sparse attention reduces this cost by restricting computation to a subset of tokens, but its effectiveness depends on efficient scoring and selection at inference time. 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 yields binary collision signals that limit ranking quality and require substantial memory to perform well. In contrast, soft LSH accumulates graded collision evidence across hash tables, preserving top-k ordering with significantly less memory. This reframes LSH from a candidate generator into a principled scoring kernel for sparse attention. Leveraging this property, SOCKET enables efficient token selection without ad hoc voting and matches or surpasses prior sparse attention methods across multiple long-context benchmarks. With a custom CUDA scoring kernel and a Flash Decode Triton backend, SOCKET achieves up to 1.5$\times$ higher throughput than FlashAttention. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2602_06283 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | SOCKET: SOft Collision Kernel EsTimator for Sparse Attention Joshi, Sahil Chowdhury, Agniva Bellinger, Wyatt Kanakamedala, Amar Singh, Ekam Le, Hoang Anh Duy Desai, Aditya Shrivastava, Anshumali Machine Learning Exploiting sparsity during long-context inference is key to scaling large language models, as attention dominates the cost of autoregressive decoding. Sparse attention reduces this cost by restricting computation to a subset of tokens, but its effectiveness depends on efficient scoring and selection at inference time. 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 yields binary collision signals that limit ranking quality and require substantial memory to perform well. In contrast, soft LSH accumulates graded collision evidence across hash tables, preserving top-k ordering with significantly less memory. This reframes LSH from a candidate generator into a principled scoring kernel for sparse attention. Leveraging this property, SOCKET enables efficient token selection without ad hoc voting and matches or surpasses prior sparse attention methods across multiple long-context benchmarks. With a custom CUDA scoring kernel and a Flash Decode Triton backend, SOCKET achieves up to 1.5$\times$ higher throughput than FlashAttention. |
| title | SOCKET: SOft Collision Kernel EsTimator for Sparse Attention |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2602.06283 |