Learning to Merge Tokens via Decoupled Embedding for Efficient Vision Transformers

1 Introduction

• Vision Transformers incur quadratic computational and memory costs with token count, motivating methods that reduce tokens while preserving performance.Token merging combines redundant tokens rather than removing them, reducing information loss relative to pruning.Source: §1 · pp. 1–2
• Intermediate ViT features serve both contextual encoding and merging, making them difficult to tailor and requiring extensive end-to-end training.This dependence also makes effective use of pre-trained models difficult and typically requires substantial data to avoid overfitting.Source: §1 · p. 2
• DTEM introduces a lightweight embedding module dedicated to token merging, decoupling merging features from intermediate ViT features.This separation allows improved merging without altering ViT parameters and supports modular optimization with pre-trained models.Source: §1 · p. 2
• DTEM learns its decoupled embeddings through continuously relaxed grouping and merging operators that make the token-merging policy differentiable.During inference, the relaxed operators are replaced with hard operators used by existing token-merging methods.Source: §1 · p. 2, Figure 1 · p. 2
• DTEM can be trained modularly by updating only the added embedding module or end-to-end by fine-tuning the full model.Modular training improves performance even with substantially smaller datasets and fewer training epochs.Source: §1 · p. 2
• DTEM consistently improves token merging across image classification, captioning, and segmentation with various ViT models.The experiments target tasks requiring different levels of representation granularity and report a better trade-off between task performance and computation cost.Source: §1 · p. 2

2 Background

• Token merging reduces N input tokens by r tokens at each Transformer block, producing N−r tokens through grouping followed by merging.The grouping identifies tokens to combine, and the merging operator pools tokens within each connected group.Source: §2 · p. 3
• The grouping operator determines which tokens merge, so performance depends strongly on its design, while reachability-matrix computation can be costly.Early clustering methods can be slow and may suffer performance drops from aggressive distribution shifts between input and merged tokens.Source: §2 · p. 3
• ToMe uses Bipartite Soft Matching to construct a sparse bipartite graph, select high-similarity edges, and merge connected tokens efficiently.It also tracks combined-token sizes and incorporates them into proportional attention.Source: §2 · p. 3
• The proportional-attention formulation applies softmax to attention scores while accounting for the sizes of combined tokens.The supplied background identifies σ as the softmax function.Source: §2 · p. 3
• Prior grouping relies entirely on similarities between intermediate ViT features because discrete clustering and matching operators block gradient flow.Consequently, improving merging requires updating intermediate features through the merging operator rather than directly learning the grouping policy.Source: §2 · p. 3

3 Method

• DTEM learns per-token embeddings used only to compute similarities and shape the merging policy, allowing improvements without altering ViT parameters.The embedding module is decoupled from the ViT forward pass and can be trained modularly.Source: §3 · p. 3
• Continuous relaxations replace discrete grouping and merging during training, enabling gradients to reach the decoupled embeddings.Soft grouping produces continuous adjacency values, while soft merging absorbs tokens proportionally to those values.Source: §3 · pp. 3–4, §3.2 · p. 4, §3.3 · p. 5

3.1 Decoupled Embedding Module

• A lightweight per-token projection produces lower-dimensional embeddings for each Transformer block, with fewer than 1% additional parameters and parallelizable computation.The output embedding dimension satisfies d′ ≪ d, and even an affine module improves merging with minimal cost.Source: §3.1 · p. 4

3.2 Soft Grouping

• Soft grouping repeatedly selects r similar token pairs using temperature-controlled softmax weights while suppressing reused nodes, producing a continuous adjacency matrix with at most one edge per node in A.As τ approaches zero, the soft selection converges toward one-hot choices; clipping constrains the total selected edges.Source: §3.2 · pp. 4–5

3.3 Soft Merging

• Soft merging continuously updates retained-token features and effective sizes, while discretized adjacency performs actual token reduction at inference.Tokens in A are absorbed into tokens in B during training, and the inference procedure uses hard merging equivalent to BSM.Source: §3.3 · p. 5, §3.4 · p. 5

3.4 Training and Inference

• Embeddings trained at a high reduction rate generalize to lower rates, and DTEM supports either frozen-backbone modular training or joint end-to-end fine-tuning.Inference discretizes the learned operators, while training alternates soft embedding updates with efficient ViT updates when needed.Source: §3.3 · p. 5, §3.4 · p. 5

4 Experiments

• DTEM consistently outperforms prior token-merging methods across image classification, captioning, and semantic segmentation.The experiments cover multiple ViT models and tasks requiring different representation granularity.Source: §4 · p. 6

4.1 Image Classification

• At 50% FLOPs reduction, DTEM improves classification performance over ToMe by +0.47% to +1.64% across DeiT-S/B and MAE-B/L models with frozen ViTs.At 35% FLOPs reduction, gains range from +0.15% to +0.47%, with less than 1% additional FLOPs at 50% reduction.Source: §4.1 · p. 6, Table 1 · p. 6

4.2 Image Captioning

• DTEM improves captioning over ToMe by +5.0 to +6.0 CIDEr for GIT-B at 31%–41% FLOPs reduction and by +2.3 to +6.0 for GIT-L.The results indicate a better trade-off between caption quality and encoder computation.Source: §4.2 · p. 8, Table 4 · p. 8

4.3 Semantic Segmentation

• DTEM improves semantic-segmentation mIoU over ToMe by +0.32 to +1.3 when reducing 25% to 50% of FLOPs.The method is evaluated with Segmenter using ViT-S on ADE20K.Source: Table 5 · p. 8, §4.3 · p. 9

4.4 Analysis

• The decoupled embedding is crucial: it learns features distinct from intermediate ViT features, while proportional attention and a shallow affine module provide effective low-cost designs.Modular training also gains +0.44% accuracy with 4000 images and converges quickly, while DTEM prioritizes merging background tokens and preserves more foreground tokens.Source: §4.4 · pp. 9–10

• DTEM learns decoupled token embeddings through continuously relaxed token merging to exploit features tailored specifically for merging.The decoupled structure resolves dependence on intermediate ViT features and supports modular training with frozen pretrained models.Source: §5 · p. 10

Image Captioning

• At reduction rates above 41% for GIT-B and 49% for GIT-L, captioning quality decreases significantly.The broader evaluation identifies these reduction ranges as challenging, resource-constrained settings.Source: §A.1 · p. 14

Classification Results

• DTEM consistently improves token merging across DeiT-T, DeiT-S/B, MAE-B/L, and higher-resolution AugReg ViT-S classification settings.The evaluations cover modular and end-to-end training, multiple reduction rates, and larger token counts.Source: §A.1 · pp. 14–15, Table 10 · p. 14, Table 11 · p. 14

Training for Longer Epochs

• With 100 training epochs, DTEM outperforms baselines on DeiT-T and is comparable to dTPS on DeiT-S while incurring much lower training costs.DTEM uses a single trained model across reduction rates and only the task loss, whereas dTPS requires multiple models and additional losses.Source: §A.1 · pp. 14–15, Table 14 · p. 15

Effect of temperature scaling

• Temperature values from 0.1 to 0.3 consistently provide gains, with an accuracy difference of 0.1% in the modular DeiT-S experiment.The temperature-scaling study trains the decoupled embedding module for 10 epochs.Source: §A.1 · p. 15, Table 16 · p. 15

End-to-End Training for Captioning

• With end-to-end GIT-base captioning at reduction rate r = 13, DTEM improves CIDEr by +1.0 to +2.0 over end-to-end ToMe.End-to-end training benefits both methods, while DTEM offers the better performance–computation trade-off.Source: §A.1 · p. 15, Table 17 · p. 16

Common details

• During training, DTEM excludes tokens with minimum effective size from relaxed merging and uses detached ViT features so only the embedding module learns from the relaxed process.In end-to-end training, embedding and ViT parameters are updated alternately, with one embedding update per nine ViT updates.Source: §A.2 · p. 16

Image Classification

• Classification uses 30-epoch AdamW training with task-specific augmentation, learning-rate searches, reduction rates, and temperature scaling across pretrained ViTs.MixUp and CutMix are omitted for DeiT-T and MAE because performance improves without them.Source: §A.2 · p. 16

Image Captioning

• Captioning trains GIT models for 20,000 iterations with beam-search evaluation and reduction rates r = 13 for GIT-base and r = 8 for GIT-large.The best validation model is evaluated on the test split.Source: §A.2 · pp. 16–17

Semantic Segmentation

• For segmentation, DTEM defines token similarity from decoupled embeddings, merges tokens before ViT components, and unmerges them afterward through duplication and residual addition.The method is applied with Segmenter and repeatedly merges tokens around attention and feed-forward components.Source: §A.2 · p. 17
• Segmenter training removes 40% of tokens at each transformer component and applies relaxed merging to the last 64 tokens selected for merging.The remaining tokens are first selected through conventional discrete merging to stabilize training.Source: §A.2 · p. 17

Pruning tokens

• Token pruning removes tokens based on task importance or attention scores, but information loss degrades performance especially at high reduction rates.Importance is commonly estimated with lightweight prediction modules or class-token attention.Source: §A.3 · p. 17

Combining tokens

• Token-merging methods combine similar or redundant tokens to reduce information loss, including soft clustering, K-medoids, matching, and importance-aware grouping.ToMe progressively merges tokens at each transformer block and can operate without training.Source: §A.3 · p. 17

Decoupled embedding

• DTEM addresses the common reliance on intermediate features by learning decoupled merging embeddings through continuously relaxed merging for inference-time computation reduction.Unlike GroupViT, which targets implicit segmentation, DTEM targets token-merging efficiency.Source: §A.3 · p. 17

Reduction profiles

• DiffRate learns how many tokens to reduce at each transformer block, whereas DTEM learns features for merging while reducing a fixed number of tokens across blocks.The paper identifies joint optimization of reduction profiles and merging features as a future direction.Source: §A.3 · p. 17

• DTEM’s effectiveness has so far been validated only on computer-vision tasks using ViTs.Source: §A.4 · p. 18
• The method improves computation time during inference but does not reduce computational cost during training.Source: §A.4 · p. 18
• Future work targets sparsifying training to improve computational efficiency across both training and inference.Source: §A.4 · p. 18

• All experiments used a GPU server with an Intel Xeon Gold 6230 CPU and 256GB RAM.Source: §A.5 · p. 18
• The server was equipped with 8 NVIDIA RTX 3090 GPUs.Source: §A.5 · p. 18
• Each NVIDIA RTX 3090 GPU provided 24GB of VRAM.Source: §A.5 · p. 18

• DTEM prioritizes merging background patches while allocating more patches to foreground objects than ToMe.Source: §A.6 · p. 18, Figure 6 · p. 19
• ToMe allocates more tokens to the background, producing a less focused allocation for foreground objects.Source: §A.6 · p. 18
• With reduction profile r = 16, the final output contains 11 remaining tokens.Source: Figure 6 · p. 19