Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Segformer++: Efficient Token-Merging Strategies for High-Resolution Semantic Segmentation

Daniel Kienzle, Marco Kantonis, Robin Schön, Rainer Lienhart

TL;DR

The paper tackles the prohibitive quadratic cost of self-attention in high-resolution semantic segmentation by introducing Segformer++, a token merging strategy adapted to the Segformer architecture. By applying merging after Spatial Reduction Attention with stage-wise rates for queries and keys/values, Segformer++ achieves substantial inference speedups (e.g., up to 61% on Cityscapes) while preserving mIoU and small-object accuracy, and it also reduces training memory and improves throughput. It presents two practical variants, Segformer++HQ and Segformer++fast, and compares them against a 2D Neighbor Merging baseline across semantic segmentation and human pose estimation tasks. The results indicate strong potential for real-time, edge-optimized deployment of transformer-based dense-prediction models on high-resolution data, with the method generalizable to other architectures that combine convolution and attention.

Abstract

Utilizing transformer architectures for semantic segmentation of high-resolution images is hindered by the attention's quadratic computational complexity in the number of tokens. A solution to this challenge involves decreasing the number of tokens through token merging, which has exhibited remarkable enhancements in inference speed, training efficiency, and memory utilization for image classification tasks. In this paper, we explore various token merging strategies within the framework of the Segformer architecture and perform experiments on multiple semantic segmentation and human pose estimation datasets. Notably, without model re-training, we, for example, achieve an inference acceleration of 61% on the Cityscapes dataset while maintaining the mIoU performance. Consequently, this paper facilitates the deployment of transformer-based architectures on resource-constrained devices and in real-time applications.

Segformer++: Efficient Token-Merging Strategies for High-Resolution Semantic Segmentation

TL;DR

The paper tackles the prohibitive quadratic cost of self-attention in high-resolution semantic segmentation by introducing Segformer++, a token merging strategy adapted to the Segformer architecture. By applying merging after Spatial Reduction Attention with stage-wise rates for queries and keys/values, Segformer++ achieves substantial inference speedups (e.g., up to 61% on Cityscapes) while preserving mIoU and small-object accuracy, and it also reduces training memory and improves throughput. It presents two practical variants, Segformer++HQ and Segformer++fast, and compares them against a 2D Neighbor Merging baseline across semantic segmentation and human pose estimation tasks. The results indicate strong potential for real-time, edge-optimized deployment of transformer-based dense-prediction models on high-resolution data, with the method generalizable to other architectures that combine convolution and attention.

Abstract

Utilizing transformer architectures for semantic segmentation of high-resolution images is hindered by the attention's quadratic computational complexity in the number of tokens. A solution to this challenge involves decreasing the number of tokens through token merging, which has exhibited remarkable enhancements in inference speed, training efficiency, and memory utilization for image classification tasks. In this paper, we explore various token merging strategies within the framework of the Segformer architecture and perform experiments on multiple semantic segmentation and human pose estimation datasets. Notably, without model re-training, we, for example, achieve an inference acceleration of 61% on the Cityscapes dataset while maintaining the mIoU performance. Consequently, this paper facilitates the deployment of transformer-based architectures on resource-constrained devices and in real-time applications.
Paper Structure (14 sections, 8 equations, 4 figures, 5 tables)

This paper contains 14 sections, 8 equations, 4 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methods
  4. Segformer Architecture
  5. Token Merging Strategies
  6. Segformer++
  7. Computational Complexities
  8. Experiments
  9. Experimental setup
  10. Inference for Semantic Segmentation
  11. Training for Semantic Segmentation
  12. Training on Human Pose Estimation
  13. Inference Speed for Different Resolutions
  14. Conclusion and Outlook

Figures (4)

  • Figure 1: Qualitative comparison of our efficient models. Even though our models are significantly faster, nearly no quality degradation can be observed. Small differences are highlighted with black arrows.
  • Figure 2: Visualization of efficient attention mechanisms
  • Figure 3: Different reduction rates $r_{q_i}$ and $r_{(kv)_i}$
  • Figure 4: Example predictions on JBD. The ground truth is visualized in black.