写在前面

  都快一个月没写博客了，事情太多，又有拖延症，得改啊 (＞﹏＜)

• 论文题目：TOKEN MERGING: YOUR VIT BUT FASTER
• 代码地址：https://github.com/facebookresearch/ToMe
• 预计提交于 ICCV 2023

一、Abstract

  本文引入一种 Token 融合的方法（Token Merging — ToMe），在无需额外训练的情况下增强现有 ViT 的性能。具体来说在 transformer 中使用一个通用且轻量化的匹配算法来逐步融合相似的 tokens。在图像、视频、音频上的性能绝佳。

二、引言

  Transformer 的应用很广，效果也很好，就是训练速度太慢。最近提出的一些方法在修剪 tokens，然而有一些缺点：减少的 tokens 可能会使得信息丢失；可能需要重新训练才能让模型更有效；大部分方法不能用于加快模型的训练；根据输入的内容来修剪 tokens，会使得批量推理存在麻烦。
  本文提出了 ToMe 来结合 tokens，而不是裁减掉这些。同时本文采用自定义的匹配算法，相对于剪枝来说更快且精度更高。此外，本文提出的方法需要或者不需要训练都可，在应用于大模型时只是会造成轻微的精度下降。在训练时如果使用 ToMe，时间会缩短，甚至在某些情况下训练时间减半。应用 ToMe 到图像、视频、音频上时能够达到 Sota 的效果。

三、相关工作

3.1 有效的 Transformer

  改进整个 Transformer 模块，改进注意力机制，剪枝注意力头的特征，注入特定域模块。本文通过融合现有ViT模型中的 tokens 来匹配特定域的模型，从而加快速度，一些时候甚至无需训练。

3.2 Token 的减少

  由于 transformer 能够使用任意数量的 token，因此一些方法在尝试剪枝的工作，然而需要训练。同时大量剪枝采用动态 token 数量的方式，但是不太实用，因为推理时没法用啊。因此为了改为动态方式，很多方法应用一个 mask 来训练而不是直接去除 token，但是这种方法又会让剪枝的作用失效。本文提出的方法可以直接应用于推理和训练，能够实现实时速度。

3.3 Token 的联合

  只有少量的方法剪枝到单个 token，但是他们的精度普遍在无训练时掉点很多。

四、Token 融合

  通过融合冗余的 tokens，不管在训练还是无训练的情况下，能够增加模型的效率。

4.1 策略

  在每一个 Transformer 的每一个块中，每一层来融合 $r$ 个 tokens。因此对于 $L$ 块来说，需要融合 $rL$ 个 tokens。更多的 $rL$ 意味着准确度更低，但是速度收益更大。此外，$rL$ 和图像内容无关。

  如上图所示，在每一个 Transformer 块中的注意力和 MLP 模块间加上 ToMe 模块（这里有个 ToMe 模块放置位置的消融实验）。

4.2 Token 相似性

  之前的方法：衡量特征之间的距离，但这并不是最优的，因为中间的特征可能包含大量不太重要的噪声。而 Transformer 本身的特性通过 QKV 自注意力机制就解决了这个问题，因为 K 使用点乘的方式总结了每一个 token 的信息。因此借鉴这种方法，本文也采用点乘相似度矩阵来决定包含相似性信息的 token。

4.3 双边软匹配

  之前聚类的方法用不了，太慢了。因此需要提出一种更有效的方案，其设计原则如有：避免任何不能进行并行化的迭代操作；逐步地进行融合操作以代替 Mask。另外一个没有用聚类方法的原因：因为聚类的效果太狠了，无脑去融合 tokens 到一个聚类中心，这可能会影响整个网络的性能，而匹配的方法则是使得大部分 tokens 留下。

  算法流程如下：

• 将 tokens 划分为两个大致数量差不多的集合 $\mathbb{A}$、 $\mathbb{B}$
• 在集合 $\mathbb{A}$ 中画一条边连接到 $\mathbb{B}$ 中最相似的那个 token
• 保留下 $r$ 个最相似的边
• 融合那些仍然相连接的边
• 拼接这两个集合

  这个做法相当于建立了一个双边图，而在 $\mathbb{A}$ 中的每个 token 仅有一条边，因此不需要计算每个 token 的相似度。实验表明这种双边软匹配的方法尽可能和随机扔掉 tokens 一样快。

4.4 追踪 Token 的尺寸

  一旦具有相同 key 的 tokens 被融合了，那么融合后的 token 可能会对 softmax 有影响，因此采用一种比例注意力的方式来弥补这一变化：
$$A=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d}}+\log s\right)$$
其中 $s$ 为包含每个 token 尺寸（token 表示的 patch 数量）的行向量。

4.5 采用融合操作的训练

  训练不是必须，但是训练可能会阻止梯度降低以及加快模型的训练。为了训练 ToMe，将 token 的融合视为池化操作并通过反向传播回传梯度。因此 ToMe 是个可以临时增加训练速度的替代方法。

五、图像实验

  在 ImageNet-1K 数据集上，使用 AugReg、MAE、SWAG、DeiT 这四种方式。所有的模型要么采用 ToMe 离线训练，要么在 MAE 和 DeiT 上应用 ToMe 训练。推理在单张 V100上进行。

5.1 设计的选项

  ViT-L/16 模型，24层。

5.2 模型间的比较

5.3 与其他方法的比较

5.4 可视化的结果

六、视频实验

七、音频实验

  上表 7。

八、结论

  本文提出 ToMe，通过逐步融合 tokens 来增加 ViT 模型的速度，在图像、视频、声音模态上效果很好。能够和其他模型结合，可以用在大模型中。

写在后面

  文章方法部分写的简洁凝练，实验做的特别多，工作量巨大。也需要一定的资本投入才能实现。剩下的附录部分等回家再给补上。溜了溜了~
  青青草原，我来了~