안녕하세요. 지난 포스팅의 [IS2D] SegNeXt: Rethinking Convolutional Attention Design for Semantic Segmentation (NIPS2022)에서는 multi-branch 스타일의 InceptionNet, Xception, ResNeXt와 같은 구조를 차용하여 CNN 모델을 설계한 SegNeXt에 대해서 소개하였습니다. 오늘은 SegNeXt에서 비교로 삼은 SegFormer라는 Transformer 기반 모델을 설명드리도록 하겠습니다.
Background
기본적으로 의미론적 영상 분할 (Semantic Segmentation)은 영상 내의 각 픽셀에 대해 어떤 카테고리 레이블에 속하는 지 예측하는 dense prediction 문제라고 볼 수 있습니다. 이를 통해 사용자는 전체 영상 내에 완벽한 이해 (complete understanding)를 제공해주죠. 이러한 문제는 자율주행, 이상탐지, 그리고 의료영상 등과 같은 정말 다양한 도메인에서 필수 기술로서 자리매김하고 있습니다.
이를 위해 제안된 다양한 모델들이 있습니다. 대표적으로 CNN-based 모델들은 inductive bias를 기반으로 설계되어 상대적으로 적은 양의 데이터로도 충분히 좋은 성능을 얻을 수 있다는 장점은 있으나 global information을 활용하지 못합니다. 대표적으로 FCN (CVPR2016)과 DeepLab Series (ICLR2015, TPAMI2017, arxiv2017, ECCV2018) 등이 있습니다. 그와 대척점인 Transformer-based 모델들은 대용량의 데이터 (ImageNet 또는 JFT-300M)으로 학습되어 self-attention 기반의 강력한 인코더 성능을 보장해주지만 고해상도 영상에서의 높은 연산량이 큰 문제가 있습니다. 대표적으로 SETR (CVPR2021)이 있었죠.
특히, SETR 같은 경우에는 Naive한 ViT를 바로 backbone 모델로 사용하고 CNN 디코더를 연결하여 feature resolution을 점진적으로 향상시키는 방법을 적용하였습니다. 하지만, ViT은 단일 스케일 특징 맵을 출력하고 고해상도 영상에서 높은 연산량을 요구하기 때문에 비효율적이라는 문제점이 있어 다양한 객체와 해상도를 다루는 semantic segmentatoin에서는 좋지 않은 선택임을 알 수 있습니다.
본 논문에서는 이러한 문제를 해결하기 위해 SegFormer라는 모델을 제안합니다. 이 모델은 SETR과 동일하게 Transformer 기반의 모델로 구성되어 있지만 효율성 (efficiency), 정확도 (accuracy) 그리고 강건도 (robustness)까지 모두 사로잡은 모델입니다. SegFormer는 본 논문에서 새롭게 제안하는 Transformer 인코더인 Mix Transformer (MiT)를 설계하여 potitional encoding이 없어 다양한 해상도에 효율적으로 적용할 수 있게 되었습니다. 또한, Light-weight All-MLP 디코더를 추가적으로 제안하여 Transformer에서 비롯된 특징 맵들을 multi-scale 단위에서 효율적으로 결합할 수 있는 방법론을 제시합니다.
SegFormer는 그림 1에서 보이는 것과 같이 ADE20K에서 대표적인 Transformer 기반 모델인 SETR보다 훨씬 높은 정확도와 효율성을 가지는 것을 볼 수 있습니다. 뿐만 아니라 이후 실험에서는 CityScapes와 COCO-Stuff에서도 높은 성능을 달성하였습니다.
SegFormer
1) Hierarchical Transformer Encoder
본 논문에서는 Mix Transformer (MiT)라는 새로운 Transformer 인코더를 설계하였습니다. MiT-B0 ~ MiT-B5까지 6개의 모델 변형 구조로 구성되어 있으며 모두 동일한 구조이며 전체적인 규모만 다릅니다. 즉, MiT-B0를 인코더로 사용하는 SegFormer는 속도는 빠르지만 정확도는 떨어지게 되고 반대로 MiT-B5를 인코더로 사용하는 SegFormer는 속도는 느리지만 정확도를 향상될 것 입니다. 따라서 본인의 어플리케이션 및 컴퓨팅 자원 환경에 맞추어 사용하면 되겠네요.
기본적으로 MiT는 ViT에 영감을 받아 거의 동일한 구조를 가지고 있습니다만 MiT의 경우에는 semantic segmentation에 보다 특화된 모델이라고 볼 수 있습니다. 제 개인적인 생각으로는 PVT와 거의 유사한 특성을 가지고 있다고 느껴집니다. 일단, ViT와 어떤 점이 다른 지 하나하나 살펴보도록 하죠.
1-1) Hierarchical Feature Representation
ViT의 가장 큰 단점은 single-scale feature map을 만들어내기 때문에 semantic segmentation과 같은 dense prediction task에 알맞지 않습니다. 왜냐하면 하나의 장면에서 동일한 객체라고 하더라도 다양한 크기와 모양으로 존재할 수 있기 때문에 이러한 변형들에도 강건하게 특징을 추출하기 위해서는 multi-scale feature map을 필수이기 때문이죠. 따라서 본 논문에서는 CNN과 유사하게 multi-scale feature를 추출하게 됩니다.
입력 영상의 해상도가 $H \times W \times 3$라고 가정할 때 각 스테이지의 출력 특징 맵인 $F_{i}$의 해상도는 $\frac{H}{2^{i + 1}} \times \frac{W}{2^{i + 1}} \times C_{i}$가 됩니다. 이는 그림 2에서도 잘 설명이 되어있네요.
1-2) Overlapped Patch Merging
다음으로 ViT의 경우에는 입력 영상을 패치로 나눌 때 겹치지 않게 패치를 나누게 됩니다. 이는 패치 간의 연속적인 정보를 무시하는 것과 동일하기 때문에 본 논문에서는 패치를 나눌 때 약간 씩 겹치게 패치를 나누는 방식을 채택합니다. 사실 이 방식은 이전에 설명드린 다양한 Transformer 방법들에서 채택하였기 때문에 깊은 설명을 필요없을 거 같습니다.
1-3) Efficient Self-Attention
$$\text{Attention} (Q, K, V) = \text{Softmax} \left( \frac{Q K^{T}}{\sqrt{d_{head}}} \right) V$$
위 수식은 일반적으로 사용되는 self-attention의 수식입니다. 일반적으로 위 수식의 복잡도는 입력 토큰의 개수 $N = H \times W$의 제곱 ($O(N^{2})$) 만큼의 복잡도가 소요되어 큰 해상도를 가진 영상에서는 사용하기 어렵다는 문제점이 있습니다. 본 논문에서는 이러한 문제를 해결하기 위해 PVT 스타일의 token reduction process를 사용하여 복잡도를 $O(\frac{N^{2}}{R})$ 만큼 줄여주었습니다. 각 스테이지에서는 $R = \{ 64, 16, 1, 1 \}$로 사용하였습니다. 이 부분에 대한 자세한 설명 역시 이전에 설명드린 PVT를 참고하시면 좋을 거 같습니다.
1-4) Mix-FFN
ViT는 패치 내의 지역적 정보를 주입하기 위해 positional embedding을 적용하였습니다. 하지만 ViT에서 사용되는 positional embedding 벡터의 길이는 고정되어 있기 때문에 다른 해상도를 가진 영상이 입력되면 interpolation해야하는 문제점이 존재하고 이는 학습된 정보가 손실되는 문제점으로 이어지게 됩니다. 따라서 이전 Transformer 모델들과 마찬가지로 FFN 내에 depth-wise convolution을 도입하여 positional embedding 없이도 convolution의 inductive bias를 도입하였습니다.
2) Lightweight All-MLP Decoder
또한, 인코더에서는 Transformer를 사용하였기 때문에 꽤나 큰 연산량이 사용되었습니다. 본 논문에서는 디코더에 만큼은 단순한 구조를 사용하여 연산량을 줄이기 위해 MLP 기반으로 설계하였습니다. 이를 수식화하면 다음과 같습니다.
단순히 각 스테이지에서 얻은 특징 맵들을 MLP를 한번씩 통과한 뒤 concat하여 MLP를 한번더 통과시켜 prediction을 수행하는 아주 단순한 구조입니다.
Experimental Results
1) Ablation Study
1-1) Model Size
1-2) MLP Decoder Dimension
1-3) Positional Embedding
1-4) MLP Decoder Design
2) Quantitative Results on ADE20K and CityScapes
3) Qualitative Results
4) Corruption Robustness
