LAPFormer: A Light and Accurate Polyp Segmentation Transformer

Colorectal cancer (CRC) is the most common cancer around the world [1]. Colonoscopy has always been recognised as the standard diagnostic for the early detection of colorectal cancer. Therefore, several deep learning methods have been proposed to aid clinical system in identifying colonic polyps. Among them, segmentation approach is significantly considered as the most appropriate way with promising result recently. However, colonoscopy has some limitations. In some previous reports, about 18% of polyps are missed from the diagnosis process [2, 3]. This is because it is an operator-driven procedure and solely dependent on the knowledge and skills of the endoscopist. With the current colonoscopy equipment, less experienced endoscopists cannot distinguish polyp regions during colonoscopy examinations [4]. More importantly, previous research has shown that increasing polyp detection accuracy by 1% reduces colorectal cancer risk by approximately 3%. Therefore, improving polyp detectability and robust segmentation tools are important in this problem.

Recently, with the vigorous development of deep learning technology, the accuracy of many classical problems has improved, including image segmentation problems. Various studies aimed to develop CADx models for automatic polyp segmentation. There are also a few studies aiming to build a specific model for polyp segmentation. HarDNet-MSEG [5] is one of them, which is an encoder-decoder architecture base on HarDNet [6] backbone that achieved high performance on the Kvasir-SEG dataset with processing speed up to 86 FPS. AG-ResUNet++ improved UNet++ with attention gates and ResNet backbone. Another study called Transfuse combined Transformer and CNN using BiFusion module [7]. ColonFormer [8] used MiT backbone, Uper decoder, and residual axial reverse attention to further boost the polyp segmentation accuracy. NeoUNet [9] and BlazeNeo [10] proposed effective encoder-decoder networks for polyp segmentation and neoplasm detection. Generally, research in changing model architecture is still a potential approach.

Among the recent deep learning architecture, Transformer based architecture has attracted the most attention. To efficient semantic segmentation, incorporating the advantages of a hierarchical Transformer encoder with a suitable decoder head has been widely researched. In this paper, we utilize a Transformer backbone as an encoder and propose a novel, light, and accurate decoder head for polyp segmentation task. The proposed decode head consists of a light feature fusion module which efficiently reduces the semantic gaps between feature from two level scales, and a feature selection module along with a feature refinement module which help handling feature from backbone better and calibrating feature comes from progressive feature fusion module before prediction.

Our main contributions are:

• •

  We propose a Light and Accurate Polyp Sementation Transformer, called LAPFormer, that integrates a hierarchical Transformer backbone as encoder.

• •

  A novel decoder for LAPFormer, which leverages multi-scale features and consists of Feature Refinement Module, Feature Selection Module to produce fine polyp segmentation mask.

• •

  Extensive experiments indicate that our LAPFormer achieves state-of-the-art on CVC-ColonDB [11] and get competitive results on different famous polyp segmentation benchmarks while being less computation complexity than other Transformer-based methods.

Semantic segmentation is one of the essential tasks in computer vision which is required to classify each pixel in the image. Recently, deep learning had an enormous impact on computer vision field, including semantic segmentation task. Many deep learning models are based on fully convolutional networks (FCNs) [12], which encoder gradually reduces the spatial resolution and captures more semantic context of an image with larger receptive fields. However, CNN still has small receptive fields, leading to missing context features. To overcome this limitation, PSPNet [13] proposed Pyramid Pooling Module to enhance global context, DeepLab [14] utilized atrous convolution to expand receptive fields.

Transformer [15] is a deep neural network architecture, originally proposed to solve machine translation task in natural language processing. Nowadays, Transformer has a huge influence on natural language processing and other tasks, including computer vision. Vision Transformer (ViT) [16] is the first model successfully apply Transformer in computer vision by dividing an image into sequences and treating each sequence as a token, and then putting them into Transformer. Following this success, PVT [17], Swin [18], and SegFormer [19] are designed as hierarchical Transformer, generating feature maps at different scales to enhance the local feature, thus improving the performance in dense prediction tasks, including detection and segmentation.

Recent advances in deep learning have helped solve medical semantic segmentation tasks effectively, including polyp segmentation. However, it remains a challenging problem due to medical image characteristics, and polyps come in different sizes, shapes, textures, and colors. U-net [20] is a well-known medical image segmentation model consisting of residual connections to preserve local features between encoder and decoder. Inspired by U-net, most models [21], [22], [23], [5], [24], [25], for medical image segmentation use an architecture containing a CNN backbone as an encoder, and a decoder tries to fuse features at different scales to generate segmentation result. PraNet [26], SFA [27] focus on the distinction between a polyp boundary and background to improve segmentation performance. CaraNet [28] designs a attention module for small medical object segmentation. Since the successful application of Transformer in computer vision, current methods [29], [30], [8] have utilized the capability of Transformer as a backbone for polyp segmentation task and yielded promising results. Other methods [31], [7], [32] use both Transformer and CNN backbone for feature extraction and combine features from two brands, thus enhancing segmentation result.

We choose MiT (Mix Transformer) proposed in [19] as our segmentation encoder for two main reasons. First, it consists of a novel hierarchically structured Transformer encoder which produces multi-level multi-scale feature outputs. Second, MiT uses a convolutional kernel instead of Positional Encoding (PE), therefore avoiding decreased performance when the testing resolution differs from training. Convolutional layers are argued to be more adequate for extracting location information for Transformer than other methods. Furthermore, MiT uses small image patches of size $4\times 4$, which are proved to favor intensive dense prediction tasks like semantic segmentation.

Assume that we have input image $X$ with spatial dimensions of $H\times W\times 3$ (represent the height, width and channel). MiT generates four different level feature $f_{i}$ with resolution $\frac{H}{2^{i+1}}\times\frac{W}{2^{i+1}}\times C_{i}$ where $i\in\{{1,2,3,4}\}$ and $C_{i+1}$ is larger than $C_{i}$. Basically, MiT has six variants with same architecture but varies in sizes, from MiT-B0 to MiT-B5.

Dense prediction is a famous task in computer vision which consists of Object Detection, Semantic Segmentation, … Object of vastly different scales all appear in the same picture. Therefore, multi-scale features are heavily required to generate good results. The most popular way to utilize multi-scale features is to construct a feature pyramid network (FPN) [33]. However, as pointed out in [34] and [35], there are big semantic gaps between feature from non-adjacent scales, features from two scale levels apart are weakly correlated, while features between adjacent scales are highly correlated.

In FPN, features from upper scale are upscaled then directly added to features from lower scale. We argue that this is sub-optimal, when fusing features, they should be further processed before directly fusing to next level. We propose Progressive Feature Fusion module, which progressively fuses features from upper scales to lower scales, therefore, reduce the information gap between the low resolution, high semantic feature maps and high resolution, low semantic ones.

Instead of fusing feature maps with addition operation, we use concatenation operation. Feature maps of all scale are upsampled to $(H/4,W/4)$. Then feature maps from upper scale are progressively fused with lower scale as follow:

where $x_{i}$ is features from scale $i$; $x_{i-1}$ is features from one scale lower than $i$; $[\dots]$ is concatenation operation; Linear is a fully-connected layer.

On conventional Encoder-Decoder architecture, after performing multi-scale feature fusion from feature pyramid, other models [20], [36], [23], [29] often perform prediction on the last output feature maps, which has the highest resolution (Fig 2(a)). Others adopt auxiliary predictions during training [37], [28] but during test time, they only make prediction from the highest resolution feature map (Fig 2(b)).

We argue that predicting on only the highest resolution feature maps is sub-optimal. This make the highest resolution feature maps to carry a lot of information, but information can be lost during feature fusion process. Adding information from high semantic branches (Figure 2(c)) before prediction can provide network with coarse polyp regions, therefore, easing the work of highest resolution feature maps.

After channel concatenation to increase information for prediction, it is crucial to emphasize feature maps that contain important details and suppress redundant information.

The data flow of Feature Selection Module is shown on Fig 3. Output feature maps from Aggregation for Prediction is scaled with a weighting vector $w$ which indicates the importance of each channel. Then we conduct channel reduction to keep the importance ones for prediction. The weighting vector $w$ is calculated as:

where $x$ is obtained by applying average pooling on output feature maps from Aggregation for Prediction mentioned above; $fc_{1},fc_{2}$ are two fully-connected layers; $\delta$ is ReLU activation function and $\sigma$ is Hard Sigmoid activation function.

Overall, the process of Feature Selection Module (FSM) can be formulated as:

where $X$ is output feature maps from Aggregation for Prediction; $w$ is weighting vector; $X_{inter}$ is intermediate feature maps; $X_{pre}$ is output of FSM module.

Low level features are important in semantic segmentation task, and its importance is especially high in polyp segmentation, because the boundary between polyp and its surroundings are not easily distinguishable.

We further enhance prediction on polyp boundary using a skip-connection. The skip-connection connects output feature maps from FSM with the lowest level feature maps in the feature hierarchy from Progressive Feature Fusion. This is a simple but surprisingly effective way which boosts the performance of our model.

Normally, feature maps of every scale from backbone pass through a lateral connection. It merges feature maps from upper scale into lower scale feature maps. Lateral connection is implemented using a $1\times 1$ Convolution. In Object Detection settings, this isn’t a problem, because features after feature fusion continue to be refined in the Detection Head.

However, in our Decoder settings, this is sub-optimal for multiple reasons:

• •

  All module in our Decoder part use $1\times 1$ Convolution. This hampers the learnability of model in the Decoder.

• •

  $1\times 1$ Convolution in lateral connection reduces number of channels in two top-most scale. This process can cause information lost.

• •

  $1\times 1$ Convolution in lateral connection also increases number of channels in two bottom-most scale. This process can cause information redundancy.

We propose to replace lateral connection with our Feature Refinement Module (FRM). Backbone feature maps from different levels are refined with FRM. It should be noted that each level has its own respective FRM. FRM is implemented using a $3\times 3$ Convolution, followed by a BatchNorm layer and ReLU activation function (Fig 4). This is very simple but harmonize well with our Encoder and Decoder architecture. Using $3\times 3$ Convolution brings two benefits:

• •

  $3\times 3$ Convolution further refines features coming from backbone, therefore, enhances the learning capability of Decoder.

• •

  Transformer Encoder lacks Local Receptive Field, result in lacking of local details. $3\times 3$ Convolution effectively emphasizes local features and cleans up noises.

For ablation study, we use MiT-B1 backbone and train model for 50 epochs average over 5 runs. All results are reported under Table I

Progressive Feature Fusion (PFF) and Aggregation for Prediction. We demonstrate the effect of Progressive Feature Fusion (PFF) in Fig 5. Visualization is done on feature maps closest to prediction head. After applying PFF, we can see that our model accurately localizes Polyp region and clears up noises from other regions. PFF effectively boosts performance of our Encoder’s across almost all datasets. Aggregating all feature maps produced by PFF helps utilize multi-scale features better and brings huge performance gain on ETIS, which consists of many small polyps. Results for both of these components are shown on the second and third row of Table I

Feature Selection Module (FSM). Applying FSM before giving prediction helps the prediction head to focus on important regions. The strength of FSM is demonstrated on Fig 6. With the same output feature maps from PFF + A, adding FSM refines the polyp segmentation region, highlights details which were not captured before. The result of FSM is shown on the forth row of Table I

Low level connection. Low level features have always been important in dense prediction tasks. For polyp segmentation, integrating low level features produces finer segmentation masks, especially the boundary regions. The importance of low level features is shown on Fig 7, adding more information from low level features refines polyp segmentation mask. This low level connection consistently boosts our model performance on almost all datasets, results are shown on the fifth row of Table I.

Feature Refinement Module (FSM). Incoperating Feature Refinement Module (FSM) increases our model’s performance on all datasets. Results after applying FSM is shown on last row of Table I. With FSM, we are able to scale our model with stronger backbone. All variations of our model are shown on Table II.

We compare our results with existing approaches on 5 benchmark datasets. Table III shows the results of SOTA methods. Our model achieves competitive performance on all datasets while being incredibly lighter than other methods. FLOPs and Params are shown on Table IV.