Reduce Vision Transformer Compute with Token Merging
Vision transformers can spend substantial computation processing many patch tokens that carry similar information. A patch covering one part of a clear sky may produce a representation close to nearby sky patches, yet ordinary self-attention continues to process each token separately.
Token merging reduces that redundancy by combining selected tokens as they move through the network. Unlike token pruning, which removes tokens, merging tries to preserve their information in a smaller set of representations. The practical goal is simple: reduce the token count in later transformer blocks while keeping task quality within an acceptable range.
This article develops a mental model for token merging, walks through a small numerical example, and explains the engineering decisions that determine whether it actually reduces latency.
Start with the token count
A vision transformer commonly splits an image into patches and maps each patch to a token vector. A classification token or other special tokens may also be present.
Imagine an image represented by eight patch tokens:
[t1, t2, t3, t4, t5, t6, t7, t8]Suppose t2 and t3 represent neighboring regions with very similar features. Keeping both may provide less additional information than keeping two tokens from unrelated parts of the image.
A merge operation can replace them with one representative:
before: [t1, t2, t3, t4, t5, t6, t7, t8]
after: [t1, t23, t4, t5, t6, t7, t8]The next transformer block now receives seven patch tokens instead of eight.
That reduction matters because transformer work depends on sequence length. Standard dense self-attention forms interactions between token positions, so its attention matrix grows quadratically with the number of tokens. Feed-forward layers and many projections also process each token, adding costs that grow roughly linearly with token count.
Reducing tokens therefore affects more than the attention matrix. The actual speedup, however, depends on the model, hardware, batch size, merge implementation, and how much of the end-to-end runtime occurs in operations that shrink with token count.
Merging is different from dropping
Token pruning asks which tokens can be removed. Token merging asks which tokens can be combined.
Consider two token vectors:
a = [1.0, 0.8]
b = [0.9, 1.0]A simple merge could average them:
m = (a + b) / 2
= [0.95, 0.90]The result no longer preserves two separate positions, but information from both vectors contributes to the new representation.
This distinction becomes important when several original tokens eventually collapse into one token. A plain average of the current token vectors can accidentally give unequal influence to the original patches.
Suppose token m12 already represents two original patches, while t3 represents one:
m12 = (t1 + t2) / 2If the next merge computes:
(m12 + t3) / 2then t3 receives half of the final weight, while t1 and t2 each receive one quarter. If the intended representation is the average of all three original patches, the correct weighted merge is:
m123 = (2 * m12 + 1 * t3) / 3Tracking a token’s represented size avoids this accidental reweighting.
Similarity decides which tokens are candidates
A merging method needs a rule for deciding which representations are redundant enough to combine. Similarity in feature space is a common signal.
For vectors x and y, cosine similarity is:
cos(x, y) = (x · y) / (||x|| ||y||)Values closer to 1 indicate directions that are more similar. The scale is useful for comparing token representations without making vector magnitude the only deciding factor.
A simplified procedure could be:
1. compute a feature vector for each mergeable token
2. compare candidate pairs
3. select high-similarity pairs without reusing a token
4. merge each selected pair
5. pass the shorter sequence to the next blockThis is a teaching model, not a production algorithm. Comparing every possible pair directly can itself become expensive. Practical methods use matching strategies designed to keep merge overhead small relative to the transformer work they remove.
The original Token Merging method, often called ToMe, uses a lightweight bipartite matching scheme rather than an unrestricted all-pairs clustering procedure. It can be applied to existing vision transformers without retraining, although applying merging during training can reduce the quality gap for a chosen merge schedule.
Preserve token mass in attention
Once merged tokens represent different numbers of original patches, attention needs another consideration.
Suppose one token represents four patches and another represents one. Treating both as equally sized units can change the effect of merging beyond merely compressing similar representations. A common approach is to carry a token-size value and incorporate that size when attention weights are computed.
One conceptual form is:
score(i, j) = q_i · k_j / sqrt(d) + log(size_j)before the softmax.
Because:
exp(score + log(size)) = size * exp(score)the size term gives a merged token influence proportional to the amount of token mass it represents, all else equal.
This idea is sometimes called proportional attention. It does not guarantee that merging is lossless. Two separate tokens can have different relationships with later queries, and a single merged key-value representation cannot preserve every interaction they would have produced independently. Token size corrects one source of distortion; it does not reconstruct information discarded by the merge.
Choose a merge schedule, not just a merge rate
The number and placement of merges define an accuracy-compute trade-off.
If a model starts with 196 patch tokens and merges 16 tokens after each of several blocks, later blocks can operate on much shorter sequences. Aggressive early merging saves more downstream work because the reduced sequence passes through more remaining layers.
That same choice can also remove distinctions before the network has formed sufficiently useful representations. Two patches that look similar in an early feature space may later need different roles for the task.
A merge schedule therefore has at least two dimensions:
- how many tokens to merge, which controls the strength of compression;
- where to merge them, which controls both downstream compute savings and the stage at which spatial distinctions disappear.
Uniformly merging the same number after every block is a reasonable baseline, not a universal optimum. Models with hierarchical stages, unusual special tokens, dense prediction heads, or task-specific spatial requirements may need a different schedule.
Treat the schedule as part of the model configuration. Record it with evaluation results so quality and performance numbers remain reproducible.
Measure end-to-end latency
A lower theoretical operation count does not automatically produce a proportional wall-clock improvement.
Token merging adds work for similarity calculation, matching, aggregation, bookkeeping, and potentially tensor rearrangement. A shorter token dimension may also interact differently with optimized kernels. On some hardware, a modest reduction in sequence length may not move execution into a meaningfully better operating region.
Benchmark the complete inference path with representative inputs. Useful measurements include:
- median and tail latency at the batch sizes the service actually uses;
- throughput under realistic concurrency;
- peak memory or accelerator memory;
- task quality at each merge schedule;
- time spent in merging compared with time removed from transformer blocks.
Avoid reporting only floating-point operation estimates. They describe an important part of the compute story, but not data movement, launch overhead, kernel efficiency, or application-level preprocessing and postprocessing.
For a latency-sensitive service, compare the unmodified model and several increasingly aggressive schedules. The useful operating point is usually a measured quality-latency trade-off rather than the schedule with the smallest token count.
Validate the task that will run in production
Classification accuracy can hide errors that matter for other vision tasks.
Merging nearby or semantically similar regions may be relatively forgiving for image classification because the final output summarizes the whole image. Dense tasks can be more sensitive. Segmentation, keypoint estimation, localization, and other spatial outputs may need distinctions between tokens that a classifier can safely compress.
The merge mechanism must also handle special tokens deliberately. A class token, register token, or architecture-specific control token should not be treated as an ordinary image patch unless the method explicitly supports that behavior.
Validation should therefore match the deployed task and data distribution. Check more than one aggregate metric when small structures, boundaries, rare classes, or fine spatial detail matter. Inspect examples where the compressed model diverges from the baseline; those cases often reveal whether the merge rule is collapsing information the task still needs.
Common mistakes
The first mistake is treating high feature similarity as proof that two tokens are interchangeable. Similarity is only a heuristic. Later layers can use small differences that the current metric considers unimportant.
The second is forgetting represented token size. Repeated unweighted averaging can make merge history affect contribution weights in unintended ways.
The third is merging too aggressively near the input simply because early reduction offers the largest theoretical saving. Early features may not yet express the semantic distinctions needed for reliable matching.
The fourth is evaluating only the final quality metric. A method that preserves accuracy but adds irregular operations can fail to improve latency on the target accelerator.
The fifth is assuming a schedule transfers unchanged across resolutions. Higher-resolution inputs usually create more patch tokens, changing both redundancy and the compute profile. A fixed number of merges can represent a very different compression ratio.
When token merging is a good fit
Token merging is worth testing when a transformer processes many redundant visual tokens and later blocks dominate enough runtime for sequence reduction to matter. It is especially attractive when modifying the architecture or training a new model is expensive, since some approaches can be attached to pretrained models and evaluated directly.
A simpler option can be better when the input resolution can be reduced without harming the task, when a smaller model already meets quality targets, or when the sequence is short enough that merge overhead dominates the savings. Token pruning can also be preferable when the application has a reliable importance signal and discarded tokens genuinely carry little useful information.
The key comparison is not merging versus an untouched large model in isolation. Compare it with the simplest alternatives that meet the same quality requirement: lower resolution, a smaller backbone, fewer layers, token pruning, or another inference configuration supported by the model.
Make compression a measured decision
Token merging turns representational redundancy into a compute trade-off. Similar tokens are combined, their represented size can be tracked, and later transformer blocks process a shorter sequence. The technique is useful because it targets work inside the network without requiring every removed token to vanish completely.
Start conservatively. Add a small merge rate, preserve special tokens, track token mass, and measure task quality together with end-to-end latency. Then increase compression only while the measured trade-off remains useful for the deployment target. That process turns token merging from a theoretical optimization into an engineering choice backed by evidence.