Vision Transformers: From Semantic Segmentation to Dense Prediction

TL;DR

This work shifts dense visual prediction from CNN-centric designs to Vision Transformer-based architectures by introducing SETR, a patch-based Transformer encoder with no spatial downsampling, and three decoder variants. To make Transformers practical for multi-scale dense tasks, the authors propose Hierarchical Local-Global (HLG) Transformers, a pyramidal backbone combining local, window-based attention with global context via window embeddings, plus parameter-sharing and depth-wise MLP enhancements. Through extensive experiments on ADE20K, Cityscapes, and COCO, SETR and especially HLG backbones achieve state-of-the-art or competitive results across semantic segmentation, object detection, and instance segmentation, with pretraining strategies (e.g., BEiT) and auxiliary losses providing significant gains. The findings demonstrate the viability and benefits of global-context learning at full receptive field for dense visual prediction, offering a scalable alternative to convolutional pyramids with strong cross-task performance.

Abstract

The emergence of vision transformers (ViTs) in image classification has shifted the methodologies for visual representation learning. In particular, ViTs learn visual representation at full receptive field per layer across all the image patches, in comparison to the increasing receptive fields of CNNs across layers and other alternatives (e.g., large kernels and atrous convolution). In this work, for the first time we explore the global context learning potentials of ViTs for dense visual prediction (e.g., semantic segmentation). Our motivation is that through learning global context at full receptive field layer by layer, ViTs may capture stronger long-range dependency information, critical for dense prediction tasks. We first demonstrate that encoding an image as a sequence of patches, a vanilla ViT without local convolution and resolution reduction can yield stronger visual representation for semantic segmentation. For example, our model, termed as SEgmentation TRansformer (SETR), excels on ADE20K (50.28% mIoU, the first position in the test leaderboard on the day of submission) and performs competitively on Cityscapes. However, the basic ViT architecture falls short in broader dense prediction applications, such as object detection and instance segmentation, due to its lack of a pyramidal structure, high computational demand, and insufficient local context. For tackling general dense visual prediction tasks in a cost-effective manner, we further formulate a family of Hierarchical Local-Global (HLG) Transformers, characterized by local attention within windows and global-attention across windows in a pyramidal architecture. Extensive experiments show that our methods achieve appealing performance on a variety of dense prediction tasks (e.g., object detection and instance segmentation and semantic segmentation) as well as image classification.

Figures (14)

• Figure 1: Schematic illustration of the proposed SEgmentation TRansformer (SETR) (a). We first split an image into fixed-size patches, linearly embed each of them, add position embeddings, and feed the resulting sequence of vectors to a standard Transformer encoder. To perform pixel-wise segmentation, we introduce different decoder designs: (b) progressive upsampling (resulting in a variant called SETR-PUP); and (c) multi-level feature aggregation (a variant called SETR-MLA).
• Figure 2: Left: Hierarchical local-global Transformers backbone architecture. Right: Successive hierarchical local-global Transformer layer.
• Figure 3: Local-Global attention mechanism. Local attention (in dash line) is applied within each window. Global attention is applied between global feature $Z_G$ and local feature $Z_L$.
• Figure 4: Left: plain local attention; Right: dilated local attention.
• Figure 5: Qualitative results on ADE20K: SETR (right column) vs. dilated FCN baseline (left column) in each pair. Best viewed in color and zoom in.