Research map / Architectures and efficiency / Attention and transformers
Accelerating Transformers with Spectrum-Preserving Token Merging
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PiToMe, a token merging method for ViTs that uses a graph-energy score to protect informative tokens and merge redundant ones, with spectral-preservation theory, evaluated off-the-shelf and retrained across classification, retrieval, VQA, and text classification.
Problem
Token merging in Transformers via bipartite soft matching (BSM) is sensitive to token-splitting strategy and merges informative tokens in deeper layers, hurting accuracy at a given speedup.
Approach
An energy score derived from graph energy over a cosine-similarity token graph (with an exponential-linear margin function) identifies large clusters of similar tokens as mergeable and isolated tokens as protected; the mergeable tokens are then split into sets A and B by sorted energy order before BSM matching.
Key results
- Off-the-shelf ViT-MAE-H on ImageNet-1k: PiToMe 86.4% vs ToMe 85.6% at the same 92.8 GFLOPs, with about 0.5% average drop from the base model.
- CLIP on Flickr30k off-the-shelf: PiToMe keeps retrieval near baseline, with about 0.3% average drop versus about 4.5% for other methods (per abstract).
- LLaVA-1.5-7B off-the-shelf at r=0.9: PiToMe VQA-v2 75.4 vs ToMe 75.2 and DiffRate 72.0, with reduced inference time.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use energy-based token protection instead of leaving tokens unprotected (ablation) | PiToMe · Image-text retrieval, r=0.925 | Rsum: 568.9 -> 573.4 | Helped |
| use energy-based token protection instead of leaving tokens unprotected (ablation) | PiToMe · Text classification, r=0.6 | accuracy: 86.99 -> 89.20 | Helped |
| use energy-ordered matching split instead of random sets (ablation) | PiToMe · Image-text retrieval, r=0.925 | Rsum: 567.9 -> 573.4 | Helped |
| use energy-ordered matching split instead of random sets (ablation) | PiToMe · Text classification, r=0.6 | accuracy: 87.17 -> 89.20 | Helped |
| set alpha to 1.0 rather than 0.5 or 0. (ablation) | PiToMe · Image-text retrieval | recall sum: 519.98 vs 518.66 vs 515.90 | Helped |
| use an adaptive margin rather than a fixed margin (ablation) | PiToMe · Image-text retrieval | accuracy | Helped |
| replace the energy score with CLS-token attention or mean attention (ablation) | PiToMe · Image-text retrieval and text classification | performance | Hurt |
| use a fixed number of removed tokens per layer instead of a remaining-token ratio (ablation) | PiToMe · Image classification and image-text retrieval | performance at matched FLOPS | Hurt |
| use PiToMe rather than DiffRate | DeiT-T · ImageNet-1k, off-the-shelf | accuracy: 69.9 -> 70.8 | Helped |
| use PiToMe rather than the full ViT | MAE-H · ImageNet-1k, off-the-shelf | accuracy: 86.9 -> 86.4 | Hurt |
| compress LLaVA-1.5-7B vision tokens with PiToMe rather than ToMe | LLaVA-1.5-7B · TextVQA, off-the-shelf | score: 41.1 -> 43.0 | Helped |
| use PiToMe rather than the original BERT | BERT · SST-2, retrained | accuracy: 91.4 -> 91.7 | Helped |
Surprising findings
- On some VQA datasets (VizWiz, ScienceQA), the compressed LLaVA model surpasses the uncompressed baseline, which the authors attribute to merging less significant tokens.
- The authors report that ToMe's fixed-k per-layer schedule is suboptimal compared with a ratio-r schedule given equal FLOPs.
Limitations
- Energy score relies on a fully connected token graph, so complexity grows with input size; sparse graphs are suggested.
- No unmerge mechanism, so extension to generative tasks such as image generation or segmentation is left open.
- The compression ratio r is fixed rather than learned differentiably.
Experimental setup
Datasets: Flickr30k MS-COCO ImageNet-1k VQA-v2 GQA VizWiz ScienceQA TextVQA MME SST-2 IMDb
Benchmarks: Flickr30k MS-COCO ImageNet-1k VQA-v2 GQA VizWiz ScienceQA TextVQA MME SST-2 IMDb
Metrics: recall@k Rsum top-1 accuracy FLOPs throughput speedup
Baselines: ToMe ToFu DiffRate DCT A-ViT Dynamic-ViT SP-ViT E-ViT PowER-BERT LTP
Builds on: Token Merging (ToMe) Bipartite Soft Matching proportional attention graph coarsening
Closest papers by meaning
- Efficient Time Series Processing for Transformers and State-Space Models through Token Merging (ICML 2025 · similarity 0.86)
- Learning to Merge Tokens via Decoupled Embedding for Efficient Vision Transformers (NeurIPS 2024 · similarity 0.86)
- Video Token Merging for Long Video Understanding (NeurIPS 2024 · similarity 0.85)
- Mixture of Tokens: Continuous MoE through Cross-Example Aggregation (NeurIPS 2024 · similarity 0.85)
- CORE: Compact Object-centric REpresentations as a New Paradigm for Token Merging in LVLMs (CVPR 2026 · similarity 0.84)
- Token Cropr: Faster ViTs for Quite a Few Tasks (CVPR 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.