🤖 AI

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

The paper proposes STEP, a hybrid token-reduction framework combining dynamic superpatch merging and early-exit pruning via a lightweight CNN policy, which significantly improves the computational efficiency and throughput of Vision Transformers on high-resolution semantic segmentation tasks with minimal accuracy loss.

Original authors: Michal Szczepanski, Martyna Poreba, Karim Haroun

Published 2026-05-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Michal Szczepanski, Martyna Poreba, Karim Haroun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a tour guide leading a massive group of tourists (the "tokens") through a giant, complex city (a high-resolution image) to find specific landmarks (semantic segmentation).

In a traditional Vision Transformer (ViT), the guide treats every single brick, leaf, and blade of grass in the city as a separate person to manage. Even if a whole park is just green grass, the guide stops to talk to every single blade of grass individually. This is incredibly slow and exhausting, especially when the city is huge (high resolution).

The paper introduces a new method called STEP (SuperToken and Early-Pruning) to make this tour much faster without losing the important details. It does this in two clever ways:

1. The "Grouping" Strategy (SuperTokens)

Instead of treating every tiny piece of the image as a separate person, STEP uses a smart assistant called dCTS. This assistant looks at the city and says, "Hey, this whole block of sky is just blue, and this whole field is just green. Let's group them together!"

The Analogy: Imagine instead of managing 4,000 individual tourists, the assistant groups 100 people standing in a park into one "Super-Group." Now, the guide only has to talk to that one group representative instead of 100 individuals.
The Result: The guide can skip over boring, uniform areas (like the sky or a blank wall) and focus their energy on the complex parts of the city (like a busy market or a building with many windows). The paper found that this alone can cut the number of people the guide has to manage by 2.5 times.

2. The "Early Exit" Strategy (Pruning)

As the tour continues, some tourists figure out exactly where they are very quickly. Maybe a tourist sees a "Pizza" sign and immediately knows, "I'm in the pizza district!" They don't need to hear the rest of the tour guide's speech.

The Analogy: STEP installs "Exit Doors" at various points along the tour path. If a tourist is 100% sure of their location, they are allowed to leave the tour early. They stop walking and stop listening, saving the guide the energy of explaining the rest of the route to them.
The Result: The guide only has to keep the "confused" or "uncertain" tourists for the full duration of the tour. The paper shows that up to 40% of the tourists can be sent home early, saving a massive amount of time and energy.

The Big Picture Results

When the researchers tested this on a supercomputer (an NVIDIA A100 GPU) with very large, detailed maps (images up to 1024x1024 pixels):

Speed: The tour became much faster. In some cases, the guide could process the city 3.4 times faster than the old method.
Efficiency: The computer didn't have to do as much math. The workload dropped by up to 4 times.
Accuracy: The best part is that the guide didn't get lost. Even with fewer people and shorter tours, the guide still found the landmarks almost as accurately as before (only a tiny drop in accuracy, less than 2%).

The Catch (The "Traffic Jam")

The paper also noticed a funny side effect. Even though the guide had fewer people to manage, the speed didn't always get faster in a straight line.

The Analogy: Imagine the guide is in a car. Even if there are fewer passengers, if the driver has to constantly stop, check a map, and decide who gets to leave the car, the car might still move slowly. The act of "checking who is ready to leave" (the pruning mechanism) creates a little bit of traffic and confusion that slows things down slightly.
The Conclusion: The method is incredibly efficient at reducing the work (math), but the process of deciding who to cut off needs to be smoother to get the maximum speed boost.

In short: STEP is like a smart tour guide who groups similar tourists together and lets the confident ones leave early. This makes exploring huge, detailed cities much faster and less tiring, while still getting the job done correctly.

Technical Summary: STEP (SuperToken and Early-Pruning) for High-Resolution Semantic Segmentation

1. Problem Statement

Vision Transformers (ViTs) have achieved state-of-the-art performance in semantic segmentation due to their ability to capture long-range dependencies. However, their practical deployment is hindered by high computational and memory costs, primarily driven by the quadratic complexity of the self-attention mechanism relative to image resolution. As input resolution increases (e.g., to 1024×1024), computational cost and memory consumption grow substantially.

Existing token reduction strategies face specific limitations in dense prediction tasks:

Patch-level methods (e.g., CTS) often rely on fixed hyperparameters (e.g., a fixed number of merged patches), leading to suboptimal performance on complex images (over-merging) or simple images (under-merging).
Token-level pruning (hard pruning) is often unsuitable for semantic segmentation because it permanently removes tokens, whereas dense prediction requires preserving spatial information for pixel-wise outputs.
Hybrid approaches often lack adaptability to varying image content or fail to address the specific scalability challenges of high-resolution inputs.

The paper posits that a balance must be struck between computational efficiency and accuracy preservation, particularly for high-resolution scenarios where standard grid-based processing is inefficient.

2. Methodology: STEP Framework

The authors propose STEP (SuperToken and Early-Pruning), a hybrid token-reduction framework designed to enhance ViT efficiency without significantly compromising accuracy. STEP operates on two complementary levels:

2.1 Dynamic CTS (dCTS): Adaptive Patch Merging

Unlike static merging strategies, STEP introduces dCTS, a lightweight, class-agnostic policy network based on EfficientNetLite0.

Mechanism: The network evaluates groups of adjacent patches (windows of $2\times2$ , $4\times4$ , $8\times8$ , and $16\times16$ ) to predict a similarity score $S \in [0,1]$ .
Decision Logic: If the similarity score exceeds a predefined threshold $\tau$ , the patches are merged into a superpatch. This process follows a coarse-to-fine hierarchy to prevent conflicts between nested groups.
Output: The result is a grid of superpatches with non-uniform sizes. These are resized to a standard resolution (16×16) via bilinear interpolation and projected into the embedding space as supertokens.
Benefit: This allows the token count to adapt dynamically to the semantic homogeneity of the image content (e.g., merging large background areas while preserving resolution at object boundaries).

2.2 Early-Pruning Mechanism

STEP integrates an early-exit strategy inspired by DToP (Dynamic Token Pruning) to halt the processing of "easy" tokens.

Architecture: The ViT encoder is divided into $S$ stages. Each stage includes an auxiliary decoding head (using the Attention-to-Mask module) that computes a confidence score for each token.
Pruning Logic: Tokens with confidence scores exceeding a threshold $\tau$ are masked and halted from propagating to subsequent layers. Only low-confidence (challenging) tokens continue through the deeper layers.
Integration: This strategy is complementary to dCTS; while dCTS reduces the initial sequence length, early pruning progressively shortens the set of active tokens as the network deepens.

3. Key Contributions

The paper outlines the following specific contributions:

Novel Hybrid Framework: The proposal of STEP, which combines dynamic patch merging (dCTS) with early token pruning to address efficiency in semantic segmentation.
High-Resolution Evaluation: Extensive experiments on high-resolution images (up to 1024×1024) using ViT-Base and ViT-Large backbones. The authors claim this is the first attempt to systematically evaluate token pruning strategies specifically in the context of high-resolution semantic segmentation.
Comprehensive Analysis: Detailed ablation studies on threshold parameters for merging, optimal placement of auxiliary pruning heads, and the trade-offs between computational complexity (GFLOPs), inference speed (FPS), and accuracy (mIoU).
Scalability Demonstration: Evidence that STEP maintains competitive accuracy while significantly reducing computational complexity and inference time, particularly as image resolution and model size increase.

4. Experimental Results

Experiments were conducted on COCOStuff10k, ADE20K, and Cityscapes datasets using NVIDIA A100 GPUs.

4.1 Performance on Standard Resolution (512×512)

dCTS alone: When applied to ViT-Large, dCTS reduces the token count by a factor of 2.5 compared to standard 16×16 patching. This yields a 2.6× reduction in computational cost and a 3.4× increase in throughput with minimal accuracy drop.
Full STEP: Applying both dCTS and early pruning achieves up to a 4× reduction in computational complexity and a 1.7× gain in inference speed, with a maximum accuracy drop of no more than 2.0%.
Configuration: A single pruning head placed at the 16th or 18th layer of ViT-Large offered a favorable trade-off between efficiency and latency compared to multi-head configurations.

4.2 Performance on High Resolution (768×768 and 1024×1024)

Scalability: As resolution increases, standard SegViT suffers dramatic increases in GFLOPs and drops in FPS. STEP configurations maintain a stable trade-off.
ViT-Large at 1024×1024: The dCTS variant alone reduced GFLOPs from 2086 to 802 (2.6× less complex) and increased FPS from 12 to 41 (3.4× faster), with only a marginal mIoU drop.
Token Pruning Dynamics: Up to 40% of tokens can be confidently predicted and halted before reaching the final encoder layer. In high-resolution settings with ViT-Large, an average of 48% of tokens are halted early.

4.3 Observations on Throughput vs. FLOPs

The authors note a critical finding: Throughput (FPS) does not scale linearly with FLOP reductions.

While GFLOPs decrease significantly, the inference speed gain is sometimes limited (e.g., a 4× GFLOP reduction does not always yield a 4× FPS increase).
Cause: The masking operations required for early pruning introduce bottlenecks due to irregular control flow, dynamic memory access patterns, and tensor shape variability. The paper suggests that hardware inefficiencies, such as unnecessary GPU-to-CPU memory transfers, may be diminishing expected performance gains.

5. Significance and Claims

The paper claims that STEP represents a significant step toward making ViTs viable for high-resolution dense prediction tasks.

First of its Kind: The authors assert this is the first work to systematically assess token pruning across different resolution levels in semantic segmentation.
Efficiency-Accuracy Balance: STEP demonstrates that it is possible to achieve substantial computational savings (up to 4×) and speedups while keeping accuracy degradation within acceptable limits (<2.0%).
Adaptability: The dCTS module proves particularly robust, adapting merging rates to image complexity without requiring fixed hyperparameters, making it suitable for diverse real-world scenes.

The authors conclude that while STEP offers clear advantages in efficiency and scalability, future work must address the hardware-level inefficiencies of the pruning mechanism itself to fully realize the potential throughput gains predicted by FLOP reductions.