[논문 리뷰] FCBFormer: FCN-Transformer Feature Fusion for Polyp Segmentation(2022)

이번에 리뷰할 논문은 FCN-Transformer Feature Fusion for Polyp Segmentation 입니다.

 

https://paperswithcode.com/paper/fcn-transformer-feature-fusion-for-polyp

Papers with Code - FCN-Transformer Feature Fusion for Polyp Segmentation

 

 

 

 

 

본 논문에서 제안하는 FCBFormer모델은 대장 내시경(Colonoscopy) 영상의 대장암(colorectal cancer) 부분의 영역을 검출하기 위해 제안된 모델입니다. 수동으로 대장 내시경에서 용종을 segmentation하는 것은 많은 시간 비용이 들었고 예전 부터 해당 문제를 Deep Learning으로 접근하는 연구들이 진행되었습니다. 또 최근의 transformer를 기반으로 semantic segmentation 모델을 사용하였을 때 좋은 performance를 내는 것은 사실이나, input image에 비해 낮은 해상도($ \frac{H} {4} \times \frac{W} {4} $)의 segmentation map을 예측한다는 한계가 존재했습니다. 

 

물론 용종을 segmentation하기 위해 많은 Fully convolution networks(FCNs) 모델 역시 제안되었으나, transformer 기반 모델에 비해 좋은 성능을 보여주진 못했습니다. 이는 아무래도 FCNs 와 달리 첫 번째 layer 부터 attention을 수행하여 global receptive field에 기초하여 feature 를 추출할 수 있었기 때문에 보다 dense prediction을 수행하는데 유리하다고 볼 수 있겠죠. 

 

본 논문에서는 transformer 로 추출한 feature가 full-size로 upsampling 될 때 FCNs의 feature를 사용하여 보완하자는 아이디어를 냈으며, 이를 통해 Fully convolutional branch(FCB) 와 transformer branch(TB)를 병렬적으로 수행하여 output을 prediction하는 FCBFormer를 제안하게 됩니다. 

 

 

 

 

 

 

 

 

FCBFormer

우선 전체적인 모델 architecture는 다음과 같습니다.

 

 

FCBFormer는 크게 TB와 FCB로 구성됩니다. 각각 input image에 대해 병렬적으로 연산하여 추출한 feature map을 concate한 후 Prediction Head(PH) 를 거쳐 최종 output 을 출력하게 됩니다. TB의 경우 $ H \times W $ 해상도의 이미지를 $ \frac{H} {4} \times \frac{W} {4} $ 해상도의 feature map으로 출력하며, FCB의 경우 input image 해상도 그대로 $ H \times W $ 의 feature map을 출력합니다. 결국 TB의 output feature map을 $ H \times W $ 으로 upsampling 해야하며, 이에 따라 거친 feature map이 출력되게 됩니다. 이런 TB의 출력을 PH가 잘 추론할 수 있도록 도와주기 위해 FCB의 출력을 같이 concate하여 PH의 input으로 구성하게 되는 것이죠. 

 

 

 

Transformer branch (TB)

FCBFormer의 한 요소인 TB 부터 살펴보도록 하죠. TB의 경우 SSFormer의 영향을 받았다고 합니다. SSFormer는 swin transformer로 구성된 encoder에서 생성된 각 계층의 feature map들을 모두 concate하여 output을 예측하는 구조를 가집니다. 이때 각 계층의 feature map을 aggregate 하는 부분을 PLD 라고 하는데 본 논문에서는 이 PLD를 좀 더 개선한 PLD+를 TB의 decoder로 구성합니다. 

 

또한 SSFormer와 달리 TB의 encoder는 Pyramid Vision Transformer (PVT) v2 로 구성합니다. 제 생각에는 PVT v2의 overlapping patch embedding을 사용하여 좀 더 local continuty가 손실되는 것을 막아 dense한 feature map을 추출하려 한 것으로 보이네요. 또한 former의 경우 global 한 feature는 잘 추출하는 반면 local한 부분에 weakness 가 존재하기 때문에 더 local한 부분에 집중하는 모습을 보이는 것 같습니다.

 

 

 

Transformer encoder

다시 TB의 구조를 보자면 우선 input image가 들어왔을 때 PVT v2-B3 에 적용됩니다. PVT v2는 pyramid feature representation 덕에 dense prediction 에 대한 feature 추론을 가능하게 하죠. 공간 차원의 크기를 유지하는 전통 transformer와 달리 다양한 multi-scale의 feature를 추출하기 때문에 dense prediction을 수행하는 것이 가능합니다. 

 

 

 

Improved pregressive locality decoder (PLD+)

TB의 encoder 부분인 PVT v2 를 통해 생성된 다양한 scale의 feature map을 Stepwise Features Aggregation(SFA)을 사용하여 다양한 scale의 정보들을 모두 aggregate합니다. 이러한 방식은 이미 semantic segmentation을 수행하고자 제안된 여러 CNN 기반 모델들이 사용했던 기법이며, SSFormer에서 처음으로 이러한 방식을 transformer 에 적용했었습니다. 

 

하지만 SSFormer에서 제안된 기존의 PLD는 단순히 각 layer의 feature를 동일한 resolution으로 up-sampling 한 후 concatenation을 수행하기 때문에 본 논문에서는 group normalisation, SiLU activation function, residual connection과 같은 현대적인 요소들의 부재와 상대적으로 적은 layer 수를 기존 PLD의 한계라고 생각합니다. 

 

 

위 figure는 SSFormer의 구조이며, 해당 구조에서 사용된 PLD를 개선하여 다음과 같이 PLD+로 구성합니다.

 

 

또한 LE는 다음과 같이 RB block 2개로 구성되어있는데 이 RB block의 첫 번째 convolution을 channel 64로 반환하도록 하여 모든 sequence layer가 64개의 channel을 반환할 수 있도록 조정하였습니다.

 

 

또한 LE의 마지막 요소인 Up-sampling을 해상도가 88x88 이 되도록 고정하여 encoder에서 출력한 모든 feature map의 해상도가 통일되도록 조정해주었습니다.

 

 

 

 

 

 

 

 

Fully convolutional branch (FCB)

다음으로 TB와 병렬적으로 수행되는 FCB를 알아보도록 하겠습니다. 

 

 

FCB는 TB와 달리 고도로 융합된 multi-scale feature를 full-size의 해상도로 출력합니다. FCB의 구조는 다음 3가지의 요소들로 구성이 되죠.

 

• Residual blocks (RBs)
• strided convolutional layers for down sampling
• nearest neighbour interpolation for upsampling
• dense U-Net style skip connections

 

TB에서 추출된 중요하긴 하나 up-sampling되어 거친 feature와 결합하여 PH가 full-size의 segmentation map을 추론할 수 있도록 도와줍니다.

 

 

 

 

 

 

 

 

Prediction Head (PH)

PH는 최종 prediction을 하기 위한 head로서 FCB와 TB의 output을 모두 합성한 full-size의 tensor를 input으로 받습니다. 

 

 

PH의 마지막 Pred block은 1x1 conv를 수행하여 class의 개수만큼 channel 수를 조절한 후 sigmoid 를 적용합니다. 

 

결국 PH는 TB에서 추출한 중요하지만 개략적인 feature 와 FCB에서 추출한 세밀한 feature 를 결합하여 최종 segmentation map을 예측하도록 수행됩니다.