ICCV 2023 | DiffRate：首个可微分Token压缩，性能SOTA

Token压缩，指的是在模型推理过程中减少token数量，其是加速Vision Transformers (ViTs)的重要手段。目前，Token压缩主要分为两种：token剪枝(如 EViT [1])和token合并(如 ToMe [2])。但现有的工作，无论是token剪枝还是token压缩，都需要人工去定义token压缩率，为不同的FLOPs定义不同的token压缩率不仅繁琐，且通常会导致次优结果。

为了解决这个问题，我们提出可微压缩率(Differentiable Compression Rate, DiffRate)，论文DiffRate : Differentiable Compression Rate for Efficient Vision Transformers已被ICCV 2023会议接收，代码、预训练模型均已在GitHub上开源。

论文链接：https://arxiv.org/abs/2305.17997

代码链接：https://github.com/OpenGVLab/DiffRate

一、DiffRate的优势

Diffrate具有如下优势：

• 此前工作将压缩率视为不可微分的超参，而DiffRate可以损失函数的梯度传递到压缩率上,以此实现可微优化。
• 基于DiffRate, token剪枝和token合并可以被很自然的结合，DiffRate可以自适应的去学习每一层的token剪枝压缩率和token合并压缩率。
• DiffRate极大改善了token压缩的性能。例如，将DiffRate学习到的token压缩率直接应用到预训练ViT-H (MAE)上，无需fine-tuning，可以将模型FLOPs压缩40%，吞吐量提升1.5x，且在ImageNet上性能损失仅有0.16%，这甚至优于此前token压缩方法fine-tuning后的性能。

二、 方法详述

Token压缩可以发生在每一个Transformer Block中。一个Transformer Block中包含有两个模块，Attention和MLP。由于token交互只发生在Attention模块，为了节约更多计算量，通常选择在Attention与MLP模块中间进行token压缩，其基本形式可以表示为下图:

但是，上式中剪枝率显然是不可微的，即梯度无法从loss传递到剪枝率，为了解决这个问题，我们提出了一项新颖的技术来通过可微搜索确定最优压缩率。

2.1 可微的离散代理(Differentiable Discrete Proxy)

为了使压缩率可微，我们的核心思想是对所有图片都保留Top-K重要的token，这既可以保证适配并行计算，也能最大限度的保留ViT模型的性能。为了实现此目的，我们提出了可微的离散代理( Differentiable Discrete Proxy, DDP)。其主要包含两个模块：根据一个metric进行token排序，构建了一个排序-剪枝-合并的token压缩流程；根据一个重参数技巧使得token压缩操作中的 Top-K 选择操作对K可微。

2.1.1 Token排序 (Token Sorting)

为了对token进行排序，需要一个重要性指标，而在在此前的工作中已经被充分研究。为此，我们直接选择EViT [1]提出的class attention作为重要性指标：

2.1.2 压缩率重参 (Compression Rate Re-parameterization)

通过以上attention masking的操作，我们实现了训练时保证梯度链的完整性，且计算结果和测试时直接丢弃token等价。

2.2 训练目标 (Training Objective)

我们通过最小化以下loss来求解最优的token剪枝和合并压缩率。

三、实验结果

压缩后若是无需fine-tuning，显然具有更好的易用性。表1展示了进行token压缩后不fine-tuning的实验结果，可以看出，相比于EViT和ToMe, DiffRate显著提高了性能。并且，在ViT-S(DeiT), ViT-B (DeiT), ViT-L (MAE), ViT-H (MAE)上，DiffRate压缩30%-40%左右的FLOPs开销，在ImageNet上性能损失皆小于0.4%。尤其是在ViT-H (MAE)上，40%的FLOPs降低仅带来了0.16%的性能损失。

表2展示了执行完token压缩并fine-tuning 30 epoch的结果。可以观察到fine-tuning 30 epochs后性能有进一步提升。指的注意的是，DiffRate不fine-tuning的情况下，性能就达到或者优于现有token压缩方法fine-tuning后的结果。

图5展示了不同token压缩设置下DiffRate搜索到的压缩率与人工定义的压缩率(EViT, ToMe)，以及10000个随机采样的压缩率的对比，可以看出，DiffRate搜索到的压缩率基本接近最优曲线。

表4对一些设计模块进行了消融实验。

(a)对重要性指标进行消融，

(b)对token压缩操作进行消融，

(c)对搜索最优压缩率使用的样本数量进行消融，

(d)对训练时长进行消融。

值得注意的是，如(c)所示，仅使用1000个样本或者2.7小时训练时间，可微压缩率即可收敛到较好的结果。表明提出的方法具备数据高效性和训练高效性。

图6展示了DiffRate token压缩过程的可视化。表示DiffRate可以成功丢弃语义无关的背景token，并融合相似的前景token，以消除冗余信息。

四、结语

该文对token压缩进行研究，提出了一个统一的token剪枝与token合并流程，并提出可微压缩率以根据计算量约束自适应决定网络每一层的token剪枝和token合并压缩率。实验结果证明了提出方法在性能上的优势，尤其是可以在不fine-tuning的情况下，优于此前需要fine-tuning的方法。

Reference

[1] Not all patches are what you need: Expediting vision transformers via token reorganizations.

[2] Token merging: Your ViT but faster.

[3] DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification

作者：陈锰钊 邵文琪

来源：公众号【OpenGVLab】