SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers
SegFormer (2021) 논문을 읽고 작성한 글입니다.
최근 Transformer 관련 논문들을 관심있게 공부하고 있는데, 이번에 segmentation task 에서 간단하고 효율적인 transformer 디자인을 새롭게 제안한 논문이 있어서 소개드리고자 합니다.
저번에 “Model architecture in medical image segmentation” 글에서 잠깐 transformer를 사용했던 모델들을 소개드리면서 아래 그림을 보여드린 적이 있었습니다.
그림에서 볼 수 있듯이 기존 모델들은 대부분 CNN 과 transformer 를 같이 사용했고, Decoder에서는 transformer를 거의 사용한 경우가 없었습니다. SegFormer 이전에도 PVT(Pyramid Visoin Transformer), Swin Transformer, Twins 와 같은 transformer-based segmentation 모델들이 있었는데, 모두 decoder에 transformer를 사용한 적은 없었습니다.
이전 연구들과 다르게 이번에 소개드릴 SegFormer 논문에서는 encoder 와 decoder 모두에 transformer를 사용하여 efficiency, accuracy, robustness를 모두 고려한 모델 디자인을 제안하였습니다. 간단히 제안한 방법 위주로 설명드리도록 하겠습니다.
1. hierarchical Transformer encoder
모델의 전체적인 구조는 위 그림에서 보실 수 있습니다. 먼저, encoder 부터 설명드리도록 하겠습니다.
hierarchical Transformer 구조가 encoder의 핵심입니다. 이전 연구들은 single-resolution feature map을 생성했었는데, SegFormer에서는 계층적인 구조로 CNN과 같이 high-resolution coarse feature와 low-resolution fine feature 모두를 생성할 수 있도록 하였습니다.
Transformer Block 구조를 보시면 Efficient Self-Attn 와 Mix-FFN 다음에 Overlap Patch Matching 을 통해 hierarchical feature map을 생성할 수 있도록 합니다. Overlap Patch Matching에서는 convolution network를 사용하여 stride, padding으로 overlap 되는 patch를 샘플링해서, local continuity를 계속 갖고 갈 수 있도록 합니다.
Efficient Self-Attn 에서는 K(key)의 dimension reduction을 통해서 complexity를 줄였다고 합니다. 그리고 Mix-FFN에서는 convolution layer을 가운데 넣어서 zero padding으로 positional encoding을 대신할 수 있도록 하였습니다. 기존의 positional encoding의 resolution은 고정되어 있다는 단점이 있는데, zero padding을 사용하면 train 과 test의 resolution이 달라도 성능에 크게 영향이 없다는 장점이 있다고 합니다.
2. lightweight All-MLP decoder
decoder 구조를 보면 굉장히 단순한 것을 볼 수 있습니다. 각 transformer를 통해 나온 multi-level feature들은 MLP Layer를 통과해서 channel dimension을 맞추고 upsample 해서 concatenate 됩니다. 그 다음, MLP 를 통과하여 최종 segmentation map이 나옵니다.
위 그림에서 일반적인 CNN-based network (DeepLabv3+)와 SegFormer의 effective receptive field 분석 결과를 비교한 것을 볼 수 있습니다. 확실히 SegFormer의 receptive field가 훨씬 큰 것을 볼 수 있습니다. hierarchical transformer encoder 구조로 이를 가능하게 했고, 그렇기 때문에 단순한 MLP 구조의 decoder로도 굉장히 좋은 성능을 낼 수 있다고 합니다.
Reference:
[1] SegFormer
Leave a comment