Prompt-based Dynamic Token Pruning for Efficient Segmentation of Medical Images

TL;DR

This paper introduces a prompt-driven token pruning method for Vision Transformers that reduces computational costs in medical image segmentation by selectively processing relevant tokens, enabling faster and more efficient analysis.

Contribution

The paper presents a novel prompt-based adaptive token pruning approach that improves efficiency without sacrificing segmentation accuracy in medical imaging tasks.

Findings

Reduces tokens by approximately 35-55%

Enhances inference speed and efficiency

Maintains segmentation accuracy with fewer tokens

Abstract

The high computational demands of Vision Transformers (ViTs) in processing a large number of tokens often constrain their practical application in analyzing medical images. This research proposes a Prompt-driven Adaptive Token ({\it PrATo}) pruning method to selectively reduce the processing of irrelevant tokens in the segmentation pipeline. The prompt-based spatial prior helps to rank the tokens according to their relevance. Tokens with low-relevance scores are down-weighted, ensuring that only the relevant ones are propagated for processing across subsequent stages. This data-driven pruning strategy improves segmentation accuracy and inference speed by allocating computational resources to essential regions. The proposed framework is integrated with several state-of-the-art models to facilitate the elimination of irrelevant tokens, thereby enhancing computational efficiency while preserving segmentation accuracy. The experimental results show a reduction of $\sim$ 35-55% tokens; thus reducing the computational costs relative to baselines. Cost-effective medical image processing, using our framework, facilitates real-time diagnosis by expanding its applicability in resource-constrained environments.