KFF: K-feature fusion token merging for vision transformer

Recently, Vision Transformer (ViT) has achieved better performance than Convolutional Neural Networks (CNNs) in various vision applications. However, they are usually more computationally expensive than CNNs and face challenges in training and inference efficiency. Token merging is an effective and training-free way to reduce model complexity. However, since few tokens are exactly the same, prevalent similarity-based merging methods are challenging to avoid feature information loss and accuracy degradation. To address this issue, we propose a novel K-feature fusion token merging algorithm that significantly reduces the similarity metric error and token merging error with almost no accuracy loss. Specifically, we first reveal that similarity measurement errors and merging strategies have a significant impact on the performance of token merging algorithms, and the currently popular K-based similarity method will cause obvious feature shifts during the merging process. Based on this observation, we present a new feature-enhanced K-feature fusion token similarity calculation method. By combining the keys (K), which summarize the information contained in each token, and the more detailed intermediate features, the error of similarity measurement is greatly reduced. Then, we design a similarity-weighted average token merging algorithm to combine tokens that is faster and more accurate than ordinary average token merging. Extensive experiments show that our approach yields better model performance when reducing comparable computational effort and improving throughput without extra training. For example, for ViT-B on ImageNet, our method reduces 49.58 % of tokens and improves throughput by 30 % with only a 0.44 % drop in accuracy.