TOKEN MERGING: YOUR VIT BUT FASTER

#Token Merging #Vision Transformers #Efficiency #ICLR 2023

摘要： 本文介绍了Token Merging (ToMe)，一种简单的方法，可以在不需要训练的情况下提高现有ViT模型的吞吐量。ToMe通过一种通用且轻量级的匹配算法逐渐结合相似的tokens，这种方法的速度与剪枝一样快，但更准确。现成的ToMe可以将ViT-L @ 512和ViT-H @ 518模型在图像上的吞吐量提高2倍，将ViT-L在视频上的吞吐量提高2.2倍，每种情况下的准确率仅下降0.2-0.3%。ToMe也可以在训练期间轻松应用，实际上将MAE在视频上的微调训练速度提高了2倍。通过训练使用ToMe，进一步最小化准确率下降，使得ViT-B在音频上的吞吐量提高2倍，仅损失0.4%的mAP。从定性上看，我们发现ToMe将对象部分合并为一个token，甚至在视频的多个帧上也是如此。总的来说，ToMe在图像、视频和音频方面的准确性和速度与最先进的技术竞争。

主要方法/架构：

• Token Merging (ToMe)：在Transformer中逐渐合并相似的tokens，使用自定义的匹配算法。
• Bipartite Soft Matching：一种高效的匹配算法，用于决定哪些tokens应该合并。
• Proportional Attention：在合并tokens后，使用比例注意力来调整softmax注意力，以考虑tokens现在可能代表多个输入patch。

实验对比数据结果：

• 图像实验：在ImageNet-1k上，使用ViT模型，ToMe在不训练的情况下提高了吞吐量，同时保持了与基线相近的准确率。
• 视频实验：在Kinetics-400视频分类任务上，ToMe在不训练的情况下提高了吞吐量，并且在训练期间进一步提高了吞吐量，同时减少了训练时间。
• 音频实验：在AudioSet-2M音频分类任务上，ToMe在不训练的情况下提高了吞吐量，并且训练后的模型在准确率上仅略有下降。

图示：

• 图1展示了Token Merging的流程，包括在Transformer块之间合并相似的patches。
• 图2展示了不同匹配算法的比较，ToMe的双分图软匹配算法几乎与随机剪枝tokens一样快，同时保持了高准确率。
• 图3展示了在ImageNet-1k上应用ToMe的模型范围，包括不同的ViT模型，以及与基线的比较。
• 图4和图6展示了ToMe在图像和视频上的可视化结果，展示了如何将对象的部分合并为一个token。

结论： ToMe通过逐渐合并tokens来提高ViT模型的吞吐量，同时保持了与最先进的技术竞争的准确性。ToMe可以自然地利用输入中的冗余，适用于任何具有冗余的模态。在图像、视频和音频上的广泛实验表明，ToMe在每个领域都能达到与最先进的技术竞争的速度和准确性。ToMe可以被视为一种“自然”的层次模型，类似于Swin或MViT，但使用的是纯Transformer块。ToMe可以与这些方法结合，创造出一种全新的架构。此外，ToMe在大型模型上表现良好，可以减少训练时间和内存使用，意味着它可以成为训练巨大模型的核心组件。未来的工作将探索这些可能性，并希望ToMe能够引领更好、更高效的Transformer的创造。