ICLR 2023 Oral | ToMe：无需训练的Token融合模型加速算法

科技

2023-05-06 15:05

©Paperweekly 原创 · 作者 | An.

单位 | 中科院自动化所

研究方向 | 计算机视觉、模型压缩

论文标题：

Token Merging: Your ViT But Faster

论文链接：

https://arxiv.org/pdf/2210.09461.pdf

代码链接：

https://github.com/facebookresearch/ToMe

动机&背景

过往的多篇工作 [1-3] 表明在 Transformer 结构中添加特定于视觉的归纳偏置能够以较少的计算量获得更好的性能表现。过往的 token 剪枝方法存在几大不足：a）修剪导致的信息丢失会损害模型性能；b）现有方法都需要重新训练模型才能保持较高的性能；c）大多现有方法不能加快训练速度；d）不同输入的 token 修剪数量不一致，导致批处理推理不可行。

与此同时，原始的 ViT 结构有着许多理想的特性：a）内部简单的矩阵乘法更适合并行计算；b）支持强大的自监督预训练技术（例如 MAE [4]）；c）没有过多的先验假设，很容易扩展到多模态任务；d）易于规模缩放，能够很好地推广到大模型。

因此，本文提出了一种无需训练的、只降低 FLOPs 的类 token 剪枝方法 ToMe（Token Merging），该方法在原始 ViT 结构，利用一种通用且快速的匹配算法渐进式地将 ViT 结构的 tokens 进行融合，以克服过往 token 剪枝方法的不足。

该算法在只损失 0.2~0.3% 精度的情况下，在最先进的 ViT-L@512 和 ViT-H@518 可实现两倍推理速度的提升。不同于过往的方法，ToMe 也可以在训练阶段获得速度收益。该方法在图像、视频和音频任务重均获得了更好的准确率和速度综合表现。

方法：Token Merging

如第一小节分析的那样，本文的目标是设计一个 token 融合模块，将其插入现有的 ViT [5] 模型中。通过合并冗余的 tokens，无需训练即可获得推理速度的提高。Token 剪枝方法的基本逻辑都是类似的，即利用某种方式找到相似的 tokens，利用某一手段去除 tokens 的冗余，本文提出的 ToMe 也不例外，图 1 是 ToMe 方法的示意图（关键模块的细节见图 1(c)），下面我将从几点差异来解析本文的方法细节。

▲ 图1. ToMe 方法的示意图

2.1 策略

1. 过往的 token 剪枝方法大多以「比例」作为逐层 token 剪枝的超参数，本文则选择了「数量」，在具有个块的模型中，ToMe 将渐进地融合个 tokens；

2. 过往的一些 token 剪枝方法采用「动态」的策略，不同的图像会剪去/融合不同数量的 tokens，这固然能带来模型准确率的提升，但也限制了模型批量推理和训练。因此本文将对不同图片输入采用相同数量的 tokens 融合策略，以确保模型在训练/推理时都可以批量处理；

3. 如图 1 所示，本文不同于过往的方法，将 token 融合模型应用于 MHA 和 MLP 之间，这一位置有助于被融合的 tokens 更好地传递信息，并使得 ToMe 可以利用注意力中的特征来决定哪些 tokens 需要融合，这两点都在后续的消融实验中得以验证。

2.2 Token 相似度

在进行相似 tokens 的融合时，我们先对“相似”下一个定义。过往方法 [6] 认为如果两个 token 的特征之间的距离很小，则意味着它们很相似。然而，现有的大规模 ViT 模型的中间特征可能都存在过参数化的问题。例如，ViT-B/16 的中间特征的通道数为 768，而每个 token 对应的原始图像的像素值也只有 31616=768 个，这意味着中间特征可能包含不显著的早上，这会导致常用的距离相似度度量受到干扰。

因此本文选择了自注意力机制中 QKV 的 K 的特征来判断相似度，利用每个 token 的 K 的惦记相似度度量（例如余弦相似度）来确定哪些包含相似信息。

2.3 二分软匹配

在获得各个 tokens 的相似度后，我们需要一种快速的方法来确定要融合哪些 tokens。过往的解决方案集中在 kmeans 聚类或图分割。但这些方法的运行效率都太低了。本文基于两个设计目标提出了一种更高效的解决方案，一是避免任何无法并行化的迭代，二是希望融合是渐进的。本文选择了匹配而非聚类，因为聚类会导致同一组的 tokens 数量不受限制，这可能对网络产生不利影响。具体算法如下（可见图 1(c)）：

将 tokens 划分为大小大致相等的 AB 两组；
从 A 中的每个 token 中选择一条与 B 中 token 相似度最大的边（二分匹配问题）；
保留整体最相似的个边；
找到连通分量并将其融合（本文默认计算特征的加权平均值）；
将两组合并在一起。

这一算法无需计算全部 tokens 间的相似性，连通分量的合并成本也可以忽略不计。后续实验表明，采用合适的 AB 分组策略即可获得几乎和随机丢弃 tokens 一样快的速度。

2.4 Token 数量跟踪

由于自注意力模块 softmax 的计算公式中涉及归一化因子，当 token 被融合时，该因子也应同步更新，所以在融合的过程中需要实时跟进 token 数量的变化。本文提出了一种既快速又有效的注意力公式，以解决这一问题：

该注意力机制称为比例注意力，公式中的是一个包含每个融合后 tokens 的大小的行向量。

2.5 在训练中应用 ToMe

从上文的介绍可知，ToMe 类似于一种池化操作，因此我们无需任何特殊的梯度技巧就可以将其应用于训练阶段。作者的实验表明，在无需修改原始训练参数的情况下，ToMe 能够有效提高训练速度并获得性能收益。

实验

1. ImageNet-1k 性能对比：如表 1 所示，ToMe 在极少的精度损失下显著提升了推理速度（约 100%），ToMe 能让 ViT-L 在速度上与较低规模的模型相当。

▲ 表1. 仅在 ImageNet-1k 上训练的模型性能对比（*表示 SimMIM 自监督预训练，†表示自主训练 baseline，‡的测速数据源于原始论文[V100]）

2. 剪枝方法对比：表 2 对比了不同加速剪枝方法的性能，ToMe 在精度和推理速度上都达到了 SOTA，同时，ToMe 还可直接应用在训练过程中，使得训练速度提高 1.5 倍，这是过往所有方法都无法实现的。

▲ 表2. 各个剪枝在 ViT-S 模型上的性能对比（灰色指在训练中应用 ToMe，蓝色指训练后的 ToMe）

3. 设计选择：表 3 分析了方法部分提到的多种设计改进的有效性，包括特征选择、距离函数、头聚合方式、融合方法、分区方式和比例注意力。实验表明，本文所提出多个设计在精度和速度的综合考量下取得了最优的效果。

▲ 表3. ToMe 消融实验（ViT-L/16，MAE，ImageNet-1k，V100 GPU，紫色为最终的默认选择）

4. 匹配算法&融合时间表对比分析：从图 2 左图中可知，融合优于修剪，匹配优于聚类，二分图匹配实现了更好的速度和精度的平衡。从右图可知，在不同时间表中，简单的 Token 常数降低就可以取得较优的结果。

▲ 图2. 左图是匹配算法的对比分析，右图是融合时间表的对比分析

5. 可视化分析：图 3 的可视化结果显示出了一个惊奇的现象——简单的 ToMe 就可以将相同的实例对应的 Token 合并在一起，无论是前景还是背景。

▲ 图3. ViT-H MAE ImageNet-1k 在训练过程中应用 ToMe 的融合结果

参考文献

[1] Xiaoyi Dong, Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, and Baining Guo. Cswin transformer: A general vision transformer backbone with cross-shaped windows. In CVPR, 2022.

[2] Haoqi Fan, Bo Xiong, Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. In ICCV, 2021.

[3] Benjamin Graham, Alaaeldin El-Nouby, Hugo Touvron, Pierre Stock, Armand Joulin, Herv ́e J ́egou, and Matthijs Douze. Levit: a vision transformer in convnet's clothing for faster inference. In ICCV, 2021.

[4] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Doll ́ar, and Ross Girshick. Masked autoencoders are scalable vision learners. In CVPR, 2022.

[5] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020.

[6] Dmitrii Marin, Jen-Hao Rick Chang, Anurag Ranjan, Anish Prabhu, Mohammad Rastegari, and Oncel Tuzel. Token pooling in vision transformers. arXiv:2110.03860 [cs.CV], 2021.

更多阅读