ASAP: Attention Sink Anchored Pruning¶
会议: NeurIPS 2026
arXiv: 2605.22372
OpenReview: https://openreview.net/forum?id=caL9VT4BtX
领域: 模型压缩
关键词: Vision Transformer, Token 剪枝, Attention Sink, 随机游走, 扩散距离
一句话总结¶
将注意力汇流(Attention Sink)重塑为马尔可夫链的概率吸引子,以汇流为几何原点在累积转移概率流形上计算扩散距离,实现单步自适应背景压缩与重要前景保留,完全无需训练即可加速达 48%。
研究背景与动机¶
视觉 Transformer(ViT)在高分辨率图像、长视频和多模态大语言模型(VLM)中面临自注意力机制带来的二次复杂度瓶颈。现有的免微调 Token 剪枝与合并方法(如 ToMe、FastV、PPT、VisPruner 等)主要依赖单层注意力分数或局部键向量的成对余弦相似度进行冗余过滤。然而,这些局部浅层度量无法捕捉深层网络中涌现的全局流形非线性结构,常常导致低注意力但高语义价值的小前景物体被误剪枝。
这一缺陷因视觉 Transformer 中普遍存在的注意力汇流(Attention Sink)现象而雪上加霜:在无意义的背景补丁或边界处,注意力权重常常异常汇集形成极端离群值。依赖局部标量注意力的剪枝方法将高注意力等同于重要性,极易陷入汇流陷阱——要么误将无用背景汇流锚定为核心保留对象,要么因汇流吸收了绝大部分注意力而使真正的前景目标权重被压低。现有部分图分析方法(如 Zero-TPrune)虽然尝试引入 PageRank,但由于在各层独立计算稳态分布,不仅计算代价高昂,依然未能突破单层局部注意力的信息瓶颈。
本文敏锐地洞察到:汇流补丁虽然缺乏语义信息,但其空间位置高度规律且天然形成在平坦无信息冗余的背景区域,具有极强的结构可预测性。因此,无需将汇流视作病态表征去刻意消除或避让,反而可以将其转化为绝佳的全局几何坐标原点。核心 idea:将 ViT 信息流建模为带自环残差的惰性随机游走,把注意力汇流视作概率质量吸引子,利用累积转移矩阵中到汇流的扩散距离直接构建单步径向聚类与加权池化,实现免训练的高效前景与背景解耦。
方法详解¶
整体框架¶
ASAP 包含三大阶段:首先,利用带有自环倾向的惰性随机游走(Lazy Random Walk)将多层自注意力矩阵与残差连接结合,生成全局累积转移矩阵;其次,通过列和触发机制动态捕捉注意力汇流涌现的关键层 \(t^*\),并将汇流作为几何锚点计算各 Token 到汇流的扩散距离;最后,基于扩散距离进行径向扩散聚类(Radial Diffusion Clustering),对紧靠汇流的背景簇执行转移加权池化(Transition Weight Pooling),保留核心前景结构(极低预算下可直接串联即插即用的前景去冗余阶段)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入 Patch 序列<br/>X = [CLS, x_1, ..., x_N]"] --> B["阶段 1:惰性随机游走信息流建模<br/>融合残差连接生成累积转移矩阵 P(t)"]
B --> C["阶段 2:注意力汇流动态检测与扩散距离度量<br/>列和触发捕获 t*,以汇流为几何锚点算 D(x_i, x_s)"]
C --> D["阶段 3:径向扩散聚类与转移加权池化<br/>划分 K 个能级簇,将背景簇 C_1 加权池化为单一 Token"]
D --> E["极端预算可选分支:即插即用扩散度量剪枝<br/>在前景簇中进行贪心二分匹配消除信息冗余"]
D -->|常规预算输出| F["输出压缩后 Token 序列<br/>大幅减轻深层 ViT 与 LLM 解码负担"]
E -->|极端预算输出| F
关键设计¶
1. 惰性随机游走累积转移矩阵:将跨层残差信息流几何化
传统 Rollout 忽略了残差连接在马尔可夫链中的数学自环特性。ASAP 将第 \(l\) 层头平均自注意力矩阵 \(A^l\) 与残差连接结合,定义为惰性随机游走(Lazy Random Walk)有效转移矩阵: $\(\tilde{A}^l = \alpha A^l + (1 - \alpha) I\)$ 其中 \(\alpha \in (0,1)\)(默认取 0.5)平衡注意力扩散与残差身份保持。对角线引入的恒等矩阵 \(I\) 保证了矩阵元素的非负性与严格正对角线,依据 Perron-Frobenius 定理消除了马尔可夫链的周期震荡,确保了稳态收敛的数值稳定性。经过 \(t\) 层后,全局累积转移矩阵为: $\(P^{(t)} = \tilde{A}^1 \times \tilde{A}^2 \times \cdots \times \tilde{A}^t\)$ 第 \(i\) 行 \(P^{(t)}_{i,*}\) 刻画了来自 Token \(x_i\) 的信息在网络前向传递 \(t\) 步后在全图分布的概率流,构成了后续计算流形几何距离的连续空间。
2. 汇流吸引子动态捕捉与单点扩散距离:降维解耦全局语义流形
在马尔可夫流中,注意力汇流表现为单调吸收概率质量的“概率吸引子”(Mass-accumulating Attractor)。ASAP 无需人工指定剪枝层,而是实时监测各列概率质量总和,在首次突破阈值 \(\tau\) 时触发剪枝,并以 [CLS] 注意力分布的最大列定位汇流坐标 \(x_s\): $\(t^* = \min \left\{ t \;\middle|\; \max_{j \ge 1} \sum_{i=1}^N P^{(t)}_{i,j} > \tau \right\}, \quad x_s = \arg\max_{j \ge 1} P^{(t^*)}_{0,j}\)$ 以汇流 Token \(x_s\) 为全图几何坐标锚点,ASAP 定义任意 Token \(x_i\) 到汇流的扩散距离为行向量间的 \(\ell_2\) 范数: $\(D(x_i, x_s) = \| P^{(t^*)}_{i,*} - P^{(t^*)}_{s,*} \|_2\)$ 若 \(D(x_i, x_s) \approx 0\),说明该 Token 的全局路由与汇流高度重合,属于冗余背景;若距离显著大于零,则意味着信息流向完全独立于背景汇流,表征独特的前景语义结构。由于只需计算各 Token 到单点锚点的欧氏距离,其计算复杂度仅为 \(\mathcal{O}(N^2)\),相比两两计算成对余弦相似度的 \(\mathcal{O}(N^2 d)\) 或特征奇异值分解的 \(\mathcal{O}(N^3)\) 具有显著的轻量优势。
3. 径向扩散聚类与转移加权池化:无监督语义分层与保形压缩
在获取标量扩散距离后,ASAP 将归一化距离 \(\tilde{D}(x_i, x_s) \in [0, 1]\) 线性划分为 \(K\) 个径向同心层级簇: $\(C_k = \left\{ x_i \;\middle|\; k-1 \le K \cdot \tilde{D}(x_i, x_s) < k \right\}, \quad k=1, \dots, K\)$ 紧贴汇流原点的第一能级簇 \(C_1\) 聚集了绝大部分背景 Token。ASAP 并不直接丢弃 \(C_1\),而是通过扩散距离 Softmax 赋予的转移权重将其加权池化为单一背景代表元 \(\mathbf{m}_1\): $\(\mathbf{m}_1 = \sum_{x_i \in C_1} w_i \mathbf{f}_i, \quad \text{其中 } w_i = \frac{\exp(D(x_i, x_s))}{\sum_{x_j \in C_1} \exp(D(x_j, x_s))}\)$ 离汇流较远的边缘背景由于承载更多残余语义上下文而获得更高权重。外层前景簇(\(C_2, \dots, C_K\))无损保留,在极端压缩场景下可通过在 \(P^{(t^*)}\) 几何空间内执行扩散度量二分匹配过滤组内冗余。
完整示例 / 训练策略¶
整个流程完全免训练(Training-free),无需任何反向传播微调。以典型 ViT 图像分类(196 个 patch,DeiT-Base,设定 \(K=6, \tau=7\))为例: 1. 前向与累积:在第 1 至 7 层进行前向计算,每次前向仅伴随一次 \(196 \times 196\) 的矩阵乘法累计 \(P^{(t)}\)(耗时小于 0.01ms); 2. 触发与锚定:在第 8 层检测到汇流列和达到 \(\tau=7\),停止累积,以 [CLS] 对应最大列锁定位于背景路面的汇流 patch \(x_s\); 3. 计算与聚类:耗时小于 0.01ms 完成各行与汇流行的残差模长计算与区间划分,约 110 个低距离背景 patch 被划分进 \(C_1\); 4. 池化压缩与后续执行:\(C_1\) 快速池化为一个 Token,保留其余约 86 个前景 Token,拼接后形成 88 个 Token 继续输入第 9 至 12 层,直接节省后半段计算 FLOPs 达 35.2%,最终推理吞吐提升 48.2%。
实验关键数据¶
主实验:ImageNet-1K 分类与 Kinetics-400 视频理解¶
在统一评估口径下,ASAP 在保持高加速比的同时展现出超越同类免训练方法的准确率保持度:
| 数据集 / 骨干架构 | 方法 | 原始 Acc (%) | 压缩后 Acc (%) | 精度变化 \(\Delta\) | FLOPs 下降 (%) | 吞吐提升 (%) |
|---|---|---|---|---|---|---|
| ImageNet-1K (DeiT-Base) | DynamicViT | 81.80 | 81.30 | -0.50 | -34.09 | +38.01 |
| ImageNet-1K (DeiT-Base) | ToMe | 81.80 | 80.68 | -1.12 | -34.09 | +38.01 |
| ImageNet-1K (DeiT-Base) | PPT | 81.80 | 80.30 | -1.50 | -32.95 | +12.61 |
| ImageNet-1K (DeiT-Base) | ASAP (Prun.) | 81.80 | 81.33 | -0.47 | -36.36 | +48.45 |
| ImageNet-1K (DeiT-Base) | ASAP (Pool.) | 81.80 | 81.68 | -0.12 | -35.23 | +48.21 |
| ImageNet-1K (LV-ViT-S) | ToMe | 83.35 | 79.80 | -3.55 | -46.97 | +34.28 |
| ImageNet-1K (LV-ViT-S) | ZeroTPrune | 83.35 | 81.50 | -1.85 | -46.97 | - |
| ImageNet-1K (LV-ViT-S) | ASAP (Ours) | 83.35 | 82.57 | -0.78 | -46.97 | +35.66 |
| Kinetics-400 (CLIP ViT-L) | ToMe | 80.78 | 78.18 | -2.60 | -34.09 | +20.0 |
| Kinetics-400 (CLIP ViT-L) | PPT | 80.78 | 75.93 | -4.85 | -32.95 | +18.9 |
| Kinetics-400 (CLIP ViT-L) | ASAP (Pool.) | 80.78 | 80.57 | -0.21 | -35.23 | +23.7 |
消融实验:关键模块与几何锚点有效性验证¶
对核心设计模块(马尔可夫累积、锚点选择、聚类参数等)的消融展示了该框架的稳健性:
| 实验配置 | 测试基准 / 模型 | 评测指标 | 结果 | 说明 / 结论 |
|---|---|---|---|---|
| 累积信息流消融 (W/ Markov) | ImageNet-1K (DeiT-B) | Top-1 Acc (%) | 81.68% | 完整多层累积转移流,背景与前景清晰解耦 |
| 单层局部注意力消融 (W/O Markov) | ImageNet-1K (DeiT-B) | Top-1 Acc (%) | 79.43% | 仅用 \(t^*\) 单层注意力计算扩散距离,精度暴跌 2.25%p |
| 随机锚点 (Random Anchor) | ImageNet-1K (DeiT-B) | Top-1 Acc (%) | 77.64% | 锚点误落在前景时,导致核心目标被逆向池化为背景 |
| 汇流锚点 (Attention Sink Anchor) | ImageNet-1K (DeiT-B) | Top-1 Acc (%) | 81.68% | 汇流作为吸引子天然处于中性背景,解耦最稳定(+4.04%p) |
| PCA 中心锚点 (PCA Center) | LLaVA-NeXT-7B (160 Tokens) | POPE / MME | 75.2% / 1564 | 统计中心被显著前景特征牵引,导致前景被错误压缩 |
| 汇流锚点 (Attention Sink) | LLaVA-NeXT-7B (160 Tokens) | POPE / MME | 81.6% / 1730 | 信息动力学吸引子优于统计学方差中心(POPE 显著领先 6.4%p) |
| 粗粒度聚类 (\(K=3, \tau=7\)) | ImageNet-1K (DeiT-B) | Top-1 Acc (%) | 79.82% | 簇划分过少导致 \(C_1\) 误吞噬了部分前景边缘 |
| 最优聚类配置 (\(K=6, \tau=10\)) | ImageNet-1K (DeiT-B) | Top-1 Acc (%) | 82.03% | 簇划分足够细致且累积更充分,性能甚至反超未经剪枝基线 |
关键发现¶
- 去噪正则化抑制多模态大模型幻觉:在 LLaVA-1.5 的 POPE 评测中,保留 128 个 Token 时的评测得分达到 86.2%,在 64 个 Token(仅保留 11% 视觉 Token)时得分仍达 85.7%,甚至全面超越了未剪枝全量基线的 85.9%。定性分析表明,消除背景汇流中弥散且不具局部特异性的高频噪声后,能显著修正由于复杂背景光影引起的误判和由于大面积单色背景导致的物体漏检。
- 动态阈值具备跨架构跨模态强泛化性:尽管 DeiT(\(L=12\))与 CLIP ViT-Large(\(L=24\))在网络层深与 Token 规模上迥异,但累积列和呈现一致的二阶段相变模式:浅层平坦(约 2 左右),随后在某一特定层急剧跃升。选择 \(\tau \in [7, 8]\) 可稳定捕捉该几何相变拐点(DeiT 位于第 8 层,CLIP 位于第 12 层),且其计算矩阵乘法的累加延迟在早期停止截断后仅为数毫秒,完全被后续网络大幅削减的序列长度收益所弥补。
亮点与洞察¶
- 变缺陷为利器(Artifact as Feature):不同于以往将注意力汇流视为病态表征而试图通过引入 Register Token、注意力截断或惩罚项进行抑制的常规思路,本文将汇流的单调收敛特性作为随机游走中的概率质量吸收器,天然构建了前景与背景分离的坐标系。
- 度量复杂度降维打击:利用扩散距离在均匀稳态假设下等价于欧氏范数的理论特性,巧妙地将传统成对相似度聚类的 \(\mathcal{O}(N^2 d)\) 降低为基于单一参考锚点的标量距离运算 \(\mathcal{O}(N^2)\),并在聚类时仅需 \(\mathcal{O}(N \log N)\) 排序分桶,避免了昂贵的特征值分解或 K-Means 迭代。
- 跨模态与动态视频泛化:在 Kinetics-400 视频任务中,虽然场景变化导致各帧注意力的汇流位置在空间中动态漂移,但 ASAP 逐帧重新计算以汇流为基准的扩散距离,依然能够稳定跟踪运动前景,展现出极强的动态适应性。
局限与展望¶
- 无清晰背景场景下的退化风险:算法的核心先验假设是 dominant sink 稳定涌现于语义中性的平坦背景区域。在极端特写图像、纯纹理密集画面或多前景完全铺满全图的特殊场景下,汇流可能被迫形成于弱前景上,此时若未进行保护,可能对局部前景细节产生细微干扰。
- 自适应聚类簇数 \(K\) 的演进:目前簇数 \(K\) 仍需人工设定(推荐 \(K \in [5, 7]\)),未来可基于流形上扩散距离的密度分布自动探测第一阶导数的跳变截断点,从而实现真正无需人工干预超参数的自适应聚类。
- 延伸至 LLM 的 KV Cache 压缩:大语言模型在生成过程中同样存在严重的注意力汇流和初始 Token 沉降现象,将累积概率流与扩散距离拓展至自回归长上下文的 KV Cache 稀疏化具有广阔前景。
相关工作与启发¶
- vs ToMe / FastV / PPT: ToMe 等采用局部特征的成对余弦相似度进行双边匹配,易受单层特征噪声干扰且计算复杂度含特征维度 \(d\);FastV 等仅依赖单层注意力标量剪枝,易把汇流误当做前景保留。ASAP 采用累积多层信息的扩散流形距离,以汇流为锚点,能准确剔除真正的冗余背景。
- vs Zero-TPrune: Zero-TPrune 虽然引入了图的加权 PageRank 稳态分布,但在每一层独立迭代求解,计算开销大且仍受限于单层注意力。ASAP 结合残差连接连续乘积构建累积转移矩阵,只在汇流形成的关键层单次执行径向聚类,更加高效且全局表征更鲁棒。
- vs 视觉 Register Token 方案: Darcet 等人通过在预训练时加入额外的 Register Token 来吸收注意力离群值,需要重新训练模型;ASAP 是完全免训练的即插即用推理加速框架,直接在现有冻结权重上榨取性能。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 逆向思维极具突破性,将注意力汇流巧妙建模为马尔可夫链的概率吸引子并作为几何锚点。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖图像分类、视频动作识别和多模态大模型 7 项基准,对比详尽,消融实验扎实。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备,从马尔可夫稳态分析到实验设计层层递进,动机非常具有说服力。
- 价值: ⭐⭐⭐⭐⭐ 免训练、低算力开销、即插即用,不仅能直接落地加速 ViT / VLM,还兼具缓解幻觉的额外功效。