Source-linked AI summary
Multi-Scale High-Resolution Vision Transformer for Semantic Segmentation
Jiaqi Gu, Hyoukjun Kwon, Dilin Wang, Wei Ye, Meng Li, Yu-Hsin Chen, Liangzhen Lai, Vikas Chandra, David Z. Pan
TL;DR
ViTs’ single-scale, low-resolution representations make semantic segmentation difficult because the task requires fine-grained, position-sensitive features. HRViT integrates high-resolution multi-branch architectures with co-optimized Transformer blocks, achieving stronger segmentation performance with improved efficiency on ADE20K and Cityscapes. It therefore demonstrates potential as an efficient vision backbone for semantic segmentation.
Problem
ViTs mainly produce single-scale, low-resolution representations, whereas semantic segmentation requires high position sensitivity and fine-grained image details.
Method
HRViT combines multi-branch high-resolution Transformer architecture with heterogeneous branches, efficient building blocks, augmented attention, and mixed-scale feedforward networks.
Results
HRViT achieves 50.20% mIoU on ADE20K val and 83.16% mIoU on Cityscapes val, with 1.78 higher mIoU, 28% fewer parameters, and 21% lower FLOPs than SoTA MiT and CSWin backbones on average.
Takeaways & Limitations
HRViT pushes the performance-efficiency Pareto frontier for semantic segmentation and demonstrates potential as a strong vision backbone.
Abstract
from arXiv · showhide
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.
1. Introduction
HRViT addresses the mismatch between ViTs’ single-scale, low-resolution representations and semantic segmentation’s need for precise, fine-grained features. It combines high-resolution multi-branch representations with co-optimized Transformer blocks to improve segmentation performance and efficiency.
- ViTs’ single-scale, low-resolution representations are poorly suited to semantic segmentation’s need for position sensitivity and fine-grained details.
- Existing multi-scale ViT backbones retain sequential classification-like topologies, limiting cross-scale interaction and high-resolution representation quality.
- HRNet-style multi-branch architectures preserve and repeatedly fuse parallel resolutions, but convolutional operations limit expressivity through small receptive fields and strong inductive bias.
- HRViT integrates high-resolution multi-branch architectures with ViTs, using cross-resolution fusion to learn multi-scale representations.
- HRViT improves scalability through Transformer-block redundancy reduction, performance-efficiency co-optimized building blocks, and heterogeneous branch designs.
- 50.20% mIoU on ADE20K val and 83.16% mIoU on Cityscapes val accompanied 1.78 higher mIoU, 28% fewer parameters, and 21% lower FLOPs than SoTA MiT and CSWin backbones on average.
2. Proposed HRViT Architecture
HRViT integrates a high-resolution multi-branch topology with Transformer blocks to learn multi-scale representations efficiently for semantic segmentation. Its heterogeneous branches and optimized attention, fusion, and feedforward components balance spatial precision, receptive field, expressivity, and computational cost.
- Architecture: HRViT extends ViTs with a multi-branch topology that supports multi-scale representation learning and repeated cross-resolution interaction.The architecture is inspired by HRNet while retaining Transformer-based processing.
- Architecture: Four progressive stages contain increasingly many parallel branches, with each module combining dense fusion, patch embedding, HRViTAttn, and MixCFN blocks.The n-th stage contains n branches, and stages may include multiple modules.
- Heterogeneous branch design: Branch designs allocate narrow windows and few blocks to high-resolution paths, a deeper large-window branch to medium resolution, and few large-window blocks to low resolution.This allocation uses high-resolution branches for fine-grained detail, the medium-resolution branch for receptive field and high-level features, and the low-resolution branch for coarse segmentation maps.
- Efficient attention: HRViTAttn uses cross-shaped local self-attention with fixed window dimensions, key-value projection sharing, a parallel depth-wise convolution path, and a diversity-enhanced shortcut.The design reduces computation and parameters while adding local feature aggregation and nonlinear expressiveness.
- Feedforward and fusion components: MixCFN inserts 3×3 and 5×5 depth-wise convolution paths between linear layers to extract multi-scale local information while reducing channel redundancy.The module expands channels, splits them into two branches, and applies the two depth-wise convolution paths.
3. Experiments
HRViT is evaluated against recent vision Transformer backbones on ImageNet-1K, ADE20K, and Cityscapes, including efficiency, ablation, and architecture comparisons. Results show improved performance-efficiency trade-offs through optimized multi-branch designs and components.
- Experimental setup: HRViT is pretrained on ImageNet-1K and evaluated for semantic segmentation on ADE20K and Cityscapes against Swin, Twins, MiT, and CSWin.The experiments use standard segmentation benchmarks and compare both accuracy and computational efficiency.
- ADE20K: 3.68% higher mIoU, 40% fewer parameters, and 8% less computation are reported for HRViT-b1 versus MiT-B1 on ADE20K val with the SegFormer head.HRViT-b3 also exceeds CSWin-S in mIoU while saving 23% parameters and 13% FLOPs.
- Cityscapes: +2.16 higher mIoU, 30.7% fewer parameters, and 23.1% less computation are achieved on average versus MiT and CSWin on Cityscapes val.HRViT-b3 exceeds MiT-b4 by +0.86 mIoU while using 55.4% fewer parameters and 30.7% fewer FLOPs.
- Ablation studies: Removing all block optimizations reduces efficiency and performance, while the complete techniques improve ImageNet accuracy by 0.73% and Cityscapes mIoU by +1.18 with 20% fewer parameters and 13% less computation.MixCFN, parallel convolution paths, additional nonlinearity and BN each provide reported accuracy or segmentation gains.
- Architecture comparisons: Vanilla HRNet-MiT and HRNet-CSWin retain multi-scale benefits but become inefficient because Transformer hardware costs outweigh their performance gains.HRViT instead combines heterogeneous branches and optimized components to reduce computation and parameters while improving representability.
4. Related Work
Prior semantic-segmentation backbones obtain multi-scale features through progressive downsampling, HR architectures, or hierarchical Transformer designs. HRViT differs by using a pure-ViT multi-branch topology with heterogeneous branches and block optimizations.
- Multi-scale representation learning: Earlier segmentation frameworks downsample feature maps and recover high-resolution features through upsampling, whereas HRNet maintains high-resolution representations with repeated cross-resolution fusion.Lite-HRNet and HR-NAS further explore lightweight cross-resolution exchange and searched branch settings.
- Multi-scale representation learning: HRFormer replaces HRNetV2 residual blocks with Swin Transformer blocks, while HRViT uses a pure-ViT backbone with a novel multi-branch topology.The paper distinguishes HRViT from direct convolution-to-attention substitution by emphasizing heterogeneous branches and block optimization.
- Multi-scale ViT backbones: Several multi-scale ViTs use hierarchical architectures to generate progressively down-sampled pyramid features, contrasting with HRViT’s parallel multi-branch design.HRViT is presented as combining high-resolution architectural benefits with self-attention.
5. Conclusion
The paper presents HRViT as an efficient multi-scale high-resolution ViT backbone for semantic segmentation. Its multi-branch architecture and branch-block co-optimization improve representation quality and efficiency, while evaluation shows gains over state-of-the-art ViT backbones.
- HRViT uses multi-branch architecture to learn multi-scale representations for semantic segmentation.The architecture supports high-resolution representations through cross-scale interaction.
- Heterogeneous branch designs and jointly optimized building blocks improve HRViT efficiency.The optimized components include efficient embedding layers, augmented cross-shaped attentions, and mixed-scale convolutional FFNs.
- The authors identify branch-block co-optimization as key to efficient high-resolution ViT integration.
- HRViT outperforms state-of-the-art ViT backbones on semantic segmentation with significant performance and efficiency improvements.
A.1. Image classification on ImageNet-1K
HRViT is pretrained for ImageNet-1K classification using an HRNetV2 classification head and a 300-epoch AdamW training schedule with cosine decay and warm-up.
- HRViT is pretrained on ImageNet-1K for image classification with an HRNetV2 classification head.Four multi-scale HRViT outputs are processed by convolutional bottleneck blocks before classification.
- The ImageNet-1K training schedule runs for 300 epochs with AdamW, cosine learning-rate decay, and 20 epochs of linear warm-up.The initial learning rate is 1e-3, with batch size 1,024 and weight decay 0.05.
A.2. Semantic segmentation on ADE20K and Cityscapes
The semantic-segmentation experiments evaluate HRViT on ADE20K and Cityscapes using dataset-specific resolutions, augmentation, optimization schedules, and inference settings.
- ADE20K contains 25K images across 150 semantic categories, while Cityscapes contains 5000 fine-annotated images across 19 categories.
- Training crops are 512×512 for ADE20K and 1024×1024 for Cityscapes, with random flipping, re-scaling, and photometric distortion.
- Segmentation models use AdamW for 160 k iterations with a poly learning-rate schedule and 1,500 warm-up steps.The initial learning rate is 6e-5, with weight decay 0.01 and batch sizes of 16 for ADE20K and 8 for Cityscapes.
- Inference uses 512×2048 images for ADE20K validation and 1024×2048 images for Cityscapes validation.Cityscapes inference uses sliding-window testing with 1024×1024 patches.
B.1. Semantic segmentation on ADE20K with UperNet head
With the UperNet segmentation head, HRViT largely retains its advantages, but the head dominates model cost and reduces the visibility of backbone computation savings.
- HRViT mostly maintains its advantages on ADE20K when evaluated with the UperNet head.
- UperNet accounts for more than 89% of computation and more than 54% of parameters, diluting backbone slimming benefits.
B.2. Different block assignment strategies
Table 8 compares block assignment strategies on HRViT-b1’s third low-resolution branch across ImageNet-1K and Cityscapes. Concentrating blocks in one module weakens cross-resolution fusion, while even assignment preserves information and detail more effectively.
- Table 8 compares different block assignment strategies on HRViT-b1’s third low-resolution branch.
- Concentrating more blocks in one module, from 6-6-6-2 to 17-1-1-1, degrades ImageNet classification accuracy and Cityscapes segmentation mIoU.
- The observed degradation occurs as benefits from cross-resolution fusion diminish with increasingly concentrated block assignments.
- The authors attribute this trend to information loss in the deep low-resolution module.
- Even block assignment improves information interaction and detail preservation.