SegFormer Review

"SegFormer: Simple and Efficient Design for
Semantic Segmentation with Transformers"

Abstract

  Transformer를 경량 MLP decoder와 통합하는 간단하고 효율적이지만 강력한 semantic segmentation framwork인 SegForemr를 제시한다. SegFormer는 두가지 특징이 있다. 1) SegFormer는 multiscale feature를 출력하는 새로운 계층적 구조인 Transformer encoder로 구성된다. 이것은 잠재적 encoding을 필요로 하지 않아, test 해상도가 훈련과 다를 때 성능 저하로 이어지는 positional code의 보간을 피할 수 있다. 2) SegFormer는 복잡한 decoder를 사용하지 않는다. 제안된 MLP decoder는 다른 layer로 부터 정보를 통합하고 강력한 표현을 랜더링하기 위해 local attention과 global attention을 결합한다. 저자는 이러한 간단하고 가벼운 설계가 Transformer에서 효율적인 segmentation의 key인 것을 보여준다. 이런 접근 방식을 확장해 SegFormer-B0에서 B5까지 다양한 모델을 설계해 이전의 모델보다 훨씬 더 나은 성능과 효율성을 도달했다. 

Introduction

  Semantic segmentation은 pixel 단위로 이미지를 분류하는 task로 fully convolutional networks (FCNs)의 등장 이래로 다양한 형태의 모델이 등장하며 발전해왔다. ImageNet challenge의 classification 문제와 semantic segmentation은 밀접한 관계가 있어 Segmentation task에서 backbone 구조를 설계한느것은 여전히 활성화된 영역으로 남아 있다. 실제로, VGG 부터 다른 backbone 까지 semantic segmentation의 성능을 크게 향상 시켰다. Backbone 구조 외에도, 다른 work은 semantic segmentatoin을 구조화된 예측 문제로 정의하고, contextual 정보를 효율적으로 capture할 수 있는 모델과 연산에 초점을 맞춰 설계한다. 이러한 것의 대표적인 예제는 dilated convolution이 있다.

  Transformer는 NLP 분야에서 처음 등장해 뛰어난 성능을 보였고 이는 vision task까지 확장되어 ViT 등 다양한 모델이 등장해 vision task에서도 뛰어난 성능을 보였다. Semantic segmentation에서는 SETR이 제안되어 이 분야에서도 transformer가 유용하게 사용될수 있음을 증명했다. SETR은 backbone으로 ViT를 채택하고 몇개의 CNN decoder와 통합해 결과를 출력한다. 좋은 성능에도 불구하고 ViT는 두가지 제한점이 있다: 1) ViT는 multi-scale feature 대신에 오직 single-scale low-resolution feature를 출력한다. 2) Large 이미지에서 연산량을 많이 필요로 한다. 이러한 한계를 해결하기 위해 Wang 등은 dense한 예측을 위해 피라미드 구조를 가진 ViT인 Pyramid Vision Transforemr(PVT)를 제안했다. PVT는 object detection과 semantic segmentation에서 ResNet에 비해 상당한 개선을 보였다. 하지만, Swin Transformer와 Twins같은 다른 방법들 처럼 PVT는 주로 transformer encoder의 설계를 고려하며 추가 개선을 위한 decoder를 무시한다.

  그래서 저자는 SegFormer를 소개한다. SegFormer는 효율성, 정확성 및 견고성을 모두 고려하는 semantic segmentation을 위한 transforemr framework이다. 이전 방법들과 달리, 이 framework는 endocer와 decoder를 모두 재설계 한다.

• 새로운 positional-encoding-free 그리고 계층적 transformer encoder를 가진다.
• Lightweight ALL-MLP decoder는 복잡하지 않고 연산량이 적은 모델로 강력한 표현을 나타내도록 설계되었다.
• 그림 1처럼, SegFormer는 semantic segmentation에서 SotA를 달성했다.

  먼저, 제안된 encoder는 훈련 데이터와 다른 해상도의 이미지를 inference 할때 interpolating positional code를 피한다. 결과적으로, encoder는 성능에 영향을 주지 않고 임의의 테스트 해상도에 쉽게 적응할 수 있다. 추가로, 계층적인 부분은 인코더가 고해상도의 fine feature와 저해상도의 coarse feature를 생성하도록 한다. 이것은 해상도를 고정시켜 단일의 저해상도 feature map을 생성하는 ViT와 대조적이다. 두번째, lightweight MLP decoder를 제안한다. Decoder의 key 아이디어는 lower layer의 attention이 local적으로 유지되는 경향이 있는 반면, high layer는 매우 non-local한 transformer 유도 기능을 활용하는 것이다. 다른 layer로 부터 정보를 통합하는 MLP decoder는 local과 global attention의 결합이다. 그 결과, 간단하고 직설적이며 강력한 표현을 render할 수 있는 decoder를 얻었다.

Method

  SegFormer는 그림 2와 같으며 두가지 주요 모듈을 가지고 있다: 1) 계층적 transformer encoder, 2) lightweight All-MLP decoder. 이미지의 크기가 H×W×3 으로 주어지면, 먼저 이를 4×4 patch 크기로 자른다. 16×16을 사용하는 ViT와 다르게, 세분화된 patch를 사용한는것이 semantic segmentation에서 선호된다. 두번째, 이 patch들은 계층적 transformer encoder로 입력되어 원본 해상도의 1/4, 1/8, 1/16, 1/32 크기의 multi-level feature를 얻는다. 이러한 multi-level features를 All-MLP decoder에 통과시켜 H/4×W/4×N 해상도를 가진 segmentation mask를 얻는다.

Hierarchical Transformer Encoder

  구조는 같지만 크기는 다른 Mix Transformer encoder (MiT)를 MiT-B0 부터 B5까지 다양하게 설계한다. PVT의 효율적인 self-attention 모듈과 계층적 구조 위에 segmentation task에서 큰 이점을 보여주는 overlapped patch merging과 positional encoding free design을 포함하는 몇가지 새로운 특징을 제안한다. 

  Hierachical Feature Representation. ViT와 달리, 제안하는 encoder는 입력이 주어졌을때, multi-level multi-scale feature를 생성한다. 이러한 feature는 semantic segmentatoin의 성능을 향상시킬수 있는 고해상도의 coarse feature와 저해상도의 fine-grained feature 둘다 제공한다. 특히, H×W×3 크기의 입력이 주어지면, H/2^i+1 × W/2^i+1 × C_i,  i는 {1, 2, 3, 4}, C_i+1은 C_i 보다 큼, 해상도를 가진, 계층적 feature를 얻기 위해 patch merging을 수행한다.

  Efficient Self-Attention. 이러한 계층적 feature 표현의 주요 병목 현상은 고해상도 feature의 긴 sequence 입력을 통한 2차 self-attention 복잡성이다. 다시 말해, 기존의 multi-head self-attention은 각 헤드의 Q, K, V가 같은 차원인 N×C, N은 H×W로 sequence의 길이, 를 가진다. Self-attention은 다음과 같다.

대신에 PVT에서 소개된 sequence 감소 process를 채택한다. 이 process는 sequence의 길이를 줄일기 위해 reduction ratio R을 사용한다.

이것으로 self-attention은 O(N^2/R) 만큼 줄어든다. R은 stage-1 부터 4까지 [64, 16, 4, 1]을 사용한다. 

  Overlapped Patch Merging. 이미지 patch가 주어지면, ViT에서 사용되는 patch merging process는 N × N × 3 patch를 1 × 1 × C 벡터로 통합한다. 이것은 쉽게 확장해 2 × 2 × C_i feature path를 1 × 1 × C_i+1 벡터로 통합하여 계층적 feature map을 얻을 수 있다. 이것을 사용해 계층적 특징을 첫번째 feature map에서 두번째 feature map으로 줄어들도록 할 수 있고, 계층의 다른 feature map에 대해서 반복할 수 있다. 이러한 process는 non-overlapping 이미지나 feature patch를 결합하기 위해 설계되었다. 그러므로, 이러한 patch 주변의 local 연속성을 유지하지 못한다. 그래서 overlapping patch merging process를 사용한다. 이를 위해 K, S, P를 정의한다. K는 patch 크기, S는 인접한 두 patch의 stride, P는 padding의 크기이다. 실험에서는 K=7, S=4, P=3, 그리고 K=3, S=2, P=1 로 설정해, non-overlapping process와 동일한 크기의 feature를 생성하기 위해 overlapping patch merging을 수행한다. ViT의 원래 patch 임베딩과 유사하게, 이 작업은 PyTroch의 nn.Conv2D로 구현될 수 있다.

  Positional-Encoding-Free Design. ViT에서 positional embeddings(PE)의 해상도는 고정되어져 있다. 따라서 테스트 해상도가 훈련과 다를경우 PE를 보간해야 한다. 해상도가 일치하지 않는것은 semantic segmentation에서 일반적이기 때문에 바람직하지 않고 이것은 정확도를 떨어뜨린다. 그래서 Mix-FFN을 도입한다. Mix-FFN은 FFN에서 3×3 Conv를 직접 사용해 leak location information에 대한 zero padding의 영향을 고려한다. Mix-FFN은 다음과 같다. 

x_in은 self-attention 모듈로 부터 나온 feature다. Mix-FFN은 3×3 conv와 MLP를 각각의 FFN에 혼합한다. 실험에서는 3×3 cov가 transformer에 충분한 positional information을 제공하는것을 볼 수 있다. 부분적으로, parameter를 줄이고 효율성을 향상 시키기 위해 depth-wise conv를 사용한다.

  CPVT 또한 3×3 conv를 사용해 다른 해상도에서 조건부 PE를 생성한 다음 feature map에 추가함으로써 이 문제를 완하한다고 언급한다. 이 작업은 semantic segmentation에서 feature map에 PE를 추가할 필요가 없다고 주장함에 따라 개념적으로 한 단계 더 나아간다. 또다른 최신 연구인 CvT는 토큰 간의 공간 관계를 모델링 하기 위해 3×3 Conv를 도입했다. Converging 설계에도 불구하고 저자의 작업은 semantic segmentation에서 훈련/테스트 해상도 불일치 문제를 처리하기 위해 PE를 완전히 제거하는 것을 목표로 하기 때문에 동기 부여와 응용 모두에서 다르다. 저자의 직관은 아래의 논문에서 시작되었지만 동리한 직관은 CvT에서 논의되지 않았다.

Md Amirul Islam, Sen Jia, and Neil DB Bruce. How much position information do convolutional neural networks encode? arXiv, 2020

Lightweight All-MLP Decoder

  SegFormer는 MLP layer로만 구성된 lightweight decoder를 통합해 다른 방법에 일반적으로 사용되는 수작업 및 계산 요구가 많은 구성요소를 방지한다. 이런 간다한 decoder를 가능하게 하는 핵심은 계층적 구조인 transformer encoder가 기존의 CNN encoder에 비해서 큰 effictive receptive field (ERF)를 가지기 때문이다.

  제안된 All-MLP decoder는 4가지 main step으로 구성된다.

1. MiT encoder로 부터 나온 multi-level feature는 MLP layer를 통과해 channel 차원으로 통합된다. 
2. Feature들은 upsmapled 되어지고 concat된다.
3. Concat된 feature를 융합하기 위해 MLP layer를 사용한다. 
4. 다른 MLP layer는 segmentation mask M을 예측하기 위해 융합된 feature를 취한다.

  Effective Receptive Field Analysis. Semantic segmentatoin에서 context 정보를 포함하는 large receptive field를 유지하는것은 중요한 이슈이다. MLP decoder 설계가 transformer에서 왜 효과적인지를 분석하고 해석하기 위해 ERF를 사용해 시각화 한다. 그림 3은 DeepLabv3+와 SegFormer의 4개의 encoder와 decoder 헤드를 시각화 한 결과이다. DeepLabv3+의 ERF는 4 단계더라도 상대적으로 작다. SegFormer의 encoder는 자연스럽게 낮은 단계에서 convolution과 유사한 local attention을 생성하는 동시에 4단계에서 context를 효과적으로 capture하는 highly non-local attention을 출력할 수 있다. 그림 3을 확대해 보면 MLP head의 ERF는 non-local attention 외에 local attention에 더 많은 관심을 가지는 4단계와 다르다.

  CNN에서의 제한된 receptive field는 receptive field를 확장하지만 모델이 무거워지는 ASPP와 같은 context 모듈에 의존해야 한다. 제안된 decoder 설계는 transformer의 non-local attention으로 부터 이점을 얻어 복잡하지 않고 더 큰 receptive field로 이어진다. 하지만, 전체 receptive field가 4단계에서 제한된 field에 의해 상한되기 때문에 동일한 decoder 설계는 CNN backbone에서 잘 작동하지 않으며 이는 표 1에 나타난다.

  더 중요한 것은 decoder 설계가 기본적으로 매우 local하고 non-local attention을 동시에 생성하는 transformer 유도 기능을 활용한다는 점이다. MLP decoder는 이를 통합함으로써 매개 변수를 몇개 더 추가하여 보완적이고 강력한 표현을 render 한다.

Relationship to SETR

  SegFormer는 SETR에 비해 더 효율적이고 강력한 여러 설계를 포함한다.

• ImageNet-1K는 pre-training에만 사용된다. SETR의 ViT는 더 큰 ImageNet-22K에서 사전 훈련된다.
• Segformer의 encoder는 ViT 보다 작고 고해상도의 coarse와 저해상도의 fine feature 둘다 capture할 수있는 계층적 구조를 가진다. 반대로 SETR의 ViT encoder는 단일 저해상도 feature map만 생성할 수 있다.
• Encoder에서 PE를 제거하는 반면, SETR은 추론 해상도가 훈련 해상도와 다를 때 정확도를 낮추는 고정된 모양의 PE를 사용한다.
• MLP decoder는 SETR에 있는 decoder 보다 작고 계산더 덜 복잡하다. 

Experiments

Experimetns Settings

  Datasets: Cityscapes, ADE20K, COCO-Stuff

  Implementation details

• Base code : mmsegmentation
• Gpu : Tesla V100 8
• Augmentation : random resize (ratio : 0.5-2.0), random horizontal filpping, random crop (512×512, 1024×1024)
• Optimizer : AdamW, Iter : 160K(ADE20K, Cityscapes), 80K(COCO-Stuff), Batch size : 16(ADE20K, COCO-Stuff), 8(Cityscapes), Learning rate : 0.00006, Lr schedule : poly(factor is 1.0), Metrics : mIoU

Ablation Studies

  Influence of the size of model. 먼저, encoder의 크기가 증가함에 따라 모델의 성능과 효율성에 대해 분석한다. 그림 1에서 보여졌듯이 이를 3개의 dataset에 대한 결과를 표1a에 나타냈다. 가장 먼저 관찰하는 것은 decoder와 encoder의 크기이다. Lightweight model에서 보여졌듯이, deocder의 크기는 오직 0.4M parameter이다. MiT-B5 encoder의 경우 decoder는 모델 전체 parameter 수의 최대 4%만을 차지한다. 성능 측면에서, 전반적으로 encoder의 크기가 커지면 모든 dataset에서 성능이 개선되었다. 가장 가벼운 모델인 B-0는 좋은 성능을 보이며 소형 및 효율적이여서 실시간 애플리케이션에서 유용하며 큰 모델인 B5는 세가지 dataset 모두 SotA를 달성해 transformer encoder의 잠재력을 보여준다.

  Influence of C, the MLP decoder channel dimension. MLP decoder에 있는 channel dimension인 C에 대해서 분석한다. 표 1b에 성능을 나타냈다. C를 256으로 설정했을때 경쟁력 있는 성능과 계산량을 보여주는것을 관찰했다. C가 증가할 수록 성능이 더 좋아지지만, model이 더 커지며 덜 효율적으로 변한다. 흥미롭게도 이 성능은 768일때가 가장 적절하다. B0와 B1은 C를 256으로 설정했으며 나머지는 C를 768로 설정했다.

  Mix-FFN vs Psitional Encoder (PE). 이 실험에서는, 제안된 Mix-FFN을 사용하기 위해 transformer encoder에서 positional encoding을 제거했을 때에 대한 효과를 분석한다. PE와 Mix-FFN이 있는 transformer encoder를 훈련하고, 두가지의 다른 해상도를 가진 Cityscapes에 test 한다. 표 1c는 이 실험의 결과를 보여준다. 표 1c는 주어진 해상도에 대해 Mix-FFN을 사용한 결과가 PE를 사용한 결과보다 더 좋은 성능을 보여준다. 게다가 Mix-FFN은 해상도가 다르더라도 비슷한 정확도를 보여준다. PE는 해상도가 달라짐에 따라 정확도가 3.3% 떨어졌지만, Mix-FFN은 오직 0.7%만 떨어진다. 이 결과로 인해 Mix-FFN을 사용하는것이 PE를 사용하는것 보다 더 좋은 성능을 낸다는것을 확인했다.

  Effective receptive field evaluation. 다른 CNN 모델에 비해 큰 ERF를 가지는 MLP decoder의 이점에 대해서 논의한다. 실험에서 이 효과를 측정하기 위해 CNN 기반의 MLP decoder와 제안한 방법의 성능을 비교한다. 비교한 결과는 표 1d에 있으며, CNN encoder 기반의 방법이 제안된 방법에 비해 성능이 급격히 낮아졌다. 직관적으로 CNN은 transformer에 비해 작은 receptive field를 가지고 있어, MLP decoder가 global reasoning에 충분하지 않다. 대조적으로 MLP decoder를 가진 제안된 transformer encoder는 좋은 성능을 보였다. 게다가, transformer encoder는 high-level feature 대신 low-level local feature와 high-level non-local feature를 결합하는데 있어 필수적이다.

  Influence of difference encoders.  Transformer encoder를 Vit, Swin 그리고 제안된 방법 3가지를 비교했고 비교한 결과는 표 3과 같이 MiT-B2가 Swin-T 보다 3.1% 더 개선되었다. 게다가, MiT-B5는 ViT 기반의 encoder 보다 parameter가 더 작은데도 불구하고 mIoU가 3+% 높아졌다. 이 실험에서  MiT encoder가 Swin 이나 ViT 보다 semantic segmentation에서 더 유용하다는것을 알 수 있다.

  Influence of difference decoders.  또한, MiT encoder를 다른 decoder에 테스트 했다. 이는 표 3에 나타나며 MLP decoder가 다른 decoder에 비해 1/8 크기 이지만 성능은 비슷하게 나와 MLP decoder가 효율적인 segmentation를 위해 중요하다.

Comparison to state of the art methods

  표 2는 ADE20K와 Cityscapes에 대한 실험 결과 이다. 표의 윗 부분은 Real-time SotA 모델에 대한 비교로 MiT-B0와 비교했으며 아랫 부분은 Non Real-time SotA 에 대한 비교이다.

ADE20K에서 SegFormer-B0는 37.4%의 mIoU, 3.8M parameter, 8.4G FLOPs의 성능으로 가장 뛰어난 성능을 보였다. 게다가 SegFormer-B5의 결과는 이전의 SotA model인 SETR보다 mIoU가 1.6% 개선되어 51.8%로 SotA 이다. Cityscpaes의 경우 SegFormer-B0가 15.2 FPS와 DeeplabV3+ 보다 mIoU가 1.3% 개선된 76.2%를 기록했다. SegFormer-B5는 mIoU가 84.0%이며, 이는 SETR 보다 1.8%개선되었으며 심지어, SETR보다 5배 빠르고 4배 작다. Cityscapes test set에서 공통적인 설정을 사용하고 검증 데이터를 훈련 데이터에 병합해 Imagenet-1K pretrain과 Mapillary Vistas를 사용해 결과를 나타낸다. 이는 표 4와 같으며 결과는 82.2% mIoU를 보여 SETR을 포함한 다른 방법보다 우수하다. Mpaillary pretrian을 사용한 결과 83.1%의 SotA 성능을 보였다.

Robustness to natural corruptions

  이 실험에서는 SegFormer의 robustness를 검증한다. Cityscapes-C로 검증했으며 DeeplabV3+와 다른 SotA 모델들과 비교한다. 실험결과는 표 5와 같다.

Conclusion

• PE가 없는 계층적 transformer encoder와 lightweight All-MLP decoder를 결합한 간단하고 semantic segmentation method에서 강력한 SegFormer를 제안
• 이전 방법에서 볼 수 있는 복잡한 설계를 방지해 높은 효율성과 성능을 제공
• SegFormer는 benchmark dataset에서 SotA 결과를 달성

Paper

https://proceedings.neurips.cc//paper/2021/file/64f1f27bf1b4ec22924fd0acb550c235-Paper.pdf