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

©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]