When Token Compression Breaks: Structural Pruning vs. Token Reduction for Robust ViT Segmentation under High Compression

Abstract: Vision Transformers (ViTs) are strong backbones for semantic segmentation, but their computational cost limits deployment. Recent token compression methods for efficient transformer-based segmentation reduce this cost by decreasing the number of tokens. However, existing evaluations primarily focus on low-to-moderate compression, leaving their behavior under aggressive compression and corrupted inputs unclear. Meanwhile, structural pruning provides an orthogonal route to efficiency by removing redundant components in the ViT architecture, but is rarely compared to token compression under a unified protocol. To bridge this gap, we benchmark representative token compression and structural pruning methods for ViT-based semantic segmentation under matched FLOPs on ADE20K and Cityscapes, together with their common-corruption variants ADE20K-C and Cityscapes-C. Our results reveal a consistent trend on both clean and corrupted inputs: token compression is highly effective at mild reductions but degrades sharply when compression becomes severe, consistent with substantial information loss from overly aggressive token reduction. In contrast, structural pruning exhibits a smoother degradation curve and is more stable at high compression. Motivated by these findings, we study a prune-then-merge pipeline that applies moderate token compression on top of a moderately pruned backbone. At comparable FLOPs, this combined strategy consistently achieves a better accuracy-robustness trade-off at high compression, offering a practical recipe for deployment-oriented ViT segmentation. Code is available at https://github.com/phatnguyencs/vit-seg-compression.

• The paper introduces a unified benchmark comparing token compression and structural pruning for ViT segmentation under aggressive efficiency constraints.
• It demonstrates that mild token compression maintains performance, while aggressive compression sharply degrades mIoU compared to the smoother decline seen with pruning.
• The proposed prune-then-merge pipeline achieves optimal accuracy–robustness trade-offs, offering practical guidelines for deployment-oriented ViT segmentation.

Compression Strategies for Efficient and Robust ViT Semantic Segmentation

Overview

This paper establishes a unified, controlled benchmark to evaluate token compression and structural pruning for Vision Transformer (ViT) semantic segmentation under aggressive efficiency constraints. Segmentation-aware token compression and architectural pruning are compared with matched FLOPs on ADE20K, Cityscapes, and their common corruption variants. Numerical results, effective-rank analysis, and qualitative visualizations reveal that while mild token compression maintains performance, aggressive token removal leads to sharp degradation, particularly in robustness. Structural pruning, by contrast, yields smoother degradation curves. A prune-then-merge (PtM) pipeline combining moderate pruning with token compression achieves superior accuracy–robustness trade-offs at high compression, offering practical guidelines for deployment-oriented ViT segmentation.

Figure 1: Compression strategies and the robustness-compute trade-off: token compression reduces spatial tokens, structural pruning reduces architectural width, and the prune-then-merge pipeline combines both; right, mIoU-compute robustness on ADE20K-C shows token compression breaking down under severe compression, pruning degrades more gracefully, and PtM dominates at fixed FLOPs.

Background: Compression in Vision Transformer Segmentation

ViT-based architectures dominate semantic segmentation benchmarks but suffer from extreme computational overhead due to dense token processing and the quadratic complexity of self-attention. Deployment therefore demands aggressive compute reduction via model compression. Two major strategies are:

• Token Compression: Adaptive reduction or merging of tokens, typically by pruning redundant patches or fusing tokens with similar features, e.g., ToMe, CTS, ALGM.
• Structural Pruning: Removal of model components such as heads, channels, or entire layers, yielding permanent model width/capacity reduction, e.g., NViT.

Evaluations previously focused only on low-to-moderate compression, leaving high-compression regime behavior—especially under corrupted inputs—underexplored. Existing work rarely systematically compares these strategies or their combinations within a rigorous, unified protocol.

Unified Benchmark: Protocol and Evaluation

The benchmark utilizes a standardized Segmenter pipeline with DeiT-B as the ViT encoder, optionally pruned or compressed. Compression is exclusively applied to the encoder, ensuring all other settings are held constant across methods. The key comparison axes:

• Token Compression Methods: General-purpose (ToMe), segmentation-aware (CTS, ALGM).
• Structural Pruning: NViT, leveraging global importance criteria.
• Prune-then-Merge (PtM): Token compression atop a moderately pruned NViT.

Performance is evaluated under matched FLOPs across ADE20K and Cityscapes, both on standard and corrupted sets (ADE20K-C/Cityscapes-C). Main metrics are clean $\text{mIoU}_{\text{clean}}$ and average corrupted $\text{mIoU}_{\text{noise}}$.

Key Findings: Regime-dependent Degradation

Figure 2: Accuracy-compute trade-offs on ADE20K (left) and Cityscapes (right): mild token compression maintains nearly baseline mIoU, but aggressive token removal results in a sharp performance drop; structural pruning shows steady, more graceful degradation; PtM delivers leading mIoU-robustness at high compression.

Across datasets and corruptions, results are regime-dependent:

• Mild Compression: Segmentation-aware token compression methods (CTS, ALGM) substantially reduce FLOPs ($\sim$1.4–1.9×) with minimal mIoU loss. For instance, CTS/ALGM retain almost all clean and corrupted performance in the mild regime.
• Aggressive Compression: All token compression methods break down, but ToMe (general merging) fails most rapidly, displaying catastrophic mIoU drops (ADE20K $\to$ $23.94/16.86$), indicative of severe spatial information loss. Segm.-aware ALGM/CTS degrade less dramatically but still lose significant accuracy and robustness.
• Structural Pruning: NViT achieves consistently smoother degradation, retaining stronger mIoU and robustness under extreme ($\sim$3.8×) compression.
• Prune-then-Merge: PtM dominates at high compression, maintaining higher mIoU and robustness than single-strategy baselines at matched compute.

Failure Analysis: Representation Collapse Under Aggressive Token Compression

The paper probes the failure modes of aggressive token suppression.

• Effective Rank Analysis: The dimensionality of token feature representations, measured by normalized effective rank, plummets under severe token compression, especially for ToMe, indicating concentrated, less discriminative features. Pruning maintains consistent representation diversity.
• Qualitative Results: Visually, aggressive token merging leads to spatial prediction incoherence—boundaries and fine structures become blurred or misclassified, as shown in local zoomed regions. NViT and PtM avoid these artifacts.

  Figure 3: Under aggressive compression on ADE20K, ALGM loses spatial coherence on fine structures and local boundaries, while NViT and Prune-then-Merge yield more faithful segmentations.

Practical Implications and Theoretical Insights

The main implication is that ViT segmentation pipelines contain both architectural and token redundancy, but these axes are not interchangeable—reducing tokens too aggressively destroys semantic integrity, especially when robustness to distributional shift matters. Pruning the model architecture provides a safer, more robust path to high compression.

Combining moderate pruning with moderate token compression—the prune-then-merge strategy—optimally exploits redundancy along both axes, yielding better accuracy and robustness trade-offs at a fixed compute budget.

On the theoretical side, the regime-dependence of robustness implicates information flow preservation and feature diversity as critical qualifiers. Aggressive token compression leads to effective rank collapse, echoing theories of representational capacity and information bottlenecks in deep architectures.

Future Directions

• Cross-Task Extension: The unified benchmark could be adapted to other dense prediction domains such as object detection and panoptic segmentation.
• Hierarchical Transformers: As hierarchical ViTs supplant flat-token architectures for high-resolution tasks, analogous benchmarks assessing multi-stage spatial compression and pruning are needed.
• Compression Co-Design: Algorithmic innovation is needed for joint, perhaps learnable, co-design of pruning and token compression, targeting explicit trade-off objectives of throughput, accuracy, and robustness.

Conclusion

This work establishes that while token compression is effective only up to moderate ratios, overcompression sharply compromises both accuracy and robustness in ViT segmentation. Structural pruning enables stable high-compression performance. The prune-then-merge pipeline leverages the complementarity of both axes, offering a practical and theoretically motivated recipe for deploying robust, efficient segmentation models. The insights motivate principled, co-designed compression approaches and provide a foundation for rigorous analysis of transformer model efficiency and robustness under deployment constraints.

(2607.02237)