Token Merging: Your ViT But Faster

Daniel Bolya; Cheng-Yang Fu; Xiaoliang Dai; Peizhao Zhang; Christoph Feichtenhofer; Judy Hoffman

Source-linked AI summary

Token Merging: Your ViT But Faster

Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, Judy Hoffman

arXiv:2210.09461v3Oct 17, 2022cs.CV

TL;DR

Large ViTs are difficult to run efficiently, while pruning can lose information and require retraining. ToMe gradually merges similar tokens with a fast matching algorithm that works during inference or training. Across images, video, and audio, it achieves competitive speed and accuracy, including substantial throughput gains with small reported accuracy drops.

Problem
Massive ViT models are troublesome to run, while token pruning can lose information, require retraining, and hinder batched inference or training.
Method
ToMe gradually combines redundant tokens in existing ViTs using a fast matching algorithm and supports both inference and training.
Results
ToMe is competitive with state-of-the-art speed and accuracy across images, video, and audio.
Takeaways & Limitations
ToMe provides a practical way to accelerate ViT inference and training across modalities without requiring model modifications.
Takeaways & Limitations
The video evaluation may compensate for information loss from token merging through its use of multiple clips.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Vanilla ViTs offer broad applicability and strong performance, but their massive computational demands motivate faster alternatives. Token Merging addresses this by combining similar tokens while preserving use across training settings and modalities.

Massive ViT models are difficult to run and reproduce efficiently.
Token pruning reduces token counts but loses information, often requires retraining, and can prevent efficient batched inference or training.
Token Merging combines tokens rather than pruning them, matching pruning speed while improving accuracy and supporting use with or without training.
ToMe can increase training speed, in some cases cutting total training time in half.
Experiments across images, video, and audio find ToMe competitive with state-of-the-art methods across modalities.

2 RELATED WORK

Related work improves transformer efficiency through faster attention, pruning, or domain-specific modules. ToMe instead targets existing ViTs by combining tokens without requiring training.

Efficient-transformer research includes faster attention, pruning heads or features, and adding domain-specific modules.
Prior token-reduction methods prune tokens in NLP and vision, but generally require training.
Dynamic pruning can vary token counts across inputs, complicating batched inference and training.
Earlier token-combination approaches lacked a reasonable speed-accuracy trade-off without training.

3 TOKEN MERGING

ToMe inserts gradual token merging into standard ViT blocks using fast similarity matching and size-aware aggregation. Its fixed reduction schedule supports efficient inference and training while preserving token information.

TOKEN MERGING: ToMe inserts a token-merging module into an existing ViT to combine redundant tokens and increase throughput without necessarily training.
TOKEN MERGING: ToMe reduces r tokens per transformer block, gradually merging rL tokens across L blocks regardless of input content.
TOKEN MERGING: Merging occurs between attention and MLP branches, allowing attention features to guide merging and information to propagate before tokens combine.
TOKEN MERGING: Token similarity is measured between keys using dot-product similarity because keys summarize token information for attention.
TOKEN MERGING: Bipartite soft matching partitions tokens, links each token in A to its most similar token in B, keeps r strongest edges, and merges connected tokens.
TOKEN MERGING: Merged tokens retain size information: proportional attention makes a token act like multiple copies of its key, while aggregation weights tokens by size.
TOKEN MERGING: ToMe can be trained as a pooling operation with ordinary backpropagation, making it a drop-in replacement for faster training.

4 IMAGE EXPERIMENTS

ToMe's image experiments show that lightweight token merging preserves strong accuracy while substantially increasing ViT throughput, both off-the-shelf and during training. Its design choices, matching strategy, and comparisons support efficient scaling across ViT models.

4.1 DESIGN CHOICES: ToMe uses attention keys, post-attention merging, cosine similarity, and head-averaged keys to select similar tokens efficiently.These choices were more accurate than feature-based alternatives and concatenated attention heads.
4.1 DESIGN CHOICES: Weighted averaging by token size and alternating bipartite assignment outperform token selection, max pooling, unweighted pooling, and sequential assignment.The ablations identify weighted merging and alternating partitioning as the preferred algorithmic choices.
4.1 DESIGN CHOICES: A constant schedule is close to optimal across 15,000 sampled schedules, while decreasing schedules become preferable at higher throughputs.The schedule comparison uses an off-the-shelf AugReg ViT-B/16 on ImageNet-1k.
4.2 MODEL SWEEP: At 2× throughput, trained MAE models reach accuracy drops of 0.4% for ViT-H, 0.6% for ViT-L, and 1.7% for ViT-B.Training with ToMe also enables roughly 2× faster epochs for large MAE models with negligible accuracy drops.
4.2 MODEL SWEEP: Re-evaluating a model with a different merging rate can improve performance, allowing one ToMe-trained model to support a broad range of evaluation scales.A ViT-L model trained with r = 5 improved from 85.7% baseline accuracy to 85.8% when re-evaluated with r = 0.
4.3 COMPARISON TO OTHER WORKS: ToMe matches or exceeds pruning baselines while avoiding extra parameters, gradient tricks, and padding-related training overhead.On DeiT-S, ToMe achieves a 1.5× training speedup, and an off-the-shelf AugReg ViT-S can match the trained DeiT models without training.

5 VIDEO EXPERIMENTS

ToMe applies unchanged to video, improving ViT-L throughput and training speed while preserving competitive accuracy. Its token merging also tracks objects or parts across multiple frames.

Experimental Setup: The video evaluation compares ToMe without training or during MAE fine-tuning against state-of-the-art Kinetics-400 models in the same flop range.Comparisons include Swin, MViTv2, Spatiotemporal MAE, and X-ViT + ATS.
Results: ToMe matches Swin-B throughput and outperforms MViTv2-L even without training.With a decreasing schedule, ViT-L MAE significantly outperforms the same-flop-count ViT-B MAE baseline with or without training.
Throughput: 2.2× throughput comes with a negligible 0.2% accuracy drop for ViT-L using a constant schedule.The same setting cuts training time in half despite synchronization overhead across 8 GPUs.
Clip Count: At lower token-merging rates, accuracy does not decrease relative to the baseline across clip counts.At r = 65, accuracy degrades by approximately 0.5% compared with the four-clip sweet spot.
Visualization: Visualization shows ToMe merging the same object or part into one token across multiple video frames.The ball remains merged across its full range of motion, while a reflected ball patch appears as an extraneous match.

6 AUDIO EXPERIMENTS

ToMe is evaluated on Audio MAE with ViT-B using spectrogram inputs and AudioSet-2M. It nearly doubles throughput while limiting the mAP decrease to 0.4%.

Results: Almost 2× throughput is achieved with only a 0.4% mAP drop on audio.The metric is mAP rather than accuracy because AudioSet-2M has class imbalance.
Experimental Setup: The audio experiment uses ViT-B from Audio MAE and evaluates on AudioSet-2M.Audio signals are rasterized as spectrograms before being fed into the ViT.
Experimental Setup: The reported baselines differ because the authors compare untrained ToMe with the original result and trained ToMe with their own trained baseline.The authors attribute the difference to training implementation differences.

7 CONCLUSION

The paper concludes that Token Merging increases ViT throughput by gradually merging tokens and exploits input redundancy across modalities. Experiments on images, video, and audio find competitive speed and accuracy relative to the state of the art.

Conclusion: ToMe increases ViT throughput by gradually merging tokens and naturally exploits redundancy in the input.The conclusion frames this property as applicable to any modality with redundancy.
Conclusion: Experiments on images, video, and audio obtain speeds and accuracies competitive with the state of the art in each case.
Future Work: ToMe is presented as a natural hierarchical model that uses pure transformer blocks.The authors suggest combining it with methods such as Swin or MViT as future work.
Future Work: The paper focuses on classification, while visualizations suggest potential for tasks such as segmentation.The authors also identify large-model training, memory reduction, and more efficient transformers as future directions.

A FULL RESULTS

The appendix provides full image, video, audio, and model-specific results, with throughput measured on a V100 and experimental baselines documented separately. It records the training and evaluation settings used across these results.

Measurement: Throughput values in the appendix are measured on a V100 and may differ slightly from the main paper’s values.All main-paper throughput results use the same evaluation machine.
Image Results: The appendix lists full AugReg and SWAG off-the-shelf results without training.The original off-the-shelf models are shown in gray, and selected models are referenced in the main paper.
DeiT Results: DeiT models are trained from scratch for 300 epochs without checkpoint or fine-tuning tricks.The class and distillation tokens are not merged, and DeiT-S with r = 11 did not finish training.
Video Results: The video appendix re-evaluates ViT-L with one spatial crop and ten temporal clips, unlike the original baseline’s three spatial crops and seven temporal clips.This makes the re-evaluated baseline slightly less accurate than the original paper’s baseline.
MAE Results: Full MAE results include evaluations with and without training, using original baseline models trained by the authors.The appendix includes the gray MAE models reported in the main paper.
Audio Results: Audio experiments use the published model for off-the-shelf evaluation but a separate implementation for training, so two baselines are listed.Proportional attention is disabled for both audio evaluation modes.

B HYPERPARAMETERS

The experiments examine whether token merging requires hyperparameter changes and report results across standard image, video, and audio training or evaluation settings. The supplied results indicate that default training settings remain effective, while some baseline comparisons have reproduction limits.

B HYPERPARAMETERS: Layer decay does not need to be changed when training with ToMe, despite potentially smaller later-layer gradient updates.ToMe gradually reduces the number of tokens, which could reduce later-layer update sizes without layer decay.
B HYPERPARAMETERS: DeiT-Ti results are omitted from the main paper because the authors could not reproduce SPViT’s baseline accuracy.The trained baseline is slightly more accurate than the original baseline in the reported DeiT comparison.
B HYPERPARAMETERS: Token merging produces speed-ups during both inference and training without changing the model or hyperparameters beyond adding token merging.The full DeiT results report standard 300-epoch ImageNet-1k training and define speed-up and accuracy drop relative to baseline.
B HYPERPARAMETERS: The reported results include off-the-shelf and trained evaluations for video and audio, with video using top-1 and top-5 accuracy and audio including trained and untrained baselines.These tables distinguish evaluation settings and baseline variants across modalities.
B HYPERPARAMETERS: Default hyperparameters already achieve the highest accuracy in the reported sweep for settings affected by token merging.The sweep covers relevant image and video hyperparameters potentially affected by token merging.

C MERGING SCHEDULE

The merging schedule controls how many tokens ToMe combines at each layer. The best schedule depends on the target throughput: constant schedules work at intermediate throughput, while decreasing schedules become preferable at higher throughput.

C MERGING SCHEDULE: For lower throughputs, the highest-accuracy schedules merge more tokens toward the end of the network.For higher throughputs, constant schedules become best before linearly decreasing schedules work well at still higher throughputs.
C MERGING SCHEDULE: At throughput 1600-1800, the best schedule is close to constant across layers.This motivates using a constant schedule in that throughput range.
C MERGING SCHEDULE: For throughput beyond 1600-1800, a decreasing schedule performs best, motivating a linearly decreasing schedule in the main paper.The figure analyzes average schedules among the most accurate random samples.

D IMPLEMENTATION

ToMe implements token reduction with bipartite soft matching and parallel merge operations. The implementation normalizes attention keys, matches alternating token sets, and applies the resulting assignments to token features and related quantities.

D IMPLEMENTATION: Bipartite soft matching takes normalized attention keys shaped as [batch, tokens, channels] and computes cross-set similarity scores.The input is split into alternating token sets A and B before matrix multiplication produces the scores.
D IMPLEMENTATION: The matching separates merged and unmerged token indices, then maps each selected source token to a destination token.The implementation stores unmerged indices in `unm_idx`, merged indices in `src_idx`, and destinations in `dst_idx`.
D IMPLEMENTATION: The merge operation gathers unmerged tokens and adds selected source tokens into destination tokens with `scatter_add`.It operates on tensors shaped [batch, tokens, channels] and returns the merged representation.
D IMPLEMENTATION: The matching function returns a reusable merge function that can operate on features, token sizes, or source patches.These operations are performed in parallel without sequential loops.

E MORE VISUALIZATION

Additional visualizations show ToMe repeatedly merging semantically related image parts and tracking objects across video frames. The visualizations trace final tokens back to their input patches while noting that merged regions need not be spatially contiguous.

E MORE VISUALIZATION: Visualization regions are reconstructed by tracing final merged tokens to original input patches and coloring each region by its average color.Random border colors distinguish different tokens.
E MORE VISUALIZATION: Merged tokens need not correspond to contiguous input regions because positional encodings provide ToMe’s only spatial signal.The visualizations therefore show token associations rather than guaranteed contiguous segments.
E MORE VISUALIZATION: ToMe repeatedly merges parts and objects into tokens across many different images.Figure 8 extends the image examples from Figure 4.
E MORE VISUALIZATION: Video visualizations show ToMe performing object tracking across clips from Kinetics-400 validation data.Figure 9 highlights an instance or part merged into one token across frames.
E MORE VISUALIZATION: Each video token represents two frames because the patch size is 2 × 16 × 16, so the visualizations plot the first frame of each pair.The authors report that this frame more closely matches the merged tokens.