Token Fusion: Bridging the Gap between Token Pruning and Token Merging

TL;DR

Token Fusion (ToFu) combines token pruning and merging techniques to reduce computational costs in Vision Transformers, maintaining accuracy and enabling efficient deployment on resource-limited devices.

Contribution

The paper introduces Token Fusion, a novel scheme that integrates token pruning and merging, along with MLERP merging to preserve feature norms, enhancing efficiency without retraining.

Findings

ToFu improves computational efficiency in ViTs.

ToFu maintains or enhances model accuracy.

ToFu sets new benchmarks in classification and image generation.

Abstract

Vision Transformers (ViTs) have emerged as powerful backbones in computer vision, outperforming many traditional CNNs. However, their computational overhead, largely attributed to the self-attention mechanism, makes deployment on resource-constrained edge devices challenging. Multiple solutions rely on token pruning or token merging. In this paper, we introduce "Token Fusion" (ToFu), a method that amalgamates the benefits of both token pruning and token merging. Token pruning proves advantageous when the model exhibits sensitivity to input interpolations, while token merging is effective when the model manifests close to linear responses to inputs. We combine this to propose a new scheme called Token Fusion. Moreover, we tackle the limitations of average merging, which doesn't preserve the intrinsic feature norm, resulting in distributional shifts. To mitigate this, we introduce MLERP merging, a variant of the SLERP technique, tailored to merge multiple tokens while maintaining the norm distribution. ToFu is versatile, applicable to ViTs with or without additional training. Our empirical evaluations indicate that ToFu establishes new benchmarks in both classification and image generation tasks concerning computational efficiency and model accuracy.