Token Sparsification for Faster Medical Image Segmentation.

Area of Science:

• Computer Vision
• Machine Learning
• Medical Imaging

Background:

• Token sparsification in Vision Transformers (ViT) accelerates classification but its application to dense prediction tasks like segmentation remains unexplored.
• Existing methods for token pruning and masked image modeling (MIM) are insufficient for segmentation, leading to training inefficiencies and poor feature restoration.

Purpose of the Study:

• To develop an effective method for performing dense prediction, specifically segmentation, using sparse tokens from Vision Transformers.
• To introduce a novel pipeline that addresses the challenges of sparse token utilization in segmentation tasks.

Main Methods:

• A three-stage pipeline: sparse encoding, token completion, and dense decoding (SCD).
• Soft-topK Token Pruning (STP) for efficient sparse encoding by predicting token importance and sampling topK tokens with approximated gradients.
• Multi-layer Token Assembly (MTA) for robust token completion by integrating sparse output and intermediate tokens.

Main Results:

• The proposed SCD pipeline with STP and MTA significantly outperforms baselines in training and inference throughput (up to 120% and 60.6% higher, respectively).
• The methods maintain high segmentation quality despite using sparse tokens.
• Empirical evidence shows naïve application of existing classification sparsification techniques fails for segmentation.

Conclusions:

• The SCD pipeline, enhanced by STP and MTA, enables efficient and accurate segmentation using sparse tokens in Vision Transformers.
• This approach offers a viable solution for accelerating dense prediction tasks without compromising performance.