Where Do Tokens Go? Understanding Pruning Behaviors in STEP at High Resolutions

Abstract Vision Transformers (ViTs) achieve state-of-the-art performance in semantic segmentation but are hindered by high computational and memory costs. To address this, we propose STEP (SuperToken and Early-Pruning), a hybrid token-reduction framework that combines dynamic patch merging and token pruning to enhance efficiency without significantly compromising accuracy. At the core of STEP is dCTS, a lightweight CNN-based policy network that enables flexible merging into superpatches. Encoder blocks integrate also early-exits to remove high-confident supertokens, lowering computational load. We evaluate our method on high-resolution semantic segmentation benchmarks, including images up to $$1024\times1024$$ , and show that when dCTS is applied alone, the token count can be reduced by a factor of 2.5 compared to the standard $$16\times$$ pixel patching scheme. This yields a $$2.6\times $$ reduction in computational cost and a $$3.4\times $$ increase in throughput when using ViT-Large as the backbone. Applying the full STEP framework further improves efficiency, reaching up to a $$4\times $$ reduction in computational complexity and a $$1.7\times $$ gain in inference speed, with a maximum accuracy drop of no more than 2.0%. With the proposed STEP configurations, up to 40% of tokens can be confidently predicted and halted before reaching the final encoder layer.