TriAttention: Efficient Long Reasoning with Trigonometric KV Compression
作者: Weian Mao, Xi Lin, Wei Huang, Yuxin Xie, Tianfu Fu, Bohan Zhuang, Song Han, Yukang Chen
分类: cs.CL, cs.CV
发布日期: 2026-04-07
💡 一句话要点
TriAttention：利用三角函数KV压缩实现高效长程推理
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 长程推理 KV缓存压缩 注意力机制 三角函数 RoPE 大型语言模型 Q/K集中
📋 核心要点
- 现有方法在RoPE后进行KV重要性估计，受限于RoPE导致查询随位置旋转，代表性查询不足，影响top-key选择和推理稳定性。
- TriAttention的核心思想是在RoPE前空间利用Q/K向量的集中特性，通过三角级数建模位置偏好，估计键的重要性。
- 实验表明，TriAttention在长文本生成任务中，在保证推理精度的前提下，显著提升了吞吐量并降低了KV缓存内存占用。
📝 摘要（中文）
大型语言模型（LLMs）中的长程推理带来了严重的KV缓存内存瓶颈。现有的KV缓存压缩方法使用最近RoPE后的查询的注意力分数来估计KV的重要性。然而，查询在RoPE过程中会随位置旋转，导致代表性查询非常少，从而导致较差的top-key选择和不稳定的推理。为了避免这个问题，我们转向RoPE之前的空间，在那里我们观察到Q和K向量高度集中在固定的非零中心附近，并且在不同位置保持稳定——Q/K集中。我们表明，这种集中导致查询优先关注特定距离的键（例如，最近的键），中心通过三角级数确定首选距离。基于此，我们提出了TriAttention，通过利用这些中心来估计键的重要性。通过三角级数，我们使用这些中心表征的距离偏好来根据键的位置对其进行评分，并且还利用Q/K范数作为重要性估计的附加信号。在具有32K token生成的AIME25上，TriAttention匹配了Full Attention的推理精度，同时实现了2.5倍的吞吐量或10.7倍的KV内存减少，而领先的基线在相同的效率下仅实现了大约一半的精度。TriAttention能够在单个消费级GPU上部署OpenClaw，否则长上下文会导致Full Attention内存不足。
🔬 方法详解
问题定义：论文旨在解决大型语言模型中长程推理带来的KV缓存内存瓶颈问题。现有方法，如基于注意力分数的KV缓存压缩，在RoPE（旋转位置编码）之后进行重要性评估，但RoPE导致查询向量随位置旋转，使得代表性查询数量不足，影响了top-key的选择，进而导致推理性能下降。
核心思路：TriAttention的核心思路是利用RoPE之前的Q和K向量的集中特性。作者观察到，在RoPE之前，Q和K向量高度集中在固定的非零中心附近，并且在不同位置保持稳定。这种集中性导致查询倾向于关注特定距离的键。通过分析这种距离偏好，可以更准确地估计键的重要性。
技术框架：TriAttention主要包含以下几个步骤：1) 在RoPE之前获取Q和K向量；2) 利用三角级数对Q和K向量的集中中心进行建模，从而捕捉不同位置的距离偏好；3) 结合距离偏好和Q/K向量的范数，计算每个键的重要性得分；4) 根据重要性得分选择top-k个键进行保留，从而实现KV缓存压缩。
关键创新：TriAttention的关键创新在于利用了RoPE之前的Q/K向量的集中特性，并使用三角级数来建模位置相关的距离偏好。与现有方法直接在RoPE之后进行重要性评估不同，TriAttention避免了RoPE带来的查询向量旋转问题，从而能够更准确地估计键的重要性。
关键设计：TriAttention的关键设计包括：1) 使用三角级数来表示位置相关的距离偏好，三角函数的系数可以通过学习得到；2) 结合Q/K向量的范数作为重要性估计的补充信号，范数可以反映向量的能量大小；3) 通过调整三角级数的阶数和Q/K范数的权重，可以灵活地控制压缩的程度和性能。
🖼️ 关键图片
📊 实验亮点
TriAttention在AIME25数据集上进行了评估，结果表明，在32K token生成任务中，TriAttention能够匹配Full Attention的推理精度，同时实现2.5倍的吞吐量提升或10.7倍的KV内存减少。相比之下，领先的基线方法在相同效率下仅能达到大约一半的精度。TriAttention使得OpenClaw能够在单个消费级GPU上部署，而Full Attention则会导致内存不足。
🎯 应用场景
TriAttention具有广泛的应用前景，尤其是在需要处理长文本序列的场景中，例如长文档摘要、机器翻译、对话生成等。通过降低KV缓存的内存占用，TriAttention使得大型语言模型能够在资源受限的设备上部署，例如消费级GPU或移动设备。此外，TriAttention还可以加速长文本推理过程，提高模型的响应速度。
📄 摘要（原文）
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.