Token Merging: Your ViT But Faster

In plain words

A vision transformer cuts a picture into many small pieces, called tokens, and works on all of them in every layer. More tokens means more computation, so the model gets slower. Earlier speed-ups mostly deleted tokens, which throws away information.

This method merges instead. In every layer it finds pairs of tokens that look alike and blends each pair into one token. The matching is cheap, so it does not slow the model down much, and the blending keeps the information instead of discarding it.

You can add it to a model that is already trained, with no retraining. You can also use it during training to make training faster. The authors test it on images, video and audio.

TL;DR

Token Merging (ToMe) is a parameter-free step inserted into each block of a ViT. It merges the r most similar token pairs per layer, using cosine similarity on attention keys §3 §4.1. Off the shelf, ViT-L/16 MAE with a constant r=8 runs at 183 im/s instead of 93 im/s, at 83.92% instead of 85.66% top-1 on ImageNet-1k app. A. Trained with ToMe, video fine-tuning takes about half the time tab. 6. All numbers are author-reported.

Problem

Token pruning removes tokens to save compute, but the lost information limits how far the count can be reduced. Most pruning methods need retraining, some add parameters, few speed up training, and methods that prune a different count per image make batched inference hard §1. Pruning 98% of tokens without training loses information, and random or attention-based pruning both fall well short of merging tab. 2 §4.1. Clustering such as k-means is accurate but slow and does not bound how many tokens merge into one group §3 §4.1. The method assumes a plain ViT whose blocks each consist of attention followed by an MLP, and it assumes a fixed token count per layer for every input §3.

Method

ToMe reduces the token count by a fixed number r per layer rather than by a ratio, so the same schedule works for every image §3 "Strategy". The merge runs between the attention and MLP branches of each block, which lets information from merged tokens flow and lets attention features decide what to merge §3 fig. 1.

1. Similarity. Tokens are compared with cosine similarity on the attention keys K, averaged over heads, not on raw features §3 "Token Similarity" §4.1 tab. 1.
2. Bipartite soft matching. Split tokens into two sets A and B of about equal size, alternating in position. Give each token in A one edge to its most similar token in B. Keep the r strongest edges. Merge those r A tokens into their B partners by size-weighted averaging. Concatenate the remaining A tokens with B §3 steps 1–5 tab. 1.
3. Size tracking. A merged token records how many input patches it represents. Merging averages tokens weighted by that size §3 "Tracking Token Size" §4.1 tab. 1.
4. Proportional attention. Attention logits get log(size) added before the softmax, which behaves as if each merged key had been repeated once per patch it represents §3 eq. 1.
5. Training. Merging is treated as a pooling operation. Gradients flow through the merged tokens as through average pooling. No Gumbel-softmax or other gradient trick is used §3 "Training with Merging".
6. Schedules. The constant schedule removes r tokens in every layer. The decreasing schedule removes 2r in the first layer and 0 in the last, interpolating linearly in between §4.1 eq. 2 eq. 3.

\mathbf{A} = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d}} + \log \mathbf{s}\right) eq. 1

Constant schedule, r per layer eq. 2. Decreasing schedule, 2r falling to 0 per layer eq. 3. The paper says both remove rL tokens in total over L layers §4.1.

# One ViT block with token merging  [§3 "Strategy"] [fig. 1]
# tokens (B, N, D), sizes (B, N, 1); r = tokens merged in this block

def tome_block(tokens, sizes, r):
    # Attention runs first; proportional attention adds log(size) [eq. 1]
    attended, keys = attention(norm(tokens), log(sizes))  # keys (B, N, C)
    tokens = tokens + attended
    # Merge between attention and MLP  [§3] [code: tome/patch/timm.py#L34-L56]
    tokens, sizes = merge_tokens(tokens, sizes, keys, r)
    tokens = tokens + mlp(norm(tokens))
    return tokens, sizes

# Bipartite soft matching  [§3 "Bipartite Soft Matching"] [app. D]
# Output has N - r tokens; the class token is never merged  [app. A]

def merge_tokens(tokens, sizes, keys, r):
    # Normalize keys so a dot product is a cosine similarity  [§4.1]
    keys = keys / norm(keys, axis=-1)                # (B, N, C)
    # Alternate positions into sets A (even) and B (odd)  [§3] [tab. 1]
    keys_a, keys_b = keys[:, 0::2], keys[:, 1::2]    # (B, N/2, C) each
    # Similarity of each A token to each B token  [§3] [code: tome/merge.py#L52]
    scores = keys_a @ transpose(keys_b)              # (B, N/2, N/2)
    # Class token (A, index 0) never wins a link  [code: tome/merge.py#L54-L55]
    scores[:, 0, :] = -inf
    # Each A token links only to its most similar B token  [§3, step 2]
    best_score, best_b = max(scores, axis=-1)        # (B, N/2) each
    # Rank A tokens by link strength, strongest first  [§3, step 3]
    order = argsort(best_score, descending=True)     # (B, N/2)
    merged_a = order[:, :r]                          # A tokens that merge
    kept_a = sort(order[:, r:])                      # kept A, in position order
    # Weight every token by its size before summing  [§4.1] [§3]
    weighted = tokens * sizes                        # (B, N, D)
    # Add each merged A token into its B partner  [§3, step 4]
    # scatter_add(target, target_index, source, source_index) adds rows
    sum_b = scatter_add(weighted[:, 1::2], best_b, weighted[:, 0::2], merged_a)
    size_b = scatter_add(sizes[:, 1::2], best_b, sizes[:, 0::2], merged_a)
    # Divide by the total size to get the size-weighted average  [§4.1]
    tokens_b = sum_b / size_b                        # (B, N/2, D)
    # Concatenate kept A, then updated B  [§3, step 5]
    # Sorted kept A keeps the cls token first  [code: tome/merge.py#L66-L68]
    kept_tokens = gather(tokens[:, 0::2], kept_a)    # (B, N/2 - r, D)
    kept_sizes = gather(sizes[:, 0::2], kept_a)      # (B, N/2 - r, 1)
    tokens_out = concat(kept_tokens, tokens_b)       # (B, N - r, D)
    sizes_out = concat(kept_sizes, size_b)           # (B, N - r, 1)
    return tokens_out, sizes_out

# Per-layer merge counts  [§4.1]
# Constant: r per layer  [eq. 2]. Decreasing: 2r down to 0  [eq. 3]

def merge_schedule(num_layers, r, decreasing):
    if not decreasing:
        return [r for layer in range(num_layers)]
    counts = []
    for layer in range(num_layers):
        # Linear from 2r to 0, rounded down  [code: tome/utils.py#L101-L105]
        fraction = layer / (num_layers - 1)
        counts.append(int(2 * r * (1 - fraction)))
    return counts

Hyperparameters

Results

All results are author-reported. Setups are given per row. Blue rows are off-the-shelf (no training); grey rows are trained with ToMe; grey text marks the authors' own baselines tab. 3 tab. 5.

Ablations tab. 1: ViT-L/16 MAE, off-the-shelf, ImageNet-1k, r=8, V100 fp32. Baseline without ToMe: 85.96% at 93.3 im/s (caption).

Bold is the default. The AugReg rows do not name a model size tab. 1.

Matching algorithm tab. 2: same setting as tab. 1.

Merging schedule fig. 2: a constant schedule is close to the best of 15,000 random schedules, especially as the total merged count grows. The best random schedules around 1600–1800 im/s are near constant; beyond that, decreasing is best §4.1 app. C.

Model sweep fig. 3 §4.2: off-the-shelf, a constant schedule gives up to 2x throughput across models. ViT-B, S and Ti lose about 4–5% at 2x. ViT-L loses 2% at 224px and 0.7% at 384px (AugReg). ViT-L at 512px and ViT-H at 518px (SWAG) lose about 0.3% without training. Trained MAE models at 2x lose 0.4% (ViT-H), 0.6% (ViT-L) and 1.7% (ViT-B). Re-evaluating a model trained with r=5 at r=0 gives 85.8% against 85.7% for the baseline §4.2.

ViT-L/16 MAE, ImageNet-1k app. A: ToMe off the shelf (blue) and trained (grey text is the authors' baseline). Drops are relative to the 85.66% row.

Full MAE sweeps for ViT-B/16 and ViT-H/14 are in app. A.

Image SOTA tab. 3: ImageNet-1k, trained only on ImageNet-1k, 224px unless noted. Throughput on a V100. Rows marked ‡ take throughput from the original papers.

* SimMIM pretraining. The paper says ToMe makes ViT-L and ViT-H as fast as a smaller tier §4.3.

ToMe vs token pruning tab. 4: ViT-S, DeiT recipe, 300 epochs. Pruning methods cannot be batched and were timed in different settings, so their throughput is not directly comparable. DynamicViT was timed by the authors on their V100.

† A-ViT is fine-tuned on DeiT rather than trained from scratch tab. 4.

Video, Kinetics-400 tab. 5 tab. 6 §5: Spatiotemporal MAE ViT-L. The input column is copied from the paper. The baseline is re-evaluated at 1 crop × 10 clips.

† Uses ImageNet-21k pretraining. Other models use Kinetics-400 only.

Audio, AudioSet-2M mAP tab. 7 §6: ViT-B from Audio MAE. The metric is mAP because of class imbalance. Off-the-shelf rows compare to the original paper's number; trained rows compare to the authors' own trained baseline.

The paper summarizes audio as almost 2x throughput for a 0.4 mAP drop §6.

Other figures: in the clip-count study, lower r values show no drop against the baseline at the same clip count. At r=65 there is about a 0.5% drop compared with the four-clip sweet spot §5 fig. 5. Visualizations show merged tokens following object parts on images §4.4 fig. 4 and single objects across frames on video §5 fig. 6.

Paper vs code

The official repository is facebookresearch/ToMe, linked from the paper's title-page footnote. The repository is licensed CC BY-NC 4.0 LICENSE. This card cites it by path and line and contains no verbatim code. Pseudocode above is written from the paper and the code's logic.

Agreements:

• Alternating partition, cosine on keys averaged over heads, size-weighted merging, merge between attention and MLP, and class-token protection all match §3 §4.1 tab. 1 merge.py timm.py.
• The distillation token is protected for DeiT, matching app. A merge.py merge.py.
• Proportional attention defaults to on for the timm patch and off for the MAE patch timm.py mae.py.

Differences:

• Proportional attention for trained MAE. app. A says trained MAE models use proportional attention; §4.1 says it is used for all but off-the-shelf MAE. The code's docstrings say to set it to False for training and for MAE timm.py mae.py. The repo has no training code, so which setting produced the trained numbers is not verifiable here.
• Per-layer cap. The code lowers r to at most half of the non-protected tokens in a layer merge.py. The paper text does not mention a cap.
• Decreasing schedule totals. §4.1 says the schedule removes rL tokens in total. The code rounds each layer's count down to an integer utils.py, so the total is generally below rL.
• MAE global pooling. The MAE patch averages the final tokens weighted by size, and divides by the original token count mae.py. The paper text does not describe this pooling. §3 says sizes weight any aggregation.
• Appendix implementation. The Python listing in app. D is older than the release. It protects only the class token, always sorts the unmerged indices, and uses a sum without size normalisation. The release protects both special tokens and normalises by size outside the merge merge.py merge.py.
• Training code. The repo has no training scripts. The image and video recipes come from the MAE and Spatiotemporal MAE releases §4.2 §5, which this card does not cover.
• Stable Diffusion. The source has a commented-out appendix on Stable Diffusion. It is not part of the published results and is not covered here.

Pitfalls

• Two ViT-L/16 MAE baselines. The caption of tab. 1 gives 85.96% at 93.3 im/s. The off-the-shelf sweep in app. A uses 85.66% at 93 im/s, and all its drops are relative to that. The paper does not reconcile the two.
• Default settings do not match the appendix. tab. 1 reports 84.25% at 182.9 im/s for off-the-shelf ViT-L/16 MAE with r=8 and the default settings. app. A reports 83.92% at 183 im/s for off-the-shelf r=8 with proportional attention off, which appears to be the same setting. The paper gives no reason for the gap.
• Model not stated in part of the ablations. tab. 1 says all ablations use ViT-L/16 MAE, but its proportional-attention rows for AugReg name no model size.
• Truncation changes the token budget. A decreasing schedule computed with the code's integer rounding removes fewer than rL tokens. Sum the per-layer counts from the schedule pseudocode for your L and r before comparing with a constant schedule at the same r.
• Silent cap on short inputs. The per-layer cap in merge.py reduces r without warning when the token count is small.
• Quadratic similarity matrix. Each layer builds a matrix of about N/2 × N/2 scores per image merge.py. The paper says it avoids comparing every pair §3, but it still compares every A token with every B token. Memory therefore grows with the square of the token count, which matters for video and high resolution.
• Proportional attention for trained MAE. Follow app. A to reproduce the paper's numbers. Following the code's docstring would change the setting. See Paper vs code.
• Throughput depends on setup. Numbers use V100 fp32 at optimal batch size §4. The pruning rows were timed differently and are not comparable tab. 4. Off-the-shelf speed-ups depend on the batch and the implementation.
• Mixed-source comparisons. ‡ rows in tab. 3 are from the original papers and were not re-measured on the authors' setup.
• Video baseline re-evaluated. The baseline is evaluated at 1x10 clips, while the original paper used 3x7 app. A. Its 84.7% is therefore not the published figure.
• Video recipe not comparable. The video sweep app. B starts from an MAE checkpoint trained for half the schedule, and its accuracies are not comparable with tab. 5.
• Schedule search on one model. The 15,000-schedule search used AugReg ViT-B/16 in fp16 §4.1. The conclusion that constant schedules are near-optimal may not transfer to other models.
• Memory claim unquantified. app. A says ToMe reduces training memory, but no memory figures are reported.
• Single re-evaluation point. The re-evaluation result (85.8% vs 85.7%) is one number in the text §4.2. It is not in a table.
• Dense prediction untested. The paper's results cover classification only. The commented-out diffusion appendix is not evidence of dense-task results.

Lineage

• Builds on: Vision Transformer §1 §3 (concept page vision-transformer). Evaluates on MAE checkpoints §4 (see MAE), DeiT app. A, Spatiotemporal MAE §5 and Audio MAE §6.
• Compared against: token pruning methods A-ViT, DynamicViT and SP-ViT §4.3 tab. 4; Token Pooling, a k-means-based merge, and the k-means and greedy matchings §2 tab. 2; the random and attention-based pruning baselines tab. 2.
• Superseded by: not stated in this paper. Not known from this wiki.