Token Merging: Your ViT But Faster
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Semantic segmentation | ADE20K (val) | mIoU51.8 | 3089 | |
| Object Detection | COCO 2017 (val) | -- | 2930 | |
| Image Classification | ImageNet-1K 1.0 (val) | Top-1 Accuracy81.69 | 2386 | |
| Object Hallucination Evaluation | POPE | Accuracy86.9 | 2056 | |
| Visual Question Answering | VizWiz | Accuracy82.75 | 1863 | |
| Visual Question Answering | TextVQA | Accuracy80.22 | 1455 | |
| Visual Question Answering | GQA | Accuracy64.49 | 1445 | |
| Visual Question Answering | VQA v2 | Accuracy76 | 1429 | |
| Instance Segmentation | COCO 2017 (val) | APm0.315 | 1304 | |
| Video Object Segmentation | DAVIS 2017 (val) | -- | 1251 |