TriAttention: Efficient Long Reasoning with Trigonometric KV Compression

作者/机构: Weian Mao 1∗, Xi Lin 3∗, Wei Huang 2∗, Yuxin Xie 1, Tianfu Fu 1, Bohan Zhuang 3, Song Han 1 2, Yukang Chen 2

A1 主要贡献

大型语言模型（LLM）中的长程推理会产生跨越数万个词元的思想链序列，导致KV缓存随之成比例增长，从而造成严重的内存瓶颈。现有的KV缓存压缩方法通过保留最重要的词元来解决此问题，其重要性是根据近期查询计算的注意力分数来估计的。然而，这些方法本质上是不稳定的。它们在应用了旋转位置编码（RoPE）之后的查询（post-RoPE queries）上操作，而这些查询会随着位置旋转，导致只有最近的查询才具有最新的方向，形成一个极小的观察窗口。由于代表性查询过少，重要键（key）可能未被检测到——在短窗口内获得低注意力的词元可能会被永久驱逐，即使它在后续变得至关重要。这对于检索头（retrieval heads）尤其具有挑战性，因为相关词元可能在变得关键之前长时间处于休眠状态。在推理任务中，这种信息丢失会破坏思维链。

为解决此问题，本文将研究视角转向了RoPE之前的空间（pre-RoPE space），并观察到一个显著现象：在大部分注意力头中，Q和K向量高度集中于一个固定的非零中心，作者将此特性称为Q/K集中。如图2(A, C)所示，这种集中现象在不同位置和上下文中保持稳定。由于pre-RoPE向量在应用位置编码前不受位置旋转的影响，这种稳定性是内在的。更重要的是，pre-RoPE向量通过RoPE公式直接与注意力相关联，使得这些中心对于评估KV重要性具有重要意义。

基于这一发现，本文首先证明了这些中心通过一个三角级数与注意力行为相关联。当Q/K高度集中时，它们可以被其中心近似；将这些中心代入RoPE公式，注意力logit简化为一个仅依赖于Q-K距离的函数——一个三角级数——形成一条“注意力-距离”曲线。如图2(D)所示，这条曲线通常在特定的Q-K距离处出现峰值，实验验证了处于这些距离的键确实在实践中获得了更高的注意力。这表明Q/K集中导致注意力偏好特定距离的键，而这些中心决定了哪些距离被偏好。

本文利用这一理解设计了TriAttention，一种KV缓存压缩方法。TriAttention对键进行评分，并仅保留得分最高的键，以解决内存瓶颈。其评分函数的核心思想是利用Q中心和三角级数来评估因距离偏好而导致的键之间的重要性差异。对于Q不那么集中的少数头，作者引入Q/K范数作为补充信号。该方法使用Q/K集中度度量自动平衡这两个组成部分。

在AIME25数学推理基准测试上，当生成32K词元时，TriAttention在保持与完整注意力（Full Attention）相当的推理准确率的同时，实现了2.5倍的吞吐量提升或10.7倍的KV内存减少（如图1所示），而主流基线方法在同等效率下仅达到约一半的准确率。TriAttention还使得OpenClaw模型能够部署在单个消费级GPU上，而完整注意力会因长上下文导致内存不足。


图1. AIME25（Qwen3-8B）上的性能权衡。(A) 在等效准确率（40.8%）下，TriAttention的吞吐量是完整注意力的2.5倍。(B) TriAttention在匹配完整注意力准确率的同时，将KV缓存内存减少了10.7倍。


图2. Q/K集中现象及其对注意力的影响。(A) 在主导频带上的Pre-RoPE Q/K向量高度集中（平均合矢量长度R很高）。(B) RoPE旋转将这些向量分散成弧形图案。在(A-B)中，叠加了三个不同的输入序列，显示这种结构在不同内容下是稳定的。(C) 这种集中现象几乎在所有头中都存在。(D) 当Q/K集中时，注意力logit可以使用三角级数准确重建（皮尔逊相关系数r = 0.72）。

A3 背景知识与关键发现

背景与相关工作

KV缓存压缩通过在固定的内存预算下仅保留KV对的一个子集来工作。核心挑战在于重要性估计：确定哪些词元将从未来的查询中获得高注意力。现有方法通过分析后RoPE（post-RoPE）表示来解决这个问题，我们将在§2.2中讨论。为了理解为什么后RoPE分析具有局限性，我们首先回顾RoPE的工作原理。

2.1. 旋转位置编码 (Rotary Position Embedding)

旋转位置编码（RoPE） 是一种通过在向量空间中进行旋转来编码位置信息的方法【索引43，Roformer: Enhanced transformer with rotary position embedding，2024，Neurocomputing】、【索引7，On the token distance modeling ability of higher rope attention dimension，2024，EMNLP】、【索引2，Round and round we go! what makes rotary positional encodings useful?，2025，ICLR】，并已成为现代LLM中主要的位置编码方式。对于一个d维向量，RoPE将其划分为d/2个二维子空间，由$f \in \{0, . . . , d/2 − 1\}$索引。第f个子空间以频率$\omega_f = \theta^{-2f/d}$（通常$\theta = 10000$）旋转；我们称之为频带f。对于频带f，RoPE在位置p处应用一个角度为$\omega_f p$的旋转：

$$\begin{aligned} \begin{pmatrix} x'_{2f} \\ x'_{2f+1} \end{pmatrix} = \begin{pmatrix} \cos(\omega_f p) & -\sin(\omega_f p) \\ \sin(\omega_f p) & \cos(\omega_f p) \end{pmatrix} \begin{pmatrix} x_{2f} \\ x_{2f+1} \end{pmatrix} \end{aligned}$$

我们将这次旋转之前的Q/K向量称为pre-RoPE，旋转之后的称为post-RoPE。


图3. 三个基于DeepSeek-R1蒸馏的LLM（包括Qwen3、Qwen2.5和Llama3）的注意力重建相关性。图中显示了所有注意力头中每个头的重建皮尔逊相关性（$\bar{r}$）的分布。红色虚线表示均值。所有模型都显示出右偏分布，均值高于0.5。

2.2. 后RoPE（Post-RoPE）压缩方法

KV缓存压缩方法 通过后RoPE表示来估计词元的重要性【索引4，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024，arXiv】、【索引13，Model tells you what to discard: Adaptive KV cache compression for LLMs，2024，ICLR】、【索引12，Ada-KV: Optimizing KV cache eviction by adaptive budget allocation for efficient LLM inference，2025，NeurIPS】、【索引3，RocketKV: Accelerating long-context LLM inference via two-stage KV cache compression，2025，ICML】。我们将它们分为启发式、基于注意力和基于范数三类。

启发式方法。早期的压缩方法使用基于启发式规则的注意力模式。例如，StreamingLLM【索引52，Efficient streaming language models with attention sinks，2024，ICLR】、【索引14，When attention sink emerges in language models: An empirical view，2025，ICLR】通过仅保留一个小的固定大小缓存（几个初始的“汇聚”词元加上一个最近词元的滑动窗口）来实现无限长度的流式输入。汇聚词元利用了初始位置无论其内容如何都会获得不成比例的注意力的观察——丢弃它们会导致注意力分数失去一个稳定的“汇聚点”，从而降低性能。虽然简单，但这种固定规则无法适应依赖于内容的重要性，这促使了更复杂方法的出现。

基于注意力的方法。这些方法使用在后RoPE Q/K上计算的注意力分数来识别重要的词元。H2O【索引54，H2o: Heavy- hitter oracle for efficient generative inference of large language models，2023，NeurIPS】在解码步骤中累积注意力分数，以识别持续获得高注意力的“重击”词元。SnapKV【索引25，SnapKV: LLM knows what you are looking for before generation，2024，NeurIPS】在局部窗口内计算注意力并聚合分数，以预测哪些词元对生成重要。Scissorhands【索引26，Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time，2023，NeurIPS】利用“重要性持久性”假设，使用历史注意力来指导驱逐。R-KV【索引5，R-KV: Redundancy-aware KV cache compression for reasoning models，2025，NeurIPS】通过来自最近查询的注意力对词元进行评分，并结合冗余检测用于推理模型。LazyEviction【索引53，Lazyeviction: Lagged kv eviction with attention pattern observation for efficient long reasoning，2025，】在观察窗口内跟踪词元重要性的复现，以延迟驱逐决策。

基于范数的方法。VATP【索引15，Attention score is not all you need for token importance indicator in kv cache reduction: Value also matters，2024，EMNLP】观察到仅凭注意力分数是不够的：注意力汇聚点获得高注意力，但其值向量的范数接近于零，对输出贡献很小。通过结合值向量的范数，VATP提供了一个更细致的重要性度量。

2.3. 后RoPE方法的局限性

基于注意力和基于范数的方法的共同局限性。这两种方法都可以被视为在已经应用了位置旋转的后RoPE空间中操作。这个共同的基础导致了不同但相关的局限性。

基于注意力的方法的局限性。其关键限制是查询随着位置旋转，这将有用的观察限制在一个微小的窗口内，并导致重要的键未被检测到。

基于范数的方法的局限性。其局限性则不同：它们仅利用向量的幅度，而忽略了方向信息。理想情况下，结合Q/K的方向可以改善重要性估计——注意力既依赖于范数，也依赖于Q和K之间的夹角。然而，在后RoPE空间中，方向与位置旋转纠缠在一起：向量的方向随位置不断变化，使得方向信息难以利用。

Pre-RoPE与三角级数

如引言所述，我们观察到在pre-RoPE空间中的Q和K向量高度集中在非零中心周围，并且这种集中在不同位置和上下文中保持稳定。这并非特定头类型的属性，而是在模型中普遍存在的现象，如图2(A, C)所示。在本节中，我们首先描述这种集中现象（§3.1），然后展示它如何使注意力模式能被一个三角级数描述（§3.2），最后通过实验验证这一点（§3.3）。

3.1. Pre-RoPE集中现象

Q/K在pre-RoPE空间的分布。我们检查了Q/K在pre-RoPE空间中的分布。对于每个头，我们识别其主导频带——即对注意力logit贡献最大的频率【索引55，Elitekv: Scalable kv cache compression via rope frequency selection and joint lowrank projection，2025，arXiv】（详见附录B.7）——并在相应的2D平面上可视化Q/K向量。

观察3.1（普遍的Q/K集中）。在pre-RoPE空间中，Q和K向量在大多数注意力头中都高度集中在非零中心周围。如图2(A)所示，这种集中在不同的词元位置和输入上下文中是稳定的。


图4. 方法概述。从左到右：离线校准计算Q分布中心；然后在推理过程中，通过结合Strig和基于范数的分量对原始注意力进行评分；最右侧的面板显示了修剪后的注意力图。我们观察到一些头表现出距离偏好——远处的键倾向于获得更高的注意力。然而，我们也发现某些键尽管离查询很远，但由于其范数较低而几乎没有受到关注。这启发了我们的两个评分分量：Strig捕捉距离偏好，而基于范数的分数识别低范数的键。在这个例子中，Strig正确地为附近的键分配了低分，而基于范数的分数则识别出最早的词元（最左侧）由于其范数低而不重要，尽管它的距离最大。它们共同准确地识别出不会被关注的词元并将其修剪。有关真实注意力图的可视化，请参见附录D。

量化集中度。为了量化这种集中度，我们使用平均 resultant 长度（Mean Resultant Length）$R = \Vert E[q] \Vert / E[\Vert q \Vert]$，它衡量向量围绕其平均方向集中的紧密程度（$R \rightarrow 1$ 表示完美集中；$R \rightarrow 0$ 表示均匀分散）。图2(C)显示，在Qwen3-8B的所有头中，绝大多数的R值接近1.0，证实了Q/K集中是普遍存在的。

广泛集中的影响。这种广泛的集中现象具有重要影响：当Q/K向量近似为常数时，注意力计算将极大简化，我们将在下文展示。

3.2. 可预测的距离偏好

从集中到三角级数。当Q/K向量高度集中时，我们可以用它们的中心来近似它们。这种近似将注意力计算转化为一个仅依赖于Q-K距离的三角级数，使得注意力模式可以仅从中心进行预测。

RoPE注意力公式。考虑RoPE注意力公式。对于位置为$p_q$的查询q和位置为$p_k$的键k，RoPE以速率$\omega_f$旋转频带f。pre-softmax logit为（见附录B）：

$$\operatorname{logit}(q, k)=\sum_{f}\left\|q_{f}\right\|\left\|k_{f}\right\| \cos \left(\omega_{f} \Delta+\phi_{f}\right)$$

其中$\Delta = p_q - p_k$是Q-K距离， $q_f, k_f \in \mathbb{C}$ 是在频带f中的pre-RoPE分量，而$\phi_f = \arg(q_f) - \arg(k_f)$是它们的相位差。

近似为距离的函数。当Q/K集中时，我们近似$q_f \approx \bar{q}_f$和$k_f \approx \bar{k}_f$（中心）。由于$\bar{q}_f$和$\bar{k}_f$是常数，logit变为仅依赖于距离的函数：

$$\begin{aligned} \begin{aligned} \operatorname{logit}(\Delta) & \approx \sum_{f} \underbrace{\left\|\bar{q}_{f}\right\|\left\|\bar{k}_{f}\right\|}_{\text {amplitude }} \cos (\omega_{f} \Delta+\underbrace{\bar{\phi}_{f}}_{\text {phase }}) \\ & =\sum_{f}\left[a_{f} \cos \left(\omega_{f} \Delta\right)+b_{f} \sin \left(\omega_{f} \Delta\right)\right] \end{aligned} \end{aligned}$$

其中系数$a_f, b_f$由Q/K中心决定。这是一个关于Q-K距离$\Delta$的三角级数。

类比傅里叶合成。尽管RoPE频率遵循几何级数而非谐波级数，但其原理类似于傅里叶合成：学习到的Q/K中心决定了系数，这些系数反过来塑造了注意力与距离的关系曲线。不同的中心产生不同的曲线——有些在小距离处达到峰值（局部注意力），有些则在大距离处达到峰值（注意力汇聚点）。在所有情况下，距离偏好都编码在Q/K中心中，并且可以通过三角级数进行预测。

3.3. 实验验证

验证假设。我们通过实验测试Q/K集中是否导致注意力遵循由三角级数描述的距离偏好。我们从Q/K中心计算级数，并检查它是否能重建实际的注意力。成功的重建证实了这种因果联系，并表明这些偏好可以从中心预测。

重建过程。我们在Qwen3-8B的所有1152个注意力头（36层×32头）上使用约10K词元的序列进行测试。为了重建，我们从校准数据集中计算pre-RoPE空间中的平均Q和K向量——表示为$E[q_f]$和$E[k_f]$——并将它们代入三角级数：

$$\hat{s}(\Delta)=\sum_{f}\|\mathbb{E}[q_{f}]\|\|\mathbb{E}[k_{f}]\|\cos(\omega_{f}\Delta+\phi_{f})$$

其中$\phi_f = \arg(E[q_f]) - \arg(E[k_f])$是平均向量之间的相位差。这得出了一个关于Q-K距离$\Delta$的预测注意力曲线（图2(D)）。

量化预测质量。为了量化预测质量，我们定义了重建相关性$\bar{r}$：预测和实际注意力logit之间的平均皮尔逊相关性。对于每个查询，我们计算其实际logit与预测曲线之间的相关性，然后在所有查询中取平均值：

$$\bar{r} = \frac{1}{N} \sum_{i=1}^{N} \rho(\mathbf{a}_i, \hat{\mathbf{s}})$$

在这里，$a_i$是查询i的实际注意力logit向量，$\hat{s}$是根据方程4预测的logit，$\rho$是皮尔逊相关性。$a_i$和$\hat{s}$都在相同的对数间隔距离$\Delta = 1, 2, 4, 8, ...$上进行评估，以确保在不同距离尺度上都有均衡的覆盖。

结果分析。图2(D)展示了一个例子。对于第一层第一个头——为了避免精心挑选——预测与实际注意力紧密吻合，达到$\bar{r} = 0.72$。在三种不同架构（Qwen3, Qwen2.5, Llama3）的所有头中，$\bar{r}$的分布峰值在0.6–0.9之间，平均值高于0.5；完整分布见图3。在许多头和架构上的高相关性证实了从Q/K中心计算的三角级数能准确预测注意力模式。我们进一步发现Q/K集中是模型的内在属性：在Qwen3-8B上，跨数学、编码和聊天领域测量的MRL值几乎相同（0.977–0.980），无论领域如何，约90%的头表现出$R > 0.95$。这在不同架构中也成立，包括MLA（附录I）。

A2 方法细节

基于§3的分析，我们提出了一种KV缓存压缩方法，该方法对键的重要性进行评分，并仅保留得分最高的键（图4）。评分函数结合了两种信号（§4.1）：捕捉距离偏好的三角级数，以及作为补充的范数信息；它们之间的权重根据Q/K的集中度进行调整（§4.2）。最后，我们描述了如何应用此评分进行KV缓存修剪（§4.3）。

4.1. 键重要性评分

在KV缓存中，键已经被缓存。令$k_f \in \mathbb{C}$表示一个键k在频带f中的pre-RoPE表示，$\Delta = p_q - p_k$是Q-K距离。

三角级数分数（Trigonometric Series Score）。不同位置的键与未来查询的距离不同。为了估计每个键将获得多少注意力，我们使用Q中心作为未来查询的代理——这由§3中观察到的Q/K集中现象证明是合理的。然后，三角级数给出了在距离$\Delta$处的预期注意力：

$$S_{\mathrm{trig}}(k, \Delta) = \sum_{f} \|\mathbb{E}[q_f]\| \cdot \|k_f\| \cdot \cos(\omega_f \Delta + \phi_f)$$

这里，$E[q_f] \in \mathbb{C}$是频带f中的Q中心，从校准数据中计算得出。由于缓存中的键是已知的，我们直接使用它们的表示$k_f$。相位$\phi_f = \arg(E[q_f]) - \arg(k_f)$是Q中心与当前键之间的角度差。

基于范数的分数（Norm-Based Score）。三角级数假设Q/K正好位于它们的中心。实际上，中心周围存在变异，这种变异会影响注意力。我们用一个基于范数的项来解释这一点：

$$S_{\text{norm}}^{(0)}(k) = \sum_{f} \mathbb{E}[\|q_{f}\|] \cdot \|k_{f}\|$$

这里，$E[\Vert q_f \Vert]$是频带f中的预期查询范数，从校准数据中计算得出。与只考虑$\Vert k \Vert$的方法不同，我们通过其预期的查询贡献来加权每个频带。

4.2. 通过集中度进行自适应加权

动机。当Q/K高度集中时，三角级数分数$S_{trig}$最准确；而对于集中度较低的少数头，基于范数的分数$S_{norm}$变得更加重要。我们使用Q/K集中度作为权重因子来自动平衡这些分量。

实现。回想一下，集中度由每个频带f的平均 resultant 长度$R_f = \Vert E[q_f] \Vert / E[\Vert q_f \Vert]$来量化。当$R_f$高时，三角级数是准确的，并且$S_{trig}$是可靠的；当$R_f$低时，中心周围有更多的变异，$S_{norm}$提供了有用的补充信息。

我们通过将每个频带乘以$(1 - R_f)$来改进基于范数的分数：

$$S_{\text{norm}}(k) = \sum_{f}(1 - R_f) \cdot \mathbb{E}[\|q_f\|] \cdot \|k_f\|$$

这可以重写为：

$$S_{\text{norm}}(k) = \sum_{f} (\mathbb{E}[\|q_f\|] - \|\mathbb{E}[q_f]\|) \cdot \|k_f\|$$

直观上，当$R_f$高时（集中度强），$(1 - R_f)$很小，因此$S_{norm}$贡献很小，$S_{trig}$占主导地位。当$R_f$低时（集中度弱），保留了完整的范数贡献。

最终分数。最终的组合分数为：

$$S(k, \Delta) = S_{\text{trig}}(k, \Delta) + S_{\text{norm}}(k)$$

表1. AIME24和AIME25上的推理性能。最佳结果以粗体显示，次佳结果带下划线。所有方法都在相同的KV缓存预算下进行比较。

考虑未来位置。一个键可能从任何未来的位置被查询，因此其重要性取决于所有未来的查询位置。我们在多个偏移量上计算$S(k, \Delta + \delta)$，并将平均重要性分数定义为：

$$\tilde{S}(k) = \frac{1}{|\mathcal{D}|} \sum_{\delta \in \mathcal{D}} S(k, \Delta + \delta)$$

其中$D = \{1, 2, 4, . . . , 2^{16}\}$是未来偏移量的集合。

4.3. KV缓存修剪

评分与保留。定义了评分函数$\tilde{S}(k)$后，我们独立地对缓存中的每个键进行评分，并仅保留前B个。在选择之前，我们解决两个实际问题：减少频繁评分的计算开销，以及处理分组查询注意力（GQA）【索引1，Gqa: Training generalized multi-query transformer models from multi-head checkpoints，2023，EMNLP】，其中多个查询头共享每个KV头。

基于窗口的修剪。在每个解码步骤中对所有键进行评分计算成本高昂。相反，我们每生成$\beta = 128$个词元就触发一次修剪：当每个区间的第128个词元生成时，如果缓存超过预算B，我们就对所有键进行评分，保留前B个，并驱逐其余的。我们称每个128词元的区间为一个窗口，这遵循了R-KV【索引5，R-KV: Redundancy-aware KV cache compression for reasoning models，2025，NeurIPS】的做法。这种批处理方法显著减少了开销。

分组查询注意力（GQA）。在GQA中，每个KV头由G个查询头共享。由于$\tilde{S}$依赖于查询统计数据（$E[q_f]$和$E[\Vert q_f \Vert]$），每个键会收到G个不同的分数。这些分数在不同头之间的尺度不同，使得直接比较不可靠。

归一化后聚合。我们应用“归一化后聚合”的策略。令$\tilde{S}^{(g)}(k)$表示使用查询头g的统计数据计算出的分数。我们首先在每个头内部进行z-score归一化：

$$\hat{S}^{(g)}(k) = \frac{\tilde{S}^{(g)}(k) - \mu_g}{\sigma_g}$$

其中$\mu_g, \sigma_g$是所有键上$\{ \tilde{S}^{(g)}(k) \}_k$的均值和标准差。然后我们通过最大值进行聚合：

$$S_{\text {final }}(k)=\max _{g \in\{0, \ldots, G-1\}}\left(\hat{S}^{(g)}(k)\right)$$

表2. MATH 500上的推理性能，KV预算为512。最佳结果以粗体显示，次佳结果带下划线。

保留策略。如果任何一个查询头认为一个键是重要的，它就会被保留。我们根据$S_{final}(k)$保留前B个键，并驱逐其余的。

A4 实验环境与结果

实验环境


图5. Qwen3-8B上的性能比较。(A–C) 在三个数学推理基准上，准确率与KV缓存预算的关系。TriAttention在所有预算水平上都持续优于R-KV。(D) 在递归状态查询基准上的内存保留能力。深度指DFS递归深度；更深的递归需要保留更多的中间状态，增加了内存压力。

实验结果

推理任务性能

我们在标准数学推理基准上评估了TriAttention。表1和表2分别展示了在AIME和MATH 500上的结果。在所有模型和数据集上，TriAttention在KV缓存压缩方法中始终取得最佳性能，非常接近Full Attention，同时优于所有基线。

为了解性能如何随压缩强度变化，我们在Qwen3-8B上评估了TriAttention和R-KV在一系列KV预算（从512到4096个词元）下的表现，如图5所示。TriAttention在所有预算水平上都持续优于R-KV，其优势在低到中等预算时最为明显。在MATH 500上，TriAttention在1024的预算下与Full Attention持平，并在更高预算下略有超过。在更具挑战性的AIME基准上，TriAttention保持了显著的领先优势：在AIME25上，预算为2048时，TriAttention达到32.9%，而R-KV仅为17.5%——相差15.4个百分点。值得注意的是，TriAttention在显著压缩KV缓存的情况下，能达到或超过Full Attention的性能：在AIME25上，它以4096的预算达到了43.3%，超过了Full Attention。

内存保留基准测试

KV缓存修剪可能导致LLM在长推理链中忘记中间状态。为衡量此效应，我们设计了一个基于递归模拟的基准测试。递归算法天然地对内存保留能力构成压力：模型必须深入嵌套调用，保留中间状态，然后回溯以产生结果。更深的递归深度需要保留更多中间状态，从而产生更大的内存压力。如果KV缓存修剪算法错误地丢弃了这些关键状态，模型将丢失回溯所需的信息，错误会级联到所有后续步骤。我们使用深度优先搜索（DFS）来实例化此递归模拟。

我们在Qwen3-8B上进行评估，步数从6到20，比较了Full Attention、R-KV和TriAttention，每个步数有80个样本。两种压缩方法的KV预算均为2048，如图5(D)所示。在低到中等内存压力下（深度达16），TriAttention的性能与Full Attention相当，甚至在深度8和12时略有超出。仅在深度超过18时，TriAttention才开始落后。这表明在典型工作负载下，Full Attention的KV缓存存在冗余，而TriAttention有效地保留了内存保留所需的基本信息。相比之下，R-KV在所有深度上都持续表现不佳。最显著的是，R-KV从深度16开始出现灾难性的准确率下降，从深度14的约61%降至深度16的31%。这种急剧下降表明R-KV的修剪严重损害了模型保留中间状态的能力。

消融研究

表3. 在Qwen3-8B上，KV预算为2048的消融研究。(A) 移除三角级数分数Strig的效果。(B) 移除基于集中度的权重的效果。(C) 跨域泛化：比较在编码数据与推理数据上进行校准，两者都在推理任务上测试。

(A) 三角级数分数

(B) 基于集中度的权重

(C) 跨域校准

吞吐量和效率

我们从两个角度评估效率：与Full Attention的比较，以及与先前压缩方法的比较。吞吐量是在单个A100 80GB GPU上，以最大批次大小，在16K解码长度上测得的每秒平均生成词元数。

表4. Qwen3-8B上的吞吐量比较。在达到与Full Attention相当准确率的KV预算下，TriAttention显示出显著的吞吐量增益。

表5. 在可比的推理准确率和内存下，R-KV和TriAttention的比较。

A5 结论

我们发现在pre-RoPE空间中存在Q/K集中现象：查询和键向量聚集在固定的中心周围，导致注意力通过三角级数遵循可预测的距离偏好。这种集中是模型的内在属性，跨领域和架构保持一致。TriAttention利用这一发现，从稳定的Q/K中心估计键的重要性，避免了后RoPE方法的不稳定性。在数学推理（AIME, MATH 500）和通用任务（LongBench, RULER）上的实验表明，TriAttention持续优于所有基线，在达到与Full Attention相当性能的同时，实现了2.5倍的吞吐量提升和10.7倍的KV内存减少。

A6 附录

A. 局限性与未来工作

局限性与未来工作。在这项工作中，我们提出了TriAttention，一种新颖的KV缓存压缩机制，它利用pre-RoPE空间中Q/K集中的内在稳定性，通过三角级数来估计键的重要性。我们在复杂推理基准上的广泛评估证实，TriAttention有效地匹配了Full Attention的准确性，同时实现了2.5倍的吞吐量增加和10.7倍的KV内存占用减少。尽管当前实现展示了显著的效率提升，但通过设计专为我们独特压缩操作优化的、硬件感知的推理内核，可以实现进一步的延迟降低。未来的一个主要改进方向是开发一个专用的、高性能的推理内核，旨在进一步加速三角级数的计算和随后的缓存修剪过程。除了这些工程优化，我们计划将我们的评估扩展到更广泛的领域，如编码和智能体任务，并研究更精细的压缩策略，包括针对特定头的预算，以进一步推动效率的边界。

B. RoPE与三角级数的联系

RoPE与三角级数的关系。在本附录中，我们建立了RoPE（旋转位置嵌入）和三角级数之间的数学联系，解释了注意力头如何利用RoPE实现依赖于距离的注意力模式。

B.1. RoPE作为复数旋转

RoPE作为复数旋转。RoPE对Query和Key向量应用与位置相关的旋转。对于对应于频率$\omega_f$的二维子空间，RoPE可以写成复数形式：

$$\tilde{q}_f(p)=q_f \cdot e^{i \omega_f p}, \quad \tilde{k}_f(p)=k_f \cdot e^{i \omega_f p}$$

其中$q_f, k_f \in \mathbb{C}$是Query和Key在频带f中的复数表示，p是位置。

B.2. RoPE注意力的一般形式

RoPE注意力的一般形式。在频带f中，位置为$p_q$的Query与位置为$p_k$的Key之间的点积为：

$$\operatorname{Re}(\tilde{q}_f(p_q) \cdot \overline{\tilde{k}_f(p_k)}) = \operatorname{Re}(q_f \overline{k_f} \cdot e^{i \omega_f(p_q - p_k)})$$

对所有频带求和，注意力logit为：

$$\langle q, k\rangle_{\Delta}=\sum_{f}\left\|q_{f}\right\|\left\|k_{f}\right\| \cos \left(\omega_{f} \Delta+\phi_{f}\right)$$

其中$\Delta = p_q - p_k$是相对位置，$\phi_f = \arg(q_f) - \arg(k_f)$是相位差。关键是，系数$\Vert q_f \Vert \Vert k_f \Vert$和相位$\phi_f$依赖于特定的Q/K向量，这些向量在不同词元之间是变化的。

B.3. 特殊情况：恒定的Q/K产生三角级数

恒定Q/K与三角级数。当pre-RoPE的Q和K向量在所有词元上都是恒定时，系数变为固定的。使用和角公式，logit简化为：

$$\langle q, k\rangle_{\Delta}=\sum_{f}\left[a_{f} \cos \left(\omega_{f} \Delta\right)+b_{f} \sin \left(\omega_{f} \Delta\right)\right]$$

其中系数现在是常数：

$$\begin{aligned} \begin{aligned} a_{f} & =\|q_{f}\|\|k_{f}\| \cos (\phi_{f}) \\ b_{f} & =-\|q_{f}\|\|k_{f}\| \sin (\phi_{f}) \end{aligned} \end{aligned}$$

这是一个关于$\Delta$的三角级数。我们注意到，RoPE频率遵循几何级数$\omega_f = \theta^{-2f/d}$（其中$\theta = 10000$），而不是经典傅里叶级数的谐波级数$\omega_n = n\omega_0$。

关键洞见。
* 频率$\omega_f$由RoPE预先确定（几何级数）。
* 当Q/K恒定时，系数（$a_f, b_f$）变为固定的常数。
* 模型可以通过学习合适的Q/K值来“合成”任意的距离依赖性注意力模式。

B.4. 从Q/K集中到可预测的距离偏好

从Q/K集中到可预测的距离偏好。当Q/K向量高度集中于其中心时——由高平均 resultant 长度R量化——预期的注意力logit可以使用这些中心来近似：

$$\mathbb{E}[\langle q, k\rangle_{\Delta}] \approx \sum_{f}\|\mathbb{E}[q_{f}]\|\|\mathbb{E}[k_{f}]\| \cos (\omega_{f} \Delta+\phi_{f})$$

其中$\phi_f = \arg(E[q_f]) - \arg(E[k_f])$是平均向量之间的相位差。

近似的准确性。当集中度高时（即$R \rightarrow 1$），这种近似变得准确。在这种情况下，Q/K中心完全决定了距离偏好曲线：不同的中心产生不同的注意力-距离曲线，在特定的Q-K距离处有峰值。关键是，这些偏好可以仅从中心预测，而无需观察实际的注意力分数——这是TriAttention利用的关键洞见。

B.5. 平均合矢量长度 (Mean Resultant Length)

平均合矢量长度（R） 是方向统计学中的一个标准度量【索引27，Directional Statistics，1999】，它量化了一组向量围绕其平均方向集中的紧密程度。对于一组单位向量$\{u_1, . . . , u_n\}$，平均 resultant 长度定义为：

$$R = \left\| \frac{1}{n} \sum_{i=1}^{n} u_i \right\|$$

对于具有不同大小的向量，如注意力头中的Q/K向量，我们将其推广为：

$$R = \frac{\|\mathbb{E}[q]\|}{\mathbb{E}[\|q\|]}$$

其中期望是在词元位置上取的。这个比率有直观的界限：
* $R = 1$: 所有向量都指向完全相同的方向（完美集中）。
* $R = 0$: 向量在所有方向上均匀分布（没有集中）。

在实践中的应用。在实践中，我们为每个频带f单独计算$R_f$：

$$R_{f}=\frac{\|\mathbb{E}[q_{f}]\|}{\mathbb{E}[\|q_{f}\|]}$$

高$R_f$表明对于频带f，三角级数近似是准确的，这证明了使用$S_{trig}$的合理性。$S_{norm}$中的权重因子$(1 - R_f)$确保了当集中度较低时，基于范数的评分贡献更多。

B.6. 重建相关性 (Reconstruction Correlation)

验证Q/K集中度的影响。为了验证Q/K集中能够实现可预测的距离偏好，我们测量了三角级数重建实际注意力模式的效果。我们将重建相关性$\bar{r}$定义如下。

定义。对于给定的注意力头，令$\hat{s}(\Delta)$表示在距离$\Delta$处的预测注意力logit，它是通过三角级数（主文本中的方程4）从Q/K中心计算得出的。对于每个查询i，令$a_i = (a_{i,1}, a_{i,2}, ...)$为其在距离$\Delta_1, \Delta_2, ...$处的实际注意力logit，令$\hat{s} = (\hat{s}(\Delta_1), \hat{s}(\Delta_2), ...)$为相应的预测。

单查询相关性。单查询相关性是皮尔逊相关系数：

$$r_{i}=\rho\left(\mathbf{a}_{i}, \hat{\mathbf{s}}\right)=\frac{\operatorname{Cov}\left(\mathbf{a}_{i}, \hat{\mathbf{s}}\right)}{\sigma_{\mathbf{a}_{i}} \sigma_{\hat{\mathbf{s}}}}$$


图A. 通过递归模拟评估内存。左图：在内存完整的情况下，所有中间状态都被保留，正确的值向上传播。右图：当一个中间状态丢失时（State2），错误会通过所有后续的返回值传播，从而破坏最终结果。

平均相关性。重建相关性$\bar{r}$是所有查询的平均值：

$$\bar{r}=\frac{1}{N} \sum_{i=1}^{N} r_{i}$$

距离采样。为了确保在不同距离尺度上的均衡覆盖，我们在对数间隔的距离$\Delta \in \{1, 2, 4, 8, 16, . . .\}$上进行采样。这可以防止近距离（数量众多）主导相关性，并确保长距离注意力模式得到充分的体现。

B.7. 主导频带选择

频带贡献不均。并非所有频带对注意力logit的贡献都相等。我们将主导频带定义为对预期注意力分数贡献最大的那些。对于每个头，我们计算频带f的预期贡献为：

$$C_f = \mathbb{E}[\|q_f\|] \cdot \mathbb{E}[\|k_f\|]$$

其中期望是在一个校准数据集上计算的。

选择与可视化。我们按$C_f$对频带进行排序，并选择前K个频带（通常K=2）进行可视化。这些主导频带占据了注意力logit幅度的绝大部分。主文本中的图表可视化了与这些主导频带相对应的2D复平面中的Q/K分布。

C. 递归状态查询基准详情

本附录提供了§5.2.2中使用的递归状态查询基准的详细信息。

C.1. 任务定义

任务描述。给定一个无向图$G = (V, E)$和一个目标步数k，模型模拟从指定节点开始的k步深度优先搜索，并报告：
* 当前节点：遍历当前所在的节点。
* 栈状态：从起始节点到当前节点的完整路径（按顺序）。
* 已访问节点：遍历过程中所有已访问的节点。


图B. 使用真实注意力图的方法可视化，对应于图4中的示意图。顶行：TriAttention的四个阶段。(1) 我们从pre-RoPE分布计算Q/K中心E[q], E[k]。(2) 使用三角级数，我们计算出基于距离偏好对键进行评分的Strig。(3) 我们添加由集中度加权的基于范数的分数Snorm，得到最终分数S̄(k)。(4) 我们保留得分最高的键并驱逐其余的。底行：真实注意力图，展示了图4中描述的场景。从左到右：显示距离偏好的原始注意力模式；捕捉对角线结构的Strig可视化；结合了范数信息的综合分数S̄；KV缓存修剪后的注意力，保留了基本模式。

C.2. 为什么使用DFS进行内存评估

DFS的适用性。DFS非常适合评估内存保留能力，因为：
* 历史依赖性：栈状态取决于完整的遍历历史——任何中间信息的丢失都会导致错误。
* 信息均匀分布：信息均匀分布在序列中，而不是集中在开头或结尾。
* 难度可控：任务难度直接随步数扩展。
* 确定性：算法是确定性的（邻居按升序选择），提供了唯一的真实答案。

C.3. 评估指标

评估指标。我们使用栈精确匹配作为主要指标，它要求完整路径按顺序完全正确。这是最严格的指标，对由KV缓存修剪引起的信息丢失最为敏感。

C.4. 为什么递归能测试内存

递归与内存测试。图A说明了为什么递归模拟能有效地测试内存保留能力。递归算法要求模型首先通过嵌套调用向下深入，然后回溯以产生结果。在回溯期间，模型必须回忆序列中早期的中间状态。如果任何状态被遗忘，错误将通过所有后续的返回值传播，从而破坏最终结果。

D. 方法可视化

图B提供了TriAttention流程的详细可视化，使用了真实的注意力图，补充了图4中的示意图。

E. 与其他基线的比较

与其他基线的比较。我们比较了TriAttention与其他KV缓存压缩方法，除了主文本中评估的SnapKV和R-KV。

AIME24上的比较。表A在AIME24上使用DeepSeek-R1-Distill-Qwen-7B在多个KV预算下与LazyEviction【索引53，Lazyeviction: Lagged kv eviction with attention pattern observation for efficient long reasoning，2025，】进行了比较，同时还包括了H2O【索引54，H2o: Heavy- hitter oracle for efficient generative inference of large language models，2023，NeurIPS】、TOVA【索引31，Transformers are multi-state rnns，2024，EMNLP】和RaaS【索引19，Raas: Reasoning-aware attention sparsity for efficient llm reasoning，2025，ACL】（结果引自LazyEviction）。TriAttention在每个预算下都优于所有方法，并在30%的KV预算下（46.7%）与Full Attention持平。

表A. 在AIME24（DeepSeek-R1-Distill-Qwen-7B）上不同KV预算下的比较。*结果引自LazyEviction。

F. 在多样化基准上的评估

泛化能力评估。为了评估在数学推理之外的泛化能力，我们在LongBench【索引21，Longbench: A bilingual, multitask benchmark for long context understanding，2024，ACL】（16个子任务，涵盖问答、摘要、对话、检索和代码；Qwen3-8B，50% KV预算）和RULER【索引18，Ruler: What’s the real context size of your long-context language models?，2024，CoLM】（检索任务，4K上下文）上进行了测试。我们与StreamingLLM【索引52，Efficient streaming language models with attention sinks，2024，ICLR】、PyramidKV【索引4，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024，arXiv】、KnormPress【索引10，Expected attention: ´ Kv cache compression by estimating attention from future queries distribution，2025，arXiv】、Ada-KV+SnapKV【索引12，Ada-KV: Optimizing KV cache eviction by adaptive budget allocation for efficient LLM inference，2025，NeurIPS】和H2O【索引54，H2o: Heavy- hitter oracle for efficient generative inference of large language models，2023，NeurIPS】进行了比较。

结果。表B展示了完整的LongBench结果。TriAttention在16个子任务中取得了最高的平均分（48.1），在16个任务中赢得了11个，并以+2.5的优势超过了Ada-KV+SnapKV（48.1 vs 45.6）。在RULER（表C）上，TriAttention以66.1的成绩优于所有基线，与SnapKV相比有+10.5的差距。我们还与H2O进行了比较，H2O需要O(n^2)的内存并且不能使用FlashAttention；在H2O能装入48GB GPU内存的12个LongBench子任务中（表D），TriAttention在12个任务中赢得了10个。

G. 未来偏移量设计消融

未来偏移量设计消融。我们的评分函数在多个未来偏移量D上评估键的重要性。我们对两个方面进行了消融：偏移量的范围/数量，以及间隔策略（表E）。

结果。将最大距离从128增加到4096，准确率从41.7%提高到48.8%（+7.1%），证实了未来偏移量是有益的。几何间隔（{1, 2, 4, ...}）显著优于线性间隔（45.8% vs. 28.7%, -17.1%），因为近距离位置需要更密集的采样。

H. 校准数据敏感性

校准数据敏感性。我们测试了对校准数据数量和质量的鲁棒性（表F）。

结果。性能在50k到960k词元的校准大小范围内保持稳定（45.4–45.8%）。同样，校准数据质量与准确率没有明显相关性：使用谷歌主页HTML（低质量）达到了46.2%，与ShareGPT聊天数据（46.7%）相当。这证实了校准期间捕获的Q/K统计数据是模型的内在属性，对校准数据的选择具有鲁棒性。

I. MLA架构验证

MLA架构验证。为了测试Q/K集中是否能推广到标准注意力架构之外，我们在使用多头潜在注意力（MLA）并拥有940个头的GLM-4.7-Flash上进行了评估。表G比较了GQA和MLA架构之间的重建质量（皮尔逊r）和方向集中度（MRL）。

结果。MLA架构显示出与GQA相当或更强的集中度和重建质量。值得注意的是，96.6%的MLA头达到了R > 0.95（而GQA为84.7%），这表明Q/K集中现象是

表B. LongBench结果（Qwen3-8B，50% KV预算）。16个子任务，涵盖问答、摘要、少样本分类、检索、计数和代码。粗体=压缩方法中的最佳。

表C. RULER检索结果（Qwen3-8B，50% KV，4K上下文）。粗体=压缩方法中的最佳。

架构通用的，甚至可能在MLA中更为显著。

J. 现实世界部署：单GPU上的智能体任务

现实部署。我们通过使用多轮智能体OpenClaw，在一个实际的部署场景中展示了TriAttention。我们在单个RTX 4090（24GB）上部署了带有AWQ INT4量化的Qwen3-32B，其中仅模型权重就消耗了大部分GPU内存，为KV缓存留下了非常有限的预算。这使得该设置对OpenClaw尤其具有挑战性，因为其在第一个请求中的提示已经超过15k词元，并且每个交互轮次都进一步扩展了上下文，因为智能体需要阅读和处理六个markdown文档以生成周报。

结果。使用完整注意力（基线），KV缓存在多轮交互中无限制增长，导致智能体在完成任务前就出现内存不足错误。而使用TriAttention，KV缓存压缩将内存使用量保持在预算范围内，使得智能体能够成功读取所有文档并生成报告。完成会话的截图如图C所示。

表D. 在H2O能装入48GB GPU内存的LongBench子任务上与H2O的比较（Qwen3-8B，50% KV）。H2O需要实例化完整的O(n^2)注意力矩阵，不能使用FlashAttention。粗体=最佳。

表E. Qwen3-8B（AIME24）上的未来偏移量消融。上：偏移范围的影响。下：间隔策略比较（17个偏移，范围[1, 65536]）。

表F. Qwen3-8B（AIME24）上的校准数据敏感性。上：校准数据大小的影响。下：校准数据质量的影响。

表G. GQA（Qwen3-8B）和MLA（GLM-4.7-Flash）架构之间的重建质量（皮尔逊r）和Q/K集中度（MRL）的跨架构比较。
(A) 重建质量（皮尔逊r）

(B) Q/K集中度（MRL）


图C. 在单个RTX 4090上使用Qwen3-32B（INT4）的OpenClaw演示。完整注意在多轮交互中内存不足，而TriAttention在GPU内存预算内完成了任务。完整视频可在我们的GitHub页面上观看：https://github.com/WeianMao/triattention 。