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

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后的性能。

图1 和此前token压缩方法的对比

方法详述

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

图2 Token压缩示意图

在此前的工作中，Token Compression可以是token剪枝或者token合并，通过人工设定token压缩率来决定多少token被剪枝/合并。如图1中所示，这种孤立的考虑token剪枝和token合并，且人工设定压缩率的方式，通常导致次优解。为了解决此问题，我们提出DiffRate，一个统一的token压缩方法，可以根据给定FLOPs约束 搜索最优的token剪枝压缩率和token合并压缩率以最小化任务损失：

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

可微的离散代理

(Differentiable Discrete Proxy)

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

Token排序

(Token Sorting)

图3 Token排序

为了对token进行排序，需要一个重要性指标，而在在此前的工作中已经被充分研究。为此，我们直接选择EViT [1]提出的class attention作为重要性指标： 其中， 表示每一个image token对class token的贡献程度，越高的值表明image token对最终输出贡献越大，即越重要。如图3所示，我们首先根据 对N个token进行排序。紧接着，遵循此前的token剪枝工作，为了移除语义无关的冗余token, 我们依照剪枝压缩率移除最不重要的 个token。剪枝之后，我们进一步进行token合并, 进一步选出 个不重要的token，依据cosine相似度，将他们融合到最相似的保留token中，融合操作直接采用求平均。基于此排序-剪枝-合并的流程，token剪枝率和token压缩率可以通过DiffRate被最优得决定。因此，DiffRate可以无缝合并token剪枝和token压缩。

压缩率重参

(Compression Rate Re-parameterization)

图4 压缩率重参

如图4所示，为了使得压缩率 可微，我们将其重参成N个离散的候选压缩率， 。其中， ， 表示最不重要的k-1个token需要被移除。每一个候选压缩率对应一个可学习概率 ，其中 。基于此，最终压缩率 可以表示成：

通过利用离散的候选压缩率，压缩率的优化问题转化为学习概率 。为了将可学习概率 融入计算图，我们首先根据候选压缩率与其对应的可学习概率求出每一个token被压缩的概率： 其中， 表示至少一个token会被保留。并且，从上式可以看出， 。即表明DiffRate对齐了不重要的token要有更大概率被压缩的事实。进一步的，我们将此被压缩概率转换成一个0-1 mask：

其中 表示第k个token被保留，反之亦然。在每一个transformer block中， 我们引入两个独立的重参数模块，以此同时学习剪枝和合并压缩率。因此，每一个block中有两个mask。针对每个token而言，其同时具备剪枝mask 和合并mask 。注意一旦一个token在前一个block中被移除，其在当前block中也应当被移除，因此，最终的mask定义为：

为了保证梯度链的传播，我们需要将此求得的mask融入计算图中。遵循DynamicViT, 在训练过程中我们使用attention masking来代替token丢弃。为了实现此目的，我们为self-attention operation 构建了一个和原始attention map大小一致的mask ：

这个mask的本质作用是防止被压缩token和其余token的交互，除了其本身，紧接着，我们利用此mask来修改self-attention操作中的SoftMax计算：

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

训练目标

(Training Objective)

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

其中 是任务约束， 将压缩后的FLOPs约束到T。在反向传播过程中，基于直通估计(straight-through-estimator, STE), 对于 的梯度可以计算如下：

实验结果

压缩后若是无需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%的性能损失。

表1 token压缩后不fine-tuning的实验结果 

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

表2 token压缩且fine-tuning 30epoch的实验结果

表示fine-tuning 30 epoch

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

图5 Token压缩率对比

表4对一些设计模块进行了消融实验。(a)对重要性指标进行消融，(b)对token压缩操作进行消融，（c）对搜索最优压缩率使用的样本数量进行消融，（d）对训练时长进行消融。值得注意的是，如(c)所示，仅使用1000个样本或者2.7小时训练时间，可微压缩率即可收敛到较好的结果。表明提出的方法具备数据高效性和训练高效性。

表4 DeiT-B上的消融实验

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

图6 可视化

结语

该文对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

Paper：

https://arxiv.org/abs/2305.17997

Code：（点击文末“阅读原文”直达开源链接）

https://github.com/OpenGVLab/DiffRate

往期推荐

Unmasked Teacher 非掩码教师: 高效训练视频基模型，代码模型已开源 | ICCV2023

Meta-Transformer：1个框架理解12种模态引发的质变与涌现（已开源）

VideoChat 成为央视听媒体大模型关键支撑技术

点击下方名片，获取通用视觉团队更多研究动态