Lossless Token Merging Even Without Fine-Tuning in Vision Transformers

2025-05-21

AI Summary

Vision Transformer (ViT) inference incurs substantial computational overhead, and existing token compression methods either require fine-tuning or suffer from representational degradation. Method: We propose Adaptive Token Merging (ATM), a fine-tuning-free, layer-adaptive token merging framework that achieves truly lossless compression. ATM introduces (i) a layer-wise adaptive similarity threshold enabling dynamic merging across layers and batches, and (ii) a matching mechanism jointly optimizing token similarity and size constraints to balance compression ratio and representation fidelity. Results: On DeiT-Tiny and DeiT-Small, ATM reduces FLOPs by over 30% with zero top-1 accuracy loss—outperforming all training-free methods and most fine-tuning-dependent approaches. This establishes a new paradigm for efficient ViT inference without compromising model integrity.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Learning on the Edge & Model Compression
• Data Mining & Knowledge Management: Data Compression

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Although Vision Transformers (ViTs) have become the standard architecture in computer vision, their massive sizes lead to significant computational overhead. Token compression techniques have attracted considerable attention to address this issue, but they often suffer from severe information loss, requiring extensive additional training to achieve practical performance. In this paper, we propose Adaptive Token Merging (ATM), a novel method that ensures lossless token merging, eliminating the need for fine-tuning while maintaining competitive performance. ATM adaptively reduces tokens across layers and batches by carefully adjusting layer-specific similarity thresholds, thereby preventing the undesirable merging of less similar tokens with respect to each layer. Furthermore, ATM introduces a novel token matching technique that considers not only similarity but also merging sizes, particularly for the final layers, to minimize the information loss incurred from each merging operation. We empirically validate our method across a wide range of pretrained models, demonstrating that ATM not only outperforms all existing training-free methods but also surpasses most training-intensive approaches, even without additional training. Remarkably, training-free ATM achieves over a 30% reduction in FLOPs for the DeiT-T and DeiT-S models without any drop in their original accuracy.

Problem

• Reduce Vision Transformers' computational overhead without fine-tuning
• Prevent information loss in token compression techniques
• Maintain model accuracy while significantly decreasing FLOPs

Innovation

• Adaptive Token Merging ensures lossless compression
• Layer-specific thresholds prevent undesirable token merging
• Token matching considers similarity and merging sizes