[논문리뷰] Segmenter: Transformer for Semantic Segmentation

참조

 

[논문리뷰] AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE(Vision Transformer)

 

들어가며

 

 

 

 

이 논문이 게재될 당시, Semantic Segmentation(SS)에서는 기본적으로 Encoder-Decoder 구조로 구성되어 있었습니다. 여기서 Encoder는 해상도를 줄여가며 Feature를 추출하고, Decoder에서는 추출한 Feature를 복원하여 pixel-level class를 예측하는 Fully Convolutional Networks(FCN)이 지배적으로 사용되고 있었습니다.

 

하지만, Encoder에서 설계된 Convolution 구조는 고정된 커널은 지역적 특성만 반영하고, 이미지내 전역적인 Feature를 추출하지 못합니다. 이러한 지역적 특성은 이미지내 전체 픽셀을 라벨링하는 SS에서 제한 사항이 있었습니다. 

 

위와 같은 제한 사항을 개선하기 위해, atrous convolution, spatial pyramid pooling 등 수용필드를 늘리려고 많은 노력을 해왔습니다(수용 필드가 커야 Global context를 학습할 수 있음). 

 

따라서, 본 논문은 수용필드가 큰 Transformer 구조를 SS의 Encoder에 도입하였습니다. 이 구조는 모든 레이어에서 Attention 연산을 전역적으로 진행하기 때문에 Global Context 학습에 용이하고, 이는 아래 그림처럼 SS에 유리한 것을 나타냅니다. 

 

 

이렇게 Encoder에서 추출된 전역적인 Features를 Decoder로 업샘플링하여 각 클래스를 예측합니다. 이 예측하는 방법은 두 가지로 나뉩니다.

 

1. 심플한 Linear Mapping
   
   
2. Class Mask로 생성된 패치 토큰들과, 학습 가능한 Class Embedding을 가진 Transformer 기반

 

대략적인 내용은 이렇고, 위 구조들의 세부 구조는 아래에서 자세히 다루겠습니다.

 

 

Segmenter

 

Segmenter는 Transformer 기반의 Encoder-Decoder로 전체 구조는 아래와 같습니다. 

 

 

입력 이미지가 Transformer Encoder를 거쳐 나오는 과정은 ViT와 동일하므로 본 글에서는 다루지 않겠습니다.

 

Decoder

ViT의 Encoder의 최종 출력은 \( z_L \in \mathbb{R}^{N*D} \)의 형태로 나타내집니다. 그리고 이 \(z_L \)이 Transformer Decoder로 입력되어 segmentation map \( s \in \mathbb {R}^{H*W*K} \) (K는 클래스의 수)가 최종적으로 출력됩니다.

 

Decoder는 전술했듯이, Linear와 Mask Transformer로 2가지 방식이 존재합니다. 

 

• Linear
   
  1. ViT Encoder의 출력인 패치 인코딩 \( z_L \in \mathbb {R}^{N*D} \)에서, 패치 레벨 \( z_{lin} \in \mathbb {R}^{N*K} \)을 예측하기 위해 Linear layer를 적용
     
     
  2. \( z_{lin} \in \mathbb {R}^{N*K} \)을 2D로 Reshape \( s_{lin} \in \mathbb {R}^{ \frac {H} {P}  * \frac {H} {P}  * K} \)
     -> P = 패치의 수
     
     
  3. 원리 이미지의 크기 \( S \in \mathbb {R}^{H*W*K} \)로 복원하기 위해, P 크기 만큼 bilinear interpolation
     
     
  4. 최종적으로 클래스에 대해 Softmax 적용
     
     
• Mask Transformer
  
  
  1. ViT Encoder의 출력과 별도로, \( K \)개의 학습가능한 클래스 임베딩 \( cls = [cls_1, ..., cls_k] \in \mathbb {R}^{K*D} \)을 생성
     -> 이 클래스 임베딩은 각 Class에 대한 마스크를 생성하기 위함임
     
     
  2. 위 그림처럼 ViT Encoder의 출력인 패치 임베딩 \( z_L \)과 클래스 임베딩 \( c \)는 결합(Concatenation)되어, Transformer Decoder로 함께 입력됨.
     -> Transformer Decoder는 ViT Encoder와 동일한 레이어로 M개 구성됨.
     
     
  3. M개의 레이어를 거쳐 출력된 패치 임베딩 \( z_M^{'} \in \mathbb {R}^{N * D} \)과 클래스 임베딩 \( c \in \mathbb {R}^{K*D} \)은 다시 분해 및 L2 정규화(D차원 기준)를 해주고, 아래와 같이 \( Masks(z_M^{'}, c) \in \mathbb {R}^{N*K} \) 를 계산
     
     

     

  4. 계산된 \( Masks(z_M^{'}, c) \)는 패치의 시퀀스 이므로 다시 2D로 Reshape \( s_{mask} \in \mathbb {R}^{\frac {H} {P} * \frac {W} {P} * K} \)해주고, 원래 해상도 만큼의 Bilinear Interpolation을 해줌 \( s \in \mathbb {R}^{H * W * K} \).
     
     
  5. 최종적으로 Layer normalization과 softmax를 취함.

 

마치며

 

 

 

위 결과는 ViT의 크기(Seg-B, Seg-L), 패치 사이즈(/16, /32), Linear or Mask를 적용하였을 때에 대한 비교실험입니다. Mask Transformer를 적용하였을 때 어느 정도 성능 향상이 이루어졌네요. 

 

 

오늘은 Segmenter 논문을 리뷰해 보았습니다.

긴 글 읽어주셔서 감사합니다.