论文总结

《TOKEN MERGING: YOUR VIT BUT FASTER》 作者为Daniel Bolya等人，来自Georgia Tech和Meta AI。

研究背景：

  • 视觉Transformer（ViTs）在计算机视觉领域发展迅速，但存在运行大规模模型困难的问题。
  • 最近出现了通过在运行时修剪标记来使模型更快的有前途的子领域，但标记修剪存在一些缺点。

主要贡献：

  • 提出Token Merging方法：通过合并相似的标记来增加ViT模型的吞吐量，且该方法在训练和推理时都能使用，训练时能提高训练速度，减少精度下降。
  • 进行广泛实验：在图像、视频和音频上进行实验，证明ToMe在所有情况下都具有竞争力，且能将物体部分合并为一个标记，在视频中能对物体进行跟踪。

相关工作：

  • Efficient Transformers：一些工作试图在自然语言处理和视觉领域创建更高效的transformers，本文专注于通过合并标记来加速现有ViT模型。
  • Token Reduction：最近有一些工作试图从transformers中修剪标记，但这些方法需要训练，且大多是动态的，不利于批处理推理或训练。
  • Combining Tokens：很少有工作合并标记，且之前的方法在速度 - 精度权衡方面不太理想。

方法介绍：

Token_Merging

  • 策略：在transformer的每个块中，逐渐合并标记以减少每层的标记数量r，且与图像内容无关。
  • Token Similarity：使用QKV自注意力中的键（K）来确定标记的相似性，通过点积相似性度量（如余弦相似性）来判断包含相似信息的标记。
  • Bipartite Soft Matching：提出一种更高效的匹配算法，通过将标记划分为两个大致相等的集合A和B，进行匹配和合并，该算法避免了迭代，且运行时间可忽略不计。
  • Tracking Token Size：使用比例注意力（proportional attention）来解决合并后标记不再代表一个输入补丁的问题，同时在聚合标记时需要根据标记大小进行加权。
  • Training with Merging：训练时将标记合并视为池化操作进行反向传播，不需要使用梯度技巧，且发现使用训练ViT的相同设置也是最优的。

图像实验：

  • 设计选择：通过消融实验确定了默认的参数设置，包括使用注意力键（K）、余弦相似性、平均加权的标记合并方式、交替分配标记的分区方式以及使用比例注意力（除了现成的MAE模型）。
  • 模型扫描：将标记合并方法应用于11个最先进的ViT模型，发现常数合并计划能使吞吐量提高2倍，且大型模型的精度下降较小。
  • 与其他工作比较：与其他最先进的模型和标记修剪方法进行比较，发现ToMe能提高ViT模型的吞吐量，使其在速度上与低层级模型相当，且精度下降较小。

视频实验：

  • 框架与结果：将标记合并方法应用于视频分类任务，使用Spatio - temporal MAE在Kinetics - 400上进行实验，结果表明ToMe能提高视频的吞吐量和精度，且在不同的合并计划下表现良好。
  • 可视化：通过可视化展示了ToMe在视频中对物体的跟踪能力，能够将同一物体或部分在多个帧中合并为一个标记。

音频实验：

  • 实验设置与结果：在音频MAE上进行实验，使用Huang等人的ViT - B模型在AudioSet - 2M上进行评估，结果表明ToMe能使基线的吞吐量翻倍，且mAP下降仅0.4%。

结论：

Token Merging（ToMe）能自然地利用输入的冗余性，在图像、视频和音频等领域获得了与最先进技术相竞争的速度和精度，可以被视为一种“自然”的分层模型，有望用于创建更好、更高效的transformers。