[Paper Review] Multi-Scale High-Resolution Vision Transformer for Semantic Segmentation

Jiaqi Gu, Hyoukjun Kwon+7 more|arXiv (Cornell University)|Nov 1, 2021

Advanced Neural Network Applications39 references4 citations

TL;DR

HRViT proposes a multi-scale, high-resolution Vision Transformer backbone for semantic segmentation by integrating a multi-branch architecture with self-attention mechanisms to enhance spatial precision and semantic richness. It achieves state-of-the-art performance with 50.20% mIoU on ADE20K and 83.16% on Cityscapes, outperforming SOTA models like MiT and CSWin by +1.78 mIoU while reducing parameters by 28% and FLOPs by 21%.

ABSTRACT

Vision Transformers (ViTs) have emerged with superior performance on computer vision tasks compared to convolutional neural network (CNN)-based models. However, ViTs are mainly designed for image classification that generate single-scale low-resolution representations, which makes dense prediction tasks such as semantic segmentation challenging for ViTs. Therefore, we propose HRViT, which enhances ViTs to learn semantically-rich and spatially-precise multi-scale representations by integrating high-resolution multi-branch architectures with ViTs. We balance the model performance and efficiency of HRViT by various branch-block co-optimization techniques. Specifically, we explore heterogeneous branch designs, reduce the redundancy in linear layers, and augment the attention block with enhanced expressiveness. Those approaches enabled HRViT to push the Pareto frontier of performance and efficiency on semantic segmentation to a new level, as our evaluation results on ADE20K and Cityscapes show. HRViT achieves 50.20% mIoU on ADE20K and 83.16% mIoU on Cityscapes, surpassing state-of-the-art MiT and CSWin backbones with an average of +1.78 mIoU improvement, 28% parameter saving, and 21% FLOPs reduction, demonstrating the potential of HRViT as a strong vision backbone for semantic segmentation.

Motivation & Objective

• Address the limitation of Vision Transformers (ViTs) in semantic segmentation due to single-scale, low-resolution feature representations.
• Enhance ViTs with high-resolution, multi-scale feature learning by integrating a multi-branch architecture inspired by HRNet.
• Improve efficiency and scalability of multi-branch ViTs through joint optimization of attention blocks, feedforward networks, and embedding layers.
• Balance performance and efficiency in ViT-based backbones for dense prediction tasks on resource-constrained platforms like AR/VR devices.
• Demonstrate that pure ViT backbones with heterogeneous multi-branch designs can surpass convolutional and sequential ViT baselines in segmentation accuracy and efficiency.

Proposed method

• Introduce a multi-branch, high-resolution architecture that maintains and fuses features across multiple scales throughout the network, enabling cross-scale interaction.
• Propose an augmented local self-attention mechanism that removes redundant keys/values and adds parallel convolution paths and nonlinearity for enhanced expressivity and efficiency.
• Employ mixed-scale convolutional feedforward networks (FFNs) to strengthen multi-scale feature extraction in the MLP blocks.
• Design an efficient HR convolutional stem and patch embedding layers to preserve low-level, fine-grained features with reduced computational cost.
• Implement heterogeneous branch designs where different branches use distinct architectural components (e.g., varying window sizes, attention types) to balance performance and efficiency.
• Apply branch-block co-optimization techniques, including redundancy reduction in linear layers and auxiliary shortcuts, to improve model efficiency without sacrificing representational power.

Experimental results

Research questions

• RQ1 Can a multi-branch Vision Transformer architecture effectively learn high-resolution, multi-scale representations for semantic segmentation?
• RQ2 How can the high computational cost of multi-branch ViTs be mitigated while preserving or improving performance?
• RQ3 What architectural components and co-optimization strategies are most effective in balancing efficiency and representation quality in high-resolution ViTs?
• RQ4 Does replacing residual blocks in HRNet with ViT blocks yield better performance than sequential ViT backbones in dense prediction tasks?
• RQ5 Can a pure ViT backbone with heterogeneous multi-branch design outperform both convolutional HRNet variants and sequential ViT models in semantic segmentation?

Key findings

• HRViT achieves 50.20% mIoU on the ADE20K validation set, surpassing the previous SOTA (MiT and CSWin) by an average of +1.78 mIoU.
• On the Cityscapes validation set, HRViT attains 83.16% mIoU, demonstrating a significant performance gain over SOTA ViT backbones.
• HRViT reduces model parameters by 28% and FLOPs by 21% compared to SOTA models, indicating substantial efficiency improvements.
• The optimal window size for attention in HRViT-b3 is 9, yielding a 0.34 mIoU improvement over window size 7 with only 0.8% more FLOPs.
• Excessive window sizes (e.g., 15) degrade performance by 0.26 mIoU and increase FLOPs by 3.7%, indicating a performance drop beyond optimal settings.
• The vanilla HRNet-ViT baseline with direct block substitution fails to scale efficiently due to high computational cost, highlighting the necessity of co-optimization techniques in HRViT.