Research map / Architectures and efficiency / Attention and transformers
Learning to Merge Tokens via Decoupled Embedding for Efficient Vision Transformers
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 DTEM, which learns a lightweight decoupled token embedding for ViT token merging via a continuously relaxed (soft top-k grouping and soft merging) process, usable modularly with frozen models or end-to-end.
Problem
Token merging in Vision Transformers relies on intermediate ViT features for similarity, which are hard to tailor for merging and require end-to-end training to improve merging; discrete grouping operators block gradient flow.
Approach
Add a lightweight per-block affine embedding decoupled from the ViT forward pass, used only to compute token similarities for merging, and train it through a differentiable soft grouping (revised differentiable top-k with clipping) and soft merging operator; at inference, replace with hard bipartite soft matching (ToMe).
Key results
- On ImageNet-1k with DeiT-S, DTEM achieves 79.85% top-1 with 37.2% FLOPs reduction (abstract).
- Modular training with frozen DeiT-S/B at 50% FLOPs reduction improves over ToMe by +0.47% to +1.64%.
- On GIT-B COCO captioning, DTEM raises CIDEr by +5.0 to +6.0 across 31% to 41% FLOPs reduction; on ADE20K Segmenter ViT-S, mIoU gains of +0.32 to +1.3 over ToMe.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace ToMe token similarities with learned decoupled embeddings | DeiT-S with frozen pretrained weights · ImageNet-1k val, 35% reduction | Acc@1: 79.12 -> 79.44 | Helped |
| replace ToMe token similarities with learned decoupled embeddings | DeiT-S with frozen pretrained weights · ImageNet-1k val, 50% reduction | Acc@1: 78.01 -> 78.99 | Helped |
| replace ToMe token similarities with learned decoupled embeddings | DeiT-B with frozen pretrained weights · ImageNet-1k val, 50% reduction | Acc@1: 77.92 -> 79.54 | Helped |
| replace ToMe token similarities with learned decoupled embeddings | MAE-B with frozen pretrained weights · ImageNet-1k val, 35% reduction | Acc@1: 82.33 -> 82.80 | Helped |
| replace ToMe token similarities with learned decoupled embeddings | MAE-L with frozen pretrained weights · ImageNet-1k val, 50% reduction | Acc@1: 84.21 -> 84.68 | Helped |
| replace ToMe token similarities with learned decoupled embeddings | GIT-B · COCO test, image captioning | CIDEr: +5.0 to +6.0 | Helped |
| replace ToMe token similarities with learned decoupled embeddings | GIT-L · COCO test, image captioning | CIDEr: +2.3 to +6.0 | Helped |
| replace ToMe token similarities with learned decoupled embeddings | Segmenter with ViT-S · ADE20K, semantic segmentation | mIoU: +0.32 to +1.3 | Helped |
| replace decoupled embeddings with self-attention keys for grouping (ablation) | DTEM, end-to-end training · ImageNet-1k classification | accuracy | Hurt |
| remove proportional attention (ablation) | DTEM soft grouping · ImageNet-1k classification | accuracy | Hurt |
| train the decoupled embedding with 4000 images | DTEM modular training · ImageNet-1k classification | accuracy: +0.44% | Helped |
| use temperature scale 0.1 instead of 0.05 (ablation) | DTEM with DeiT-S · ImageNet-1k val, 50% reduction | Acc@1: 78.41 -> 78.87 | Helped |
Surprising findings
- Naively applying soft grouping/merging to self-attention keys as in ToMe degrades performance, showing the decoupled embedding is necessary.
- Modular training with only the embedding module (under 1% extra parameters) matches or beats end-to-end ToMe fine-tuning in low-data regimes (under 10% of ImageNet).
Limitations
- Validated only on computer vision tasks with ViTs; applicability to other modalities such as text is untested.
- Does not reduce computational cost during training.
- Uses a fixed reduction rate per block rather than learning the reduction profile.
Experimental setup
Datasets: ImageNet-1k COCO captions ADE20K
Benchmarks: ImageNet-1k COCO captions ADE20K
Metrics: Acc@1 GFLOPs throughput (im/s) BLEU-4 METEOR CIDEr SPICE mIoU
Baselines: ToMe EViT DynamicViT dTPS Gumbel Softmax
Builds on: ToMe Bipartite Soft Matching differentiable top-k DynamicViT
Closest papers by meaning
- Video Token Merging for Long Video Understanding (NeurIPS 2024 · similarity 0.89)
- ImagePiece: Content-aware Re-tokenization for Efficient Image Recognition (AAAI 2025 · similarity 0.88)
- vid-TLDR: Training Free Token Merging for Light-weight Video Transformer (CVPR 2024 · similarity 0.88)
- Token Cropr: Faster ViTs for Quite a Few Tasks (CVPR 2025 · similarity 0.87)
- CORE: Compact Object-centric REpresentations as a New Paradigm for Token Merging in LVLMs (CVPR 2026 · similarity 0.87)
- Subobject-level Image Tokenization (ICML 2025 · similarity 0.87)
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.