RaaS: Reasoning-Aware Attention Sparsity for Efficient LLM Reasoning

• URL: http://arxiv.org/abs/2502.11147v2
• Date: Fri, 30 May 2025 11:43:48 GMT
• Title: RaaS: Reasoning-Aware Attention Sparsity for Efficient LLM Reasoning
• Authors: Junhao Hu, Wenrui Huang, Weidong Wang, Zhenwen Li, Tiancheng Hu, Zhixia Liu, Xusheng Chen, Tao Xie, Yizhou Shan,
• Abstract summary: Large Language Models (LLMs) have demonstrated strong capabilities across various domains.<n>These algorithms struggle with the "impossible trinity" of accuracy, time, and memory.<n>We propose a new algorithm RaaS that identifies milestone tokens and retains their KV vectors until they are no longer needed.
• Score: 14.409253716114213
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Large Language Models (LLMs) have demonstrated strong capabilities across various domains, with recent advancements in challenging reasoning tasks such as mathematics and programming. However, solving reasoning tasks often requires an LLM to generate long sequences, incurring $O(N)$ time and memory complexities per token, where $N$ is the current sequence length. To reduce complexities, existing sparsity-based algorithms propose to retain Key-Value (KV) vectors, the intermediate representations of only the most critical tokens. However, these algorithms struggle with the "impossible trinity" of accuracy, time, and memory. For example, the state-of-the-art algorithm, Quest, achieves high accuracy with $O(L)$ time but $O(N)$ memory ($L$ is the cache budget, $L \ll N$). To address the "impossible trinity", in this paper, we identify a new attention pattern during the decode stage of reasoning tasks, where milestone tokens (analogous to lemmas in mathematical proofs) emerge, are utilized, and then become unimportant afterward. Based on this pattern, we propose a new algorithm RaaS that identifies milestone tokens and retains their KV vectors until they are no longer needed, achieving high accuracy with $O(L)$ time and $O(L)$ memory complexities.

Related papers

• Active Learning of Symbolic Automata Over Rational Numbers [4.921484415160938]
  Angluin's $L*$ algorithm learns finite-state automata (DFAs) in time when provided with a minimally adequate teacher.<n>We extend $L*$ to learn symbolic automata whose transitions use predicates over rational numbers, i.e., over infinite and dense alphabets.<n>Our proposed algorithm is optimal in the sense that it asks a number of queries to the teacher that is at most linear with respect to the number of transitions, and to the representation size of the predicates.
  arXiv  Detail & Related papers  (2025-11-15T18:04:36Z)
• Compression Barriers for Autoregressive Transformers [0.8331054243801623]
  Key limitation of autoregressive Transformers is the large memory needed to cache all previous key-value embeddings. We show that any algorithm requires $Omega(dcdot ed)$ space and prove, using tight bounds on covering numbers, that SubGen, proposed by Zandieh, Han, Mirrokni and Karbasi, matches this bound.
  arXiv  Detail & Related papers  (2025-02-21T21:37:52Z)
• HashAttention: Semantic Sparsity for Faster Inference [91.54218318798603]
  HashAttention is a principled approach casting pivotal token identification as a recommendation problem.<n>It efficiently identifies pivotal tokens for a given query in this Hamming space using bitwise operations.<n>It can reduce the number of tokens used by a factor of $1/32times$ for the Llama-3.1-8B model with LongBench.
  arXiv  Detail & Related papers  (2024-12-19T02:34:15Z)
• FLARE: Faithful Logic-Aided Reasoning and Exploration [50.9814063216852]
  We introduce a novel approach for traversing the problem space using task decompositions.<n>We use the Large Language Models to plan a solution, soft-formalise the query into facts and predicates using a logic programming code.<n>Our method allows us to compute the faithfulness of the reasoning process w.r.t. the generated code and analyse the steps of the multi-hop search without relying on external solvers.
  arXiv  Detail & Related papers  (2024-10-14T19:39:11Z)
• HSR-Enhanced Sparse Attention Acceleration [19.776342074253435]
  We introduce a novel approach to accelerate attention computation in Large Language Models (LLMs) We leverage the inherent sparsity within attention mechanisms, both in conventional Softmax attention and ReLU attention. Our method only introduces provably negligible error for Softmax attention.
  arXiv  Detail & Related papers  (2024-10-14T05:18:02Z)
• A Training-free Sub-quadratic Cost Transformer Model Serving Framework With Hierarchically Pruned Attention [43.211427581302715]
  We propose Hierarchically Pruned Attention (HiP) to increase context length in large language models.<n>HiP reduces the time complexity of the attention mechanism to $O(T log T)$ and the space complexity to $O(T)$, where $T$ is the sequence length.<n>We show that HiP significantly reduces both prefill and decoding latencies, as well as memory usage, while maintaining high-quality generation with minimal degradation.
  arXiv  Detail & Related papers  (2024-06-14T08:32:45Z)
• One Pass Streaming Algorithm for Super Long Token Attention Approximation in Sublinear Space [11.735802740426294]
  Attention computation takes both the time complexity of $O(n2)$ and the space complexity of $O(n2)$ simultaneously. We introduce a new algorithm that only reads one pass of data in a streaming fashion. Notably, our algorithm exhibits exceptional memory-efficient performance with super-long tokens.
  arXiv  Detail & Related papers  (2023-11-24T18:35:00Z)
• Efficiently Learning One-Hidden-Layer ReLU Networks via Schur Polynomials [50.90125395570797]
  We study the problem of PAC learning a linear combination of $k$ ReLU activations under the standard Gaussian distribution on $mathbbRd$ with respect to the square loss. Our main result is an efficient algorithm for this learning task with sample and computational complexity $(dk/epsilon)O(k)$, whereepsilon>0$ is the target accuracy.
  arXiv  Detail & Related papers  (2023-07-24T14:37:22Z)
• Simplifying and Understanding State Space Models with Diagonal Linear RNNs [56.33053691749856]
  This work disposes of the discretization step, and proposes a model based on vanilla Diagonal Linear RNNs. We empirically show that, despite being conceptually much simpler, $mathrmDLR$ is as performant as previously-proposed SSMs. We also characterize the expressivity of SSMs and attention-based models via a suite of $13$ synthetic sequence-to-sequence tasks.
  arXiv  Detail & Related papers  (2022-12-01T18:53:06Z)
• Near-Optimal Regret Bounds for Multi-batch Reinforcement Learning [54.806166861456035]
  We study the episodic reinforcement learning (RL) problem modeled by finite-horizon Markov Decision Processes (MDPs) with constraint on the number of batches. We design a computational efficient algorithm to achieve near-optimal regret of $tildeO(sqrtSAH3Kln (1/delta))$tildeO(cdot) hides logarithmic terms of $(S,A,H,K)$ in $K$ episodes. Our technical contribution are two-fold: 1) a near-optimal design scheme to explore
  arXiv  Detail & Related papers  (2022-10-15T09:22:22Z)
• Time and Query Complexity Tradeoffs for the Dihedral Coset Problem [0.19731444261635428]
  Dihedral Coset Problem (DCP) in $Z_N$ has been extensively studied in quantum computing and post-quantum cryptography. Ettinger-Hoyer algorithm is known to solve the DCP in $O(log(N))$ queries. We introduce the first algorithm to improve in the linear queries regime over the Ettinger-Hoyer algorithm.
  arXiv  Detail & Related papers  (2022-06-29T05:30:54Z)
• Learning a Latent Simplex in Input-Sparsity Time [58.30321592603066]
  We consider the problem of learning a latent $k$-vertex simplex $KsubsetmathbbRdtimes n$, given access to $AinmathbbRdtimes n$. We show that the dependence on $k$ in the running time is unnecessary given a natural assumption about the mass of the top $k$ singular values of $A$.
  arXiv  Detail & Related papers  (2021-05-17T16:40:48Z)
• RNNs can generate bounded hierarchical languages with optimal memory [113.73133308478612]
  We show that RNNs can efficiently generate bounded hierarchical languages that reflect the scaffolding of natural language syntax. We introduce Dyck-($k$,$m$), the language of well-nested brackets (of $k$ types) and $m$-bounded nesting depth. We prove that an RNN with $O(m log k)$ hidden units suffices, an exponential reduction in memory, by an explicit construction.
  arXiv  Detail & Related papers  (2020-10-15T04:42:29Z)

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.