[论文解读] Token Merging: Your ViT But Faster

Daniel Bolya, Cheng-Yang Fu+4 more|arXiv (Cornell University)|Oct 17, 2022

Generative Adversarial Networks and Image Synthesis被引用 62

一句话总结

Token Merging (ToMe) 通过使用快速二部匹配过程将相似标记合并，从而加速 Vision Transformers，在图像、视频和音频任务的吞吐量提高到大约 2x（视频为 2.2x），且精度损失很小，且可在有无训练的情况下使用，覆盖图像、视频和音频任务。

ABSTRACT

We introduce Token Merging (ToMe), a simple method to increase the throughput of existing ViT models without needing to train. ToMe gradually combines similar tokens in a transformer using a general and light-weight matching algorithm that is as fast as pruning while being more accurate. Off-the-shelf, ToMe can 2x the throughput of state-of-the-art ViT-L @ 512 and ViT-H @ 518 models on images and 2.2x the throughput of ViT-L on video with only a 0.2-0.3% accuracy drop in each case. ToMe can also easily be applied during training, improving in practice training speed up to 2x for MAE fine-tuning on video. Training with ToMe further minimizes accuracy drop, leading to 2x the throughput of ViT-B on audio for only a 0.4% mAP drop. Qualitatively, we find that ToMe merges object parts into one token, even over multiple frames of video. Overall, ToMe's accuracy and speed are competitive with state-of-the-art on images, video, and audio.

研究动机与目标

• 在不进行大规模再训练或架构修改的前提下，推动更快的 ViT 推理与训练。
• 引入一个简单的 token-merging 模块，通过减少每个 transformer 层的标记数量来提高吞吐量。
• 实现对图像、视频和音频的适用性，同时保持有竞争力的准确性。
• 提供一个轻量级、可并行的合并算法，适用于现成模型。

提出的方法

• 在每个 transformer 块的 attention 与 MLP 分支之间插入 token merging 步骤。
• 使用 attention keys 的点积相似度来定义 token 相似性，以选择要合并的候选项。
• 使用快速的双分性软匹配通过按 token 大小加权平均特征来合并 token。
• 采用比例注意力以考虑合并后的 token 代表多个输入补丁的情况。
• 可选地通过将合并视为池化并在合并的 token 上反向传播来训练 ToMe。
• 探索恒定和递减的合并计划，以控制跨层移除的 token 总量。

实验结果

研究问题

• RQ1 一个轻量级的 token-merge 模块是否可以在不重新训练的情况下显著提高 ViT 吞吐量，同时保持最小的精度损失？
• RQ2 应如何匹配和合并 token 以在图像、视频和音频模态中保持信息内容？
• RQ3 合并计划和特征选择对不同 ViT 模型和预训练方案的加速和准确性有何影响？
• RQ4 ToMe 在推理阶段和训练阶段以及跨多模态场景中是否都有效？

主要发现

• ToMe 可以在图像上将最先进的 ViT-L@512 与 ViT-H@518 的吞吐量大约翻倍，准确率损失约为 0.2–0.3%。
• 在视频上，ToMe 在没有训练的情况下使 ViT-L 的吞吐量约为 2.2x，且精度损失只有 0.2–0.3%。
• 在视频的 MAE 微调中，如用 ToMe 训练，训练速度提升可达约 2x。
• 在音频方面，使用 MAE 预训练的 ViT-B 上，ToMe 约实现 2x 吞吐量提升，mAP 下降约 0.4%。
• ToMe 将跨帧的对象部件合并，并可在视频中执行部件跟踪，同时合并背景与前景信息且不会造成较大精度下降。
• ToMe 在图像、视频和音频领域与最先进的方法相竞争，且不需要额外的训练技巧或新参数。