提出新方法RaaS,让大模型推理更高效且省内存。
RaaS: Reasoning-Aware Attention Sparsity for Efficient LLM Reasoning
- 识别推理过程中的关键里程碑令牌,只保留其记忆向量。
- 实现高准确率,同时保持时间与内存复杂度均为O(L)。
- 适合需要高效长序列推理的数学、编程等任务场景。
大语言模型在数学和编程等复杂推理任务中表现突出,但生成长序列时面临O(N)的时间和内存开销(N为当前序列长度)。现有基于稀疏性的方法虽能保留关键令牌的键值向量,却难以兼顾准确性、时间和内存,陷入“不可能三角”。例如,当前最优方法Quest虽有高准确率,但仅实现O(L)时间复杂度,仍需O(N)内存(L为缓存预算,远小于N)。本文发现推理解码阶段存在一种新注意力模式:里程碑令牌(类比数学证明中的引理)被使用后即不再重要。基于此,提出RaaS算法,动态识别并保留这些里程碑令牌的键值向量,直至其不再需要,从而在保证高准确率的同时,实现O(L)时间与O(L)内存复杂度。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。