Token Merging: Your ViT But Faster
Authors: Daniel Bolya
Cheng-Yang Fu
Xiaoliang Dai
Peizhao Zhang
Christoph Feichtenhofer
Judy Hoffman
Venue: ICLR 2023 notable top 5% (2023)
Keywords: token merging
token pruning
inference speed
training speed
throughput
off-the-shelf
fine tuning

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.

Published as a conference paper at ICLR 2023
TOKEN MERGING: YOUR VIT BUT FASTER
Daniel Bolya 12+
Cheng- Yang Fu²
Xiaoliang Dai?
Peizhao Zhang?
Christoph Feichtenhofer
Judy Hoffmann
Georgia Tech
Meta Al
jubolya, jedyeeatcchhh..ch.hh. Jegatchh.edeeeeee
{chengyangfu, xiao feeeeeetet com
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 general and light-weight matching algorithm that is as fast as pruning while being more accurate. Off-the-shell ToMe can 2x the throughput of state-of of- VIT-L @ 512 and VIT-H @ 518 models on images and 2.2x the throughput of ViT-L on video with only 0.2-0.3% accuracy drop 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 's accuracy and speed are competitive with state-of-the-al on images, video, and audio.
INTRODUCTION
The introduction of transformers (Vaswani et al., 2017) from NLP vision with Vision Transformers (ViTs) by Dosovitskiy et al. (2020) has rapidly advanced the field of computer vision. However, unlike in NL.,, vision has been since dominated by domain-specific transformer hybrids like Swin (Liu al., 2021; Dong et al., 2022) using vision-specific attention, MVIT (Fan al., 2021; Li et al., 2022) using vision-specific pooling, or VeV (Graham et al., 2021) using vision-specific conv modules. The reason for this trend is simple: efficiency. Adding vision-specific inductive biases enables transformer hybrids to perform better with less compute.
Yet, vanilla ViTs still have many desirable qualities: they consist of simple matrix multiplications, making them faster than their raw flop count would suggest; they support powerful sulf-supervised pre-training techniques such MAE (He al., 2022) that can put up state-of-the being fast to train; given their lack of assumptions about the data, they can be applied with little or no changes across many modalities (Feichtenhofer et al., 2022; Huang et al., 2022); and they scale well with massive amounts of data (Zhai et al., 2021; Singh et al., 2022), recently obtaining up to 90.94% 1 on ImageNet (Wortsman et al., 2022).
results while
However, running these massive models can be troublesome, and reproducing these results with faster architecture would be difficult. A promising subfield of ViTs have recently emerged where, due to the input-agnostic nature of transformers, tokens can be pruned at runtime to enable a faster model (Rao et al., 2021; Yin et al., 2022; Meng et al., 2022; Liang et al., 2022; Kong et al., 2022). Yet, token pruning has several disadvantages: the information loss from pruning limits how many tokens you can reasonably reduce; current methods require re-training the model to be effective (some with extra parameters); most cannot be applied to speed up training; and several prune different numbers of tokens depending on the input content, making batched inference infeasible.
In this work, present Token Merging (ToMe) to combine tokens, rather than prune them. Because of our custom matching algorithm, our method is as fast as pruming while being more accurate. Moreover, our method works with or without training, which unlocks its use huge models with minimal accuracy drop. Using ToMe during training, we observe actual increases in training speed, in some cases cutting the total training time in half. And we apply ToMe without any modifications to images, video, and audio and find it to be competitive with the SotA in all cases.
Our contributions are as follows: we introduce technique to increase the throughput and real-world training speed of ViT models, both with and without training (Sec. 3) and thoroughly ablate our
Work done during an internship Meta AI. Code at http::/gi thub. ccwm//acabookeeeercccccccco

---

Published as a conference paper at ICLR 2023
design choices (Sec. 4.1); we perform extensive experiments images with several ViT models (Sec. 4.2) and compare to state-of -hete--art- art in architecture design and token pruning methods (Sec. 4.3); we then repeat these experiments for both video (Sec. 5) and audio (Sec. 6) and find ToMe works well across modalities; and we visualize our results and find ToMe merges parts of objects on images (Fig. 4) and objects over their entire range of motion on video (Fig. 6). We hope ToMe can enable the creation of more powerful, faster ViT models.
RELATED WORK
Efficient Transformers. Several works have attempted create more efficient transformers in both NLP and Vision. Some focus on faster attention (Choromanski et al., 2020; Shen et al., 2021; Dao et al., 2022; Wang et al., 2020; Bolya et al., 2022), some attempt to prune heads or features (Meng et al., 2022; Voita et al., 2019; Michel et al., 2019), and some attempt to infuse domain- specific modules (Mehta & Rastegari, 2021; Graham et al., 2021; Liu et al., 2021: 2022a; Dong et al., 2022). In this paper, we focus on speeding up existing ViT models by merging tokens to match the speed-aceuracy trade-off of more complicated domain-specific models, sometimes without training.
Token Reduction. Since transformers can operate with any number of tokens, several recent works have attempted to prune the tokens from transformers in both NLP (Goyal et al., 2020; Kim & Cho, 2020; Kim et al., 2021; Lassance et al., 2021) and Vision (Meng et al., 2022; Yin al., 2022; Kong et al., 2022; Song et al., 2022; Rao et al., 2021; Fayyaz et al., 2022; Yu & Wu, 2021). However, these methods require training, while our method can be used without training. Moreover, most pruning works are dynamic, i.e., the number of tokens varies between images or sentences. While this benefits accuracy it limits practicality, as samples with different numbers of tokens can no longer be batched. To solve this, most pruning papers apply a mask during training rather than remove tokens, which negates the speed-up from pruning. Our method, on the other hand, can be applied during both inference and training, achieving real-world speed-ups in either case.
Combining Tokens. While plenty of works prune tokens, very few combine them. Kong et al. (2022) and Liang et al. (2022) combine what they prune into single token. Group ViT (Xu et al., 2022), while not focused on efficiency, groups tokens using cross-attention for semantic segmentation. TokenLeerner (Ryoo al., 2021) uses an MIP to reduce the number of tokens. LIT (Pan et al., 2022) learns deformable token merging layers for pooling between stages. Token Pooling (Marin et al., 2021) is the most similar to our token merging but uses slow kmeans-based approach' that doesn't 't work on an off-the-shelf model". Until now, no approach has been successful in offering reasonable trade-off when combining tokens without training.
TOKEN MERGING
Our goal is to insert token merging module into an existing VIT (Dosovitskiy et al., 2020). By merging redundant tokens, we hope to increase throughput, while not necessarily having to train.
Strategy. In each block of transformer, we merge tokens to reduce by per layer. Note that r is a quantity of tokens, not a ratio. Over the L blocks in the network, we gradually merge rL tokens. Varying r gives a speed-accuracy trade-off, fewer tokens means lower accuracy but higher throughput. Importantly, we reduce r tokens regardless of the image's 's content. Some pruning methods dynamically vary the number of tokens (e.g., Kong et al. (2022)). This increases accuracy but is generally impractical, as it prevents batched inference or training without padding tokens.
As shown in Fig. 1, we apply our token merging step between the attention and MLP branches of each transformer block. This is also in contrast to prior works, which tend to place their reduction method at the beginning of the block instead. Our placement allows information to be propagated from tokens that would be merged and enables us to use features within attention to decide what to merge, both of which increase accuracy (see Tab. la).
Token Similarity. Before merging similar tokens, we must first define what "similar" means. While it may be tempting to call two tokens similar if the distance between their features is small (as in Marin al. (2021)), this is not necessarily optimal. The intermediate feature space in modern transformers is overparameterized. For instance, ViT-B/16 has enough features to completely encode the rgb pixel
'Their throughput is only 1. 1.255 the baseline because their method can't 't be parallelized. In their appendix, they show drops of 10-40% accuracy when combining tokens without training.
2

---

Published as a conference paper at ICLR 2023
8) Gradually merge tokens in esch block
b) Transtormer Block Token Merging
ToMe
Image
Patchfied
Block 11
Block 22
Block 32
6)
Bipartite Soft Matching
O0ooo
8
Step Assign
Step Draw one edge takans S and thair most similar tokoo Setie.
Step Keep the top similar edges.
Step 4: Merge connected
Step Concatenate together.
Tolens to S
Set
Figure I: Token Merging. (a) With ToMe, similar patches are merged in each transformer block: for example, the dog's ' fur is merged into single token. (b) ToMe is simple and be inserted inside the standard transformer block. (c) Our fast merging algorithm, see Appendix D for implementation.
values of each token (16 16 * 768). This means that the intermediate features have the potential to contain insignificant noise that would confound our similarity calculations.
Luckily, transformers natively solve this problem with QKV self-attention (Vaswani al., 2017). Specifically, the keys (K) already summarize the information contained each token for use in dot product similarity. Thus, we dot product similarity metric (e.g., cosine similarity) between the keys of each token to determine which contain similar information (see Tab. la, lb).
Bipartite Soft Matching. With token similarity defined, we need a fast way to determine which tokens to match in order to reduce the total number by There are several potential solutions to this problem, such as kmeans clustering (Lloyd, 1982) or graph cuts (Boykov et al., 2001). But we perform this matching L times within the network on potentially thousands of tokens, so its runtime has to be absolutely negligible. This is very much not the case for most iterative clustcring algorithms.
Thus, propose more efficient solution. Our design goals are follows: 1.) we want to avoid anything iterative that cannot be parallelized and 2.) we want the changes merging makes to be gradual. The latter is why we focus matching and not clustering, clustering places no bounds on how many tokens can be merged into one group (which may adversely affect the network). whereas matching leaves most of the tokens unmerged. Our algorithm is follows (visualized in Fig. I):
1. Partition the tokens into two sets A and of roughly equal size.
2. Draw one edge from each token in to its most similar token in B.
3. Keep the r most similar edges.
4. Merge tokens that are still connected (e.g.. by averaging their features).
5. Concatenate the two sets back together.
Because this creates a bipartite graph and each token in A has only one edge, finding connected
components in step is trivial. Moreover, we don't need to compute similarity between every pair of tokens which, if we choose A and B carefully, isn't isn't a problem for accuracy (see Tab. le). In fact, this "bipartite soft matching" is nearly as fast as just dropping tokens randomly (see Tab. 2) and takes only a few lines of code to implement (see Appendix D).
Tracking Token Size. Once tokens are merged, they no longer represent one input patch. This can change the outcome of softmax attention: if we merge two tokens with the same key, that key has less effect in the softmax term. We can fix this with simple change, denoted proportional attention:
QK
A softmax
+ +og
Vd
where 8 is row vector containing the size of each token (number of patches the token represents). This performs the same operation as if you'd dd have s copies of the key. We also need to weight tokens by any time they would be aggregated, like when merging tokens together (see Tab. Id).
Training with Merging. Each component thus far has been designed to be able to add token merging to already trained ViT model. Training with ToMe isn't necessary, but may be desirable to
3

---

Published as a conference paper at ICLR 2023
feature
acc
im/s
function
acc
im/s
aggregate
acc
im/s
Xpre
83,02
186.8
eucl
84.26
182.5
concat
84.32
180.3
83.70
182.8
cosine
84.25
182.9
mean
84,25
182.9
K 84.25
182.9
dot 82.78
183.0
Q
84.04
182.8
softmax
82.00
183.0
83.80
182.9
(a) Feature Choice. The matrix accurately summarizes the infor- mation within tokens.
(b) Distance Function. Cosine similarity is the best choice for speed and accuracy.
(c) Head Aggregation. Averag- ing over the attention heads is
bit less accurate, but faster.
method keep one
acc
im/s
order
acc
im/s
prop
im/s
81.01
185,4
sequential
81.07
183.0
mae
84,25
182.9
max pool
83.50
184.6
alternating
84.25
182.9
83.84
180.9
avg pool
83.57
183.8
random 83.80
181.7
augreg augreg
82. 15
182.8
weighted avg 84.25
182.9
83.51
180.8
(d) Combining Method. Averag- ing tokens weighted by their size, (see Eq. 1), ensures consistency.
Partition Style. Alteratinn and B ensures tokens are com- pared in subsequent layers.
(f) Proportional Attn. Without MAE pretraining, off-the-shelf models require prop attn.
Table 1: Token Merging ablation experiments using ViT-L/16 from MAE (He et al., 2022) on ImageNet-Ik Ik evaluated off- shefh without training, using 8. The baseline model without ToMe obtains 85.96% acc at 93.3 im/s. For each ablation, we report Top-1 accuracy (acc) and fp32 model throughput (im/s) V100 GPU. Our default settings are marked in purple.
reduce accuracy drop to speed up training. To train, we simply treat token merging as a pooling operation and backprop through the merged tokens as if we were using average pooling. We don't t find need
any gradient tricks such as Gumbel softmax (Jang et al., 2017) in token pruning (e.g., Kong et al. (2022)). And in fact, we find that the same settings used in training a vanilla ViT are also optimal here (see Appendix B). Thus ToMe is drop-in replacement increase training speed.
IMAGE EXPERIMENTS
We perform several experiments ImageNet-Ik (Deng et al., 2009) using ViT models trained in four different ways: AugReg (Steiner et al., 2022), MAE (He et al., 2022), SWAG (Singh et al., 2022), and DeiT (Touvron et al., 2021). For all experiments, our method or, the case of MAE and DeiT, trained with our method applied. All throughputs are measured during inference V100 GPU with optimal batch size and fp32 unless noted otherwise.
either run the model off-the-shelf with
4.1 DESIGN CHOICES
In Tab. 1, we ablate the design choices made in our approach. For each ablation, we start from our default parameters marked in purple. Unless otherwise noted, on an off-the-shelf ViT-L/16 MAE model without training (ace: 85.96%, im/s: 93.3). and merge with r 8 which gradually removes 98% of tokens over the 24 layers of the network.
Token Similarity. The tokens' features (X) are not the best in terms of performance (Tab. Ia). Moving the merging operation after attention (X vs. Xpre) and using the attention keys (K) is significantly more accurate. Then, cosine similarity is best measure token distance as shown in Tab. 1b. Finally, average K over the attention heads instead of concatenating them (Tab. 1c) for efficiency.
Algorithmic Choices. After deciding what tokens merge, we combine them by averaging weighted by token size, (see Eq. 1). In Tab. 1d, this outperforms keeping just the token in B. max pooling, or unweighted average pooling. Then, our bipartite matching algorithm requires splitting the input into two disjoint sets. Because we concatenate the sets afterward, we find that assigning tokens by alternating to work the best (Tab. le). Filling A and then filling (sequentially) performs the worst.
Propertional Attention. Once merged, tokens represent more than one input patch. We address this with proportional attention (Eq. 1), which we ablate in Tab. If. Surprisingly, proportional attention necessary for supervised models (e.g., AugReg, SWAG, DeiT), but not for MAE models.
4

---

Published as a conference paper at ICLR 2023
style
algorithm
acc
im/s
prune
random
79.22
184,4
84
prune
attn-based
79.48
183.8
Schecllle
merge
kmeans (2 iter)
80.19
169.7
82
Randoo
merge
kmeans (5 iter)
80.29
147.5
81
Tokens Morged
0%
merge
greedy matching
84.36
179.4
merge
bipartite matching
84.25
182.9
79
96%
Table 2: Matching Algorithm. Different match- ing algorithms with the same settings Tab. 1. Our bipartite algorithm is almost as fast as ran- domly pruning tokens, while retaining high ac- curacy. Matching is more suited for this setup than pruning or clustering.
78
1000
1200
1400
1600
1,000
IP16 Throughput (im/s))
Figure 2: Token Merging Schedule. Our de- fault constant merging schedule is close opti- mal when compared to 15k randomly sampled merging schedules on an AugReg ViT-B/16.
Since this discrepancy disappears after training, this is likely because MAE already removes tokens during pretraining. Nevertheless, we use proportional attention for all but off-the-shelff MAE models.
Comparing Matching Algorithms. In Tab. 2 we compare our bipartite matching to different token reduction algorithms, both pruning and merging. Pruning fast, but with 98% of the tokens removed overall, important information is lost. This is true for both pruning randomly and pruning based on what isn't attended to (Kim et al., 2021). ). In contrast, merging tokens only loses information when dissimilar tokens are merged. Thus, it's important to correctly select similar tokens to merge.
At first, kmeans (Lloyd, 1982) may seem like the obvious choice, but on top of being slow it's only slightly better than pruning. While it may minimize reconstruction error, kmeans allows large number of tokens to match to the same cluster, which increases the probability of dissimilar tokens being merged. Marin et al. (2021) study several faster clustering algorithms based on kmeans, but they are unable to obtain a better than 10% accuracy drop in their setting without training.
Instead, want a matching algorithm that only merges the most similar tokens. We could do this greedily by merging the most similar pair of tokens and then repeating without replacement times. This is accurate but sequential and thus could get slow with large r. Our bipartite matching has the accuracy of this greedy approach and the speed of pruning while having constant runtime w.r.t. to r.
Selecting a Merging Schedule. By default, we merge tokens with a constant schedule, i.e. r per layer. To evaluate the optimality of this design we randomly sample total of 15,000 merging schedules. For each schedule, we test its accuracy and fp16 throughput on ImageNet- Ik val using an off-the-shelf AugReg ViT-B/16 model. In Fig. 2, we plot the results of this experiment and find that constant schedule close to optimal, especially as the total tokens merged increases. We further analyze the best random samples (see Appendix C) and find that a linearly decreasing schedule works well at throughputs up to ~3x. Thus, we also define a "decreasing" schedule that removes 2r tokens in the first layer and tokens in the last layer, linearly interpolating for the rest. This also removes rL tokens, but is faster because more are removed early:
Constant Schedule
per layer
denoted r.
(2)
Decreasing Schedule
2x - per layer
denoted r.
(3)
4.2
MODEL SWEEP
In Fig. 3.
apply our token merging method to 11 SotA off-the-shelf VIT models from various sources. For each model, we vary r with a constant schedule to construct throughput vs. accuracy curves, starting with (no merging baseline) to r such that we run out of tokens to merge. We evaluate each model theethe-shelf. That is, by default we don't ' train; we just download the model and change a few lines code. Models are evaluated on 224px images unless otherwise noted.
Supervised Models. Both AugReg (Steiner et al., 2022) and SWAG (Singh et al., 2022) are VIT models pretrained on a large supervised (or weakly supervised) pretraining dataset and fine-tuned on ImageNet- 1k. AugReg covers optimal ViT training up until ViT-L/16, while SWAG pushes the limits of ImageNet by training huge models with massive image sizes. We apply our method off she-lf on AugReg models in Fig. 3a and SWAG models in Fig. 3b.
5

---

Published as a conference paper at ICLR 2023
85
Modell
VIT-1/16 @ 384
Accuracy
B0
VIT-U16
VIT-3/16
VIT-S/16
VII-1v16
Mode
- Off-the-Sheir
32
64
128
256
512
1024
2048
4096
Throughput (Im/s)
(a) AugReg Models. A collection of ImageNet-21k pretrained models (Steiner al., 2022).
89
Model
86
Accuracy
88
VIT--114 516
VITHL/16 & 512
VIT 8/16 384
Accurscy 199)
84
VIT-H114
Mode
VIT-L/16
66
Off the Shaff
82
WIT-1/16
Mode
80
Otfithe Shelf
Trained
84
4
16
32
64
128
256
32
126
256
512
Throughput (ImUS)
Throughput (im/s]
(b) SWAG Models. Massive weakly supervised mod- els pretrained on 3.6B images (Singh al., 2022).
(c) MAE Models. Self-supervised models pretrained on ImageNet-Ik Ik (He et al., 2022).
Figure 3: Model Sweep. We apply ToMe to several state of the art ViT models off-the-shelf. i.c. without training, varying r to produce fp32 throughput vs. accuracy curves on ImageNet- lk.
that constant schedule gives up 2x the throughput no matter the model. And even though we'r re compressing 96-98% of the tokens in each, the largest models have barely any accuracy drop: while ViT-B, S, and Ti all have around 4-5% accuracy drop 2x speed, VIT-L only suffers 2% drop on 224px images and 0.7% drop on 384px images with AugReg, Similarly, with SWAG models, ViT-L on 512px images and ViT-H on 518px images both have small 0.3% accuracy drop without training. Note this trend is not just because larger models have more tokens, since we always reduce the number of tokens by 96-98%. Instead, we think this is because large models are deeper and thus allow for more gradual change in features, which lessens the impact of merging.
Immediately, we
Self-Supervised Models. MAE (He et al., 2022) is a self-supervised pretraining method for ViT with models pretrained and fine-tuned on ImageNet-1K k. In Fig. 3c we apply our method both off-the-shelf and trained by fine-tuning the public pretrained checkpoints. When fine-tuning, we find that we can the original training recipes. We don't have to compensate for fewer tokens in later layers (see Appendix B). likely because our method is already tuned to imitate a model without merging.
And as expected, in Fig. 3c, we see the same trends as before: except this time, with training we can bring the error down to 0.4% for ViT-H, 0.6% for ViT-L, and 1.7% for ViT-B at 2 x throughput. Our approach actually implicitly addresses an issue in MAE: because MAE removes tokens during pretraining, its epoch times are ~4x faster than training a supervised model. However, normal fine-tuning uses all the tokens and doesn't 't have this benefit. Our token merging method fixes this issue and allows for roughly ~2 faster epochs at negligible accuracy drops for large models. This suggests that one could train even bigger models with token merging than was possible before.
Re-evaluating. Note that, while in Fig. 3c we train a new model for each value rf this isn't actually necessary. Instead, we can take a model trained with one value of r and re-evaluated with another. In fact, it's possible to actually improve performance by doing so. For instance, the baseline VYT-L model we train in Fig. 3c gets 85.7% accuracy. If we re-evaluate our = 5 trained model with r = O, we obtain 85.8% accuracy. Thus, it's feasible to speed up training with ToMe and not apply during evaluation to produce the same or better results. This means that, while the result of applying ToMe in Fig. 3b and Fig. 3c are similar to e.g. scaling the model size, you only have train one model with ToMe to create any model with a large range of scales.
4.3 COMPARISON TO OTHER WORKS
In this section, we compare our trained token merging models to other state-of-the the art works on ImageNet- 1k, both in terms of the overall vision space well other token reduction methods.

---

Published as a conference paper at ICLR 2023
model
input
gflops
im/s
inference
train
En--B5
456
83.6
9.9
169
method
acc
gilops
im/s
speed
VIT-B MAR
224
83.6
17.6
309
DerT-S
79.8
4.6
930
1x
Swin-B"
224
84.0
15.4
2781
A-VIT
78.6
2.9
1x
CSWin-B B
224
84.2
15.0
250
Dynamic VIT
79.3
2.9
1505
1x
MViTv2-B
224
84.4
10.2
2531
SP-VIT
79.3
2.6
1x
ToMe
ToMe Delt
79.4
2.7
1552
1.5x
VIT-L MAE
224
84.2
22.3
241
ToMe
AugRog
79.3
2.7
1550
VIT-L MAB
224
85.1
31.0
183
VIT-L MAL
224
85.71
61.6
93
Table 4: ToMe vs. pruning methods on ViT-S trained from scratch with DeiT. We time Dy- namic on our V100, but the others cannot be batched and evaluate throughput in different settings. Pruning methods require token padding during training, and thus don't 't improve training speed, while ToMe at r = 13 trains 1.5x faster than the baseline DeïT-S. We also obtain the same result without training on an off sffetf AugReg VIT-S model. blue indicates ToMe applied without training while gray indicates ToMe applied during training. A-ViT is not trained from scratch and performs slightly better than Dynamic ViT in its setting, See Appendix A for DeïT-Ti results.
EM--6
528
84.0
19.0
961
MViTv2-L
224
85.3
42.1
81
ToMe
ViT-H
224
86.1
72.6
81
VIT-H MAE
224
86.5
92.9
63
VIT-H MAE
224
86.91
167.4
35
VW-n
224
85.7
118.1
49
Table 3: Advancing a Tier. Comparison to
SoTA models trained only on ImageNet- 1k. Our method allows for the use of more complicated models for the same tier of throughput. models use SimMIM self-supervised pretraining. base- line models trained by us without ToMe for com- parison. im/s from original papers (V100).
Comparison to State of the Art. In Tab. 3, we compare our MAE fine-tuned models to state-of-the- art models trained on ImageNet- 1 without extra data: EfficientNet (Tan & Le, 2019), Swin (Liu et al., 2021), V2 (Liu et al., 2022a), CSWin (Dong et al., 2022), and MViTv2 (Li et al., 2022). All throughputs are on a single V100. Note that we use MAE pretraining which is not supported for all transformers, but provides accuracy improvement for some like Swin/Swin V2. Thus we also include SimMIM (Xie et al., 2022) pre-trained Swin and Swin V2 models for comparison.
Nevertheless, token merging with ToMe improves the throughput of ViT models such that ViT-L and ViT-H become comparable in speed to models of a lower tier, without scaling the number of features. Thus we display results of ViT "advancing a tier" in Tab. 3. More testing is need to see whether applying ToMe and model scaling at the same time would produce even better results.
Comparison to Token Pruning. In Tab. 4, we compare ToMe to token pruning works that use DeîT-S training?: A-VIT (Yin et al., 2022), Dynamic (Rao et al., 2021), and SP-VIT (Kong et al., 2022) with throughput measured on V100. Even though don't 't use gradient tricks such as gumbel softmax (Jang et al., 2017), add extra parameters, or use additional training tricks, we can already match the performance and exceed the throughput of existing much more complicated token pruning works. Moreover, most token pruning works are forced to use padding tokens or attention masking during training, negating the benefits of pruning in the first place. Our method, on the other hand, doesn't It suffer from this issue and we observe a 1.5x training speedup with DeiT. Interestingly, after 300 epochs the DeiT models have similar accuracy drop to our MAE trained ViT-I. (Appendix A). But we actually don't 't need to train at all: if we take an off- aff-the-shelf AugReg ViT-S model and apply the same merging schedule, we can match the performance of the DeiT models without training.
4.4 VISUALIZZATIONS
In Fig. 4, we show the input patches belonging to each merged token at the end of the network. We find that applying ToMe results in token merging that resembles something not unlike part segmentation (Chen et al., 2014). In the second image, the husky has different tokens for its legs, body, and face. The monkey in the 3rd image has different tokens for its hand, body, face, eyes, and mouth while the orange it's holding gets its own token despite it representing just one input
'This comparison is difficult as many token pruning works use different training strategies, some even claiming improvement in accuracy without valid baseline. A-V fine-tunes on top of DeiT, while Dyn starts DeiT training from existing checkpoint. We, the other hand, train from scratch.
7

---

Published as a conference paper at ICLR 2023
Figure 4: Image Visualizations. Results of merging on ImageNet- Ik val using a ViT-HMAF model trained with ToMe. Patches with the same inner and border color are merged together. Unlike pruning, ToMe can merge similar parts of the image whether they're re in the foreground or background.
patch. In cases where there multiple instances of the same class like the dung beetles in the fourth image and the Boston terriers in the last image, the same parts from all instances get merged together. Notably unlike pruning, ToMe is able merge a large number of tokens both in the background and the foreground without losing information. See more results and methodology in Appendix E.
VIDEO EXPERIMENTS
This framework of MAE plus token merging a powerful strategy across several domains. Because of its high redundancy, one of the most promising directions is video. Thus, we apply our token merging approach on Spatiotemporal MAE (Feichtenhofer et al., 2022) for video classification on Kinetics-400 (Kay et al., 2017), both by simply applying our method off-the-shelf without training and by applying our method during MAE fine tuning with the default training recipe like we did for images. Note that nothing in our method needs to change for video: we use the same code for both.
Results. In Tab. 5, we show the results of applying our method off-the-sheds and during MAE fine-tuning using VIT-L from Spatiotemporal MAE compared the relevant -tate-of-the-art on Kinetics-400 classífication: Swin (Liu et al., 2022b) pretrained on ImageNet-21k, MViTv2 (Li et al., 2022) pretrained with MaskFeats (Wei et al., 2022). and Spatiotemporal MAE the baseline. We also include token pruning work, X-VIT + ATS (Fayyaz et al., 2022), for completeness.
Amazingly, ToMe applied ViT-L with a constant schedule can match the throughput of Swin-B while performing better than MViTv2-L, even when evaluated without training. Moreover, with decreasing schedule, ViT-L MAE significantly outperforms the baseline ViT-B MAT MAF model with the same flop count with or without training, meaning ToMe is better than model scaling here. Training not necessary for a constant schedule, but does help with decreasing schedule.
ras
Throughput. In Tab. 6, display the throughput and training time of our method applied to VIT-L. With a constant schedule, we can increase throughput by 2.2x for a negligible 0.2% accuracy drop. Moreover, this setting cuts training time in half, even with the overhead of syncing across 8 gpus.
Clip Count. Because each forward pass only sees up to 2 seconds of vidco, it's standard practice to evaluate video recognition models with multiple clips. In Tab. 5, we evaluate with multiple clips (1 spatial crop, 10 temporal crops). We don't it factor the number of clips into flop count because this a hyperparameter every method can tune, usually resulting in only small differences as long as a minimum number of clips are used (i.e., 4 in this case). Thus, we choose the same number of clips as other models to compare. However, this might compensate for the information loss from token merging. In Fig. 5, we test if this is the case by sweeping over the number of clips for our method compared to the baseline ViT-L model. For r = 65, we see some degradation compared to the 4 clip sweet-spot (~0.55%) but for lower values, there's decrease compared to the baseline.
Visualization. We visualize the final tokens for each input patch over multiple frames of video in Fig. 6 using our trained ViT-L MAE model. Just like ToMe performs primitive part segmentation on images, it is actually able to perform primitive part tracking on video. The same object or part is merged into one across multiple frames of video like the ball in Fig. 6. Note that the extraneous red patch in the third frame is the reflection of the ball in the glass. More results in Appendix E.
Pes
8

---

Published as a conference paper at ICLR 2023
model
input
acc
gflops
model
clip/s
ft time (8gpu)
VIT ATS
16x224
80.0
259 X x
VIT-L MAL
7.3
1.0x
263 hrs
1.0x
ViT-B MAE
16x224
81.3
180 3 x
VIT-L
16.3
2.2x
136 hrs
0.5x
Swin-B!
32x224
82.7
282 x X
VIT- 765
MAE
24.9
3.4x
ToMe
MAE VIT-L
16x224
82.5
184x1x 10
Table 6: Video Throughput. Inference speed and fine-tuning time for method applied to video. ToMe halves training time for almost free. estimated based on 5 days of training.
VIT-L
MAE
16x224
83.2
184x1x 10
VIT-L
16x224
84.7
598 10
Swin-L
32x224
83.1
604 x 3 X
MViTv2-L
16x224
84.3
377
10
85
ToMe
84
ViT-L
16x2242
84.5
325 x
83
Setting
Baseline
ViT-L
MAE
16x224
84.5
281 x11x1
82
82
VIT-L
16x224
84.4
281x1x10
Table 5: Video Results. Results for our
Daseline
method without training ( (blue) or applied dur- ing MAE fine-tuning (gray) compared to SoTA Kinetics-400 in the same flop range. f uses ImageNet-211 pretraining while others only use Kinetics-400. Original baseline model in gray.
Oll-the-hhtt!
16
Total Clips
Figure 5: Varying Clips. ToMe closely matches the bascline's accuracy even with fewer clips.
model
mAP
gflops
sample/s
VIT-B BOKE
47.3
48.6
103
VTT-B VTT-BMAE MAE
46.2
36.3
136
VIT-B MIKE
43.1
24.7
200
VIT-B MALE
46.4"
48.6
103
VIT-B MAL
46.3
36.3
136
ViT-B MMKE 46.0
P20=
24.7
200
Table 7: Audio Results. ViT-B fine-tuned from
audio MAE pretraining on AudioSet-2M. ToMe can double the throughput of the baseline with only a 0.4% mAP drop. due training implemen- tation differences, the baseline we train performs worse than the original bascline.
Figure 6: Video Visualization. ViT-L res- with ToMe merges the same object across multiple frames of video. Here, the ball is merged into one token over its entire range of motion (red).
6 AUDIO EXPERIMENTS
We perform experiments on an Audio MAE (Huang al., 2022), where a spectrogram of the audio signal is rasterized and then fed into a ViT model. We use the ViT-B model from Huang et al. (2022) and evaluate on AudioSet-2M (Gemmeke et al., 2017). Results. Note, the metric reported is mAP instead of accuracy because of class imbalance. Due to training implementation differences, the bascline model we train has lower mAP than originally reported in Huang et al. (2022). Thus in Tab. 7, we compare ToMe without training to the original number, and ToMe with training to trained bascline. Regardless, on audio we obtain an almost 2x throughput increase with an mAP drop of only 0.4%. Full results for this experiment are in Appendix A.
7 CONCLUSION
In this work, we introduced Token Merging (ToMe) to increase the throughput of ViT models by gradually merging tokens. ToMe naturally exploits redundancy in the input, allowing its use for any modality with redundancy. the paper we explore extensive experiments on images, video, and audio, obtaining speeds and accuracies competitive with the state-of-the-a in each case.
ToMe can be viewed as a "natural" hierarchical model, similar to Swin or MViT but using pure transformer blocks. ToMe could be combined with these methods to create an entirely new type of architecture. Similarly, we focus on classification but our visualizations show potential on tasks like segmentation. Finally, ToMe works well on large models across domains and cuts down training time and memory usage, meaning it could be a core component of training huge models. We leave these as topics for future work and hope ToMe can lead to the creation of better, more efficient transformers.

---

Published as a conference paper at ICLR 2023
REFERENCES
Daniel Bolya, Cheng Yang Fu, Xiaoliang Dai, Peizhao Zhang, and Judy Hoffman. Hydra attention: Efficient attention with many heads. arXiv:2209.07484 [cs. cVs.CVJ 2022.
Y. Boykov, O. Veksler, and R. Zabih. Fast approximate energy minimization via graph cuts. TPAMI, 2001.
Xianjie Chen, Roozbeh Mottaghi, Xiaobai Liu, Sanja Fidler, Raquel Urtasun, and Alan Yuille. Detect what you can: Detecting and representing objects using holistic models and body parts. In CVPR, 2014.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv:2009. 14794 [Cs, 2020.
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. arXiv:2205. [cs.LG], 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
Xiaoyi Dong. Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, and Baining Guo. Cswin transformer: A general vision transformer backbone with cross-shaped windows. In CVPR, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x words: Transformers for image recognition scale. In ICLR, 2020.
Haoqi Fan, Bo Xiong, Karttikeya Mangalam, Yanghao Li. Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. In ICCV, 2021.
Mohsen Fayyaz, Soroush Abbasi Kouhpayegani, Farnoush Rezaci Jafari, Eric Sommerlade, Hamid Reza Vaezi Joze, Hamed Pirsiavash, and Juergen Gall. Ats: Adaptive token sampling for efficient vision transformers. In ECCV, 2022.
Christoph Feichtenhofer, Haoqi Fan. Yanghao Li, and Kaiming He. Masked autoencoders as spatiotemporal learners. NeurlPS, 2022.
Jort F Gemmeke, Daniel PW Ellis, Dylan Freedman, Aren Jansen, Wade Lawrence, Channing Moore, Manoj Plakal, and Marvin Ritter. Audio set: An ontology and human-labeled dataset for audio events. In /CASSP, 2017.
Saurabh Goyal, Anamitra Roy Choudhury, Saurabh Raje, Venkatesan Chakaravarthy, Yogish Sab- harwal, and Ashish Verma. Power-bert: Accelerating bert inference via progressive word-vector climination. /CML, 2020.
Benjamin Graham, Alaaeldin El-Nouby, Hugo Touvron, Pierre Stock, Armand Joulin, Hervé Jégou, and Matthijs Douze. Levit: vision transformer in convnet's clothing for faster inference. In ICCV, 2021.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. CVPR, 2022.
Po- Yao Huang, Hu Xu, Juncheng Li, Alexei Baevski, Michael Auli, Wojciech Galuba, Florian Metze, and Christoph Feichtenhofer. Masked autoencoders that listen. NeurIPS, 2022.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softman ICLR, 2017.
Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanarasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, et al. The kinetics human action video dataset. arXiv: 2017.
10

---

Published as a conference paper at ICLR 2023
Gyuwan Kim and Kyunghyun Cho. Length-adaptivy transformer: Train once with length drop, anytime with search. arXiv:2010.07003 [cs. (cs.CL], 2020.
Schoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. Learned token pruning for transformers. arXiv:2107.0000 109000001000 [es. CLJ, 2021.
Zhenglun Kong, Peiyan Dong, Xiaolong Ma, Xin Meng, Wei Niu, Mengshu Sun, Bin Ren, Minghai Qin, Hao Tang, and Yanzhi Wang. Spvit: Enabling faster vision transformers via soft token pruning. In ECCV, 2022.
Carlos Lassance, Maroua Maachou, Joohee Park, and Stéphane Clinchant. A study on token pruning for colbert. 1rX1v:2112.065400 [cs. CL), 2021.
Yanghao Li, Chao- Yuan Wu, Haoqi Fan, Karttikeya Mangalam, Bo Xiong, Jitendra Malik, and Christoph Feichtenhofer. Mvitv2: Improved multiscale vision transformers for classification and detection. In CVPR, 2022.
Youwei Liang, Chongjian Ge, Zhan Tong, Yibing Song, Jue Wang, and Pengtao Xie. Not all patches are what you need: Expediting vision transformers via token reorganizations. ICLR, 2022.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao. Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In CVPR, 2022a.
Ze Liu, Jia Ning, Yue Cao. Yixuan Wei, Zheng Zhang, Stephen Lin, and Han Hu. Video swin transformer. In CVPR, 2022b.
Stuart Lloyd. Least squares quantization in pem. IEEE transactions on information theory, 1982.
Dmitrii Marin, Jen-Hao Rick Chang, Anurag Ranjan, Anish Prabhu, Mohammad Rastegari, and Oncel Tuzel. Token pooling in vision transformers. (:Xiv:2110.03880 CVJ, 2021.
Sachin Mehta and Mohammad Rastegari. Mobilevit: light, friendly vision transformer. In ICLR, 2021.
and mobile-
Lingchen Meng, Hengduo Li, Bor- Bor-Chun Chen, Shiyi Lan, Zuxuan Wu, Yu- Yu-Gang Jiang, and Ser-Nam Lim. Adavit: Adaptive vision transformers for efficient image recognition. In CVPR. 2022.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? NeurIPS, 2019.
Zizheng Pan, Bohan Zhuang, Haoyu He, Jing Liu, and Jianfei Cai. Less is more: Pay less attention in vision transformers. In AAAI, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In NeurIPS. 2019.
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. NeurIPS, 2021.
Michael Ryoo, AJ Piergiovanni, Anurag Arnab, Mostafa Dehghani, and Anelia Angelova. Token- learner: Adaptive space-time tokenization for videos. In NeurlPS, 2021.
Zhuoran Shen, Mingyuan Zhang. Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities. In WACV, 2021.
Mannat Singh, Laura Gustafson, Aaron Adcock, Vinicius de Freitas Reis, Bugra Gedik, Raj Prateck Kosaraju, Dhruv Mahajan, Ross Girshick. Piotr Dollár, and Laurens van der Maaten. Revisiting weakly supervised pre-training of visual perception models. In CVPR, 2022.
11

---

Published as a conference paper at ICLR 2023
Zhuoran Song. Yihong Xu, Zhezhi He, Li Jiang, Naifeng Jing, and Xiaoyao Liang. Cp-vit: Cascade vision transformer pruning via progressive sparsity prediction. arXiv preprint arXiv:2203.04570 2022.
Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. TMLR, 2022.
Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In ICML, 2019.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. ICML, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurIPS, 2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In ACL, 2019.
Sinong Wang. Belinda Z Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. arXiv:2006.0477 04768 [cs.LG], 2020.
Chen Wei, Haoqi Fan, Saining Xie, Chao- Wu, Alan Yuille, and Christoph Feichtenhofer. Masked feature prediction for self-supervised visual pre-training. In CVPR, 2022.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In /CML, 2022.
Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. In CVPR, 2022.
Jiarui Xu, Shalini De Mello, Sifei Liu, Wonmin Byeon, Thomas Breuel, Jan Kautz, and Xiaolong Wang. Groupvit: Semantic segmentation emerges from text supervision. In CVPR, 2022.
Hongxu Yin, Arash Vahdat, Jose Alvarez, Arun Mallya, Jan Kautz, and Pavlo Molchanov. A-ViT: Adaptive tokens for efficient vision transformer. In CVPR, 2022.
Hao Yu and Jianxin Wu. A unified pruning framework for vision transformers. a1X11:2111.15127 [cs.CV], 2021.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In CVPR, 2021.
12

---

Published as a conference paper at ICLR 2023
A
FULL RESULTS
Results for plots and tables in the main paper. For all results, im/s indicates throughput and "speed" indicates improvement over the baseline. All throughputs are measured on a V100, but the actual values may differ little from the main paper the model may have been benchmarked on different machine. However, all results in the main paper use the same machine for throughput evaluation.
A.1 IMAGES
For each ImageNet- Ik model, display full results here.
A.1.1
AUGREG MODELS
Full results listed in Tab. 8. We make no special changes to any of these models.
model
acc
drop
im/s
speed
model
acc
drop
im/s
speed
ViT-B/16
84.57
0.00
309
1.00
VIT-S/16
81.41
0.00
953
ViT-B/16
84.61
0.03
311
1.01
ViT-S/16
81.37
-0.04
967
1.01
ViT-B/16
84.52
-0.05
322
1.04
VIT-S/16
81.35
-0.06
1001
1.05
ViT-B/16
84.39
-0.18
334
1.08
ViT-S/16
81.30
-0.10
1036
1.09
VIT-B/16
84.39
-0.18
346
1.12
VIT-S/16
81.24
-0.17
1072
1.12
ViT-B/16
84.24
-0.33
360
1.17
VYT-S/16
81.12
-0.29
1114
1.17
ViT-B/16
84.17
-0.40
373
1.21
VIT-S/16
81.02
-0.38
1151
1.21
ViT-B/16
83.99
-0.58
391
1.27
VIT-S/16
80.94
-0.46
1202
1.26
ViT-B/16
83.94
-0.63
406
VIT-S/16
80.78
-0.63
1247
ViT-B/16
83.73
-0.84
425
1.37
VIT-S/16
80.53
-0.88
1302
1.37
ViT-B/16
10
83.37
-1.21
443
1.44
VIT-S/16
10
80.33
-1.08
1364
1.43
ViT-B/16
83.28
-1.30
464
1.50
VIT-S/16
80.06
-1.35
1424
1.49
VIT-B/16
12
82.86
-1.72
488
1.58
VIT-S/16
12
79.60
-1.81
1487
1.56
ViT-B/16
82.60
-1.98
511
VIT-S/16
79.30
-2.10
1564
ViT-B/16
14
82.04
-2,54
540
1.75
ViT-S/16
14
78.89
-2.52
1646
1.73
ViT-B/16
15
81.39
-3.18
571
1.85
VIT-S/16
15
78.14
-3.26
1738
1.82
ViT-B/16
80.38
-4.20
603
1.95
VIT-S/16
16
77.01
-4,40
1836
1.93
VIT-1/16
85.82
0.00
95
1.00
VIT-T//66
75.50
0.00
2558
1.00
ViT-L/16
85.80
-0.02
100
1.05
VIT-TV/16
75.39
-0.10
2471
0.97
VIT-/116
85.70
-0.12
107
1.13
VIT-TV/16
75.40
-0.09
2551
1,00
ViT-L/16
85.58
-0.24
VIT-TW/6
75.34
-0.15
2651
VIT-L/16
85.37
-0.45
125
1.32
VIT-TV/16
75.27
-0.23
2735
1.07
ViT-L/16
85.17
-0.65
137
1.44
VIT-Ti/16
75.18
-0.32
2849
1.11
ViT-L/16
84.71
-1.11
VIT-Ti/16
75.06
-0.44
2944
1.15
ViT-L/16
7
84.26
-1.56
167
1.76
VIT-TV/6
74.88
-0.62
3074
1.20
ViT-L/16
83.55
-2.27
186
1.96
VIT-Ti/16
74.76
-0.74
3191
1.25
VIT-TV/16
74.50
-1.00
3331
1.30
3T4
86.92
0.00
28
1.00
VIT-T//66
10
74.26
-1.24
3469
1.36
ViT-L/16@384
86.87
-0.04
31
1.12
VIT-Ti/16
73.76
-1.73
3629
1.42
ViT-L/16@388 384
10
86.85
-0.07
36
1.28
VIT-Ti/16
12
73.53
-1.97
3792
1.48
ViT-L/16@384
15
86.75
-0.17
42
1.50
VIT-Ti/16
13
73.04
-2.46
3980
1.56
ViT-L/16@384
20
86.53
-0.39
50
1.78
VIT-TV/16
14
72.65
-2.84
4175
1.63
VIT-L/16@384
23
86.14
-0.78
56
2.00
VIT-T//66
15
71.80
-3.70
4396
1.72
VIT-TV/6
16
70.79
-4.71
4610
1.80
(a) VIT-B, VIT-L, and VIT-L 384px.
(b) VIT-S and VIT-TI.
Table 8: Full AugReg Off-the-Shelf Results. Results without training. The original off-the-shelf models are listed in gray. We include the blue model in Tab. 4.
A. SWAG MODELS
Full results listed in Tab. 9. Again, make special changes for these models.
13

---

Published as a conference paper at ICLR 2023
model
drop
im/s
speed
VIT-B/16 @ 384
85,30
0.00
85.7
1.00
VIT-B/16 @ 384
85.27
-0.03
88.6
1.03
VIT-B/16 @ 384
10
85.21
-0.09
94.6
1.10
VIT-B/16 @ 384
15
85.18
-0.13
101.7
1.19
VIT-B/16 384
85.09
-0.22
109.1
1.27
VIT-B/16 @ 384
25
85.03
-0.27
118.3
1.38
VIT-B/16 @ 384
30
84.98
-0.33
127.7
1.49
VIT-B/16 @ 384
35
84.90
-0.40
139.6
1.63
VIT-B/16 @ 384
40
84.89
-0.41
152.8
1,78
VIT-B/16 @ 384
84.59
-0.71
167.7
1.96
VIT-L/16 @ 512
88.06
0.00
12.8
1.00
VIT-L/16 @ 512
88.02
-0.04
13.7
1.06
VIT-L/16 @ 512
10
87.98
-0,08
14.7
1.14
VYT-L/16 512
87.95
-0.11
16.1
1.25
VIT-L/16 @ 512
20
87.96
-0.10
17.5
1.36
VIT-L/16 @ 512
25
87.87
-0.19
19.3
1.51
VIT-L/16 512
87.89
1.66
VIT-L/16 @ 512
35
87.82
-0.24
23.7
1.84
VIT-L/16 @ 512
40
87.80
40,26
26.3
2.05
VIT-H/14 @ 518
88.55
0.00
4.7
1.00
VYT-H/14 @ 518
88.53
-0.02
5.1
1.07
VIT-H/114 @ 518
10
88.44
-0.11
5.5
1.16
VIT-H/14 @ 518
15
88.49
-0.06
6.0
1.26
VIT-H/14 @ 518
88.46
-0.09
1.38
VIT-H/14 @ 518
25
88.39
-0.16
7.2
1.51
VYT-H/14 @ 518
30
88.34
-0.21
7.9
1.67
VIT-H/14 @ 518
35
88.35
-0.19
8.8
1.85
VIT-H/14 @ 518
40
88.25
-0.30
9.8
2.06
Table Full SWAG Off-the-Shelf Results, Results are without training. models are listed in gray. We mention the blue models in the abstract.
The original off-the-shelf
A.1.3 MAE MODELS
We evaluate MAE models both off the shelf and trained with ToMe in Tab. 10. For off-the-shelf evaluation we disable proportional attention as noted in Sec. 4.1, but we enable for the trained models. Note that we compare to baselines we trained ourselves, which may slightly underperform the official baselines (for ViT-L). When training, we fine-tune from official pretrained weights and use the original training recipe. Unlike prior work, we intend for ToMe to replace standard training, not augment it, in order to receive the benefits of faster training times and less memory usage.
A.1.4 DEIT MODELS
We present DeiT results in Tab. 11. For DeiT, we train from scratch with the default training recipe for 300 epochs. Unlike other token pruning works, don't use any tricks such starting from an existing checkpoint or fine-tuning. Note that for merging, in addition to not merging the class token, don't 't merge the distillation token. In Tab. 11, don't train for all values of just the baseline r 0 and those between 8 and 16. r = for DeiT-S didn't finish training.
A.2
VIDBO
We run the ViT-L model from Feichtenhofer et al. (2022) off the shelf. In Tab. 12, we show the results of this experiment by sweeping over For each setting, we evaluate with I spatial crop and 10 temporal clips. Note that the original baseline is evaluated with spatial crops and 7 temporal clips, while we re-evaluated it with 1 10. Thus, the baseline has slightly lower accuracy than the original paper. Like with images, for these off-thor -he-hhe- MAE pretrained models we don't use proportional attention.
14

---

Published as a conference paper at ICLR 2023
model
drop
im/s
speed
model
drop
im/s
speed
VIT-B/16
0
83.62
0.00
305
1.00
VIT-B/16
0
83.62
0.00
309
1.00
ViT-B/16
83.55
-0.07
1.01
ViT-B/16
83.60
-0.02
1.01
ViT-B/16
2
83.50
-0.12
317
1.04
ViT-B/16
2
83.52
-0.10
322
1.04
VIT-B/16
3
83.44
-0.18
329
1.08
ViT-B/16
3
83.49
-0.13
334
1.08
ViT-B/16
83.39
-0.23
1.12
ViT-B/16
83.43
-0.19
346
1.12
VIT-B/16
5
83.36
-0.26
355
1,16
ViT-B/16
5
83.43
-0.20
360
1.17
ViT-B/16
6
83.22
-0.40
368
1.21
ViT-B/16
83.39
-0.24
373
1.21
VIT-B/16
7
83.01
-0.61
384
1.26
ViT-B/16
7
83.31
-0.31
391
1.27
VIT-B/16
8 82.93
-0.69
400
1.31
ViT-B/16
8
83.16
-0.46
406
1.31
ViT-B/16
82.69
-0.93
ViT-B/16
83.20
-0.42
ViT-B/16
10 82.52
-1.10
436
1,43
ViT-B/16
10
83.01
-0.61
443
1.44
VIT-B/16
11 82.18
-1.44
457
1.50
ViT-B/16
11
82.94
-0.68
464
1.50
ViT-B/16
81.92 12
-1.70
479
1.57
ViT-B/16
12
82.65
-0.97
488
1.58
VIT-B/16
13
81.41
-2.21
504
1.65
VIT-B/16
13
82.55
-1.07
511
1.65
ViT-B/16
14
80.85
-2.77
532
1.74
ViT-B/16
14
82.35
-1.28
540
1.75
VIT-B/16
15
80.01
-3.61
561
1.84
ViT-B/16
15
82.12
-1.50
571
1.85
VIT-B/16
16
78.75
-4.87
592
1.94
VIT-B/16
16
81.91
-1.71
603
1.95
VIT-L/16
0
85.66
0.00
93
1.00
VIT-L/16
0
85.66
0.00
93
1.00
ViT-L/16
85.63
-0.03
1.06
ViT-L/16
85.79
0.13
1.05
ViT-L/16
85.59
-0.07
105
1.13
ViT-L/16
2
85.69
0.03
105
1.12
VIT-L/16
3
85.51
-0.15
114
1.23
ViT-L/16
85.39
-0.27
123
1.33
ViT-L/16
85.54
-0.12
123
1.32
ViT-L/16
5
85.26
-0.40
134
1.45
ViT-L/16
5
85.59
-0.07
134
1.44
ViT-L/16
6
85.03
-0.63
147
1.59
VIT-L/16
6
85.37
-0.28
147
1.58
ViT-L/16
84.55
-1.11
1.76
VIT-L/16
7
85.23
-0.43
1.75
VIT-L/I6
8
83.92
-1.74
183
1.97
VIT-L/I6
8
85,05
-0.61
183
1.96
VIT-L/16
8
84.16
-1.50
241
2.58
VIT-H/14
0
86.88
0.00
35
1.00
VIT-H/14
0
86.88
0.00
35
1.00
VIT-H/14
86.86
-0.02
VIT-H/14
86.76
-0.12
VIT-H/14
2
86.82
-0.06
40
1.14
ViT-H/14
2
86.73
-0.15
40
1.14
VIT-H/14
86.80
-0.08
43
1.24
VIT-H/14
3 86.79
-0.09
43
1.24
ViT-H/14
86.69
-0.19
1.34
ViT-H/14
86.82
-0.06
VIT-H/14
5
86.52
-0.36
51
1,47
VIT-H/14
5
86.74
-0.14
51
1,47
ViT-H/14
6
86.31
-0.58
1.62
ViT-H/14
6
86.51
-0.37
1.62
ViT-H/14
7
85.94
-0.94
63
1.81
ViT-H/14
7
86.47
-0.41
63
1.81
ViT-H/14
7%
86.06
-0.82
81
2.34
(a) Applied without training.
(b) Applied during MAE fine-tuning.
Table 10: Full MAE Results. Results are with and without training. The original baseline models we train are listed in gray. We include the gray models Tab. 3.
A.3 AUDIO
Full results for our audio experiments can be found in Tab. 13. We used the model from Huang et al. (2022) to evaluate off-the-shelf. However, for training we train with our own implementation that's different from the paper. For this reason, in Tab. 13, we list two different baselines (one from the original paper, and the other trained by us). In this case, we don't use proportional attention during off off-the-shelf evaluation training.
B HYPERPARAMETERS
In Tab. 14, perform limited hyperparameter search parameters that would be affected by applying ToMe: layer decay, drop path, and the number of epochs.
Layer decay reduces learning rate based the layer of the network. Since ToMe gradually reduces the number of tokens, the size of gradient updates in later layers might already lower without layer decay. However, find that it's necessary change that parameter.
15

---

Published as a conference paper at ICLR 2023
drop
gflops
speedup
drop
grlops
speedup
0
79.96
4.61
0
71.84
1.26
79.68
-0.28
3,43
1.34x
71.69
-0.15
0.93
1.35x
9 79.69
-0.27
3.28
1,40x
9
71.63
-0.21
0.89
1.41x
10
79.59
-0.37
3.14
1,47x
10
71.24
-0.59
0.85
1.47x
71.40
-0.44
0.81
1.55x
12
79.49
-0.47
2.85
1.62x
12
71.35
-0,48
0.78
1.62x
13
79.42
-0.54
2.71
1.70x
13
71.27
-0.57
0.74
1.71x
14
79.24
-0.72
2.57
1.79x
14
71.16
-0.67
0.70
1.80x
15
79.21
-0.75
2.43
1,90x
15
71.02
-0.82
0.66
1.90x
16
79.13
2.30
2.00x
70.74
-1.09
0.63
2.01x
(a) DeiT-S. We include the gray result to compare against token pruning methods, though other speed- accuracy trade-offs may be more desirable. Note the baseline we train is slightly more accurate than the original.
(b) DeïT-Ti. We omit DeiT-Ti results from the main
paper because we are not able to reproduce the base- line accuracy in SPVIT (Kong al., 2022).
Table 11: Full DeiT Results. Results from rumning standard 300 epoch DeiT-S and DeiT-Ti training run with the official codebase on ImageNet-Ik: Note that make no changes to hyperparameters or the model other than to add token merging. Not only is our method much simpler than other token reduction approaches, but this speed-up is observed both during inference and training. Speed-up and accuracy drop is relative to the baseline (gray).
model
r
top-1
top-5
gflops
clips/s
speed
ViT-L 16x4
0
84.66
96.47
598
7.3
1.00
ViT-L 16x4
10
84.70
96.42
545
8.0
1.10
VIT-L 16x4
20
84.82
96.43
493
9.0
1.23
ViT-L 16x4
30
84.82
96.40
443
10.1
1.38
VIT-L 16x4
40
84.86
96.40
394
11.4
1.56
ViT-L 16x4
45
84.82
96.43
371
12.3
1.68
ViT-L 16x4
50
84.69
96.45
348
13.1
1.79
VIT-L 16x4
55
84.60
96.32
325
14.1
1.92
ViT-L 16x4
60
84.51
96.30
303
15.1
2.06
VIT-L 16x4
64
84.41
96.33
286
16.0
2.19
ViT-L 16x4
65
84.51
96.33
281
16.3
2.22
VIT-L 16x4
65%
82.49
95.65
184
24.9
3.39
Table 12: Full Video Off-the-Shelf Results. Results are without training. The original model is listed gray. We include the blue models Tab. 5. Evaluation is 10 (this is not factored into the flop count). We include both top-1 and top-5 accuracy Kinetics-400.
mAP
drop
sample/s
speed
mAP
drop
sample/s
speed
47.29
0.00
103.3
1.00
46.43
0.00
103.3
1.00
47.16
-0.13
108.7
1.05
46.21
-0.23
108.7
1.05
10
46.92
-0.37
115.6
1.12
10
46.36
-0.08
115.6
1.12
15
46.60
-0.69
125.8
1.22
15
46.17
-0.27
125.8
1.22
20
46.21
-1.08
136.2
1.32
20
46.29
-0.14
136.2
1.32
45.70
-1.59
148.9
1.44
46.09
-0.34
148.9
1.44
30
45.00
-2.29
162.7
1.58
30
46.09
-0.34
162.7
1.58
35
44.18
-3.11
181.3
1.76
35
46.12
-0.31
181.3
1.76
43.09
-4200
199.9
1.94
45.96
-0.47
199.9
1.94
(a) VIT-B on AudioSet-2M without training.
(b) ViT-B on AudioSet-2M with training.
Table 13: Full Audio Results. Results are with and without training. The original baseline model (left) and the baseline we train (right) are listed in gray. We include the blue and gray models in Tab. 7.
Drop path randomly drops out entire attention and MLP blocks with some probability. This has the effect of regularizing layers so that they don't rely single block. Because the K matrix
16

---

Published as a conference paper at ICLR 2023
from blocks that could be dropped out, we test the value of this parameter. Again, we find this not necessary to change.
We also perform the same experiments video except with just layer decay and the number of epochs, testing whether ToMe requires increasing the number epochs (due to seeing fewer tokens overall). And again, the default parameters work the best.
layer decay
drop path
layer decay
epochs
0.55
0.05
81.59
0.875
50
82.83
0.55
0.20
81.31
0.825
50
83.52
0.65
0.05
81.73
0.750
50
83.57
0.65
0.10
81.83
0.750
75
83.35
0.65
0.20
81.76
0.750
100
82.33
0.75
0.05
81.68
0.75
0.10
81.77
0.75
0.20
81.82
0.85
0.05
81.35
0.85
0.10
81.54
0.85
0,20
81.46
(a) Image Fine-Tuning We sweep over ViT-B/16 MAE fine-tuning ImageNet-Ik with r 16 and find that the default parameters from the official code release work the best.
(b) Video Fine-Tuning, We test some additional pa- rameters for ViT-L spatiotemporal MAE fine-tuning on K400 with 65. The defaults also work best here. Accuracy is for 3x5 evaluation. Note that this fine-tuning started from a MAE pre- trained model trained for half its schedule, these numbers aren't comparable the numhers in Tab. 5.
Table 14: Training Hyperparameters don't 't need to be updated when training with token merging. We perform a sweep over relevant hyperparmaters that might be affected by token merging for images and video. The default setting, marked in purple, already has the highest accuracy.
C MERGING SCHEDULE
Paun Thraarpuu 954
Mess Throuppeet 1603
Nox
Moxi Thasgggen
Troighostt554
Mex Throuphpt
Net
1502
Figure 7: Merging Schedule. The average the top 100 merging schedules for the experiment in Fig. 2. For each throughput range, we find the highest accuracy schedules and average their number of tokens merged per layer. For lower throughputs, merging more tokens at the end better. For higher throughputs, constant schedule becomes the best. Then for even higher throughputs, linearly decreasing schedule works well.
In Fig. 7, we plot the average number of tokens merged in each layer for the most accurate random samples in Fig. 2. Around throughputs of 1600-1800, the best schedule is close to constant, which is why constant is close to optimal in this range. For throughputs beyond that, however, decreasing schedule is best. For this, reason we define a linearly decreasing schedule in addition to a constant schedule in the main paper.
17

---

Published as a conference paper at ICLR 2023
D IMPLEMENTATION
The following is an implementation of our "bipartite soft matching" in PyTorch (Paszke et al., 2019):
def setrritt__of__mtteeeee
(k: torch. Tensor, r: int) torch. Tensor: **" Input is k from attention, size [batch, tokens, channels]. "*" k k norommddm==11 keepdim=True)
a, = k[..., ::2, :], k[..., 1::2, :]
scores e b. transpose (-1, -2)
scores[.... 0, :] -math. inf # don't merge cls token
node_max, node_idx scores max (dim=-1) edge_idx node_max meaxxargeee argsort (dim=-1, descending-Truett.... None]
unm_idx edge_ idx ..... ri, :] # Unmerged Tokens
src_idx edge_ide idx [.... :r, :] Merged Tokens
dst_idx node_idx Nonel. gather (dim=-2, index=src_idx)
unm_idx unm_idx sort
Sort cls token back to idx 0
def merge (x: torch Tensor) -> torch. Tensor: ..n Input sre,
shape [batch, tokens, channels]. ..*
x(..., ::2, :l, x(..., 1::2, :]
n, t1, src.
unm src. gather (dim=-2, index= unm_xiunm_idx idx. expand (n, t1 r, c))
src src gather (dim=-2, index=sro_ide expand (n, r, c)) dst dst ada dst_idx expand (n, r, c), src)
return torch. cat ([unm, dst], dim=-2)
return merge
This returns a lambda function that can be applied to any matrix or vector (i.e. to merge features, to calculate token size, calculate source patches). Note how this done all at once parallel- there are no sequential loops.
E MORE VISUALIZATION
To create the visualizations in Fig. 4 and Fig. 6, we follow each final merged token back to its original input patches. Then for each token, we color its input patches with the average color in that region. To make sure different tokens are distinct from each other, we also assign each token random border color. Note that tokens do not necessarily represent contiguous input regions. The only spatial signal ToMe has comes from the position encodings.
Fig. we present several more examples of merging on images as continuation of Fig, ToMe's propensity for part and object segmentation appears time and time again across many different images.
In Fig. 9, we also display more results of ToMe performing object tracking on video. Note that in (Feichtenhoferr et al., 2022), each token represents more than one frame. Namely, the patch size is 2 x 16 x 16 and thus frames of video correspond to each token. We plot the first frame from the two, because find that more closely matches the merged tokens.
18

---

Published as a conference paper at ICLR 2023
5:3
Figure 8: More visualization on images. Continuation of Fig. 4.
19

---

Published as a conference paper at ICLR 2023
Figure 9: More visualization on video. Continuation of Fig. 6. In each clip, we highlight an instance or part being merged one token across frames (red). Clips are from Kinetics-400 val.
20


Reviews:
[{"review_id": "llM4mrezNH3", "reviewer": "Reviewer_MRoC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a simple method to merge tokens within ViTs with or without training, increasing the throughput by up to 2x across various image, video, and audio tasks. ToMe inherits the spirit of simplicity and can server as a drop-in replacement to increase the training speed or inference speed.", "summary_of_the_review": "In short, the idea is interesting and work for ViTs. A lot of visualizations also intuitively explain and validate the idea. Thus I recommend to accept this work.\n\nAs for its generalizability, more experiments on both small models and large models can be considered as the cost is not so high.", "strength_and_weaknesses": "Strength:\n\n* The idea is simple yet effective and can benefit both ViT training and inference a lot.\n* The authors clearly have deep insights about the (larger or smaller) ViT designs. The introduction of the literature is clear and informative.\n* Using matching instead of clustering algorithm makes sense to me, results also demonstrate the expected throughput achievements.\n\nWeakness:\n* More direct comparisons are desired, e.g., in Tab. 3, you compare ViT-L with other ViTs. How about comparing them apple-to-apple. E.g., ViT-L w/ ToMe vs. ViT-L (attached in appendix; better to remove ahead), MViTv2 w/ ToMe vs. MViTv2, Swin w/ ToMe vs. Swin, etc.\n* It may be easier to improve the throughput for large ViT models without hurting the model accuracy a lot. How about smaller ones, e.g., LeViT?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "lUqRF8DVN9", "reviewer": "Reviewer_vh1A", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper works on improving the efficiency of a ViT model by merging tokens each layer. Specifically, the authors first uniformly split ViT tokens into two sets and merge (i.e., replace two tokens with their mean feature) the top k similar pairs across sets in each layer. The merge runs in negligible time and can work on test time without retraining. Experiments on image, video, and audio data show the proposed method can reduce ~50% tokens within <0.5% accuracy drop.", "summary_of_the_review": "This paper has a simple and effective idea on an important problem, and has great application values. All the technical contributions are well motivated and well ablated. The state-of-the-art comparisons to existing efficient transformer methods (Table 4) are currently a little bit below expectation, but this does not undermine the significance of this work, especially given that the model can work without retraining. My current recommendation is accept, and am happy to raise to strong accept if the authors address my concerns during the discussion session.", "strength_and_weaknesses": "Strengths\n\n+ The proposed method of greedily-merging tokens is very elegant and makes a lot of sense to me, especially that the method can work without retraining. This makes the proposed method easy to use for many off-the-shelf transformer-based models.\n\n+ The paper shows results on three different modalities: images, videos, and audio, and all displayed competitive results. This shows the proposed method is general and robust.\n\n+ The ablation studies in Table 1 are comprehensive, and show the proposed method is robust under different hyer-parameters.\n\n+ The qualitative examples in Figure 4 are interesting. Please make sure they are not cherry-picked.\n\n+ The paper is well motivated, well-written, and easy to understand.\n\nWeaknesses\n\n- The comparison to existing efficient transformers in Table 4 are not as exciting as expected, especially given that DynamicViT is also very simple. Most of the advances are in training speed instead of on speed-accuracy trade-off.\n\n- It will be good to include other efficient transformer methods for videos and audio to provide more context (if there is any) about the numbers.\n\n- [Minor] last line of the code in Appendix D. Should it be `return merge(k)`.?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HbmcYtYIKRX", "reviewer": "Reviewer_eArF", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": " - This work proposes to reduce the computational redundancy in vision transformers (ViT) with Token Merging, without significantly changing the original architecture. \n\n - With the introduced ToMe module, similar tokens are gradually combined with a simplified matching algorithm. The author also perform studies on the comparison with pruning.\n\n - Extensive experiments on image and video recognition are conducted, across a variety of transformer families. ", "summary_of_the_review": " - Reducing the computational burdens in vision transformers is a non-trivial task. Given the simplicity and effectiveness of ToMe, I believe it's beneficial in further research.", "strength_and_weaknesses": "#### **Strength**\n\n - The problem, reducing the computational burden in ViTs, has been an important issue in developing and deploying transformers. \n\n - While redundancy in self-attention has already been investigated before, there are few solutions except for pruning-based methods, which require additional attention. The simple strategy in this work shows promising results, even without being trained. \n\n - The authors perform concrete studies and experiments across different tasks and architectures, showing the reliability ToMe.\n\n#### **Weakness**\n\n - This method could be viewed as a parallel implementation of the matching algorithm. We expect more discussion with previous slower variants, such as the clustering algorithm in [1].\n\n - The authors mentioned that the method could be used in training to reduce training complexity. One more common use case here is to adopt such idea to recent MIM-based methods such as MAE, where longer pre-training epochs are required.  However, the authors have only used it in fine-tuning MAE pre-trained checkpoints.\n\n[1] Token Pooling in Vision Transformers. arxiv 2110.03860.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "0NpuyHhVTs", "reviewer": "Reviewer_Gshq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents a token merging strategy for efficient inference of ViTs. A Bipartite Soft Matching method is proposed. It divides all the tokens into two parts, and draws an edge from each token in the first part to its most similar token in the second part according to the cosine similarity between pairs of tokens. The model efficiency can then be decided by how many similar tokens should be merged. \n\nDifferent from most previous works that aim at compressing ViTs with training, the approach presented in this paper can be conducted in the inference stage directly. Thorough experiments are performed on images, video, and audio, which reflect the effectiveness of the proposed approach.", "summary_of_the_review": "The novelty of this paper is clear but not that significant to match the score of 8. The experiments are thorough. I think this paper deserves a score of 6.", "strength_and_weaknesses": "Strength\n- The motivation of this paper is clear. For a large ViT model, there is no need to keep all the tokens involved in the decision making as some of them may contain useless information. Discarding some of the tokens that are similar to others can efficiently reduce the inference time.\n\n- The proposed approach is simple and easy to follow.  It can be added in either the training or inference stage. For users who want to use large-scale models but are with limited computational resources, this paper provides an appropriate tool to conquer this issue.\n\n- Experiments show that the performance does not drop much when r is small but the speedup is clear. Moreover, compared to other ViT based models, this paper also shows its advantages.\n\n- Experiments on all images, video, and audio show that the proposed approach indeed behaves well.\n\n- I really like the visualizations in Figure 4. It seems that for regions with similar textures or patterns, they can be merged for efficient inference.\n\nWeaknesses\n- Like the author said, the vanilla ViT shows great potential in visual recognition especially after the emerging of MIM-based methods, like MAE and the proposed approach works well with the vanilla ViT. However, in some cases, a pyramid structure of ViT is required especially for downstream tasks, like semantic segmentation. It would be better if the approach could be utilized in ViTs, like Swin Transformer.\n\n- From the experimental results, we can see that for large models, increasing the value of r does not lead to too much performance drop. However, for small-sized or even tiny-sized models, the performance drops much.  Though the reason is clear but I still think it should be mentioned explicitly.\n\n- From Figure 3(a), it seems that after compressing the ViT-L model, the trade-off between the performance and latency is not as good as ViT-B without compression.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%