PaperlayerView exact revision on arXiv

Source-linked AI summary

SegViT: Semantic Segmentation with Plain Vision Transformers

Bowen Zhang, Zhi Tian, Quan Tang, Xiangxiang Chu, Xiaolin Wei, Chunhua Shen, Yifan Liu

arXiv:2210.05844v2cs.CV

TL;DR

Plain ViT segmentation methods have not fully exploited attention for mask prediction. SegViT uses learnable class-token similarities to generate masks and achieves strong benchmark results, while its Shrunk structure reduces computation by up to 40%.

  • Problem

    Plain ViT segmentation methods typically learn pixel-level representations, while the potential of their attention mechanism for segmentation remains not fully explored.

  • Method

    SegViT uses an Attention-to-Mask module that converts similarities between learnable class tokens and spatial feature maps into category masks, with QD and QU forming a Shrunk structure.

  • Results

    SegViT achieves 55.2% mIoU on ADE20K, 65.3% mIoU on PASCAL-Context, and 50.3% mIoU on COCO-Stuff-10K, while the Shrunk variant reaches 55.1% ADE20K mIoU at 373.5 GFLOPs.

  • Takeaways & Limitations

    Attention maps can serve as mask predictions in a plain ViT segmentation paradigm, while the Shrunk structure retains competitive performance with reduced computation.

Abstract

from arXiv · show

We explore the capability of plain Vision Transformers (ViTs) for semantic segmentation and propose the SegVit. Previous ViT-based segmentation networks usually learn a pixel-level representation from the output of the ViT. Differently, we make use of the fundamental component -- attention mechanism, to generate masks for semantic segmentation. Specifically, we propose the Attention-to-Mask (ATM) module, in which the similarity maps between a set of learnable class tokens and the spatial feature maps are transferred to the segmentation masks. Experiments show that our proposed SegVit using the ATM module outperforms its counterparts using the plain ViT backbone on the ADE20K dataset and achieves new state-of-the-art performance on COCO-Stuff-10K and PASCAL-Context datasets. Furthermore, to reduce the computational cost of the ViT backbone, we propose query-based down-sampling (QD) and query-based up-sampling (QU) to build a Shrunk structure. With the proposed Shrunk structure, the model can save up to $40\%$ computations while maintaining competitive performance.

1 Introduction

SegViT explores plain ViTs for semantic segmentation by converting attention similarities into category masks, while a Shrunk structure reduces computation. It reports strong benchmark performance and competitive accuracy at substantially lower cost.

  • Semantic segmentation requires pixel-level classification, but FCNs have difficulty explicitly modeling correlations among spatial locations because of limited receptive fields.
  • SegViT learns class tokens and identifies local patches with higher similarity instead of relying on the previous per-pixel classification paradigm.
  • The ATM module uses class tokens as queries, spatial feature maps as keys and values, and Sigmoid-transformed similarity maps as segmentation masks.Softmax cross-attention simultaneously updates the class tokens, while mask generation adds negligible computation.
  • The ATM decoder is presented as a new semantic-segmentation paradigm that uses spatial information in attention maps to generate category masks.
  • 55.2% mIoU is achieved on ADE20K, while PASCAL-Context reaches 65.3% mIoU and COCO-Stuff-10K reaches 50.3% mIoU with new state-of-the-art performance.
  • 40% lower computation is achieved by the Shrunk structure while maintaining competitive performance; its ADE20K variant reaches 55.1% mIoU at 373.5 GFLOPs versus 637.9 GFLOPs.The structure combines query-based down-sampling and query-based up-sampling around the ViT backbone.

2 Related Work

Related work develops semantic segmentation from per-pixel convolutional predictions toward attention-based and token-based approaches. Plain ViT decoders must also address the backbone’s uniform feature resolution and the need for high-resolution prediction maps.

  • Semantic segmentation methods use dilated convolutions, spatial pyramid pooling, or attention to enlarge context and capture long-range dependencies.
  • Recent methods decouple per-pixel classification using a fixed number of learnable tokens, global classification, and binary matching to permit overlapping feature maps.These methods still commonly calculate feature maps statically and require feature-merge modules such as FPN.
  • Figure 1 contrasts Softmax attention maps with Sigmoid masks derived from the same query-key similarity map.
  • Vision Transformer backbones provide attention-based alternatives to convolutional networks, while hierarchical variants add pooling and down-sampling for easier decoder integration.
  • High-resolution feature maps are important for segmentation, but plain non-hierarchical transformers keep feature-map resolution unchanged across layers.

3 Method

SegViT uses attention between learnable class tokens and plain ViT feature maps to generate class-specific masks, cascading ATM decoders across backbone layers. Its Shrunk structure combines query-based down- and up-sampling to reduce computation while retaining competitive performance.

  • 3.1 Encoder: Plain ViT preserves the token count across layers, providing feature maps at a fixed spatial resolution for the segmentation pipeline.The backbone reshapes the image into tokens and repeatedly applies transformer layers without changing their number.
  • 3.2 Decoder: ATM applies cross attention between learnable class embeddings and ViT feature maps, using class tokens as queries and feature maps as keys and values.The attention module inherits a transformer decoder structure and updates the class-token queries sequentially across backbone layers.
  • 3.2 Decoder: Similarity maps between class queries and spatial keys are passed through Sigmoid to produce one mask per class, while updated class tokens yield class probabilities.The masks have shape N × L and can be reshaped to N × H/P × W/P; inference combines mask groups with class probabilities.
  • 3.2 Decoder: SegViT cascades ATM modules across different ViT layers and combines their outputs to exploit low-level and higher-level feature information.The method uses the output tokens from one ATM decoder as class embeddings for the next decoder, while masks are summed orderly and supervised with classification and mask losses.
  • 3.2 Decoder: The Shrunk structure uses query-based down-sampling and up-sampling to reduce feature resolution and recover information lost by down-sampling.QD changes the resolution from 1/16 to 1/32, while QU uses low-level 1/16-resolution features to produce a 1/16-resolution output.
  • 3.2 Decoder: 40% of overall computational cost can be reduced while performance remains competitive when the Shrunk structure combines QD with QU.Applying QD alone decreases performance, whereas QU recovers the performance and preserves resolution and low-level feature information.

4 Experiments

Experiments evaluate SegViT on three semantic-segmentation benchmarks and ablate ATM, multi-layer inputs, and Shrunk components. SegViT reaches strong accuracy with reduced computational cost, while QD and QU jointly retain performance in the Shrunk structure.

  • Results on ADE20K: SegViT achieves 55.2% mIoU on ADE20K with ViT-Large, while its Shrunk version reaches 55.1% at 373.5 GFLOPs.The Shrunk version is reported as much cheaper than the ViT-Large backbone alone at 612.3 GFLOPs.
  • Results on COCO-Stuff-10K: 50.3% mIoU on COCO-Stuff-10K exceeds StrucToken by 1.2%, while the Shrunk version achieves 49.4% at 224.8 GFLOPs.The Shrunk version has computational cost similar to a dilated ResNet-101 backbone but higher performance.
  • Results on PASCAL-Context: SegViT reaches 65.3% mIoU over 59 classes and 59.3% over 60 classes on PASCAL-Context, outperforming ViT-backbone methods with less computational cost.The two metrics differ by whether the background class is included.
  • Ablation Study: 3.1% performance improvement results when mask and classification supervision are separated, compared with ATM trained using regular cross-entropy loss.Applying ATM with regular cross-entropy loss provides a 0.5% boost over the SETR-naive baseline.
  • Ablation Study: +1.3% improvement comes from adding lower-layer feature maps, with three layers empirically selected for best performance.The ablation reports further gains when more layers are included.
  • Ablation Study: QD alone decreases performance by 2.7%, but QU recovers it; using QD and QU jointly retains performance while reducing computational cost.QD uses attention-based nonlinear downsampling between keys and queries, while QU preserves resolution and supplies low-level feature information.

5 Conclusion

The paper presents SegViT, which uses attention maps for semantic-segmentation masks with plain ViT backbones. It also introduces Shrunk to reduce computation, while noting a GPU-memory limitation from global attention.

  • 5 Conclusion: SegViT derives category mask predictions during attention calculation through the Attention-to-Mask module.The method uses spatial information in attention maps for semantic segmentation.
  • 5 Conclusion: Shrunk reduces computational cost by 40% while maintaining competitive performance across semantic-segmentation benchmarks.The conclusion describes both SegViT and Shrunk as structures for semantic segmentation with ViT backbones.
  • 5 Conclusion: Global attention can consume substantial GPU memory, potentially restricting applicability on some devices.This is identified as a limitation of the proposed structures.

A Appendix

The appendix provides additional evaluation results to demonstrate SegViT’s performance.

  • A Appendix: Additional evaluation results are presented to demonstrate the performance of SegViT.

A.1 Illustration of the accuracy vs. computational cost

On ADE20K, SegViT offers the best performance among ViT-backbone methods while improving the accuracy–computational-cost trade-off. The Shrunk version substantially reduces computation while retaining competitive performance.

  • A.1 Illustration of the accuracy vs. computational cost: SegViT achieves the best performance and a better accuracy–computational-cost trade-off among methods using ViT backbones.The comparison is measured on the ADE20K validation split at single-scale inference.
  • A.1 Illustration of the accuracy vs. computational cost: The SegViT Shrunk version dramatically reduces computation while retaining competitive performance.

A.2 More Visualization Results

SegViT produces satisfactory segmentation results across indoor and outdoor scenes on ADE20K, COCO-Stuff-10K, and PASCAL-Context under single-scale inference.

  • SegViT achieves competitive segmentation results on ADE20K, COCO-Stuff-10K, and PASCAL-Context.The evaluation uses single-scale inference across indoor and outdoor scenes.

A.3 More Ablation Study Results

Ablations on a ViT-Base backbone compare decoder feature merging and loss choices, showing that ATM-based token merging improves performance over FPN.

  • 50.6% performance is obtained with ATM-based token merging, improving over FPN's 46.7% on a plain ViT-Base backbone.The comparison evaluates decoder feature merge types and loss types.
  • Plain ViT backbones maintain feature-map resolution, while their final-layer maps contain the richest semantic information.This explains why FPN, which is useful for hierarchical backbones, is less effective for plain ViTs.
  • Figure 6 reports competitive segmentation results on COCO-Stuff-10K.
  • Figure 7 reports competitive segmentation results on PASCAL-Context with 60 classes.
  • Table 9 presents ablation results for decoder methods, feature merge types, and loss types using ViT-Base throughout.
Loading 2210.05844v2…