Learning to Merge Tokens in Vision Transformers

TL;DR

The paper addresses the computational burden of Vision Transformers by introducing PatchMerger, a simple learnable module that merges a variable number of input tokens into a fixed small set of output tokens between encoder blocks. This enables substantial compute savings (40–60% for large models) while maintaining comparable upstream and downstream performance after fine-tuning, across ViT and V-MoE variants. The authors provide extensive experiments on JFT-300M and ImageNet, including up to 10-shot and full finetuning transfers, and show that strategic placement and token-count choices yield favorable Pareto frontiers. They also discuss connections to attention mechanisms and propose future work on adaptive per-input compute for even greater efficiency.

Abstract

Transformers are widely applied to solve natural language understanding and computer vision tasks. While scaling up these architectures leads to improved performance, it often comes at the expense of much higher computational costs. In order for large-scale models to remain practical in real-world systems, there is a need for reducing their computational overhead. In this work, we present the PatchMerger, a simple module that reduces the number of patches or tokens the network has to process by merging them between two consecutive intermediate layers. We show that the PatchMerger achieves a significant speedup across various model sizes while matching the original performance both upstream and downstream after fine-tuning.

Figures (9)

• Figure 1: (Left) Overview of the PatchMerger module with $N$ input tokens and $M$ output ones. (Right) Placement and token shapes when adding a PatchMerger between two encoder blocks.
• Figure 2: Patch Merger for ViT. Performance versus total training FLOPs. Colors represent different ViT variants, markers represent either standard ${\hbox{$\bullet$}}{}$ ViT or ${\hbox{$\blacksquare$}}{}$ Merger ViT. The brown Merger ViT (top-right most) corresponds to H/11. Lines show the Pareto frontier of Merger ViT (solid) and ViT (dashed). Figure \ref{['fig:vit_upstream_fewshow_runtime']} in the Appendix shows performance versus runtime for these models.
• Figure 3: Patch Merger for V-MoE. Performance versus total training FLOPs. Colors represent different V-MoE variants, markers represent either standard ${\hbox{$\mathbf{x}$}}{}$ V-MoE or $\varhexagonblack$ Merger V-MoE. Lines show the Pareto frontier of Merger V-MoE (solid) and V-MoE (dashed). Figure \ref{['fig:vmoe_upstream_fewshow_runtime']} in the Appendix shows performance versus runtime for these models.
• Figure 4: Full finetuning on ImageNet. Performance versus total training FLOPs. Colors represent different ViT and V-MoE variants (including H/11 for Merger ViT), markers represent either standard ${\hbox{$\bullet$}}{}$ ViT, ${\hbox{$\blacksquare$}}{}$ Merger ViT, ${\hbox{$\mathbf{x}$}}{}$ V-MoE or $\varhexagonblack$ Merger V-MoE. Lines show the Pareto frontiers in each case. Figure \ref{['fig:downstream_runtime']} in the Appendix shows performance versus runtime for these models.
• Figure 5: Placement Comparison. The discontinuous line corresponds to placing the merger in the middle of the network as done in previous sections --after the 6th, before the 7th encoder block, in this case. The number of output tokens is 8 in all cases.