Context-Aware Token Selection and Packing for Enhanced Vision Transformer

TL;DR

This paper introduces SPA, a dynamic token selection and packing algorithm for vision transformers that improves efficiency and accuracy by selecting informative tokens contextually, leading to better performance and reduced computational costs.

Contribution

The paper presents SPA, a novel context-aware token selection and packing method that enhances vision transformer efficiency and accuracy over existing sparse attention approaches.

Findings

0.6 mAP improvement in object detection

16.4% reduction in computational costs

Superior performance across diverse datasets

Abstract

In recent years, the long-range attention mechanism of vision transformers has driven significant performance breakthroughs across various computer vision tasks. However, the traditional self-attention mechanism, which processes both informative and non-informative tokens, suffers from inefficiency and inaccuracies. While sparse attention mechanisms have been introduced to mitigate these issues by pruning tokens involved in attention, they often lack context-awareness and intelligence. These mechanisms frequently apply a uniform token selection strategy across different inputs for batch training or optimize efficiency only for the inference stage. To overcome these challenges, we propose a novel algorithm: Select and Pack Attention (SPA). SPA dynamically selects informative tokens using a low-cost gating layer supervised by selection labels and packs these tokens into new batches, enabling a variable number of tokens to be used in parallelized GPU batch training and inference. Extensive experiments across diverse datasets and computer vision tasks demonstrate that SPA delivers superior performance and efficiency, including a 0.6 mAP improvement in object detection and a 16.4% reduction in computational costs.