Task-Adaptive Feature Transformer with Semantic Enrichment
for Few-Shot Segmentation

The goal of few-shot segmentation is to meta-train the model so that it can adapt to unseen classes and perform segmentation with only a few labeled samples from new classes. For meta-training and evaluation, training set $D_{train}$ and test set $D_{test}$ with no-overlapping categories are used. In both meta-training and evaluation phases, episodes are composed with classes randomly selected from $D_{train}$ and $D_{test}$, respectively. For 1-way $N$-shot segmentation, each episode consists of a support set $S=\{(\mathbf{x}_{n},\mathbf{y}_{n})\}_{n=1}^{N}$ with $N$ image-label pairs and a query set $Q=\{(\bar{\mathbf{x}}_{m},\bar{\mathbf{y}}_{m})\}_{m=1}^{M}$ with $M$ image-label pairs, where $\mathbf{x}_{n}$ and $\bar{\mathbf{x}}_{m}$ represent images and $\mathbf{y}_{n}$ and $\bar{\mathbf{y}}_{m}$ indicate the corresponding binary masks for the given class. In each episode, the model adapts to the given class using $S$, and then predicts the segmentation masks from the query images $\{\bar{\mathbf{x}}_{m}\}_{m=1}^{M}$. At the end of each episode, the predicted segmentation masks are compared with the query masks $\{\bar{\mathbf{y}}_{m}\}_{m=1}^{M}$. In the meta-training phase, a loss is computed based on the difference between the predicted masks and the query masks and used to update the model parameters. In the evaluation phase, the model performance is evaluated from the comparison between the predicted masks and the query masks.

TAFT is a plug-in module which enables few-shot segmentation for an existing segmentation network composed of an encoder and a decoder. By transforming the features from the encoder, TAFT provides an effective task-conditioning to the segmentation network. As seen in Figure 1, when an episode with $S$ and $Q$ is given, TAFT computes prototypes $\{\mathbf{c}_{fg},\mathbf{c}_{bg}\}$, which are the pixel averages of the embedded features from the support images for foreground and background classes as done in (Wang et al., 2019; Zhang et al., 2019b, 2020). Unlike in previous few-shot segmentation methods, however, TAFT also makes use of another set of class references, $\{\mathbf{r}_{fg},\mathbf{r}_{bg}\}$. The reference vectors $\mathbf{r}_{k}$’s are randomly initialized and meta-trained with the encoder $f$ and the decoder $g$ of the segmentation network. While the prototypes $\mathbf{c}_{k}$’s are apparently driven by and thus depend on the current task’s input images, $\mathbf{r}_{k}$’s are completely decoupled from the current task. Given the prototypes $\mathbf{c}_{k}$ and references $\mathbf{r}_{k}$, TAFT constructs a linear transformation matrix $\mathbf{P}$ such that the transformed prototype $\mathbf{Pc}_{k}$ is brought near $\mathbf{r}_{k}$ for all $k$.

By using the reference vectors, we can utilize a stable task-agnostic feature space instead of a highly varying feature space from the images of the unseen class. We observe that while the prototypes $\mathbf{c}_{k}$ may change significantly from one task to next as inputs vary, the references $\mathbf{r}_{k}$ change only by little from one episode to next. As such, $\mathbf{r}_{k}$ offers a stable feature space, while the linear transformation $\mathbf{P}$, which is constructed anew in every episode, provides quick task-conditioning. See Figure 5 of Appendix for visual illustration of the stability of the references $\mathbf{r}_{k}$. The decoder makes use of the query features transformed to the space where $\mathbf{r}_{k}$ reside. In this sense, the stability of $\mathbf{r}_{k}$ helps the decoder to develop a steady and reliable strategy to generate the segmentation mask.

Figure 1 visualizes 1-way 1-shot segmentation. The reference vectors $\{\mathbf{r}_{fg},\mathbf{r}_{bg}\}$ for foreground and background have the length matching the number of channels in the high-level feature. Encoder $f$ extracts the high-level feature $\mathbf{h}=f(\mathbf{x})$ from the support image $\mathbf{x}$. The downsized labels $\tilde{\mathbf{y}}_{fg}$ and $\tilde{\mathbf{y}}_{bg}$ of the same size as the high-level feature are generated by average pooling. Although we can use a simple resizing when generating the downsized labels, resizing often leads to binary hard labels with all foreground or all background when the object in the image is too large or too small; this is problematic in prototype computation. On the other hand, average pooling generates soft labels with continuous values between 0 and 1, and we can compute both foreground and background prototypes for all cases. Using the downsized soft label $\tilde{\mathbf{y}}_{fg}$ and feature $\mathbf{h}$, the foreground prototype $\mathbf{c}_{fg}$ is computed as

where $\mathbf{h}^{(i,j)}$ and $\tilde{\mathbf{y}}_{fg}^{(i,j)}$ denote the pixels in $\mathbf{h}$ and $\tilde{\mathbf{y}}_{fg}$, respectively. Likewise, using soft label $\tilde{\mathbf{y}}_{bg}$ and feature $\mathbf{h}$ the background prototype $\mathbf{c}_{bg}$ is computed as

where $\tilde{\mathbf{y}}_{bg}^{(i,j)}$ are the pixels in $\tilde{\mathbf{y}}_{bg}$. For $N$-shot setting with $N$ support images and support labels, the prototypes $\mathbf{c}_{fg},\mathbf{c}_{bg}$ are computed as the means of the sample prototypes computed by equations (1) and (2). Given the prototypes $\{\mathbf{c}_{fg},\mathbf{c}_{bg}\}$ and reference vectors $\{\mathbf{r}_{fg},\mathbf{r}_{bg}\}$, we compute the prototype matrix $\mathbf{C}$ as $\left[\frac{\mathbf{c}_{fg}}{\lVert\mathbf{c}_{fg}\rVert},\frac{\mathbf{c}_{bg}}{\lVert\mathbf{c}_{bg}\rVert}\right]$ and reference matrix $\mathbf{R}$ as $\left[\frac{\mathbf{r}_{fg}}{\lVert\mathbf{r}_{fg}\rVert},\frac{\mathbf{r}_{bg}}{\lVert\mathbf{r}_{bg}\rVert}\right]$.

We can construct the transformation matrix $\mathbf{P}$ by finding a matrix such that $\mathbf{PC}=\mathbf{R}$. In general, $\mathbf{C}$ is not a square matrix and does not have an inverse. One way to find a reasonable $\mathbf{P}$ is to compute $\mathbf{P}=\mathbf{R}\mathbf{C}^{+}$, where $\mathbf{C}^{+}$ is the pseudo-inverse of $\mathbf{C}$ computed as $\{\mathbf{C}^{T}\mathbf{C}\}^{-1}\mathbf{C}^{T}$. This gives the least square fit between $\mathbf{PC}$ and $\mathbf{R}$. Note that in 1-way segmentation the matrix $\mathbf{C}^{T}\mathbf{C}$ is $2\times 2$, and the inversion of this matrix requires very little computation. Matrix $\mathbf{P}$ depends on $\mathbf{C}$, which is driven by the encoder input representing the current task. Thus, it can be said that task-conditioning is achieved via application of the linear transformation using $\mathbf{P}$. Given the high-level feature from a query image, TAFT transforms it to the task-agnostic feature pixel-by-pixel using $\mathbf{P}$. The pixel-wise feature transformation can be easily done using a $1\times 1$ convolution layer with weight $\mathbf{P}$. With this task-adaptive feature transformation, the pixels of the task-agnostic feature get to settle close to the corresponding reference vectors, and the decoder can easily distinguish the foreground and background pixels in the task-agnostic feature. Note that the reference vectors are the only learned part in TAFT. The transformation matrix is not trained but rather computed for each task.

While the proposed TAFT module allows the model to adapt to a new task, the quality of the segmentation prediction is highly dependent on the quality of the high-level features from the encoder. If the encoder fails to provide the informative high-level features, it is impossible to obtain the useful task-agnostic features and therefore the decoder fails to predict the exact segmentation masks. To improve the quality of the high-level features from the encoder, we propose a semantic enrichment (SE) module that utilizes the following two additional components.

First, we adopt a pixel-wise self-attention module on the top of the encoder. Using the self-attention module, we can correlate pixels in the feature with the other feature pixels in the same feature and capture the context among the feature pixels. For the self-attention, we utilize a multi-head scaled dot-product attention proposed in (Vaswani et al., 2017). To apply the self-attention, we first flatten each feature $X\in\mathbb{R}^{(H\times W\times d)}$ into a one-dimensional feature sequence $X_{seq}\in\mathbb{R}^{(HW\times d)}$. Then, we process the feature sequence through the attention layer and obtain the attended feature sequence having the same size with the input sequence. Finally, we reshape the sequence into a two-dimensional attended feature $X_{attn}\in\mathbb{R}^{(H\times W\times d)}$. By processing the feature using the self-attention module, the high-level features are refined considering the context among the feature pixels, and the quality of semantic information in the high level features is improved.

Our second component in SE module is an auxiliary decoder that conducts multi-class semantic segmentation. While the labels of the 1-way segmentation episodes contain only the information about the target class, there exist multiple objects of various classes in the corresponding training images in general. To make use of the various classes in the training images, we utilize an auxiliary decoder which predicts the segmentation masks of training classes from the features of the training images. For the supervision of the auxiliary decoder, we generate auxiliary segmentation labels for the training images from the original segmentation labels, by making the all test classes as the background. Using the auxiliary decoder and the auxiliary labels, we obtain an auxiliary loss to train the auxiliary decoder and the encoder network. By training the encoder for the various classes in the image via auxiliary loss, the encoder learns to extract the more general features with the richer semantic information which is useful for few-shot segmentation.

Since the attention module operates on the high-level feature reduced to 1/16 of the original image, the pixel-wise attention does not require significant additional computation. Note that the auxiliary decoder is used only during training and not utilized for evaluation.

Algorithm 1 shows processing of a given episode for the TAFT-SE modules combined with a segmentation algorithm during meta-training (which is the same as the inference process except for the loss computation, the auxiliary decoder and parameter update part). The prototypes are computed using the downsized labels and the pixel-wise attended encoder features from the support set (line 2 to 10). Then the transformation matrix $\mathbf{P}$ is constructed using the prototypes and reference vectors (line 11 to 12). For training the reference vectors, we utilize a regression loss $L_{R}$. Using the reference vectors and task-agnostic features, we generate the downsized predictions and compare them with downsized labels (line 16). In prediction, we use the pixel-wise inner product with the reference vectors and normalize the prediction scores for each pixel using the softmax activation function. Since the downsized labels contain continuous values in $[0,1]$, the mean-squared-error (MSE) loss for regression is utilized as the regression loss and computed over all pixels of downsized predictions. Using $L_{R}$, the reference vectors are trained so that the pixels in task-agnostic features can be easily distinguished. Using the task-agnostic features, the decoder predicts the segmentation masks (line 17). The segmentation loss $L_{S}$ is computed between the labels and the predicted segmentation masks (line 18). The segmentation loss is a cross-entropy loss computed over all pixels in images, which is generally used for semantic segmentation. Through meta-learning using $L_{S}$, the decoder learns to generate segmentation masks using the task-agnostic features where each pixel is close to the corresponding reference vector. At the same time, the auxiliary decoder predicts the segmentation masks and the auxiliary segmentation loss is computed using the auxiliary labels. The auxiliary loss is also a cross-entropy loss computed over whole pixels as the segmentation loss (line 19). The auxiliary loss helps the encoder learn to extract more general features.

In this paper, we test the TAFT-SE modules in conjunction with Deeplab V3+. Deeplab V3+ consists of an encoder, a decoder, and the Atrous Spatial Pyramid Pooling (ASPP) module. For segmentation, the high-level feature containing the semantic information is extracted from the image by the encoder. The low-level feature including the shape information is also extracted from the middle layer of the encoder. The high-level feature is then processed by the ASPP module to capture the multi-scale information, and the decoder network generates the segmentation mask using the ASPP output feature and the low-level feature.

Figure 2 illustrates how TAFT-SE modules are plugged in with Deeplab V3+. The attention module is applied on the top of encoder and refines the high-level features. TAFT operates on the high-level feature and generates the task-agnostic features. The task-agnostic features are then further processed by the ASPP module, which is included as a part of the decoder in the figure. Then, the decoder generates the segmentation masks using the low-level features and the ASPP-processed task-agnostic features together. At the same time, the auxiliary decoder takes high-level features and low-level features as input, and predicts the multi-class segmentation masks. Note that in Deeplab V3+, TAFT-SE does not process the low-level feature, since the low-level feature mostly contains the shape information which can be considered general.

Figure 3 displays some qualitative results of TAFT-SE combined with Deeplab V3+. We visualize the 1-shot segmentation results of aeroplane, bus, dog and sheep classes from each split of PASCAL-$5^{i}$. The images in the first row are support images, with the support labels shown together at the bottom left of the support images. The images in the second and row show the query images with the prediction results. We can see that TAFT-SE on Deeplab V3+ successfully segments the objects from the query images using only a single support sample in each class. More qualitative results can be found in Appendix.

Semantic segmentation is a computer vision problem that aims to predict the label for each pixel of the image. The fully-convolutional network (FCN) of (Long et al., 2015) utilizes the convolutional neural network (CNN) architectures designed for classification such as VGG or ResNet by convolutionalization. The convolutionalization process transforms the fully connected layers as $1\times 1$ convolution layers. In this way, the network can preserve the location information while predicting the class of each pixel. UNet of (Ronneberger et al., 2015) modifies the FCN architecture to utilize the multi-scale features to obtain better segmentation performance with less data. PSPNet of (Zhao et al., 2017) proposes a Pyramid Pooling Module (PPM) to effectively aggregate context information over different regions to predict the detailed segmentation masks. Deeplab of (Chen et al., 2017a) and Dilation of (Yu & Koltun, 2016) introduce a dilated convolution which is able to expand the receptive field without a parameter increase. The advanced versions of Deeplab (Chen et al., 2017b, 2018) propose the Atrous Spatial Pyramid Pooling (ASPP) module based on dilated convolution to utilize the multi-scale receptive field. Although the segmentation methods work well when trained with ample data, they cannot be easily adapted to a previously unseen class with only a few data samples, since they are designed to be trained with large datasets. On the other hand, the proposed TAFT-SE modules can be easily combined with an existing segmentation method to enable few-shot segmentation.

Few-shot learning aims to develop a general classifier that can quickly adapt to the unseen task with only a few labeled samples. In carrying out few-shot classification, two main strategies have been developed. One is metric-based few-shot learning (Oreshkin et al., 2018; Snell et al., 2017; Vinyals et al., 2016; Yoon et al., 2019), the goal of which is to learn a mapping to the metric space where samples from the same category cluster together while those from different categories are kept far apart.

Another is the optimization-based meta-learner (Finn et al., 2017; Jamal & Qi, 2019; Nichol & Schulman, 2018; Ravi & Larochelle, 2017; Rusu et al., 2018). This approach aims to train a meta-learner which in turn trains an actual learner. The meta-learner supports learning of the actual learner so that the learner can adapt well to a new task using only a small number of updates based on a few samples.

Among the few-shot learning algorithms, TapNet of (Yoon et al., 2019) is most closely related to our TAFT module in that stand-alone meta-learned reference vectors $\boldsymbol{\Phi}$, which are similar in spirit to $\{\mathbf{r}_{fg},\mathbf{r}_{bg}\}$ of the present method, are employed and linear transformation is used. Despite the similarity, there are some key differences in both philosophical viewpoints as well as design methodologies. In TAFT, the reference vectors are viewed as a more stable form of prototypes residing in a task-agnostic space, which serves as the destination for linear transformation. TAFT brings the embedded features all the way to this task-agnostic space, where inference and reference updates take place. This tends to make the references in TAFT more stable than those in TapNet, helping reliable generation of segmentation masks in the decoder and subsequent training of decoder modules. As for specific design methodology, TapNet’s linear transformation attempts to align in-class pairs of vectors while at the same time distancing out-class vectors as much as possible, to maximize classification accuracy. TapNet employs linear nulling of errors for this purpose. In comparison, for TAFT, the least square fit criterion is more appropriate for linear transformation, as it is well-suited for segmentation.

OSLSM of (Amirreza Shaban & Boots, 2017) is a first algorithm adopting the few-shot learning strategy to semantic segmentation. OSLSM learns to generate the conditioning parameters for element-wise scaling and shifting. Co-FCN of (Rakelly et al., 2018) generates a globally pooled prediction predicts the segmentation mask by fusing it with query features in segmentation branch. Recent methods utilize the prototype idea of (Snell et al., 2017), to use the information in the support set efficiently. SG-One of (Zhang et al., 2020) computes the prototype by the masked-average pooling and compute cosine-similarity map to guide segmentation. Following SG-One, many methods utilize the prototype computed by the masked-average pooling for few-shot segmentation (Liu et al., 2020a, b; Nguyen & Todorovic, 2019; Wang et al., 2019; Yang et al., 2020; Zhang et al., 2019b). For example, CANet of (Zhang et al., 2019b) predicts the mask via dense comparison between query feature and prototype, and iteratively updates it using optimization module. PANet of (Wang et al., 2019) introduces a novel regularization method of prototype alignment for more effective meta-training. PMMs of (Yang et al., 2020), PPNet of (Liu et al., 2020b), ASGNet of (Li et al., 2021) and SCL of (Zhang et al., 2021a) utilize multiple prototypes to leverage the semantic information of different regions of support images. Instead of using prototypes, PGNet of (Zhang et al., 2019a) and DAN of (Wang et al., 2020) adopt the graph attention mechanism to model the relationship between support feature pixels and query feature pixels. PFENet of (Tian et al., 2020) utilizes the prior mask computed from pixel-wise cosine similarity between query feature and masked support feature instead of prototype. PFENet also introduces a feature enrichment module which is able to utilize the multi-scale information with hierarchical relations across different scales. CWT of (Lu et al., 2021) utilizes the encoder and decoder pretrained with training classes and adopts a self-attention mechanism to adapt the classifier weights to a given task. MMNet of (Wu et al., 2021) adopts a learnable memory embeddings to memorize the middle-level meta-class embeddings during meta training and utilize them for few-shot learning.

Our TAFT is similar to the works of (Wang et al., 2019; Zhang et al., 2019b, 2020) in the sense that the prototypes are utilized to represent the pixels. The prior works directly utilize the prototypes that vary widely from task to task in comparing the pixels of the feature with the prototypes. However, the proposed TAFT method does not rely on comparison between the prototypes and the pixels of feature. Instead, TAFT utilizes feature transformation converting task-specific features into fairly stable, task-agnostic features. While the prototype-based approaches require sample-specific comparison between the prototypes and pixels of the feature for each image sample, TAFT does not require any sample-specific process; it generates the transformation matrix $\mathbf{P}$ and use it for every query image. This enables an efficient parallel computation to predict the segmentation masks of different samples at the same time.

Our SE module is similar to CWT of (Lu et al., 2021) in terms of adopting a self-attention mechanism. While CWT conducts self-attention between features and classifier weights to adapt the model to the given task, our attention module performs self-attention among the pixels from the same feature to enrich the semantic information in the feature. Our SE module and CWT are also similar in terms of conducting semantic segmentation for all training classes during training. However, while CWT utilizes the encoder and decoder pretrained for semantic segmentation task in few-shot adaptation, our SE module conducts the multi-class segmentation only to assist the training of the encoder through the auxiliary loss.

PASCAL-$\mathbf{5^{i}}$ is used for our experiments. PASCAL-$5^{i}$ is a dataset based on PASCAL VOC 2012, proposed by (Amirreza Shaban & Boots, 2017) for few-shot segmentation. The 20 classes of VOC 2012 are divided into 4 splits, and each split contains 5 disjoint classes. In our experiments, the 5 classes in one of 4 splits are selected as the test classes, and the remaining classes are used as the training classes. The training samples of training classes and the test samples of test classes are used for training and testing, respectively. The details of class split can be found in (Amirreza Shaban & Boots, 2017). Data augmentation is not applied.

COCO-$\mathbf{20^{i}}$ is also utilized for our experiment. It is constructed from the MS-COCO dataset. The 80 classes of MS-COCO are divided into 4 splits, and each split contains 20 disjoint classes. We utilize the class split suggested in (Nguyen & Todorovic, 2019); the details of the split can be found in (Nguyen & Todorovic, 2019). Again, data augmentation is not applied.

We use the mean Intersection-over-Union(mIoU) score as the evaluation metric. As done in prior works, we compute the foreground IoU for each class, and use the averaged per-class foreground IoU as mIoU. We also report the evaluation result with another metric, the Foreground-Background IoU(FBIoU) score suggested in (Rakelly et al., 2018). The FBIoU score is computed as the average of foreground IoU and background IoU computed over all test images.

TAFT is implemented as a single $1\times 1$ convolutional layer between encoder and decoder. The weights of this layer reflects the elements of the linear transformation matrix P, so that every pixels of feature can be transformed by P.

For the self-attention in SE, we adopt a scale-dot product attention layer with a single head. In the attention module, the dimensions of query-key-value in self-attention are the same as the input feature dimension. The auxiliary decoder has the same shape with the original decoder, except the last convolutional layer that conducts the pixel-wise classification; the auxiliary decoder has more output channels for more classes.

For training TAFT-SE on Deeplab V3+, we use stochastic gradient descent with the learning rate of 0.01 and a momentum of 0.9 for both PASCAL-$5^{i}$ and COCO-$20^{i}$ datasets. For the encoder, we utilize ResNet-50 and ResNet-101 networks. We initialize the encoder with the ImageNet pretrained ResNet-50 or ResNet-101, so we apply a 10 times smaller learning rate to the encoder as done in Deeplab V3+.

We modified the strides in the ResNet encoders so that the low-level features and high-level features are downsized by a factor of 4 and 16, respectively. Since PASCAL-$5^{i}$ include many small objects, we use the optimized atrous rates of $[1,4,7,11]$ in the ASPP module.

For all experiments, $3\times 10^{4}$ episodes are used for meta-training and the learning rate is decayed by a factor of 10 after training $2\times 10^{4}$ episodes. The weight decay with the rate of $3\times 10^{-4}$ is applied for regularization. Training is done using episodes with 12 queries.

For every experiment, we use the images and labels resized to $512\times 512$, considering the crop size used in Deeplab V3+. In evaluation, 30,000 episodes are used. For PASCAL-$5^{i}$ experiments, the multi-scale input test with scale factors $[0.7,1,1.3]$ is used as done in prior works of (Zhang et al., 2019a, b). In evaluation, the support and query samples are scaled with a specific ratio, and the corresponding predictions are scaled again into the original ratio. The final prediction is done by averaging 3 predictions from different scales.

We compare TAFT-SE on Deeplab V3+ with prior approaches. In Table 1, we display the mean mIoU and FBIoU scores for both 1-shot and 5-shot segmentations for PASCAL-$5^{i}$. For both ResNet-50 and ResNet-101 backbones, our 1-shot mean mIoU scores fall somewhat lower than previous SOTA methods. However, for 5-shots, TAFT-SE modules do indeed extend Deeplab V3+’s few-shot capability to beyond what is possible with the SOTA performance by prior methods by a fair margin. Accordingly, the $\Delta$ value, which represents the ability of the few-shot segmentator to learn from an increasing number of shots, is also the largest except for PPNet in ResNet-50 experiments, and the largest in ResNet-101 experiments. In fact, as shown in Table 7 of Appendix, when compared with PFENet (which gives a SOTA 1-shot mIoU performance in ResNet-101 experiments) across a wider range of 1, 3, 5, 7 and 10 shots, our method already starts to outperform PFENet from 3 shots and on, by a significant margin.

For FBIoU scores, our TAFT-SE on Deeplab V3+ with ResNet-50 backbone falls short of SOTA performance on 1-shot. But for 5-shot, our extension yields the best FBIoU score. In ResNet-101 experiments, TAFT-SE extension exhibits the best 1-shot and 5-shot scores and the best $\Delta$ value, highlighting its ability to learn quickly from an increasing number of examples. The higher FBIoU score means that the proposed method achieves a higher background IoU than prior works. We conjecture that the usage of background pixel information results in higher background IoU and FBIoU scores, as in PANet of (Wang et al., 2019) and PPNet of (Liu et al., 2020b). While most prior works only utilize the foreground prototype or foreground pixel information, TAFT-SE utilizes the prototypes and reference vectors for both foreground and background.

Table 2 displays the mean mIoU scores for COCO-$20^{i}$. Again, for both ResNet-50 and ResNet-101 backbone, our TAFT-SE extension on Deeplab V3+ gives the best 5-shot performance with large margins.

Table 3 shows the results of our ablation experiments with PASCAL-$5^{i}$ dataset. For the all ablation experiments, we use Deeplab V3+ with the ResNet-50 backbone as a base segmentator, with the same hyperparameter setting. We can see that the attention module and the auxiliary loss consistently improve the performances. For both 1-shot and 5-shot, the auxiliary loss provides more gains compared to the attention module. In Appendix, we show the performance of TAFT-SE extension on FCN of (Long et al., 2015), another well-known segmentation model.