ToSA: Token Merging with Spatial Awareness

TL;DR

ToSA introduces a token merging method for Vision Transformers that integrates spatial information via depth images, leading to better scene structure preservation and improved efficiency in visual tasks.

Contribution

ToSA is the first token merging approach that combines semantic and spatial cues using depth images, enhancing ViT acceleration and accuracy.

Findings

Outperforms previous token merging methods on multiple benchmarks.

Reduces runtime of Vision Transformers significantly.

Improves scene structure preservation during token merging.

Abstract

Token merging has emerged as an effective strategy to accelerate Vision Transformers (ViT) by reducing computational costs. However, existing methods primarily rely on the visual token's feature similarity for token merging, overlooking the potential of integrating spatial information, which can serve as a reliable criterion for token merging in the early layers of ViT, where the visual tokens only possess weak visual information. In this paper, we propose ToSA, a novel token merging method that combines both semantic and spatial awareness to guide the token merging process. ToSA leverages the depth image as input to generate pseudo spatial tokens, which serve as auxiliary spatial information for the visual token merging process. With the introduced spatial awareness, ToSA achieves a more informed merging strategy that better preserves critical scene structure. Experimental results demonstrate that ToSA outperforms previous token merging methods across multiple benchmarks on visual and embodied question answering while largely reducing the runtime of the ViT, making it an efficient solution for ViT acceleration. The code will be available at: https://github.com/hsiangwei0903/ToSA