SegFormer (2021) 논문 리뷰

SegFormer의 논문 제목은 SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers다. (링크)

 

저자는 Enze Xie, Wenhai Wang, Zhiding Yu, Anima Anandkumar, Jose M. Alvarez, Ping Luo다.

 

Github: 링크

 

Abstract

SegFormer는 Transformers와 가벼운 mulilayer perceptron (MLP) 디코더 구조를 하나로 통합한 효율적인 모델이다. 2가지 특징을 지닌다. 1) SegFormer는 새로운 트랜스포머 인코더 구조를 통해서 multiscale 피쳐를 생성한다. 이는 positional encoding을 필요로 하지 않아서 다양한 해상도에 대한 positional codes에 대한 interpolation이 불필요하다. 2) 복잡한 디코더 구조를 피한다. 하지만 MLP 디코더를 사용해서 global and local attention을 혼합한다. 다양한 스케일의 SegFormer는 제안하는데 B0 부터 B5까지가 있다. 베스트 모델은 SegFormer-B5다. B0가 가장 작고 B5가 가장 크다. 

 

3. Method

 

SegFormer의 구조는 위 Figure 2와 같다. 

 

주어진 이미지의 크기는 H x W x 3인데 이를 4 x 4 크기의 패치들로 나눈다. ViT의 16 x 16과 다르다. 더 작은 패치의 크기는 dense prediction task에 더 유리하다. 이를 hierarchical Transformer encoder에 입력으로 집어 넣어서 multi-level features { 1/4, 1/8, 1/16, 1/32 } of original image resolution을 얻는다. 이를 All-MLP decoder에 보내서 H / 4 x W / 4 x $N_{cls}$ ㅎ상도의 segmentation mask를 예측한다. $N_{cls}$는 카테고리의 개수다. 

 

3.1. Hierarchical Transformer Encoder

일련의 Mix Transformer encoderss (MiT)를 제시하는데 MiT-B0 부터 MiT-B5로, 구조는 모두 같지만 크기만 다르다. B0가 가장 작고 B5가 가장 크다. 

 

Hierarchical Feature Representation

Patch merging을 통해서 hierarchical features map $F_i$를 얻는다. $F_i$의 해상도는 H / $2^{i+1}$ x W / $2^{i+1}$  x $C_i$, 이때 $i \in$ {1, 2, 3, 4}다. 그리고 $C_{i+1}$ 은 $C_i$보다 더 크다. 

 

Overlapped Patch Merging

주어진 이미지 패치에 대해서 ViT에서 쓰인 patch merging process를 사용해서 N x N x 3 패치를 1 x 1 x C 벡터로 통합한다. 이는 2 x 2 x $C_i$의 피쳐들을 1 x 1 x $C_i$로 통합하여 hierarchical feature maps로 이어진다. 이를 통해서 $F_1$ (H / 4 x W / 4 x $C_1$)을 $F_2$ (H / 8 x W / 8 x $C_2$)로 축소할 수 있다. 이를 지속적으로 반복해서 hierarchy를 달성한다. 위 과정에서는 본래 non-overlapping image or feature pathces다. 하지만 저자들은 Local continuity를 유지하기 위해서 overlapping patch merging process를 제시한다. Path size K, stride S, padding size P를 제시한다. K = 7, S = 4, P =3과 K = 3, S = 2, P = 1을 이용해서 overlapping patch merging을 수행하는데 이는 non-overlapping process와 같은 크기의 피쳐를 생성한다. Local continuity는 서로 연관있는 이미지 픽셀들끼리의 연속성을 의미한다. 

 

 

Efficient Self-Attention

원래의 Self-Attention의 computational complexity는 $O(N^2)$다. 저자들은 Pyramid vision transformer: A versatile backbone for dense prediction without convolutions 논문 (링크) 에서 제시된 방법을 도입해서 계산 복잡도를 줄인다.

 

Reduction ratio R을 도입해서 시퀀스의 길이를 아래와 같이 줄인다.

 

$\hat{K}$ = Reshape (N / R, C R) (K)

$K$ = Linear(CR, C)(\hat{K})

 

$K$는 줄어들 대상이 되는 시퀀스의 길이다. 위 Reshape의 뜻은 $K$를 N / R x (C R) 크기로 다시 변환한다. Linear($C_{in}, C_{out}$)으로 in-channel dim와 out-channel dim을 정의하기에, CR 크기를 C로 변환한다. 따라서 K는 N / R x C 가 된다. 따라서 self-attention의 computational complexity는 $O(N^2)$에서 $O(N^2 / R)$로 변한다. 저자들은 R을 [64, 16, 4, 1]오 설정한다. 

 

 

Mix-FFN

기존의 ViT는 positional encoding (PE)를 사용하지만 이는 고정된 해상도에만 사용할 수 있다. 따라서 저자들은 CVPT Conditional positional encodings for vision transformers 논문 (링크)에서 제시한 방법대로 3 x 3 Conv로 data-driven PE를 만든다. 저자들은 PE가 semantic segmentation에서는 사실 반드시 필요한 사항은 아니라고 주장한다. 대신 위치 정보를 누설하기 위해서 zero padding을 고려하는 Mix-FFN을 제안한다. 그 식은 아래의 (3)이다.

 

 

저자들은 depth-wise conv를 활용해서 파라미터 수를 줄이고 효율성을 향상시켰다.

 

3.2. Lightweight All-MLP Decoder

MLP로만 이루어진 디코더를 제시하는데 이는 보다 단순한 구조를 지니며, 전통적인 CNN의 인코더들 보다 더 큰 크기의 effective receptive field (ERF)를 갖는다.  

 

자세한 구조는 아래 식 (4)와 같다. 

 

마지막의 M은 predicted mask다. 

 

Effective Receptive Field Analysis

 

 

Figure 3를 보면 전통적인 구조의 DeepLabv3+ 보다 SegFormer가 가장 깊은 단계인 Stage-4에서 더 큰 receptive fields를 갖는걸 확인할 수 있다.

 

Stage-2에서는 local attention을, Stage-4에서는 non-local attention까지 포함하는 것을 확인할 수 있다. 

 

MLP Head를 비교해 보면 SegFormer가 보다 응집되어 있는 혀앹의 강한 local attention을 가짐을 알 수 있다.

 

3.3.  Relationship to SETR

SETR Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers 논문 (링크)와의 비교 항목이다.

 

• SETR은 ImageNet-22K와 ViT를 지니지만, SegFormer는 ImageNet-1K를 사전학습에 사용했다.
• SegFormer의 인코더는 계층적 구조이며, ViT 보다 작다. 또한 저해상도와 고해상도의 피쳐를 모두 학습한다. 반면에 SETR의 ViT 인코더 구조는 단일한 저해상도만 생성한다.
• Positional Encoder를 제거해서 다양한 해상도에 대응 가능하지만, SETR은 그렇지 않다.
• MLP로만 이루어진 디코더는 SETR의 여러개의 3x3 conv로 구성된 무거운 디코더 보다 가볍다. 또한 계산적 오버헤드도 무시가능하다. 

 

4. Experiments

Cityscapes, ADE20K, COCO-Stuff를 활용해서 성능을 평가한다. mmsegmentation으 코드 베이스를 활용해서 학습했다.

8개의 Tesla V100을 사용했다. ImageNet-1K로 인코더를 사전학습하고 디코더를 랜덤하게 초기화했다. 데이터 증강은 random resize with ratio 0.5-2.0, random horizontal flipping를 적용했으며, random cropping to 512 x 512, 1024 x 1024, 512 x 512를 Cityscapes, ADE20K, COCO-Stuff에 대해서 각각 적용했다. AdamW로 옵티마이저를 설정했다. 

 

 

 

Table 1은 Ablation study고, Table 2는 다른 모델들과의 성능 비교다.

 

Table 1을 보면, 모델의 크기가 커질 수록, MLP의 차원인 C가 커질수록, Mix-FFN을 PE 대신 사용했을 경우, 그리고 ResNet과 ResNeXt의 CNN 구조 대신 트랜스포머 인코더를 사용했을 때 더 좋은 성능임을 확인할 수 있다.

 

Table 2를 보면 SegFormer가 대체로 파라미터 수는 더 적지만 더 좋은 성능을 달성했음을 알 수 있다. 

 

 

끝으로 B0 부터 B5까지의 모델에 대한 구체적인 사양을 담은 Table 6로 마무리한다.

 

 

 

Semantic Segmentation 결과로 마무리한다.