ClusterAttention：一种无训练双向注意力加速方法

# ClusterAttention：双向注意力的免训练加速方法  
来源：https://arxiv.org/html/2608.26965  
*独立研究员 2026年8月27日*  

###### 摘要  
本文介绍ClusterAttention，一种通用的双向注意力层免训练加速方法。现有的稀疏注意力方法要么依赖于输入结构（如语言中的顺序或图像中的空间邻近性），要么使用缓慢的聚类过程并通过多次前向传播分摊计算成本。ClusterAttention则采用快速递归聚类方法，根据每个注意力头中键和查询的几何形状自适应地生成有效聚类。该方法允许任意设置聚类大小。我们利用这一特性，将所有聚类大小固定为2的幂，使得块稀疏注意力在GPU上每个键查询交互的延迟与密集注意力相当。我们还推导了稀疏注意力输出误差的表达式，解释了一个反直觉的实验发现：紧密聚类可能比随机聚类导致更大的误差。随后，我们推导了通过质心补偿被排除聚类后的误差，表明该误差随聚类收紧而减小。我们将这种补偿机制集成到方法中。在大规模表格数据上，ClusterAttention将TabPFN-3\[5 (https://arxiv.org/html/2608.26965#bib.bib4)\]加速了2到6倍，同时保留至少99%的密集注意力精度。据我们所知，这是首个可成功应用于非结构化输入和单次前向传播场景的免训练方法。在Wan 2.1-14B T2V\[16 (https://arxiv.org/html/2608.26965#bib.bib15)\]视频生成任务中，与专为该领域开发的领先方法SVOO\[9 (https://arxiv.org/html/2608.26965#bib.bib9)\]相比，ClusterAttention在无需离线校准的情况下，输出更接近密集注意力且加速比更高（1.8倍 vs 1.4倍）。111由于时间和预算限制，当前评估有限。代码将发布于 https://github.com/SpoketKasper/ClusterAttention，其中包含完整的预印本历史记录。  

## 1 预备知识  
### 1.1 稀疏注意力的动机  
双向注意力是现代Transformer模型中常见的操作，无需因果结构。其应用包括视觉Transformer、视频生成模型、文本嵌入模型、基因组模型和表格数据模型。然而，该操作的主要缺点是计算成本与token数量呈二次方关系，因为每个token需关注所有token（包括自身）。但许多连接往往较弱，不会显著影响操作输出。通过廉价剪枝这些交互，仅对每个token执行其与重要子集的注意力操作的方法称为**稀疏注意力方法**，而完整的注意力计算则对应称为**密集注意力**。本研究中，决定哪些查询关注哪些键及其方式的过程称为**路由**。当特定键查询交互以单一分辨率token直接发生时，称为执行**逐token注意力**。在给定注意力连接数下保留尽可能多注意力质量的路由被称为**先知路由**。可在推理时作为预训练注意力层简单修改使用的方法称为**免训练方法**，因其无需进一步训练（尽管仍可能结合修改进行训练）。在某些场景中，如表格数据\[5\]、高分辨率图像（例如病理学）\[18\]、视频生成\[7\]和基因组数据\[11\]，token数量庞大。通常的解决方案包括限制token数量、接受高延迟，或根本改变注意力计算。另一种使这些场景计算可行的途径是使用稀疏注意力，精确近似双向注意力而非改变其核心计算。本研究针对这一方向，其高层流程将所有键查询间的注意力替换为以下步骤：  
1. 对键和查询进行聚类  
2. 评估簇间得分，为每个查询簇选择键簇子集  
3. 在选定簇上执行注意力  
4. 可选：补偿未选中的簇  

若聚类、评分和选择的开销相对于节省的注意力计算较小，且稀疏注意力的单次键查询交互延迟不显著慢于密集注意力，则该方法可提供有效加速。若使用补偿，其增加的延迟也不应超过节省的计算量。  

### 1.2 稀疏注意力的误差分析  
本节分析稀疏注意力输出与密集注意力输出间的误差。定义\(v_i\)为token \(i\)的值，\(w_i\)为其从当前查询获得的权重，则该查询的注意力输出为：  
\[ o=\sum_{i}w_{i}v_{i} \quad (1) \]  
若选择键值子集\(S\)进行注意力计算，其输出为：  
\[ o_{S}=\frac{\sum_{i\in S}w_{i}v_{i}}{\sum_{i\in S}w_{i}}=\frac{\sum_{i\in S}w_{i}v_{i}}{w_{S}} \quad (2) \]  
其中定义\(w_{S}:=\sum_{i\in S}w_{i}\)。设排除token集为\(\bar{S}\)，其注意力输出为\(o_{\bar{S}}\)，则有：  
\[ o=w_{S}o_{S}+(1-w_{S})o_{\bar{S}} \quad (3) \]  
因此误差为：  
\[ o-o_{S}=(w_{S}-1)o_{S}+(1-w_{S})o_{\bar{S}}=(1-w_{S})(o_{\bar{S}}-o_{S}) \quad (4) \]  
该分解表明最小化注意力误差需关注两个因素：一是**遗漏注意力质量**（稀疏注意力方法的常见焦点），二是包含集与排除集注意力输出（即加权平均值）的偏差。当键值协方差较高时，选择与查询点积最高的键的稀疏选择会自然导致第二项偏差较大。本文发现一个反直觉现象：随机分配的簇（基于查询簇与键簇质心进行选择）可优于针对特定键集高注意力质量召回率设计的聚类方法（即使误差表达式中第一项较小），因为随机簇自然诱导较小的第二项偏差。这在视觉Transformer DINOv2\[12\]和表格数据Transformer TabPFN-3\[5\]中均有体现，见第4.1节和第4.2节。  

解决第一项因素需要紧密的查询簇和键簇，而解决第二项因素则更复杂。然而，我们可通过修改计算使簇紧密度直接降低误差。具体而言，若通过未选中簇的键值质心与查询的交互将其纳入计算，则可推导出（见附录A）补偿后的误差为：  
\[ o-\hat{o}_{S}=\frac{1}{\hat{d}_{S}}\sum_{c\in C_{\bar{S}}}\|c\|\left(\delta_{c}(\bar{v}_{c}-o)+\mathrm{Cov}_{c}(w,v)\right) \quad (5) \]  
其中\(\hat{o}_{S}\)为经质心补偿的稀疏注意力输出，\(C_{\bar{S}}\)为被排除的簇集，\(\delta_{c}\)为softmax指数中Jensen不等式引起的注意力质量误差。当大部分注意力质量被密集覆盖时，可证明\(\hat{d}_{S}\)的低估很小（实验也验证了这一点）。将其近似为1，可简化为：  
\[ o-\hat{o}_{S}\approx\sum_{c\in C_{\bar{S}}}\|c\|\left(\delta_{c}(\bar{v}_{c}-o)+\mathrm{Cov}_{c}(w,v)\right) \quad (6) \]  
第一项称为**Jensen项**，第二项称为**协方差项**。显然，Jensen项的第二因子不易调整（除非通过极松散的簇使质心移向值均值），而第一因子可通过在键空间紧密的簇最小化。协方差项随键和值的分布扩展而增大，因此通过在键值空间均紧密的簇可最小化。本文将这种包含未选中簇的方法称为**条纹均值补偿**，因其通过包含token均值进行补偿，产生视觉上的条纹注意力矩阵。在DINOv2上的实验表明，即使聚类时未考虑值，Jensen项仍比协方差项大约一个数量级（但测试尚不充分）。这是通过密集计算人为移除分量并检查输出误差测量的。  

## 2 相关工作  
本节首先简要介绍一组类似工作，然后说明ClusterAttention与它们的异同。  

**SpargeAttn** 是一种免训练稀疏化注意力方法\[20\]，依赖输入数据结构（如有意义的排序或图像/视频空间的空间填充曲线）创建聚类。本文使用行主序聚类实现，因未找到官方空间填充曲线形式实现，且论文表明两者差异不大。SpargeAttn选择大小为2的幂的簇以适配GPU分块，并测量簇内方差：若超过阈值，则将簇纳入所有计算（因基于质心的路由可能不准确）。  

**Clustered Attention** 方法按头对查询聚类，使用每个查询簇的质心表示簇内所有查询\[15\]。Vyas等人将该方法应用于预训练模型及新模型训练。他们使用基于查询的局部敏感哈希和汉明空间K-Means的快速聚类，并通过为每个查询簇计算其权重最高的\(k\)个键的密集注意力来改进近似。  

**AdaCluster** 是一种免训练方法，按头分别聚类键和查询\[13\]。Tan等人指出键和查询在注意力中的不同作用，并据此提出不同聚类方法：键在欧氏空间使用定制多阶段K-Means聚类，查询在普通K-Means聚类前进行归一化以实现基于角度的聚类（更紧凑）。键簇到查询簇的分配使用TensorQuest（Quest\[14\]的改进）。该方法明确针对视频扩散Transformer，但也可能适用于其他双向注意力领域。  

**SVOO** 是一种免训练方法，执行离线逐层稀疏性分析\[9\]。Luo等人指出视频扩散Transformer中的稀疏性与输入无关，而是每层的固有属性。他们认为若两个键对所有查询产生相似的注意力逻辑值，则从查询视角看它们相似；若两个查询与所有键的点积相似，则从键视角看它们相似。基于此，他们设计了一种迭代优化随机初始化键查询簇的聚类方法。键簇到查询簇的分配使用簇质心点积。键和查询每\(N=20\)个扩散步骤重新聚类。该方法同样明确针对视频生成，但可能更广泛适用。本文使用固定预算（1024个键簇和256个查询簇），这是作者推荐的配置，因其测试表明最优簇数与token数相关性较小。需注意的是，本文在非目标应用（视频生成）及更大token数范围内测试此配置，可能并非最优。  

ClusterAttention融合了上述工作的多个思想：与AdaCluster类似考虑键查询的不同作用，与SVOO类似关注键查询交互，并采用与SpargeAttn类似的适配分块的簇大小。总体上，此类稀疏注意力方法均遵循类似模式：聚类键和/或查询→利用簇选择跳过哪些键查询交互→执行稀疏注意力操作。其差异在于各部分的具体实现。ClusterAttention使用考虑键查询交互的空间变换，文献中未发现类似应用。它还沿主成分递归分裂以生成预设大小的簇，实现快速聚类且单次交互延迟不显著慢于密集注意力。文献中未发现将基于主成分的聚类方法应用于注意力设置的案例，也未发现除基于输入结构聚类（如SpargeAttn）外的其他固定大小聚类工作。  

## 3 方法  
### 3.1 概述  
ClusterAttention包含三部分：键和查询的**聚类**、注意力矩阵中待处理块的**分配**以及**注意力**计算执行。聚类对每个注意力头单独进行，且键和查询分别处理。分配和注意力同样按头执行。本节复杂度计算针对单头。聚类采用递归分裂方法，将键或查询集投影到近似...