Learning to Merge Tokens via Decoupled Embedding for Efficient Vision Transformers

TL;DR

Vision Transformers incur high costs from self-attention on many tokens. DTEM introduces a decoupled embedding learned through a continuously relaxed token merging process to tailor token similarities for merging, enabling modular or end-to-end training while maintaining accuracy. Across ImageNet-1k classification, COCO captioning, and ADE20K segmentation, DTEM achieves substantial FLOPs reductions (e.g., 37.2% for DeiT-small) with competitive or improved accuracy, and ablations confirm the importance of the decoupled embedding and the soft grouping design. This work offers a practical, scalable path to efficient ViTs by decoupling merging features from contextual features and learning merging-aware representations directly.

Abstract

Recent token reduction methods for Vision Transformers (ViTs) incorporate token merging, which measures the similarities between token embeddings and combines the most similar pairs. However, their merging policies are directly dependent on intermediate features in ViTs, which prevents exploiting features tailored for merging and requires end-to-end training to improve token merging. In this paper, we propose Decoupled Token Embedding for Merging (DTEM) that enhances token merging through a decoupled embedding learned via a continuously relaxed token merging process. Our method introduces a lightweight embedding module decoupled from the ViT forward pass to extract dedicated features for token merging, thereby addressing the restriction from using intermediate features. The continuously relaxed token merging, applied during training, enables us to learn the decoupled embeddings in a differentiable manner. Thanks to the decoupled structure, our method can be seamlessly integrated into existing ViT backbones and trained either modularly by learning only the decoupled embeddings or end-to-end by fine-tuning. We demonstrate the applicability of DTEM on various tasks, including classification, captioning, and segmentation, with consistent improvement in token merging. Especially in the ImageNet-1k classification, DTEM achieves a 37.2% reduction in FLOPs while maintaining a top-1 accuracy of 79.85% with DeiT-small. Code is available at \href{https://github.com/movinghoon/dtem}{link}.

Figures (10)

• Figure 1: Comparison of our method with conventional token merging. Contrary to prior works that merge tokens directly based on intermediate features in ViT, our method leverages a decoupled embedding to extract features tailored for token merging. The embedding module is trained via continuous relaxation of grouping and merging operators, i.e., soft grouping and merging, respectively, that allow differentiation.
• Figure 2: Classification results with LV-ViT-S. $^*$ indicates the results with off-the-shelf frozen pretrained model.
• Figure 3: Ablation study on decoupled embedding module design: (a) decoupled embedding dimension and (b) number of hidden layers.
• Figure 4: Image captioning evaluation results when token merging is applied. We report with caption evaluation metrics: BLEU-4 (B@4), CIDEr (C), METEOR (M) and SPICE (S). Reduction represents the decreases in FLOPs within the ViT encoder, and # indicates the number of tokens passed to language decoder.
• Figure 4: Image classification results on data and train efficiency: (a) dataset size and (b) training epochs. In the experiments, DTEM is modularly trained on DeiT-S model, while ToMe undergoes end-to-end training.