SASFormer: Transformers for Sparsely Annotated Semantic Segmentation

Sparsely annotated semantic segmentation aims to address the segmentation problem using minimally labeled visual areas. What’s the point[2] brought point-level labels to the semantic segmentation problem for the first time, which boost the performance by introducing objectiveness priors. After that, ScribbleSup[4] extends point-level labels into the scribbling ones in order to bridge the performance gap between fully annotated methods and unlabeled ones. ScribeSup constructs a graphical model to transmit information from scribbles to unknown pixels, hence providing semantic guidance for unknown pixels. Since then, approaches for sparse annotated semantic segmentation have emerged and developed. In an attempt to enhance performance at the lowest feasible cost, researchers have begun investigating the association between labeled and unlabeled regions. For example, Tang et al. [7] offered a variety of regularized losses for sparsely supervised segmentation based on dense CRF and kernel cut. BPG [9] tries to integrate semantic characteristics and textural information by constructing a prediction refinement network, whilst a boundary regression network is introduced to facilitate performance by generating clearly defined semantically separate parts. Recently, SPML [11] tackles this challenge by developing a semi-supervised metric learning approach with four unique forms of attraction and repulsion relationships TEL [1] presents a tree energy loss, in which minimal spanning trees are constructed to represent low-level and high-level pair-wise affinities. The majority of these techniques either use multi-stage processes for progressive inference or include intricate frameworks for supplementary training. While in this work, we make the first attempt to model relationships between distinct areas and offer category recommendations to unlabeled regions based on vision transformer.

The Transformer was initially proposed to model long-term dependencies in natural language processing tasks [15]. In 2020, Alexey Dosovitskiy introduced the pure Transformer architecture, which achieved remarkable results in image classification [16]. Subsequently, the Transformer has been widely employed in various computer vision tasks, such as object detection [17, 18, 19], semantic segmentation [12, 20] and video processing [21, 22]. Recently, researchers have started to utilize the long-term dependencies mechanism of the Transformer to capture the relationship between image categories and local image features, aiming to optimize weakly supervised problems [23]. However, these methods mainly focus on utilizing the correlation between class tokens and patch tokens, while few works consider the optimization and utilization of the correlation among different patch tokens. Our work takes an early attempt to delve deeply into the issue of imprecise correlation among different patch tokens in the Transformer block and proposes a simple yet effective solution.

As depicted in Fig 2, given an image $I\in\mathbb{R}^{W\times H\times 3}$, we partition it into $p\times p$ resolution patches. These patches are fed into hierarchical transformer encoder to provide multilevel characteristics. We pass multi-level features into the decoder block in order to obtain segmentation probabilities $\mathbf{P}$ with resolution $M\times C$, where $M$ is the patch number and $C$ is the number of categories. On the one hand, we build multi-level patch attention maps with transformer layers in encoder blocks during the training phase. In addition, segmentation probabilities are interpolated into resolutions corresponding to the width of patch attention maps. Patch attention maps are then multiplied by interpolated segmentation probabilities, yielding propagated segmentation probabilities $\mathbf{Y}$. For labeled regions of interest, we perform standard cross entropy loss function $L_{seg}$ to supervise segmentation probabilities with ground truth labels. For unlabeled regions, we introduce an affinity loss function $L_{aff}$ to match segmentation probabilities to propagated segmentation probabilities. During testing, segmentation probabilities are immediately assigned to the class with the maximum probability at each pixel, resulting in the final segmentation map. The overall loss function is defined as follows:

$$L=L_{seg}+\alpha*L_{aff}$$

(1)

We introduce Segformer [12] as our backbone in order to describe our method more conveniently. More results of different backbones can be referred to Sec 4.3. Hierarchical transformer encoder contains $L=4$ encoder blocks. Each encoder block contains $N_{l}$ consecutive transformer layers. We get the patch attention map with the efficient self-attention of transformer layer, which is defined as follows:

$$\mathbf{A}_{l,n}=softmax(\frac{\mathbf{Q}_{l,n}\mathbf{K}_{l,n}^{\mathbf{T}}}{\sqrt{D}})$$

(2)

where $\mathbf{Q}_{l,n}\in\mathbb{R}^{M_{l}\times D}$ and $\mathbf{K}_{l,n}\in\mathbb{R}^{M^{\prime}_{l}\times D}$ are the query and key representations of $n$-th transformer layer in $l$-th encoder block, respectively. $M_{l}\in\{\frac{W*H}{4},\frac{W*H}{8},\frac{W*H}{16},\frac{W*H}{32}\}$ denotes resolution size in different encoder blocks. $M^{\prime}_{l}$ denotes resolution size from $M_{l}$ after sequence reduction process[12] to be more efficient. $D$ indicates dimension of patch embeddings. $\mathbf{T}$ is the transpose operator.

Efficient attention map $\mathbf{A}_{l,n}\in\mathbb{R}^{M_{l}\times M^{\prime}_{l}}$ records dependency of each patch token of $n$-th transformer layer in $l$-th encoder block. We aggregate these attention maps over transformer layers in each decoder block as the following:

$$\mathbf{A}_{l}=\frac{1}{N}\sum_{n=1}^{N}\mathbf{A}_{l,n}$$

(3)

where $\mathbf{A}_{l}\in\mathbb{R}^{M_{l}\times M^{\prime}_{l}}$ means patch attention map in $l$-th encoder block.

To better comprehend what information patch attention maps capture at various levels, we randomly draw a point with a red pentagram in the inputs and analyze the dependency between the reference point and other pixels on patch attention maps of various encoder blocks, as depicted in the baseline of Fig 3. $\mathbf{A}^{i,j}_{l,n}$ means the amount of $j$-th input token embedding contributes to $i$-th output token of $n$-th transformer layer in $l$-th encoder block. Likewise, $\mathbf{A}^{i}_{l}$ means the aggregated amount of each input token embedding contributes to $i$-th output token in $l$-th encoder block. We construct $\mathbf{A}^{r}_{1}$, $\mathbf{A}^{r}_{2}$, $\mathbf{A}^{r}_{3}$ and $\mathbf{A}^{r}_{4}$ to illustrate dependency between the reference point and each pixel in the first, second, third, forth encoder block, respectively.

It can be seen that $\mathbf{A}_{1}$ and $\mathbf{A}_{2}$ prefer to assign larger weights to patches that have the same color and texture as the reference point. For instance, $\mathbf{A}^{r}_{1}$ and $\mathbf{A}^{r}_{2}$ of the first input emphasize the white color of the bird, while the second input highlights the dark brown color of the material. $\mathbf{A}_{3}$ and $\mathbf{A}_{4}$ are more concerned with semantic consistency, although $\mathbf{A}_{4}$ is more class-specific. Patch attention maps establish correlations between pixels in color space and high-level characteristics, facilitating the propagation of segmentation information from labeled to unlabeled pixels. However, transformer tends to distract attention from targets to background area[18], which can be clearly seen in those of baseline($\mathbf{A}_{2}$) and baseline($\mathbf{A}_{3}$). We hypothesize that the inability of attention maps to capture dependence, which impacts the transfer of information across pixels, is partly manifested in distractions from irrelevant background.

For each encoder block, we interpolate segmentation probabilities $\mathbf{P}$ into the resolution of $M^{\prime}_{l}$ and $M_{l}$, yielding $\mathbf{P}_{l}^{{}^{\prime}}$ and $\mathbf{P}_{l}$ respectively. Then we transfer segmentation predictions from labeled pixels to unlabeled pixels by multiplying patch attention map $\mathbf{A}_{l}$ and interpolated segmentation probabilities $\mathbf{P}_{l}^{{}^{\prime}}$, which is described as follows:

$$\mathbf{Y}_{l}=\mathbf{A}_{l}\otimes\mathbf{P}_{l}^{{}^{\prime}}$$

(4)

where $\mathbf{Y}_{l}$ is propagated segmentation probabilities, $\otimes$ denotes matrix multiplication.

$\mathbf{Y}_{l}$ and $\mathbf{P}_{l}$ are normalized along the category dimension, which is defined as:

$$\mathbf{Y}^{*(m_{l},c)}_{l}=\frac{\exp(\mathbf{Y}^{(m_{l},c)}_{l})}{\sum_{c=1}^{C}{\exp(\mathbf{Y}^{(m_{l},c)}_{l})}}\qquad 1\leq m_{l}\leq M_{l}$$

(5)

$$\mathbf{P}^{*(m_{l},c)}_{l}=\frac{\exp(\mathbf{P}^{(m_{l},c)}_{l})}{\sum_{c=1}^{C}{\exp(\mathbf{P}^{(m_{l},c)}_{l})}}\qquad 1\leq m_{l}\leq M_{l}$$

(6)

where $m_{l}$ and $c$ denote patch index and class index. $\mathbf{Y}^{(m_{l},c)}_{l}$ and $\mathbf{P}^{(m_{l},c)}_{l}$ mean the position of row $m_{l}$, column $c$ of matrix $\mathbf{Y}_{l}$ and $\mathbf{P}_{l}$.

The affinity loss $L_{aff}$ attempts to maximize the similarity between propagated probabilities and the interpolated segmentation probabilities over all encoder blocks. Formally, the affinity loss is defined as:

$$L_{aff}=\frac{1}{L}\sum_{l=1}^{L}{\|\mathbf{Y}_{l}^{*}-\mathbf{P}_{l}^{*}\|_{1}}$$

(7)

where $\|\|_{1}$ and L denote L1 regularization term and encoder block number, respectively.

In the process of segmentation training, we obtain additional bonus information, as illustrated in Fig 3. Under the constraint of consistency between propagated and interpolated segmentation probabilities, pair-wise dependencies between patch tokens with different labels on patch attention maps are simultaneously deteriorated, while pair-wise dependencies between patch tokens with the same label are enhanced. By comparing to the baseline in Fig 3, it is evident that the background noise in the patch attention map has been suppressed in our proposed SASFormer.

The point-wise annotation [2] and scribble-wise annotation [4] of Pascal VOC 2012 dataset [24] are utilized for point-supervised and scribble-supervised settings, respectively.

We adopt the segformer with pre-trained weights on ImageNet as our backbone for experiments. Our model is trained for 80k training epochs with initial learning rate 0.001, batch size 16 and input resolution 512x512. The SGD optimizer is used with momentum 0.9 and weight decay 1e-4. Random horizontal flip, random brightness in [-10, 10], random resize in [0.5, 2.0] and random crop are all employed in data augmentation. In practice, we set $\alpha$ in Eq.1 1.2 for scribble setting and 0.2 for point setting. All experiments are conducted on PyTorch with four 32GB V100 GPUs. * indicates that we reproduce the results using the provided codes, unless otherwise specified.

Tab 1 and Tab 2 present the experimental results on the Pascal VOC 2012 for point-wise and scribble-wise annotations, respectively. In Tab 1 for point-wise annotations, the baseline based on segformer-B4 employs the cross-entropy loss function only can reach 64.2% mIoU. TEL [1] improves mIoU performance by 5.13%. Our proposed SASFormer outperforms TEL by 4.43% with the same backbone, yielding a mIoU of 72.78%. As for scribble-wise annotations in Tab 2, the baseline framework achieves a mIoU of 73.3%. SASFormer produces the best performance with the mIoU of 79.49%, surpassing the baseline by 6.19%. It demonstrates that SASFormer achieves remarkable performance compared to all other state-of-the-art approaches, which adopt a single-stage training strategy without any additional data. The qualitative results for point-wise annotations and scribble-wise annotations showcase that not only does our technique capture object outlines more accurately, but it also minimizes category misidentification, as shown in Fig 4.