Segmenter: Transformer for Semantic Segmentation

Idea

Apply Vision Transformer (ViT) for semantic segmentation, unlike DPT - Vision Transformers for Dense Prediction, this paper also uses mask Transformer as the decoder.

The class masks make the decoder focuses on specific foreground areas, and could help in extracting more relevant and localized features for segmentation

Architecture

Encoder

Adopt flattening and linear projection as the embedding method. then added with position embedding. Just like the original vision transformer paper

Decoder

The sequence of path encodings ${\mathbf{z}}_L\in {\mathbb{R}}^{N\times D}$ is decoded to a segmentation map $\mathbf{s}\in {\mathbb{R}}^{H\times W\times K}$ where $K$ is the number of classes.

The decoder learns to map patch-level encodings coming from the encoder to patch-level class scores. Next these patch-level class scores are upsampled by bilinear interpolation to pixel-level scores.

Linear

A point-wise linear layer is applied to the patch encodings ${\mathbf{z}}_L\in {\mathbb{R}}^{N\times D}$ to produce patch-level class logits ${\mathbf{z}}_{\mathrm{lin}}\in {\mathbb{R}}^{N\times K}$. The sequence is then reshaped into a 2D feature map ${\mathbf{s}}_{\mathrm{lin}}\in {\mathbb{R}}^{H\times W\times K}$. A Softmax is then applied on the class dimension to obtain the final segmentation map

Mask Transformer

Introduce a set of $K$ learnable class embeddings $\text{cls}=\left[{\mathrm{cls}}_1,\dots ,{\mathrm{cls}}_K\right]\in {\mathbb{R}}^{K\times D}$. Each class embedding is initialized randomly and assigned to a single semantic class.

The class embedding $\text{cls}$ are processed jointly with patch encoding ${\mathbf{z}}_L$ by the decoder.

Mask transformer generates $K$ masks by computing the scalar product between L2-normalized patch embeddings ${\mathbf{z}}_M^{\prime }\in {\mathbb{R}}^{N\times D}$ and class embeddings $\mathbf{c}\in {\mathbb{R}}^{K\times D}$ output by the decoder

$$\mathrm{Masks}({\mathbf{z}}_M^{\prime },\mathbf{c})={\mathbf{z}}_M^{\prime }{\mathbf{c}}^T$$

These these mask sequences are reshaped into a 2D mask to form ${\mathbf{s}}_{\mathrm{mask}}\in {\mathbb{R}}^{H/P\times W/P\times K}$ and bilinearly upsampled to the original image size to obtain a feature map $\mathbf{s}\in {\mathbb{R}}^{H\times W\times K}$.