# Token Sparsification for Faster Medical Image Segmentation

arXiv: https://arxiv.org/abs/2303.06522 · HTML: https://mlsanity.com/paper/2303.06522/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Token pruning and merging for ViTs, one of 12 directions in Efficient LLM inference and serving.
- Builds on 23 of its 29 references that are in the map.
- First posted on arXiv 11 Mar 2023.

## Placed in

- Efficient LLM inference and serving / [Token pruning and merging for ViTs](https://mlsanity.com/inference/token-pruning-and-merging-for-vits/)

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [Vision Transformers for Dense Prediction](https://mlsanity.com/paper/2103.13413/) (2021)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [Deep High-Resolution Representation Learning for Visual Recognition](https://mlsanity.com/paper/1908.07919/) (2019)
- [Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers](https://mlsanity.com/paper/2012.15840/) (2020)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [Exploring Plain Vision Transformer Backbones for Object Detection](https://mlsanity.com/paper/2203.16527/) (2022)
- [The Medical Segmentation Decathlon](https://mlsanity.com/paper/2106.05735/) (2021)
- [Self-Supervised Pre-Training of Swin Transformers for 3D Medical Image Analysis](https://mlsanity.com/paper/2111.14791/) (2021)
- [Attention Gated Networks: Learning to Leverage Salient Regions in Medical Images](https://mlsanity.com/paper/1808.08114/) (2018)
- [Benchmarking Detection Transfer Learning with Vision Transformers](https://mlsanity.com/paper/2111.11429/) (2021)
- [Reparameterizable Subset Sampling via Continuous Relaxations](https://mlsanity.com/paper/1901.10517/) (2019)
- [Differentiable Patch Selection for Image Recognition](https://mlsanity.com/paper/2104.03059/) (2021)
- [Domain Adaptive Relational Reasoning for 3D Multi-Organ Segmentation](https://mlsanity.com/paper/2005.09120/) (2020)
- [Differentiable Subset Pruning of Transformer Heads](https://mlsanity.com/paper/2108.04657/) (2021)
- [D-Former: A U-shaped Dilated Transformer for 3D Medical Image Segmentation](https://mlsanity.com/paper/2201.00462/) (2022)

## Built on by (most cited first, 4 of 4)

- [Enhancing Efficiency in Vision Transformer Networks: Design Techniques and Insights](https://mlsanity.com/paper/2403.19882/) (2024)
- [BATS: Resource-Efficient Volumetric Segmentation with Boundary-Aware Mixed-Resolution Tokens](https://mlsanity.com/paper/2607.26829/) (2026)
- [Prompt-based Dynamic Token Pruning for Efficient Segmentation of Medical Images](https://mlsanity.com/paper/2506.16369/) (2025)
- [Reg4Pru: Regularisation Through Random Token Routing for Token Pruning](https://mlsanity.com/paper/2602.02163/) (2026)

## Abstract (arXiv)

Can we use sparse tokens for dense prediction, e.g., segmentation? Although token sparsification has been applied to Vision Transformers (ViT) to accelerate classification, it is still unknown how to perform segmentation from sparse tokens. To this end, we reformulate segmentation as a sparse encoding -> token completion -> dense decoding (SCD) pipeline. We first empirically show that naively applying existing approaches from classification token pruning and masked image modeling (MIM) leads to failure and inefficient training caused by inappropriate sampling algorithms and the low quality of the restored dense features. In this paper, we propose Soft-topK Token Pruning (STP) and Multi-layer Token Assembly (MTA) to address these problems. In sparse encoding, STP predicts token importance scores with a lightweight sub-network and samples the topK tokens. The intractable topK gradients are approximated through a continuous perturbed score distribution. In token completion, MTA restores a full token sequence by assembling both sparse output tokens and pruned multi-layer intermediate ones. The last dense decoding stage is compatible with existing segmentation decoders, e.g., UNETR. Experiments show SCD pipelines equipped with STP and MTA are much faster than baselines without token pruning in both training (up to 120% higher throughput and inference up to 60.6% higher throughput) while maintaining segmentation quality.