TLDR: STEP (SuperToken and Early-Pruning) is a new method that makes Vision Transformers (ViTs) more efficient for semantic segmentation, especially with high-resolution images. It combines dynamic patch merging (dCTS) to create larger “supertokens” from similar image regions and an early-pruning mechanism that halts processing for confidently predicted tokens. This reduces computational cost by up to 4x and increases inference speed by up to 1.7x with a minimal accuracy drop, making ViTs more practical for demanding tasks.
Vision Transformers (ViTs) have become a cornerstone in advanced computer vision tasks like semantic segmentation, where the goal is to classify every pixel in an image. While incredibly powerful, ViTs come with a significant drawback: high computational and memory costs, especially when dealing with high-resolution images. This is where a new approach called STEP, or SuperToken and Early-Pruning, steps in to make these powerful models more efficient without sacrificing much accuracy.
Making ViTs More Efficient with STEP
STEP is a clever hybrid framework designed to reduce the number of “tokens” that a ViT has to process. In simple terms, ViTs break down an image into small patches, and each patch becomes a token. Processing too many tokens is computationally expensive. STEP tackles this in two main ways: dynamic patch merging and early token pruning.
Dynamic Patch Merging with dCTS
Imagine an image with a large, uniform area, like a clear sky. Does a ViT really need to process every tiny patch in that sky individually? Probably not. This is the idea behind dynamic CTS (dCTS), a core component of STEP. Instead of treating every 16×16 pixel patch as a separate token, dCTS intelligently merges adjacent patches that are semantically similar into larger “superpatches.” These superpatches then become “supertokens.” This process is guided by a lightweight network that assesses the similarity of neighboring patches. If a group of patches likely belongs to the same semantic class (e.g., all part of the sky), they are merged. This significantly reduces the initial number of tokens the ViT has to handle, leading to immediate savings in computation and memory.
Early Token Pruning
The second part of STEP involves “early-pruning.” As tokens move through the ViT’s layers, some become easier to predict than others. STEP incorporates an early-exit mechanism, inspired by a method called DToP, which identifies these “easy-to-predict” tokens at intermediate layers. Once a token reaches a certain confidence level in its prediction, it’s “halted” and removed from further processing in deeper layers. This means only the more challenging or uncertain tokens continue through the entire network, further reducing the computational load.
Real-World Impact and Performance
The researchers evaluated STEP on high-resolution semantic segmentation benchmarks, including images up to 1024×1024 pixels, using popular datasets like COCOStuff10k, ADE20K, and Cityscapes. The results are quite impressive:
- When dCTS is used alone, the number of tokens can be reduced by a factor of 2.5 compared to standard methods. This translates to a 2.6x reduction in computational cost and a 3.4x increase in processing speed (throughput) when using a large ViT model.
- Applying the full STEP framework pushes efficiency even further, achieving up to a 4x reduction in computational complexity and a 1.7x gain in inference speed. This is achieved with a maximum accuracy drop of no more than 2.0%, which is a small trade-off for such significant efficiency gains.
- Up to 40% of tokens can be confidently predicted and halted before reaching the final layers of the network.
While STEP offers substantial computational savings, the researchers noted that the inference speed (FPS) doesn’t always increase proportionally to the reduction in computational operations (GFLOPs). This is likely due to hardware inefficiencies introduced by the dynamic masking operations, which can cause bottlenecks. Further optimization in how these dynamic operations are handled on GPUs is an area for future work.
Also Read:
- Boosting Video Encoding Efficiency with ResidualViT
- ToMA: Practical Speedups for High-Fidelity Image Generation
A Step Forward for Efficient Vision Transformers
STEP represents a significant advancement in making Vision Transformers more practical for real-world applications, especially those requiring high-resolution image processing like autonomous driving or medical imaging. By intelligently merging patches and pruning confident tokens early, STEP allows ViTs to maintain strong performance while drastically cutting down on computational and memory demands. This research paves the way for more scalable and efficient semantic segmentation models. You can read the full research paper here.
