提出解耦嵌入机制,让ViT更高效地合并视觉令牌。
Learning to Merge Tokens via Decoupled Embedding for Efficient Vision Transformers
- 用独立模块学习专用于合并的嵌入表示
- 在ImageNet上减少37.2%计算量,准确率仍达79.85%
- 可模块化训练,兼容现有ViT模型
近期视觉变压器(ViTs)的令牌压缩方法依赖中间特征进行令牌合并,限制了合并性能且需端到端训练。本文提出解耦令牌嵌入合并(DTEM),通过连续松弛的合并过程学习一个与ViT前向传播解耦的轻量级嵌入模块,专门提取用于合并的特征。该结构摆脱了对中间特征的依赖,支持仅训练解耦嵌入或端到端微调。DTEM在图像分类、图文生成和分割任务中均实现稳定提升,尤其在ImageNet-1k分类任务中,使用DeiT-small时将FLOPs降低37.2%,同时保持79.85%的顶级准确率。代码已公开。
原文摘要 · Abstract (English)
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}.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。