HypTM: Efficient Vision Transformers With Hyperbolic Token Merging

Abstract

Vision Transformers (ViTs) have shown tremendous potential in various computer vision applications attribute to the attention mechanism, which comes at the expense of increased computational costs. Therefore, token compression methods such as pruning and merging are proposed to facilitate the deployment of ViTs on resource-constrained devices. However, existing approaches assess tokens in Euclidean space, neglecting their complex geometric and hierarchical relationships. In this paper, we propose a novel Hyperbolic Token Merging (HypTM) method by exploring the intrinsic properties of images in hyperbolic space. Our HypTM is inspired by the insight that the contribution of distinct tokens to the overall image varies across the hierarchy. To this end, we develop a hyperbolic-driven token assessment algorithm that dynamically adjusts token significance according to hierarchical structure. Additionally, we design a hierarchy enhancement merging strategy to advance correlations among tokens within the same hierarchy. Without any retraining, our HypTM achieves notable improvements over the state-of-the-arts (SOTAs), while maintaining comparable FLOPs (floating point operations). For the DeiT-S model on ImageNet dataset, we achieve a 41.3% reduction in FLOPs with only a 0.48% decrease in Top-1 accuracy, which is even 0.09% higher than that of the meticulously fine-tuned comparison approach.