Token merging

Overview

A transformer's cost grows with its token count, and in a vision transformer the count is set by the image resolution and patch size. Token merging reduces that count inside the network. In each block, it finds tokens that are similar and replaces each matched pair with one token, so the next block runs on fewer tokens 2210.09461 §1 2210.09461 §3.

The shared idea is that merging keeps information that pruning throws away. Averaging a pair of tokens keeps a share of both, while deleting a token removes it entirely. The paper argues that this is why merging can reach a higher token reduction at the same accuracy 2210.09461 §1. Its comparison at the same setting supports this, with merging ahead of random and attention-based pruning 2210.09461 tab. 2.

Token merging is one family of token reduction. The paper's comparison set is pruning methods, which remove tokens, and clustering methods, which merge them more slowly 2210.09461 §2 2210.09461 tab. 2. See the token pruning page for removal-based methods. Only one carded merging method exists so far, so this page covers its design space.

Variants

The carded method, ToMe, has four design choices that the paper ablates: how many tokens to merge per layer, how to match them, how to combine merged tokens, and whether to change attention. The paper also contrasts its matching with other matching rules and with a second schedule.

Bipartite soft matching with a fixed count (ToMe)

Split the tokens into two sets of about equal size by alternating positions. Each token in the first set links to its most similar token in the second set, using cosine similarity on the attention keys averaged over heads. The r strongest links merge, and the rest of the tokens pass through unchanged 2210.09461 §3 2210.09461 §4.1 2210.09461 tab. 1. The count r is a fixed number per layer, not a ratio, so every image in a batch loses the same number of tokens 2210.09461 §3.

The ablation uses ViT-L/16 MAE with r = 8 and reports 84.25% top-1 at 182.9 im/s off the shelf, against 85.96% at 93.3 im/s for the baseline without merging 2210.09461 tab. 1. The matching is the authors' main choice, so the other matching rules below are compared against it. Card: Token Merging: Your ViT But Faster.

Other matching rules

The paper compares two clustering-style rules with the same r = 8 setting. K-means with 2 iterations reaches 80.19% at 169.7 im/s, and with 5 iterations 80.29% at 147.5 im/s. Greedy matching reaches 84.36% at 179.4 im/s, which is slightly more accurate than ToMe and slower 2210.09461 tab. 2. Greedy matching is therefore the nearest alternative in the paper, and the choice between the two is a trade-off between about 0.1 points of accuracy and about 3.5 im/s of throughput at this setting 2210.09461 tab. 2. The paper does not give a cost model for either rule beyond these measurements.

Constant and decreasing schedules

A constant schedule removes r tokens in every layer. A decreasing schedule removes 2r in the first layer and falls linearly to 0 in the last. Both remove about rL tokens in total over L layers 2210.09461 §4.1 2210.09461 eq. 2 2210.09461 eq. 3. On a 15,000-schedule search, the best schedules below about 1600 to 1800 im/s are close to constant, and above that range the decreasing schedule is best 2210.09461 fig. 2 2210.09461 app. C.

Trained with merging, the decreasing schedule at r = 8 reaches 84.2% at 241 im/s, and the constant schedule at r = 8 reaches 85.1% at 183 im/s 2210.09461 tab. 3. The choice is therefore a trade-off between throughput and accuracy at the same r.

Off the shelf or trained with merging

Merging can be added to a trained model with no weight change. Off the shelf, ViT-L/16 MAE at r = 8 loses 1.74 points against its baseline, and at r = 4 it loses 0.27 points 2210.09461 app. A. The paper's model sweep reports up to 2x throughput with a constant schedule across several models, with losses that vary by model size 2210.09461 §4.2 2210.09461 fig. 3.

Training with merging recovers part of the loss. ViT-L/16 MAE trained with merging at r = 8 loses 0.61 points at 183 im/s 2210.09461 app. A. The paper's training recipes are the MAE defaults, which the authors did not change for merging 2210.09461 app. B. Merging is also used to speed up video and audio fine-tuning 2210.09461 §5 2210.09461 §6.

When to use which

Compare with pruning at the same r before committing. The paper's comparison is at one setting 2210.09461 tab. 2, and the throughput of pruning methods was measured in other settings in its tab. 4, so do not compare those numbers across methods.

Implementation sketch

Pseudocode of the common core, one block at a time. Tokens and sizes are the per-token inputs, and r is the number of merges in this layer.

# Merge r similar token pairs inside one transformer block  [arxiv:2210.09461 §3]
# tokens (N, D), sizes (N,) with the number of input patches per token
# keys (N, C): attention keys, averaged over heads  [arxiv:2210.09461 §4.1]

def merge_block(tokens, sizes, keys, r):
    # Normalize keys so a dot product is a cosine similarity  [arxiv:2210.09461 §4.1]
    keys = keys / norm(keys)
    # Alternate positions into two sets A and B  [arxiv:2210.09461 tab. 1]
    key_a, key_b = keys[0::2], keys[1::2]
    # Link each A token to its most similar B token  [arxiv:2210.09461 §3]
    score, link = max(key_a @ key_b.T, axis=-1)
    # Keep the r strongest links  [arxiv:2210.09461 §3]
    order = argsort(score, descending=True)
    merged, kept = order[:r], order[r:]
    # Combine by size-weighted average  [arxiv:2210.09461 tab. 1]
    weighted = tokens * sizes
    sum_b = scatter_add(weighted[1::2], link[merged], weighted[0::2][merged])
    size_b = scatter_add(sizes[1::2], link[merged], sizes[0::2][merged])
    # Concatenate kept A tokens with updated B tokens  [arxiv:2210.09461 §3]
    return concat(tokens[0::2][kept], sum_b / size_b)

Proportional attention adds log(size) to the attention logits, so a merged token counts as the patches it represents 2210.09461 §3 2210.09461 eq. 1. The paper uses it for trained models but not for off-the-shelf MAE, see the contradictions below. The class token is protected from merging 2210.09461 app. A. Schedules change only the value of r per layer 2210.09461 eq. 2 2210.09461 eq. 3.

Contradictions

• Two baselines for ViT-L/16 MAE. The caption of the ablation table gives 85.96% at 93.3 im/s 2210.09461 tab. 1. The appendix sweep gives 85.66% at 93 im/s, and its drops are relative to that row 2210.09461 app. A. The paper does not reconcile the two.
• Same setting, different numbers. The ablation reports 84.25% at 182.9 im/s for off-the-shelf ViT-L/16 MAE at r = 8 with default settings 2210.09461 tab. 1. The appendix reports 83.92% at 183 im/s for what looks like the same setting, with proportional attention off 2210.09461 app. A. The paper gives no reason for the gap.
• Proportional attention for trained MAE. The appendix says trained MAE models use it 2210.09461 app. A. The main text says it is used except for off-the-shelf MAE 2210.09461 §4.1. The two readings give different settings for the trained rows.
• Video baseline. The video comparison re-evaluates the baseline at one crop and ten clips, not the 3x7 protocol of the original paper, so its 84.7% is not the published figure 2210.09461 app. A 2210.09461 tab. 5.

Papers

• Token Merging: Your ViT But Faster: merges r similar token pairs per block with bipartite soft matching and size-weighted averaging, with off-the-shelf and trained results on images, video and audio.