Token系列：Joint Token Pruning and Squeezing Towards More Aggressive Compression of Vision Transformers

写在前面

  新的一周又开始了，冲冲冲~
  这篇论文和之前面一篇博客（Vit 中的 Token 改进版本：Token Mreging: Your Vit But Faster 论文阅读笔记）中对 Transformer 的 Token 改造有异曲同工之妙，接下来一起看看吧~

• 论文地址：Joint Token Pruning and Squeezing Towards More Aggressive Compression of Vision Transformers
• 代码地址：https://github.com/megvii-research/TPS-CVPR2023
• 收录于：CVPR 2023
• 2023 每周一篇博文，主页更多干货内容持续输出，欢迎关注~

一、Abstract

  ViTs 模型受限于计算成本而很难部署到实际应用上，之前的方法通过裁剪 tokens 来实现精度和成本之间的权衡，但裁剪策略一定会导致信息丢失。本文进行的量化实验揭示了裁剪 tokens 对性能的影响是显而易见的。于是提出 Token Pruning & Squeezing module (TPS) 模块来更高效率地压缩 ViTs。具体来说，首先采用裁剪操作，得到需要保留和裁剪的 tokens 子集，其次，通过无向的最近邻匹配和基于相似度的融合步骤，压缩裁剪的 tokens 信息到保留住下来的 tokens 内。实验效果很好，并且本文提出的方法对于 token 裁剪过程中的错误也有很高的鲁棒性。

二、引言

  Transformers 模型仍然需要相当高的计算预算，主要在于计算和内存成本随 tokens 长度呈平方关系。一些方法提出裁剪 tokens 策略来平衡成本与性能，此外类似知识蒸馏等技术则进一步降低了性能代沟。但随着裁剪比例增加，性能下降无可避免，因为丢掉了一些关键信息和辅助的上下文，特别是当保留下来的 tokens 数量接近 10 时。此外，一些激进的裁剪策略可能会导致目标不完整或者背景上下文信息丢失，如下图所示：

  EViT、Evo-ViT 提出将这些裁剪掉的 tokens 聚合为一个 token，但忽略了这些被裁剪 tokens 的不一致性，从而导致特征的坍缩或者更加激进的 token 裁剪。
  本文认为被裁剪的 tokens 信息值得更好的利用，于是本文进行了一个实验，结果如下图所示：

  相比于 dynamicViT，保留策略能够在原始结构上带来额外的性能补充，而且随着裁剪程度加深，性能也更好。
  为留存下裁剪 tokens 的信息，本文提出 Joint Token Pruning & Squeezing (TPS) 模块来缓解激进的 ViTs 压缩方法。TPS 模块利用特征分离机制将裁剪的 tokens 压缩至保留下来的 tokens 内，如下图 c 所示：

  首先，基于 tokens 的打分结果，TPS 模块将输入的 tokens 划分为两个互补的子集：保留集、裁剪集。之后，利用一种无向的最近邻匹配算法来分离每个 token 到相邻的保留 token 内，而这些 token 称之为主 token，同时那些未被选择的 token 保留独一性。当固定住保留 tokens 的基数时，便能在推理时执行并行计算。此外引入两个弹性的变体，内部 block：dTPS，外部 block：eTPS，本质是两个能即插即用在传统 ViTs 和混合 ViTs 的模块。
  在 ImageNet1K、iNaturalist 2019 数据集上进行的实验表明了模型的效率、弹性、鲁棒性：首先根据不同的 token 裁剪、重组织策略表明 TPS 的性能，之后与 SOTA 的 Transformer 进行性能比较，最后将 TPS 整合进不同的 ViTs，此外还评估随机 token 选择策略下的 TPS 鲁棒性。总体贡献如下：

• 提出 TPS 及其变体 dTPS、eTPS，保留丢弃的 tokens 信息，并促进更好的 ViTs 压缩算法；
• 大量实验表明有效；
• 大量实验表明 TPS 的弹性，相比于 token 裁剪和重组更加鲁棒。

三、相关工作

原始 ViTs

  DeiT、LV-ViT、T2T-ViT、PS-ViT。

混合 ViTs

  Swin Transformer、CVT、PVT。

Token裁剪

  虽有 DynamicViT、AdaViT、A-ViT、ATS，但主流的深度学习框架并不太能支持这些推理时会发生 token 长度变化的模型。且最大的缺点在于裁剪的信息损失会使得精度下降，并且限制更激进的 token 裁剪。于是 Evo-ViT、EViT、SPViT 将裁剪的 tokens 塌缩进一个 token 来重新组织数量，这一操作记为 token 重组。尽管 token 重组缓解了一部分信息损失，但还是导致性能下降。SPViT、IA-RED2 在训练计划上下功夫，但引入了太多的超参数且优化困难。而本文引入一种新的方法 joint token Pruning & Squeezing，仅需要简单地微调预训练模型。

四、方法

4.1 动机

  实验验证被丢弃的 tokens 中蕴含有用的信息。

  从上图中可以看到，随着裁剪变得激进，性能下降更加明显。但当交换裁剪的 token 和保留的 token 时，一些被裁剪的 token 能够解决部分错误样本。此外，随着裁剪密度增加，附加的性能增益也在上升。这表明随着裁剪比例增加，专有信息越来越重要。于是促使本文将裁剪的 tokens 同化到保留的 tokens 上，从而减少信息的丢失。

4.2 Token 裁剪

   本文引入两种变体：dTPS、eTPS，使得 TPS 能够和任意的裁剪技术相融合。

  如上图所示，dTPS 采用 dynamicViT 中可学习的 token 得分预测头，并通过直通式 Gumbel-Softmax 对二值化掩码进行采样来达到可微分的目的，而 eTPS 采用 EViT 中的类别 token 注意力值来衡量 token 的重要性。在推理阶段，采用固定比例 $\rho$ 的 Top-K 操作来实施 token 选择策略，同时确保 token 的尺寸不变以完成并行推理。用 $S^r$ 表示保留下来的 tokens，用 $S^p$ 表示将被裁剪的 tokens。

4.3 Token 压缩

  为避免生成额外的 token，将裁剪的 tokens 注入相似的保留 tokens 内。方法：应用一种无向的最近邻匹配算法：以多对一的方式将 $S^p$ 映射到 $S^r$，之后利用一种基于相似度的融合方式来同化这些新信息。

匹配

  给定两个 token 子集：$S^p$、$S^r$，相应的 tokens 索引分别为 $I^r$、$I^p$。用 $c_{i,j}$ 表示所有 $i\in I^p$，$j\in I^r$ 的相似度矩阵。对于每一个裁剪的 token $x_i\in S^p$，从 $S^r$ 中找到其最近邻 token $x_{\ast }^{host}\in S^r$ 作为主 token：
$$x_{\ast }^{\text{host }}=\underset{{\mathbf{x}}_j\in S^r}{\mathrm{argmax}}{c}_{i,j}$$注意 token 匹配步骤是无向的，即 $S^p$ 和 $S^r$ 之间无向匹配，同时多个裁剪的 token 也能共享主 token，每个保留下来的 token 也能作为主 token。之后在矩阵 $M\in {\mathbb{R}}^{N^p\times N^r}$ 中记录每个匹配值：
$$m_{i,j}=\{\begin{array}{ll}1,& x_j\text{ is the host token of }{x}_i\\ 0,& \text{ otherwise }\end{array}$$其中 $N^p$ 和 $N^r$ 分别表示两个子集的 tokens 数量。
  一种高效计算相似度矩阵的方式是采用 $S^p$ 和 $S^r$ 之间的余弦相似度：
$$c_{i,j}=\frac{{\mathbf{x}}_i^T{\mathbf{x}}_j}{\Vert x_i\Vert \Vert x_j\Vert },\text{ for }i\in I^p,j\in I^r$$在匹配步骤中，相似度矩阵 $c_{i,j}$ 可以直接生成，因此不需要额外的参数。

融合

  EViT 采用 tokens 重要性得分来对聚合的 tokens 重新加权，而本文采用基于相似度的加权计划。通过 mask $M$ 来控制两个子集中所有 tokens 的融合步骤，来确保仅有主 tokens 和裁剪的 tokens 混合。具体来说，通过融合原始特征和裁剪的 tokens 特征来更新保留的 tokens $x_j$：
$$y_j=w_j{x}_j+\sum _{x_i\in S^p}{w}_i{x}_i$$其中 $w_i$ 为每个被裁减 token $x_i\in S^p$ 的权重，$w_j$ 为保留的 token 权重，$y_j$ 为更新的 token 权重。融合的权重 $w_i$ 取决于 mask 值 $m_{i,j}$ 和相似度矩阵 $c_{i,j}$：
$$w_i=\frac{\exp (c_{i,j})m_{i,j}}{{\sum }_{x_i\in S^p}\exp (c_{i,j})m_{i,j}+\mathrm{e}}$$保留下来的 token 总是有着最大的融合权重 $w_j$，因为 $x_j$ 与其自身的相似度才等于 $1$：
$$w_j=\frac{e}{{\sum }_{x_i\in S^p}\exp (c_{i,j})m_{i,j}+e}$$于是未成为主 token 的那部分的保留 token 保持不变，而那些裁剪的 tokens 将被压缩进入主 tokens 来代替原有的 tokens。
  由于匹配和融合策略确保了被处理的 tokens 数量等于保留下来的 tokens 数量，于是推理过程能够维持并行处理。

4.4 TPS 嵌入到混合 ViTs 中

  对于常规扁平的 Transformer block，TPS 能够轻松嵌入其中来减少 token 的数量。而当 block 包含一些涉及到空间结构的操作，例如卷积、池化等，则 TPS 需要做轻微调整。例如在 PVT 模型中，将 TPS 嵌入每阶段的第一个 block 中，来生成决策 $D$。对于注意力层，则减少输入 token 的维度以及后续的 query $Q$。如果在其内部应用空间压缩层，则丢弃的 token 用 $0$ 来代替从而维持其空间结构。

五、实验

数据集和评估指标

  Baseline：DynamicViT、EViT；指标：Top-1 精度、GFLOPs；输入尺寸：$224\times 224$；数据集：ImageNet1K、iNaturalist 2019。

实验细节

  采用 DeiT 中相同的数据增强方式。AdamW 优化器，余弦学习率计划。训练设置与损失函数同 Baseline，除基础学习率设为 $\frac{batchsize}{1024}\times 2.5\times 10^{-4}$，原因是原始的设置不稳定，特别是当采用激进的裁剪策略时。

5.1 主要结果

与 Baseline 的比较

  

视觉上的比较

  

与 SOTA 的比较

  

在更多 Backbones 上的拓展

  

细粒度的视觉分类

  

5.2 消融实验

训练的 Epochs

  

特征类型

  

相似度矩阵

  

5.3 鲁棒性实验

  

六、结论和限制

  本文提出 Token Pruning & Squeezing (TPS) 来更加激进地压缩 ViTs。相比于 token 裁剪和重组，TPS 能够避免更大的性能下降。实验结果表明了方法的有效性、弹性、鲁棒性。
  仍有一些限制：首先一些混合 ViTs 中的结构化空间操作限制了 TPS 直接整合到这些模型中；其次微调预训练模型的过程可能会被更加先进的训练方法替代，从而减少总体训练时间。

补充材料

可视化、两种变体的细节；前向过程；Block 的位置；参数；性能；

TPS 嵌入混合 ViTs；PVT 训练、推理；CvT 训练、推理；详细的实验设置；

ImageNet-1K 分类；TPS 嵌入 DeiT；TPS 嵌入 LV-ViT；TPS 嵌入 PS-ViT；

TPS 嵌入 PVT；TPS 嵌入 CvT；iNaturalist 2019 分类；TPS 嵌入 DeiT；

TPS 采用更大的模型和更大的输入尺寸；TPS 在不同保存比率下的表现；

关于 TPS 设计中的更多消融实验；

写在后面

  论文后面附上了大量的补充材料，足以见作者们花了大量的时间完成这项工作。这篇文章读来通俗易懂，创新点也是值得称赞，最后的补充材料完美弥补了可能认为的工作量不足问题，是篇上乘顶会文章。