RaaS: Reasoning-Aware Attention Sparsity for Efficient LLM Reasoning
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2025
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866908385918582784 |
|---|---|
| author | Hu, Junhao Huang, Wenrui Wang, Weidong Li, Zhenwen Hu, Tiancheng Liu, Zhixia Chen, Xusheng Xie, Tao Shan, Yizhou |
| author_facet | Hu, Junhao Huang, Wenrui Wang, Weidong Li, Zhenwen Hu, Tiancheng Liu, Zhixia Chen, Xusheng Xie, Tao Shan, Yizhou |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_11147 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | RaaS: Reasoning-Aware Attention Sparsity for Efficient LLM Reasoning Hu, Junhao Huang, Wenrui Wang, Weidong Li, Zhenwen Hu, Tiancheng Liu, Zhixia Chen, Xusheng Xie, Tao Shan, Yizhou Machine Learning Artificial Intelligence 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. |
| title | RaaS: Reasoning-Aware Attention Sparsity for Efficient LLM Reasoning |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2502.11147 |