Rényi Entropy: A New Token Pruning Metric for Vision Transformers

作者: Wei-Yuan Su, Ruijie Zhang, Zheng Zhang 年份: 2026 会议: arXiv 分类: 剪枝与稀疏化

论文笔记:Rényi Entropy: A New Token Pruning Metric for Vision Transformers

元信息

项目内容
机构University of California, Santa Barbara
日期March 2026
项目主页-
对比基线EViT, TCA, PruMerge, VisPruner
链接arXiv / Code

一句话总结

提出基于 Rényi 熵 的无训练 token 重要性度量 Col-Ln,通过列向 ℓn-范数替代不可靠的 [CLS] 注意力,从第一层就能准确识别重要 token。

核心贡献

揭示 [CLS] 指导剪枝的缺陷: 在 ViT 早期层中 [CLS] token 的注意力尚未成熟,导致重要前景 token 被错误丢弃,甚至不如随机剪枝

提出 Col-Ln 度量: 基于 Rényi 熵 推导出的无参数度量,利用 patch 间的集体共识而非单一 [CLS] 代理来衡量 token 重要性

广泛验证: 在 ViT 和 LVLM 上全面超越 SOTA 方法,同时可作为现有 [CLS] 方法的修正机制

问题背景

要解决的问题

ViTMulti-Head Self-Attention 复杂度为 O(N2)O(N^2),高分辨率输入下推理成本高昂

Token Pruning 是加速推理的关键技术,但现有方法依赖 [CLS] token 的注意力权重来估计 patch 重要性

现有方法的局限

语义不成熟: 在早期层(如 L0-L3),[CLS] token 随机初始化,尚未聚合足够的全局信息,其注意力图具有误导性

任务偏置: [CLS] token 受最终任务头监督,注意力分布偏向训练目标相关特征,泛化性差

Table 1 实证:在 ViT-Small 上,[CLS] 方法在多个 keep rate 下都不如随机剪枝

本文的动机

信息丰富的 token 会从多个其他 token 获得集中的注意力(低熵 = 高共识),而冗余 token 获得的注意力分散(高熵 = 低共识)

Rényi 熵 量化列方向的注意力集中度,可以从第一层就可靠地识别重要 token

方法详解

核心思想

Col-Ln 放弃以 [CLS] 为中心的全局代理方式,转而采用集体共识机制

  • 输入: 注意力矩阵 ARN×N\mathbf{A} \in \mathbb{R}^{N \times N}
  • 核心: 对每个 token jj,计算其被所有其他 token 关注的列向注意力分布Rényi 熵
  • 判断: 低熵 → 高共识 → 重要 token;高熵 → 低共识 → 冗余 token
  • 输出: 保留 Rényi 熵最低的 Top-K token

核心模块

模块1: Rényi 熵重要性度量

设计动机: 利用 Rényi 熵 的数学性质将熵最小化等价转化为 ℓn-范数最大化,避免直接计算对数

具体实现:

  • 计算注意力矩阵每列的 Rényi 熵,度量 token jj 被其他 token 关注的集中程度
  • 利用 Rényi 熵与 ℓn-范数的单调关系,将最小熵选择转化为最大 ℓn-范数选择
  • 阶参数 nn 控制对高概率事件的敏感度,nn 越大越能隔离达到强共识的 token

模块2: Col-Ln 修正机制(Correcting)

设计动机: 作为现有 [CLS] 方法的补充,“拯救”被 [CLS] 错误标记为不重要的关键 token

具体实现:

  • 通过 rescue ratio cc 分配预算:kcls=r(1c)k_{\text{cls}} = \lfloor r \cdot (1-c) \rfloor 个 token 由 [CLS] 选择
  • 剩余 kcol=rkclsk_{\text{col}} = r - k_{\text{cls}} 个 token 从 [CLS] 未选中的候选中按 Col-Ln 排序选取
  • 两组合并后保留,缓解早期层的过早信息丢失

关键公式

公式1: Rényi 熵定义

Hn(Columnj)=11nlog(i=1N(Ai,j)n)H_n(\text{Column}_j) = \frac{1}{1-n} \log\left(\sum_{i=1}^{N}(A_{i,j})^n\right)

含义: 度量 token jj 收到的列向注意力分布的集中程度

符号说明:

  • Ai,jA_{i,j}: token ii 对 token jj 的注意力权重
  • nn: 阶参数,n>1n > 1 时对高概率事件更敏感
  • NN: token 总数

公式2: ℓn-范数等价形式

Hn(Columnj)=n1nlog(A:,jn)H_n(\text{Column}_j) = \frac{n}{1-n} \log(\|\mathbf{A}_{:,j}\|_n)

含义: 将 Rényi 熵改写为 ℓn-范数的对数形式

符号说明:

  • n\|\cdot\|_n: ℓn-范数
  • n1n\frac{n}{1-n}: 当 n>1n > 1 时为严格负系数

公式3: 最小熵等价于最大 ℓn-范数

Bottom-K(Hn(Columnj),j)    Top-K(A:,jn,j)\text{Bottom-}K(H_n(\text{Column}_j), j) \iff \text{Top-}K(\|\mathbf{A}_{:,j}\|_n, j)

含义: 由于系数为负且对数单调递增,最小化 Rényi 熵等价于最大化列 ℓn-范数

符号说明:

  • Bottom-K\text{Bottom-}K: 选择最小的 KK
  • Top-K\text{Top-}K: 选择最大的 KK

公式4: Col-Ln 重要性分数

Sj=(i=0NAi,jn)1nS_j = \left(\sum_{i=0}^{N}|A_{i,j}|^n\right)^{\frac{1}{n}}

含义: 每个 token jj 的 Col-Ln 重要性分数,即注意力矩阵第 jj 列的 ℓn-范数

符号说明:

  • SjS_j: token jj 的重要性分数
  • nn: 范数阶数(论文中使用 n=4n=4

公式5: 修正模式预算分配

kcls=r(1c),kcol=rkclsk_{\text{cls}} = \lfloor r \cdot (1-c) \rfloor, \quad k_{\text{col}} = r - k_{\text{cls}}

含义: 将保留预算 rr 按 rescue ratio cc 分配给 [CLS] 和 Col-Ln 两个选择器

符号说明:

  • rr: 总保留 token 数
  • cc: rescue ratio(论文中使用 c=0.8c = 0.8
  • kclsk_{\text{cls}}: 由 [CLS] 注意力选择的 token 数
  • kcolk_{\text{col}}: 由 Col-Ln 从剩余候选中选择的 token 数

关键图表

Figure 1: Visual Comparison / Token 剪枝视觉对比

Figure 1: Visual Comparison of Token Pruning{:width 600}

说明: 对比 [CLS] 注意力和 Col-Ln 在初始层(L₀–L₅)的 token 剪枝效果。[CLS] 在早期层频繁误删前景 token,Col-Ln 从第一层起就能准确保留语义重要的 token。

Figure 2: Method Diagram / 方法对比图

Figure 2: Method Comparison{:width 600}

说明: 左侧为传统 [CLS] 方法,依赖 [CLS] 行注意力选 token;右侧为 Col-Ln 方法,计算每列的 ℓn-范数作为重要性分数,利用所有 patch 的集体共识。

Figure 3: Training Loss Comparison / 训练损失对比

Figure 3a: Training Loss (0,3,6){:width 600}

Figure 3b: Training Loss (3,6,9){:width 600}

说明: EViT 框架下的训练损失对比。(a) 在 (0,3,6) 调度下,Col-Ln 的初始损失显著低于 [CLS],说明从一开始就保留了更多关键信息。(b) 在 (3,6,9) 调度下差异缩小,符合 [CLS] 在深层更可靠的预期。

Figure 4: Heatmap Visualization / 注意力热力图

Figure 4: Heatmap Attention Visualization{:width 600}

说明: 对比 [CLS] 注意力(上方)和 Col-Ln 重要性分数(下方)在 L₀–L₅ 的热力图。[CLS] 在早期层注意力分散且噪声大,Col-Ln 从第一层就集中于前景目标。

Figure 5: Layer-wise Attention / 不同阶数的逐层注意力

Figure 5: Layer-wise attention comparison{:width 600}

说明: ViT-Base 上 [CLS] 注意力与不同阶数 nn 的 Col-Ln 的逐层可视化。高阶(如 n=4n=4)的 Col-Ln 在所有层都能产生更集中、更有区分度的重要性分布。

Table 1: EViT on ViT-Small / [CLS] 失败实证

Keep rate (rr)[CLS]Random
0.766.173.2
0.874.077.6
0.979.080.0

说明: 在 (0,3,6) 调度下,[CLS] 在所有 keep rate 上都不如随机剪枝,实证确认其在早期层的灾难性失败。

Table 2: Early-Layer Pruning / 早期层剪枝

Modelpp[CLS]OursGFLOPs
ViT-S/16-81.3881.384.6
480.9881.144.2
880.1280.593.7
1279.0679.943.3
1677.2278.782.9
2073.8776.652.4
2466.5471.402.0
ViT-B/16-84.5484.5417.6
483.8784.3315.9
883.0783.9814.2
1282.0583.4712.6
1680.5382.7311.0
2077.9381.019.3
2472.0277.297.7
ViT-L/16-85.8385.8361.6
884.8284.9848.1
1682.1582.9434.8
2465.5472.1621.5

说明: 仅在前 6 层剪枝。Col-Ln 在所有模型和剪枝强度下均优于 [CLS],在最激进的 p=24p=24 下优势分别达 +4.86%(ViT-S)、+5.27%(ViT-B)、+6.62%(ViT-L)。

Table 3: All-Layer Pruning / 全层剪枝

Modelpp[CLS]OursGFLOPs
ViT-S/16-81.3881.384.6
281.2181.294.3
480.9281.084.0
680.5080.763.7
879.9980.393.4
1079.2279.763.1
1278.1078.632.9
ViT-B/16-84.5484.5417.6
284.1984.4816.5
483.8884.3015.3
683.5184.1714.2
882.9883.8713.1
1082.2783.3912.0
1281.4482.7010.9
ViT-L/16-85.8385.8361.6
285.6185.6853.8
485.2785.3446.1
684.4084.6438.5

说明: 全层逐层剪枝。Col-Ln 在所有配置下一致优于 [CLS],ViT-B 在 p=12p=12 时优势达 +1.26%。

Table 4: OOD Robustness / 分布外鲁棒性(TCA 框架)

MethodLayersIAVRSAvgGFLOPs
CLIP-ViT-B/16-68.3750.2061.8677.5548.2461.2417.6
TCA3,6,965.4647.4259.1466.0443.9156.3911.9
Ours3,6,965.3149.2759.1372.1045.6458.39-
TCA0,3,662.2638.9755.7160.4341.0951.699.2
Ours0,3,662.2540.4755.5164.3341.0952.73-

说明: ImageNet OOD 变体上的鲁棒性。Col-Ln 在 keep rate 0.7 下将平均 OOD 准确率从 56.39% 提升至 58.39%(+2.0%),无需微调。

Table 5: Cross-dataset Classification / 跨数据集分类

MethodLayersC101OPSCOFF101AS397DTDESUFCAvg
CLIP-ViT-B/16-92.9889.1359.3271.3886.1124.3365.4845.5748.1668.8365.80
TCA3,6,990.8785.1559.3271.1779.2923.8262.1044.9256.8968.9764.25
Ours3,6,992.7485.9957.6470.8581.2922.9863.1443.9758.1169.3464.61
TCA0,3,684.6782.1252.1369.5975.9522.0559.0243.0346.9362.3359.78
Ours0,3,689.9082.4749.6668.5377.4122.1760.9343.0348.9367.0161.00

说明: 10 个下游数据集上的跨数据集分类。Col-Ln 平均准确率 64.61% 超过 TCA 的 64.25%。

Table 6: LVLM Pruning (PruMerge) / 大视觉语言模型剪枝

MethodrrVQAv2GQAVizWizSQATextVQAPOPEMMEMMBMMB-CNMM-VetAvgRel.%
LLaVA-1.5 Full57676.762.054.269.558.285.9150564.658.129.763.4100.0
PruMerge25670.855.756.068.054.773.7136461.856.426.859.293.4
Ours25675.259.554.467.955.485.9147163.758.529.262.398.3
PruMerge19269.054.355.967.254.271.3129859.353.026.657.891.2
Ours19274.658.955.067.854.785.2143463.357.029.961.897.5

说明: 在 LLaVA-1.5-7B 上使用 PruMerge 框架。r=256r=256 时 Col-Ln 保持原模型 98.3% 性能(vs PruMerge 的 93.4%),POPE 指标完全无损。

Table 7: Early-Layer Correcting / 早期层修正

Modelpp[CLS]OursGFLOPs
ViT-S/16-81.3881.384.6
480.9881.034.2
880.1280.653.7
1279.0679.933.3
1677.2278.812.9
2073.8776.532.4
2466.5471.332.0
ViT-B/16-84.5484.5417.6
483.8784.3415.9
883.0784.0214.2
1282.0583.5412.6
1680.5382.6811.0
2077.9381.029.3
2472.0277.287.7
ViT-L/16-85.8385.8361.6
884.8284.8548.1
1682.1582.9934.8
2465.5472.6921.5

说明: 修正模式下的早期层剪枝结果,Col-Ln 有效修正 [CLS] 的错误判断。

Table 8: All-Layer Correcting / 全层修正

Modelpp[CLS]OursGFLOPs
ViT-S/16-81.3881.384.6
281.2181.324.3
480.9281.064.0
680.5080.803.7
879.9980.463.4
1079.2279.833.1
1278.1078.702.9
ViT-B/16-84.5484.5417.6
284.1984.4616.5
483.8884.3215.3
683.5184.1514.2
882.9883.8913.1
1082.2783.5312.0
1281.4482.8210.9
ViT-L/16-85.8385.8361.6
285.6185.6953.8
485.2785.2846.1
684.4084.7138.5

说明: 全层修正模式下 Col-Ln 同样一致优于纯 [CLS] 方法。

Table 9: EViT on ViT-Small / EViT 框架对比(ViT-Small)

Keep RatePruning Layers[CLS]OursRandomGFLOPs
--81.481.481.44.6
0.73,6,978.979.077.33.0
0,3,666.175.073.22.3
0.83,6,980.580.579.53.5
0,3,674.078.677.62.9
0.93,6,981.381.380.64.0
0,3,679.080.580.03.7

说明: 在 (0,3,6) 调度下 Col-Ln 以 +8.9% 的巨大优势超越 [CLS](75.0% vs 66.1%)。

Table 10: EViT on ViT-Base / EViT 框架对比(ViT-Base)

Keep RatePruning Layers[CLS]OursRandomGFLOPs
--84.584.584.517.6
0.73,6,983.083.081.011.6
0,3,669.779.977.88.8
0.83,6,983.984.182.713.3
0,3,676.782.181.411.3
0.93,6,984.584.583.915.4
0,3,681.883.983.414.3

说明: ViT-Base 上 (0,3,6) 调度的优势更为惊人:keep rate 0.7 时 Col-Ln 79.9% vs [CLS] 69.7%(+10.2%)。

Table 11: Trend Analysis / 趋势分析(ViT-Small)

Keep RatePruning Layers[CLS]OursRandomGFLOPs
--81.481.481.44.6
0.70,3,666.175.073.22.3
1,4,776.377.074.82.5
2,5,877.978.375.92.8
3,6,978.979.077.33.0
4,7,1079.679.878.33.3

说明: 随剪枝调度向深层移动,[CLS] 逐渐改善但 Col-Ln 始终保持最优,证明其在所有层深度的鲁棒性。

Table 12: LVLM Pruning (VisPruner) / VisPruner 框架

MethodrrVQAv2GQAVizWizSQATextVQAPOPEMMEMMBMMB-CNMM-VetAvgRel.%
LLaVA Full57676.762.054.269.558.285.9150564.658.129.763.4100.0
VisPruner25676.060.453.968.457.786.4145463.457.231.162.798.9
Ours25676.060.454.069.157.986.5148363.458.032.163.299.7
VisPruner19275.259.454.668.757.485.7145962.357.230.562.498.4
Ours19275.259.454.668.757.886.0148263.157.933.363.099.4

说明: 在 VisPruner 框架下 r=256r=256 时保持原模型 99.7% 性能,接近无损压缩。

Table 13: EViT Fine-tuning / EViT 微调结果

ModelPruning Layers[CLS]OursRandomGFLOPs
ViT-S0,3,677.279.377.32.3
3,6,981.081.179.43.0

说明: 微调 30 epochs 后,(0,3,6) 调度下 Col-Ln 仍比 [CLS] 高 +2.1%,且 [CLS] 与 Random 持平(77.2 vs 77.3)。

Table 14: EViT on DeiT-Small / DeiT-Small 推理

Keep RatePruning Layers[CLS]OursRandomGFLOPs
--79.879.879.84.6
0.73,6,978.578.577.53.0
0,3,664.775.174.52.3
0.83,6,979.379.378.53.5
0,3,672.177.777.32.9
0.93,6,979.779.779.34.0
0,3,677.479.178.93.7

说明: 在 DeiT-Small 上同样验证了 Col-Ln 的优越性,(0,3,6) keep rate 0.7 时 +10.4% 优于 [CLS]。

Table 15: EViT on DeiT-Base / DeiT-Base 推理

Keep RatePruning Layers[CLS]OursRandomGFLOPs
--82.082.082.04.6
0.73,6,980.680.678.83.0
0,3,671.376.675.72.3
0.83,6,981.381.380.33.5
0,3,677.079.579.02.9
0.93,6,981.781.881.24.0
0,3,680.281.180.83.7

说明: DeiT-Base 上结果一致,[CLS] 在 (0,3,6) 下严重退化。

Table 16: Hyperparameters / 超参数

HyperparameterValue
Backbone Modelvit_small_patch16_augreg
Input Resolution224×224
Total Batch Size2048
OptimizerAdamW
Learning Rate (Base)2×1052 \times 10^{-5}
Min Learning Rate2×1062 \times 10^{-6}
Weight Decay1×1061 \times 10^{-6}
LR SchedulerCosine decay
Training Epochs30
Warmup Epochs0
Pruning Start Epoch0
Keep Rate (rr)0.7
Rescue Ratio (cc)0.8

说明: EViT 微调实验的完整超参数配置。

实验

数据集

数据集规模特点用途
ImageNet-1K1.28M 训练 / 50K 验证1000 类ViT 分类评估
ImageNet OOD 变体5 个变体(I/A/V/R/S)分布外鲁棒性TCA OOD 评估
10 个下游数据集多样化跨域泛化TCA 跨数据集评估
VQAv2, GQA, VizWiz 等 10 个多模态视觉语言理解LVLM 评估

实现细节

Backbone: ViT-S/16, ViT-B/16, ViT-L/16, DeiT-S, DeiT-B, CLIP-ViT-B/16

优化器: AdamW,学习率 2×1052 \times 10^{-5}

Batch Size: 2048

训练轮数: 30 epochs(微调实验)

关键超参: 范数阶数 n=4n=4,rescue ratio c=0.8c=0.8

可视化结果

Figure 1 和 Figure 4 的可视化清晰展示:[CLS] 在早期层注意力分散、噪声大,经常关注背景区域;Col-Ln 从第一层就能精准定位前景目标

Figure 3 的训练损失曲线表明 Col-Ln 从训练初期就保留了更多关键信息

批判性思考

优点

理论优雅: 从 Rényi 熵出发推导出 ℓn-范数,数学上清晰且计算高效

无需训练: 完全无参数,可直接插入任何 ViT/LVLM 架构

双重使用模式: 既可作为独立替代方案(Col-Ln Pruning),也可作为修正机制(Col-Ln Correcting),灵活性强

实验全面: 覆盖 ViT-S/B/L、DeiT-S/B、CLIP-ViT、LLaVA,以及 EViT/TCA/PruMerge/VisPruner 四个框架

在最重要的场景(早期层激进剪枝)中优势最为显著

局限性

仅评估 ImageNet 分类和 VL 任务: 缺少检测/分割等密集预测任务的验证,这些任务对空间信息更敏感

阶参数 nn 的选择: 论文固定 n=4n=4,缺乏对不同 nn 值的系统性消融实验

Rescue ratio cc 固定为 0.8: 不同模型/任务可能需要不同的 cc

深层剪枝优势缩小: 在 (3,6,9) 等深层调度下,Col-Ln 与 [CLS] 差距不大,说明该方法的核心价值主要在早期层

缺少推理延迟测量: 仅报告 GFLOPs,未提供实际 wall-clock time

潜在改进方向

自适应选择不同层的最优阶参数 nn

将 Col-Ln 与 Token Merging 结合,既剪枝又合并

扩展到密集预测任务(如 DETR、SegFormer)

分析 Col-Ln 在不同注意力模式(如 FlashAttention)下的兼容性

可复现性评估

  • 代码开源(GitHub: Wayne0758/SparseAttention)
  • 预训练模型
  • 训练细节完整(Table 16)
  • 数据集可获取

关联笔记

基于

EViT: 核心 [CLS]-based token pruning 框架

Rényi 熵: 理论基础,信息论中的广义熵

对比

TCA: test-time adaptation 框架,验证 OOD 鲁棒性

PruMerge: LVLM token 剪枝+合并框架

VisPruner: VLM 视觉 token 剪枝框架

方法相关

Token Pruning: 核心任务

Multi-Head Self-Attention: Col-Ln 直接作用于注意力矩阵

ViT: 目标架构

硬件/数据相关

ImageNet: 主要评估数据集

速查卡片

Rényi Entropy: A New Token Pruning Metric for Vision Transformers

  • 核心: 用 Rényi 熵推导的列向 ℓn-范数替代不可靠的 [CLS] 注意力来度量 token 重要性
  • 方法: 计算注意力矩阵每列的 ℓn-范数(n=4n=4),选择范数最大的 token 保留
  • 结果: ViT 早期层剪枝优于 [CLS] 高达 +10.2%;LVLM 上保持 98-99% 原模型性能
  • 代码: GitHub

笔记创建时间: 2026-04-01