技术摘要：头自主性 (Autonomy-of-Heads, AoH)

问题陈述

大语言模型（LLM）的长文本推理目前受限于注意力机制的二次方计算成本以及键值（KV）缓存线性增长的内存需求。现有的解决方案（如稀疏注意力机制和 KV 压缩方法）通常依赖于依赖于输入的信号来决定保留哪些 Token 或注意力头。这些信号包括运行时注意力分数（例如 Heavy-Hitter）、观察窗口、校准提示词或学习到的门控机制。虽然这些方法行之有效，但它们需要运行时观察、额外的训练或校准过程，这使得部署成本高昂，且依赖于特定的输入行为。

作者提出了一个问题：仅凭冻结权重本身是否可以为确定注意力头的功能提供有用的先验知识，而无需依赖运行时分数、校准提示词或额外的训练？

方法论：头自主性 (AoH)

AoH 是一种无需数据且无需训练的方法，它直接从冻结的查询-键（Query-Key）投影的谱几何（spectral geometry）中识别出“检索头”（需要全局上下文）和“流式头”（通过局部/汇点 Token 即可满足）。

1. 理论基础

该方法为每个头 $h$ 定义了一个核注意力算子：
$$M_h = W_K^{h\top} W_Q^h \in \mathbb{R}^{d_{model} \times d_{model}}$$
该矩阵代表了特定于头 $h$ 的冻结查询-键匹配方向。作者假设 $M_h$ 的谱几何编码了该头的功能角色：

• 集中谱（Concentrated Spectra）： 表示存在少量占主导地位的查询-键匹配方向，这与需要全局搜索特定上下文的检索头相关。
• 弥散谱（Diffuse Spectra）： 表示缺乏占主导地位的全局匹配方向，这与专注于汇点（sink）Token 和近期窗口的流式头相关。

2. 有效秩作为分类器

AoH 使用 $M_h$ 的**有效秩（Effective Rank, ER）**作为衡量头功能的权重空间度量。ER 是通过归一化后的 $M_h$ 奇异值的香农熵的指数计算得出的：
$$\text{eff\_rank}(h) = \exp\left(-\sum_{k=1}^{r_h} \hat{\sigma}_k \log \hat{\sigma}_k\right)$$

• 低 ER： 暗示存在集中谱 $\rightarrow$ 检索头（分配完整的全局注意力）。
• 高 ER： 暗示存在弥散谱 $\rightarrow$ 流式头（分配受限的汇点及近期窗口注意力）。

3. 高效计算

构建完整的 $d_{model} \times d_{model}$ 矩阵 $M_h$ 计算成本很高。作者推导出了一个高效的 $d_{head}$ 维计算方法。通过利用 Sylvester 行列式定理，可以从一个较小的代理矩阵 $C_h$ 中计算出 $M_h$ 的非零奇异值：
$$C_h = (W_Q^h W_Q^{h\top})(W_K^h W_K^{h\top}) \in \mathbb{R}^{d_{head} \times d_{head}}$$
这使得计算复杂度从 $O(d_{model}^2)$ 降低到 $O(d_{head}^2 \cdot d_{model})$，使得预计算头标签变得可行。

4. 部署策略

• 头分类： 在推理之前，根据 ER 对头进行排序。前 $k$ 个头（低 ER）被指定为检索头；其余为流式头。
• KV 缓存管理：
  • 检索头： 保留完整的全局 KV 状态。
  • 流式头： 仅保留固定大小的“汇点”缓存（例如前 128 个 Token）和一个“近期”窗口（例如最后 256 个 Token）。
• 实现： 该方法与分组查询注意力（GQA）兼容，通过将头级分数聚合到组级（使用均值聚合）来实现。它能够集成到标准的块状预填充（chunked prefill）和 FlashAttention-2 中，无需专门的内核。

核心贡献

1. 无需数据的头选择： 引入了 AoH，这是一种直接从冻结的查询-键几何中识别检索头和流式头的方法，无需校准数据或运行时观察。
2. 有效秩分类器： 开发了一个基于核注意力矩阵 $M_h$ 的有效秩分类器，并开发了一种通过 $d_{head}$ 维代理矩阵计算其谱的高效算法。
3. 稀疏注意力策略： 将 AoH 实现为一种简单的、无需训练的稀疏注意力策略，在显著降低延迟和内存的同时保持了准确性。

实验结果

作者在 LongBench 上对多个模型（Qwen2.5-7B, Qwen3-8B, Llama3.1-8B-Instruct）进行了评估。

• 准确度： 在 50% 稀疏度（50% 的头作为检索头，50% 作为流式头）下，AoH 平均保留了 96.5% 的全注意力性能。
  • 在 Qwen3-8B 上，AoH 仅比全注意力低 0.32 分。
  • 在 Llama3.1-8B-Instruct 上，差距为 1.10 分。
  • AoH 始终优于 H2O、SnapKV、Quest 和 DuoAttention 等基线方法，也优于随机选择和反向选择头的消融实验。
• 效率：
  • 延迟： 在 256K 上下文长度和 75% 稀疏度下，与全注意力相比，AoH 在 Llama3.1-8B 上实现了高达 3.24 倍的预填充加速和 9.14 倍的解码加速。
  • 内存： 在 75% 稀疏度下，KV 缓存内存占用减少高达 3.98 倍；在 256K Token 且 50% 稀疏度下，减少了 50.0%。

意义与主张

论文声称头也“知道自己知道什么”；它们的功能角色（检索 vs 流式）被编码在其冻结权重的谱几何中。AoH 提供了一个简单、可扩展且无需训练的先验，用于稀疏注意力。

• 独立性： 与以往的工作不同，AoH 不依赖于输入相关的信号、学习到的门控或校准提示词。它使得稀疏注意力策略可以从第一个 Token（预填充阶段）就开始应用，而不是在运行时观察之后。
• 兼容性： 该方法与现有架构（包括 GQA）和标准推理引擎（FlashAttention）兼容，具有实际部署价值。
• 鲁棒性： 作者证明了有效秩排序在不同上下文长度下是稳定的，并且即使在考虑 RoPE（旋转位置编码）旋转时也保持高度一致。

作者总结道，AoH 为提高长上下文推理效率提供了一条可行路径，且无需额外的训练开销或复杂的运行时启发式算法。