用三角函数压缩注意力键值,让大模型长文本推理更快更省内存。
TriAttention: Efficient Long Reasoning with Trigonometric KV Compression

- 基于查询与键向量在位置旋转前的集中特性,用三角函数建模关键位置偏好。
- 在AIME25上32K token生成时,准确率媲美全注意力,吞吐提升2.5倍或显存降10.7倍。
- 适合需要长上下文推理但资源受限的场景,如单卡部署大型模型。
大语言模型进行长序列推理时,键值缓存(KV cache)内存瓶颈严重。现有压缩方法依赖旋转位置编码(RoPE)后的注意力得分来估计重要性,但因查询向量随位置旋转导致代表性查询稀少,造成关键键选择不佳且推理不稳定。为此,我们转向RoPE前的空间,发现查询(Q)和键(K)向量在固定非零中心附近高度集中,且位置间保持稳定——即Q/K集中性。该集中性导致查询倾向于关注特定距离的键(如最近键),中心决定偏好距离并通过三角函数系列体现。基于此,我们提出TriAttention,利用这些中心信息构建位置偏好评分机制,并结合Q/K范数作为重要性估计的补充信号。在AIME25数据集上32K token生成任务中,TriAttention达到与全注意力相当的推理准确率,同时实现2.5倍更高吞吐或10.7倍的KV内存减少;而领先基线在相同效率下准确率仅达一半。该方法使OpenClaw模型可在单个消费级显卡上部署长上下文推理,否则将因显存不足而失败。
原文摘要 · Abstract (English)
Extended reasoning in large language models (LLMs) creates severe KV cache memory bottlenecks. Leading KV cache compression methods estimate KV importance using attention scores from recent post-RoPE queries. However, queries rotate with position during RoPE, making representative queries very few, leading to poor top-key selection and unstable reasoning. To avoid this issue, we turn to the pre-RoPE space, where we observe that Q and K vectors are highly concentrated around fixed non-zero centers and remain stable across positions -- Q/K concentration. We show that this concentration causes queries to preferentially attend to keys at specific distances (e.g., nearest keys), with the centers determining which distances are preferred via a trigonometric series. Based on this, we propose TriAttention to estimate key importance by leveraging these centers. Via the trigonometric series, we use the distance preference characterized by these centers to score keys according to their positions, and also leverage Q/K norms as an additional signal for importance estimation. On AIME25 with 32K-token generation, TriAttention matches Full Attention reasoning accuracy while achieving 2.5x higher throughput or 10.7x KV memory reduction, whereas leading baselines achieve only about half the accuracy at the same efficiency. TriAttention enables OpenClaw deployment on a single consumer GPU, where long context would otherwise cause out-of-memory with Full Attention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。