ToSA: Token Merging with Spatial Awareness

2025-06-24

AI Summary

Existing ViT acceleration methods perform token merging solely based on feature similarity, neglecting the sparse yet geometrically critical spatial structure inherent in early ViT layers. To address this, we propose ToSA (Token Merging with Semantic and Spatial Awareness), the first approach to explicitly incorporate pseudo-spatial tokens—inspired by deep image generation—into the bottom layers of ViTs. ToSA jointly clusters tokens using both semantic features and explicit spatial position relationships, enabling spatially aware, adaptive token merging. This preserves scene geometry while improving the fidelity and rationale of information compression in early layers. Experiments demonstrate that ToSA significantly outperforms state-of-the-art token merging methods across multiple visual and embodied question-answering benchmarks, achieving both reduced inference latency and enhanced task performance.

Technology Category

• Computer Vision: Multi-modal Vision
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.

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: Vertical and domain-specific search
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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

Problem

• Improves token merging in ViTs by adding spatial awareness
• Reduces computational costs while preserving scene structure
• Enhances performance on visual and embodied question answering

Innovation

• Combines semantic and spatial awareness for merging
• Uses depth image to generate pseudo spatial tokens
• Preserves scene structure while reducing computational costs