Inference-time sparse attention with asymmetric indexing

• URL: http://arxiv.org/abs/2502.08246v1
• Date: Wed, 12 Feb 2025 09:39:54 GMT
• Title: Inference-time sparse attention with asymmetric indexing
• Authors: Pierre-Emmanuel Mazaré, Gergely Szilvasy, Maria Lomeli, Francisco Massa, Naila Murray, Hervé Jégou, Matthijs Douze,
• Abstract summary: Self-attention in transformer models is an incremental associative memory that maps key vectors to value vectors. Standard partitioning methods yield poor results in this context. We introduce SAAP (Self-Attention with Asymmetric Partitions), which overcomes these problems.
• Score: 23.305984099821618
• License:
• Abstract: Self-attention in transformer models is an incremental associative memory that maps key vectors to value vectors. One way to speed up self-attention is to employ GPU-compliant vector search algorithms, yet the standard partitioning methods yield poor results in this context, because (1) keys and queries follow different distributions and (2) the effect of RoPE positional encoding. In this paper, we introduce SAAP (Self-Attention with Asymmetric Partitions), which overcomes these problems. It is an asymmetrical indexing technique that employs distinct partitions for keys and queries, thereby approximating self-attention with a data-adaptive sparsity pattern. It works on pretrained language models without finetuning, as it only requires to train (offline) a small query classifier. On a long context Llama 3.1-8b model, with sequences ranging from 100k to 500k tokens, our method typically reduces by a factor 20 the fraction of memory that needs to be looked-up, which translates to a time saving of 60\% when compared to FlashAttention-v2.

Related papers

• Squeezed Attention: Accelerating Long Context Length LLM Inference [64.11145320159126]
  We propose Squeezed Attention as a mechanism to accelerate LLM applications where a large portion of the input prompt is fixed. We use K-means clustering offline to group the keys for the fixed context based on semantic similarity and represent each cluster with a single centroid value. We then compute exact attention using only these important keys from the fixed context, thereby reducing bandwidth and computational costs.
  arXiv  Detail & Related papers  (2024-11-14T18:54:19Z)
• RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval [24.472784635757016]
  RetrievalAttention is a training-free approach to both accelerate attention computation and reduce GPU memory consumption. We show that RetrievalAttention achieves near full attention accuracy while only requiring access to 1--3% of the data.
  arXiv  Detail & Related papers  (2024-09-16T17:59:52Z)
• Training-Free Exponential Context Extension via Cascading KV Cache [49.608367376911694]
  We introduce a novel mechanism that leverages cascading sub-cache buffers to selectively retain the most relevant tokens. Our method reduces prefill stage latency by a factor of 6.8 when compared to flash attention on 1M tokens.
  arXiv  Detail & Related papers  (2024-06-24T03:59:17Z)
• Breaking the Attention Bottleneck [0.0]
  This paper develops a generative function as attention or activation replacement. It still has the auto-regressive character by comparing each token with the previous one. The concept of attention replacement is distributed under the AGPL v3 license at https://gitlab.com/Bachstelzecausal_generation.
  arXiv  Detail & Related papers  (2024-06-16T12:06:58Z)
• ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition [3.659659889927316]
  ChunkAttention is a prefix-aware self-attention module for large language models. It can detect matching prompt prefixes across multiple requests and share their key/value tensors in memory at runtime. Experiments show that ChunkAttention can speed up the self-attention kernel by 3.2-4.8$times$ compared to the state-of-the-art implementation.
  arXiv  Detail & Related papers  (2024-02-23T09:29:19Z)
• Similarity search in the blink of an eye with compressed indices [3.39271933237479]
  Graph-based indices are currently the best performing techniques for billion-scale similarity search. We present new techniques and systems for creating faster and smaller graph-based indices.
  arXiv  Detail & Related papers  (2023-04-07T23:10:39Z)
• CITADEL: Conditional Token Interaction via Dynamic Lexical Routing for Efficient and Effective Multi-Vector Retrieval [72.90850213615427]
  Multi-vector retrieval methods combine the merits of sparse (e.g. BM25) and dense (e.g. DPR) retrievers. These methods are orders of magnitude slower and need much more space to store their indices compared to their single-vector counterparts. We propose conditional token interaction via dynamic lexical routing, namely CITADEL, for efficient and effective multi-vector retrieval.
  arXiv  Detail & Related papers  (2022-11-18T18:27:35Z)
• ClusTR: Exploring Efficient Self-attention via Clustering for Vision Transformers [70.76313507550684]
  We propose a content-based sparse attention method, as an alternative to dense self-attention. Specifically, we cluster and then aggregate key and value tokens, as a content-based method of reducing the total token count. The resulting clustered-token sequence retains the semantic diversity of the original signal, but can be processed at a lower computational cost.
  arXiv  Detail & Related papers  (2022-08-28T04:18:27Z)
• Learning Tracking Representations via Dual-Branch Fully Transformer Networks [82.21771581817937]
  We present a Siamese-like Dual-branch network based on solely Transformers for tracking. We extract a feature vector for each patch based on its matching results with others within an attention window. The method achieves better or comparable results as the best-performing methods.
  arXiv  Detail & Related papers  (2021-12-05T13:44:33Z)
• You Only Sample (Almost) Once: Linear Cost Self-Attention Via Bernoulli Sampling [38.34914626128062]
  We show that a Bernoulli sampling attention mechanism based on Locality Sensitive Hashing (LSH) decreases the quadratic complexity of such models to linear. We evaluate our algorithm on the GLUE benchmark with standard 512 sequence length where we see favorable performance relative to a standard pretrained Transformer.
  arXiv  Detail & Related papers  (2021-11-18T14:24:34Z)
• IRLI: Iterative Re-partitioning for Learning to Index [104.72641345738425]
  Methods have to trade between obtaining high accuracy while maintaining load balance and scalability in distributed settings. We propose a novel approach called IRLI, which iteratively partitions the items by learning the relevant buckets directly from the query-item relevance data. We mathematically show that IRLI retrieves the correct item with high probability under very natural assumptions and provides superior load balancing.
  arXiv  Detail & Related papers  (2021-03-17T23:13:25Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.