Token merging reduces data for split vision transformer inference
Consommé: Conditioned Token Merging for Split Inference
Computer Vision and Pattern Recognition
Summary
Performing split inference with vision transformers can be expensive because sending intermediate data between devices uses a lot of bandwidth. The authors study how to reduce the number of data tokens sent while using a fixed pretrained model. They found that a special summary token (CLS) can be compressed and used to decide which image patches to keep. Their method, Consommé, efficiently merges tokens in one step, making split inference cheaper and faster without losing much accuracy.
What this means in practice
- •For edge device developers: Reduce communication overhead in split vision transformer inference using token merging with minimal accuracy loss.
- •For mobile app developers: Deploy efficient vision transformer models that maintain accuracy while lowering data transmission between phone and server.
Authors
Man Ngo, Jan S. Rellermeyer
Abstract
Split inference partitions a Vision Transformer between a resource-constrained edge device and a server. However, transmitting intermediate tokens across the communication channel is costly. Deploying fine-tuned backbones to every endpoint becomes expensive as well. We study token reduction for split inference under a frozen backbone, where only a small adapter is shipped. We observe that the CLS token plays a dual role: it acts as the global summary the receiver needs for classification, and it serves as a signal on which patches to keep. We propose Consommé, a method utilizing both characteristics. There, an information bottleneck compresses the CLS token into a compact latent which is transmitted in place of the full CLS token and conditions a learned patch gate merging patches into exactly $k$ tokens via soft assignment. It reaches any budget in a single step, thus escaping the per-layer cap of bipartite matching (BPM). On DeiT-S/T, it leads all methods at aggressive token budgets and remains operable at shallow splits where BPM cannot operate, while matching fine-tuned methods at gentle budgets at a fraction of their deployment cost.