Adaptive MLP Pruning for Large Vision Transformers

Adaptive MLP Pruning for Large Vision Transformers (AMP)

一句话总结

提出 label-free 信息熵准则替代 one-hot cross entropy 做 Taylor pruning 重要性评估，结合二分搜索自适应确定每层 MLP 压缩比，在 CLIP 和 DINOv2 等大 ViT 上实现约 40% 参数/FLOPs 削减且 zero-shot 性能近无损。

核心问题

大型 ViT 中 MLP 模块占参数量 ~81%（如 EVA-CLIP-E），是压缩的核心靶点。传统 Taylor 剪枝用 one-hot CE loss 计算重要性，忽略非目标类别的预测信息 → 重要性评分不准确。

关键公式

信息熵准则（替代 one-hot CE）

实例间相似度：

$s_{ij} = \frac{z_i^{\text{cls}} \cdot z_j^{\text{cls}}}{\|z_i^{\text{cls}}\| \cdot \|z_j^{\text{cls}}\|}$

预测概率（softmax over batch）：

$p_{ij} = \frac{\exp(s_{ij}/\tau)}{\sum_{j'=1}^{B}\exp(s_{ij'}/\tau)}$

信息熵：

$\mathcal{E} = -\frac{1}{B}\sum_{i=1}^{B}\sum_{j=1}^{B} p_{ij} \cdot \log p_{ij}$

优势：(1) 无需标签（label-free）→ 适用于 DINOv2 等无发布损失函数的模型；(2) 利用全部预测信息而非仅正确类别。

Taylor 重要性评分

$\mathcal{I}_k = \left|\sum_{n=1}^{N}\hat{h}_k^{(n)} \cdot \nabla_{\hat{h}_k^{(n)}}\mathcal{C}\right|$

其中 $\mathcal{C}$ 使用信息熵 $\mathcal{E}$ 替代传统 CE loss。

自适应二分搜索

搜索范围 $[M_{\min}, M_{\max}] = [0, M_0]$（$M_0$ = 原始隐藏维度）

设定熵增量阈值 $\Delta E$

每次测试 $M_t = (M_{\min} + M_{\max}) / 2$

熵增量 $\leq \Delta E$ → 继续剪；否则回退

每层 6 次迭代，从最后一个 block 开始逐层处理

知识蒸馏恢复

$\mathcal{L}_{\text{distill}} = \frac{1}{C}\|z^{\text{cls}} - \hat{z}^{\text{cls}}\|^2 + \frac{1}{N \cdot C}\|z^{\text{patch}} - \hat{z}^{\text{patch}}\|^2$

剪枝只影响隐藏层维度，输出维度不变 → 可直接做特征对齐蒸馏。

关键图表

Figure 1: AMP 方法概览

左：信息熵 vs one-hot CE 的重要性评估对比。右：自适应二分搜索确定每层压缩比。

Figure 2: 信息熵 vs One-hot CE

One-hot CE 只看目标类别预测；信息熵利用所有预测可能性 → 更准确的重要性评分。

实验结果

Table 1: Zero-Shot 分类（无蒸馏 / 有蒸馏）

蒸馏后完全恢复甚至超越原始性能。

Table 3: 与其他剪枝方法对比（无蒸馏）

AMP 在无蒸馏场景下超越其他方法 42.7个百分点。

消融实验

二分搜索的贡献是压倒性的（+46.5pp）。

对我们工作的启示

1. MLP 是 ViT 的参数大户：与 “Grow, Don’t Overwrite” 的发现一致（MLP 扩展效果最好）
2. 信息熵作为容量指标：可以迁移到模型增长——信息熵低的层可能容量不足，需要增长
3. 自适应压缩比：避免人工调参，二分搜索思路可用于自适应增长决策
4. Label-free 评估：对无监督/自监督模型的结构变化（增长/剪枝）很有价值

局限性

只在 zero-shot 场景验证；fine-tuning 后差距可能缩小

信息熵阈值 $\Delta E$ 的选择仍需手动

未在 LLM 上验证（仅视觉 Transformer）

蒸馏恢复需要 ImageNet-1K 训练集（虽然只占 LAION-2B 的 0.06%）

相关概念

Taylor pruning

卷积滤波器剪枝

CLIP