Source-linked AI summary
TopFormer: Token Pyramid Transformer for Mobile Semantic Segmentation
Wenqiang Zhang, Zilong Huang, Guozhong Luo, Tao Chen, Xinggang Wang, Wenyu Liu, Gang Yu, Chunhua Shen
TL;DR
TopFormer addresses the computational burden that limits Vision Transformers for dense prediction on mobile devices. It uses a token pyramid and semantics injection to build scale-aware hierarchical features, achieving better segmentation results than MobileNets with lower latency, including a 5% mIoU gain over MobileNetV3 on ADE20K.
Problem
Vision Transformers' full-attention mechanism requires computational resources beyond many mobile and embedded devices, limiting their use for dense prediction tasks such as semantic segmentation.
Method
TopFormer pools tokens from multiple scales into compact scale-aware semantics and fuses them with corresponding tokens to augment hierarchical representations for dense prediction.
Results
TopFormer obtains better results than MobileNets with lower latency on an ARM-based device, while its base model achieves 5% higher mIoU than MobileNetV3 on ADE20K.
Takeaways & Limitations
The tiny TopFormer enables real-time segmentation on an ARM-based mobile device with competitive results and offers a favorable accuracy–computational-cost trade-off.
Takeaways & Limitations
Window partitioning remains time-consuming on mobile devices, while reducing token numbers can sacrifice recognition accuracy.
Abstract
from arXiv · showhide
Although vision transformers (ViTs) have achieved great success in computer vision, the heavy computational cost hampers their applications to dense prediction tasks such as semantic segmentation on mobile devices. In this paper, we present a mobile-friendly architecture named \textbf{To}ken \textbf{P}yramid Vision Trans\textbf{former} (\textbf{TopFormer}). The proposed \textbf{TopFormer} takes Tokens from various scales as input to produce scale-aware semantic features, which are then injected into the corresponding tokens to augment the representation. Experimental results demonstrate that our method significantly outperforms CNN- and ViT-based networks across several semantic segmentation datasets and achieves a good trade-off between accuracy and latency. On the ADE20K dataset, TopFormer achieves 5\% higher accuracy in mIoU than MobileNetV3 with lower latency on an ARM-based mobile device. Furthermore, the tiny version of TopFormer achieves real-time inference on an ARM-based mobile device with competitive results. The code and models are available at: https://github.com/hustvl/TopFormer
1. Introduction
TopFormer targets mobile semantic segmentation by combining CNN efficiency with Transformer semantics while avoiding the computational and latency limitations of existing approaches.
- Vision Transformers’ full attention has quadratic token complexity, creating heavy computational demands for mobile and embedded dense prediction.
- Windowed attention reduces computation but can remain time-consuming on mobile devices, while token reduction may sacrifice recognition accuracy.
- TopFormer combines a lightweight CNN-based Token Pyramid Module with a Transformer-based Semantics Extractor for mobile-friendly dense prediction.
- Scale-aware semantics are split across channels and fused with corresponding tokens to create hierarchical features for segmentation.
- Experiments cover ADE20K, Pascal Context, COCO-Stuff, and COCO object detection, with latency measured on an off-the-shelf ARM-based computing core.
- 5% higher mIoU than MobileNetV3 is reported for TopFormer-Base on ADE20K, with lower latency on an ARM-based mobile device.
2. Related Work
Related work spans lightweight Vision Transformers, efficient CNNs, and mobile semantic segmentation, motivating a Transformer that can outperform MobileNets with lower segmentation latency.
- Light-weight Vision Transformers: Lightweight Vision Transformers reduce token counts or combine CNNs and Transformers, but many remain computationally heavy or lack mobile latency advantages.
- Light-weight Vision Transformers: MobileViT and Mobile-Former improve classification performance over MobileNets, yet reported mobile latency does not show corresponding advantages.
- Efficient Convolutional Neural Networks: Efficient CNNs use inverted bottlenecks, depth-wise convolutions, channel shuffling, or permutation operators to reduce computation and improve information flow.
- Mobile Semantic Segmentation: Mobile segmentation methods pursue efficiency through multi-scale inputs, cascade networks, lightweight backbones, and cross-level feature aggregation.
3. Architecture
TopFormer builds a token pyramid, extracts scale-aware global semantics from pooled multi-scale tokens, injects them into local tokens, and decodes the augmented hierarchy.
- Architecture: The architecture contains a Token Pyramid Module, Semantics Extractor, Semantics Injection Module, and Segmentation Head.
- Token Pyramid Module: The Token Pyramid Module uses stacked MobileNet blocks to process images into tokens across multiple scales.
- Semantics Extractor: Tokens from different scales are average-pooled to a small common resolution, concatenated by channel, and passed to the Vision Transformer.
- Semantics Extractor: The Vision Transformer applies global self-attention to pooled multi-scale tokens, producing full-image receptive fields and scale-aware semantics.
- Semantics Injection Module: The Semantics Injection Module combines local tokens with global semantics to alleviate their semantic gap before fusion.
- Segmentation Head: The segmentation head upsamples tokens to a common size, sums all scales element-wise, and applies two convolutional layers.
- Architecture and Variants: TopFormer-Tiny, TopFormer-Small, and TopFormer-Base provide variants with different complexity targets; the last three scales feed SIM and the segmentation head.
4. Experiments
The experiments evaluate semantic segmentation across public datasets, analyze implementation and competing methods, study component effects through ablations, and test object-detection generalization.
- Experiments cover several public datasets and compare implementation details and semantic-segmentation results with other approaches.
- Ablation studies analyze the effectiveness and efficiency of different network components.
- Object-detection experiments assess the method’s generalization ability beyond semantic segmentation.
4.1. Semantic Segmentation
TopFormer is evaluated for mobile semantic segmentation across multiple datasets, with ablations examining its token pyramid, semantics extractor, semantic injection, and computational profile. Results show strong accuracy–efficiency trade-offs, including competitive or superior performance against CNN- and ViT-based baselines and real-time inference for the tiny model.
- Experimental Setup: TopFormer is evaluated on ADE20K, PASCAL Context, and COCO-Stuff using mIoU, with latency measured on an ARM-based computing core.Full-precision models are converted to TNN before latency measurement.
- ADE20K Results: 3.8% higher accuracy gives TopFormer-small an advantage over LR-ASPP with comparable latency, while TopFormer-tiny matches performance using 0.6G versus 1.3G computation.The LR-ASPP comparison uses a MobileNetV3-Large-reduce backbone and reports 81 ms latency for that baseline.
- ADE20K Results: 1.2% higher accuracy and fewer computations distinguish TopFormer-small from HR-NAS-B, while TopFormer-base matches SegFormer with 1.8 GFLOPs versus 8.4 GFLOPs.HR-NAS-B reports 2.2 GFLOPs and 34.9 mIoU; SegFormer reports 37.4 mIoU and 3.8M parameters.
- ADE20K Results: 32ms inference enables TopFormer-tiny to perform real-time segmentation on an ARM-based mobile device after resizing inputs to 448 × 448.The paper reports a slight performance drop at this input resolution.
- Ablation Studies: Ablations report about 10% mIoU gain from SASE and 2.4% mIoU gain from MHSA, with SASE outperforming PSP and ASPP at much lower computation cost.These experiments support the scale-aware semantics design for mobile devices.
- Cross-Dataset Results: TopFormer-base is 8% more accurate than a MobileNetV3 model with comparable computation on COCO-Stuff, while Pascal Context results are best among the compared CNN- and ViT-based approaches.Pascal Context evaluation considers both 59 categories and 60 categories including background.
4.2. Object Detection
TopFormer is evaluated for object detection on COCO using RetinaNet with different backbones to produce feature pyramids.
- COCO evaluation uses 118K training images, 5K validation images, and 20K test images.
- All models are trained on COCO train2017 and evaluated on the val2017 split.
- RetinaNet is used as the object detection method, with different backbones producing feature pyramids.
5. Conclusion and Limitations
TopFormer combines CNN and ViT advantages to balance accuracy and computational cost, while its tiny version enables real-time ARM-based inference. The authors identify minor object-detection improvements as the major limitation.
- TopFormer combines CNN and ViT advantages to achieve a trade-off between accuracy and computational cost.
- The tiny TopFormer version provides real-time inference on an ARM-based mobile device with competitive results.
- Minor improvements on object detection are identified as TopFormer’s major limitation.
- Future work will target improved object detection and applications in dense prediction.
6. ImageNet Pre-training
The paper evaluates TopFormer for ImageNet classification using a classification architecture that converts global semantics into class scores, with results reported in Table 12.
- The classification architecture appends average pooling and a Linear layer to global semantics to produce class scores.
- ImageNet classification uses ImageNet pre-trained parameters for initialization.
- For 224 × 224 input images, the Semantics Extractor receives target-resolution input tokens at 1 32×32 of the input size.
- Classification results are reported in Table 12.
7. Network Structure
The detailed network structures show that the ViT-based Semantics Extractor accounts for most parameters, despite the Token Pyramid Module having the most layers.
- The Token Pyramid Module contains the most layers in the detailed network structures.
- The ViT-based Semantics Extractor accounts for the vast majority of model parameters.
8. The Performance on Cityscapes
On Cityscapes, TopFormer is evaluated with full- and half-resolution inputs, with the full-resolution variant improving mIoU over a MobileNetV2-based L-ASPP while using less computation.
- Evaluation setup: The experiments directly feed full-resolution and half-resolution inputs into trained segmentation models for testing.This evaluation examines the accuracy–computation trade-off at large input resolutions.
- Evaluation setup: Full-resolution testing uses 1024 × 2048 inputs, while half-resolution testing uses 512 × 1024 inputs.These variants are denoted Ours(f) and Ours(h), respectively.
- Results: 2.6% higher mIoU than MobileNetV2-based L-ASPP is achieved by TopFormer with full-resolution input and lower computation.The full-resolution variant is denoted Ours(f).
- Architecture: The proposed architecture uses a Token Pyramid Module and a Semantics Extractor whose configurations include convolutional blocks and Transformer depth and head counts.The architecture configurations are specified for 512 × 512 inputs.