Token Merging (ToMe)

Token merging (ToMe) is a method for speeding up transformers, especially the Vision Transformer (ViT), by combining similar tokens inside the network so that each block processes fewer of them. Images contain much redundancy, such as patches of sky or road that look almost the same, and averaging such tokens together removes computation while keeping their information. ToMe works on already trained models without retraining and can also be applied during training. Daniel Bolya and colleagues at Georgia Tech and Meta AI introduced it in Token Merging: Your ViT But Faster, an oral paper at ICLR 2023. Off the shelf, it doubled the throughput of ViT-L at 512 pixels and ViT-H at 518 pixels on images with an accuracy drop of only 0.2% to 0.3%, and it sped up a ViT-L video model 2.2 times with a similar drop.

How Token Merging Works#

A ViT splits an image into patches, one token each, and every block runs self-attention, whose cost grows with the square of the token count, followed by an MLP, whose cost grows linearly. ToMe inserts a merging step between the attention and the MLP of every block and removes a fixed number of tokens, r, each time, so a 24-block model with r = 8 ends with 192 fewer tokens than it started with. The merge uses bipartite soft matching, chosen because it is fast and needs no iterative clustering:

1. Split: Divide the tokens alternately into two sets, A and B.
2. Match: Connect each token in A to its most similar token in B, measured by the cosine similarity of the attention keys averaged over heads.
3. Keep the best edges: Keep the r most similar pairs.
4. Merge: Average each kept pair into one token, weighted by how many original patches each token already represents, and concatenate the two sets again.

Because merged tokens stand for several patches, ToMe adds the logarithm of each token's size to the attention scores, called proportional attention, so a token built from many patches carries its share of attention. The whole procedure is differentiable, which lets models be trained with ToMe for extra speed.

This PyTorch example applies bipartite soft matching to Meta's DINOv2-small ViT from Hugging Face Transformers on a 448 x 448 image of a bus, merging r tokens after each of its 12 blocks, and compares the resulting image embedding with the one from the unmerged model:

from urllib.request import urlopen

import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoImageProcessor, AutoModel

processor = AutoImageProcessor.from_pretrained("facebook/dinov2-small")
model = AutoModel.from_pretrained("facebook/dinov2-small")
image = Image.open(urlopen("https://ultralytics.com/images/bus.jpg")).resize((448, 448))
pixels = processor(images=image, do_resize=False, do_center_crop=False, return_tensors="pt").pixel_values


def merge(x, r):
    """Bipartite soft matching: merge the r most similar pairs between alternating token sets A and B."""
    cls, a, b = x[:, :1], x[:, 1::2], x[:, 2::2]
    similarity = F.normalize(a[0], dim=-1) @ F.normalize(b[0], dim=-1).T
    best, match = similarity.max(dim=-1)  # each A token's most similar B token
    order = best.argsort(descending=True)
    merged, kept = order[:r], order[r:]
    index = match[merged][None, :, None].expand(1, r, x.shape[-1])
    b = b.scatter_reduce(1, index, a[:, merged], reduce="mean")  # average merged A tokens into their B matches
    return torch.cat([cls, a[:, kept], b], dim=1)


with torch.no_grad():
    for r in (0, 32, 64):
        x, counts = model.embeddings(pixels), []
        for layer in model.encoder.layer:
            counts.append(x.shape[1])
            x = merge(layer(x), r)  # merge r tokens after every transformer block
        embedding = model.layernorm(x)[:, 0]  # CLS token: the image embedding
        reference = embedding if r == 0 else reference
        similarity = F.cosine_similarity(embedding, reference).item()
        print(f"r={r}: mean tokens per block {sum(counts) / 12:.0f}, CLS similarity to unmerged {similarity:.3f}")

The unmerged model processes 1025 tokens in every block, 1,024 patches plus the CLS token. Merging 32 tokens per block lowers the average to 849 while the image embedding keeps a cosine similarity of 0.983 to the original, and merging 64 per block lowers it to 673 at a similarity of 0.887, showing how r trades compute for fidelity.

Where Token Merging Is Used#

Token Merging for Fast Stable Diffusion applies the same idea to the transformer blocks of the Stable Diffusion denoising network, where it can remove up to 60% of the tokens, generating images up to twice as fast and with up to 5.6 times less memory, with no extra training. Beyond images, the ToMe paper doubled the throughput of an audio ViT-B trained with merging, at a 0.4% drop in mAP, and sped up masked autoencoder fine-tuning on video by up to 2 times. Researchers apply similar merging and pruning to the visual tokens of vision-language models, where LLaVA-1.5 spends 576 tokens of the language model's context on every image.

In the Ultralytics package, Meta's Segment Anything Model (SAM) encodes each image at 1024 x 1024 pixels as 4,096 patch tokens, and in the MobileSAM comparison its ViT-H image encoder takes 452 ms of the 456 ms SAM needs per image on a GPU, the share of compute that token reduction methods target.

Token Merging vs. Related Concepts#

• Token pruning: Methods such as DynamicViT and EViT discard tokens judged unimportant, which loses their information and usually requires training. ToMe merges tokens instead, so the information is averaged in, and it works on models trained without it.
• Tokenization: Creates the initial tokens from patches or text before the network runs. ToMe reduces their number inside the network.
• Flash attention and sparse attention: Make attention cheaper or faster for a given number of tokens. ToMe lowers the number of tokens, which speeds up the MLP layers too, and can be combined with them.
• Mixture of Depths (MoD): Lets a learned router decide which tokens skip a block in language models. ToMe permanently fuses redundant tokens using a fixed similarity rule.
• Lower input resolution: Shrinks the token count uniformly across the image. ToMe merges mostly in redundant regions, such as background, and keeps detailed regions at full resolution.