[SegFormer] SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

저자: Enze Xie, Wenhai Wang, Zhiding Yu
발행년도: 2021년
인용수: 3500회
논문 링크: https://arxiv.org/abs/2105.15203
arXiv ID: 2105.15203

---

SegFormer 논문 읽고 정리해봤다

Vision Transformer가 나온 이후로 semantic segmentation에 트랜스포머를 어떻게 잘 활용할 수 있을까 고민이 많았다. 특히 CNN 기반의 모델들이 워낙 잘 되어 있어서, 굳이 트랜스포머로 바꿔야 하나 싶은 생각도 들었다.

그런데 기존 ViT 기반 방법들을 보면 좀 아쉬운 점들이 있었다. 일단 positional encoding에 너무 의존적이라 입력 해상도가 바뀌면 성능이 확 떨어진다. 그리고 디코더 부분도 너무 복잡하게 설계되어 있어서, 과연 이게 효율적인가 하는 의문이 들었다.

그래서 SegFormer 논문을 읽어봤는데, 접근 방식이 정말 깔끔했다.

문제를 바라보는 관점이 달랐다

저자들이 지적한 핵심 문제는 두 가지였다. 첫째, positional encoding의 고정된 해상도 문제다. ViT에서 사용하는 절대적 위치 인코딩은 학습할 때 본 해상도에서만 잘 동작한다. 테스트할 때 다른 해상도의 이미지가 들어오면? 성능이 뚝 떨어진다.

둘째, 복잡한 디코더 구조다. 기존 방법들은 CNN의 Feature Pyramid Network처럼 복잡한 디코더를 설계했다. 트랜스포머 인코더가 이미 충분히 강력한데, 디코더까지 복잡하게 만들 필요가 있을까?

사실 이 부분 읽으면서 "맞아, 왜 다들 디코더를 그렇게 복잡하게 만들었지?" 하는 생각이 들었다.

해결책이 놀라울 정도로 간단했다

SegFormer의 핵심 아이디어는 hierarchical Transformer encoder와 lightweight All-MLP decoder의 조합이다.

처음엔 "에이, 그냥 계층적 구조랑 MLP 디코더 쓴다고 뭐가 달라지겠어" 싶었는데, 자세히 보니 디테일이 장난 아니었다.

Hierarchical Transformer Encoder의 비밀

일반적인 ViT는 이미지를 16x16 패치로 나누고, 모든 레이어에서 동일한 해상도를 유지한다. 반면 SegFormer는 CNN처럼 점진적으로 해상도를 줄여나간다.

Stage 1: H/4 × W/4 (패치 크기 4×4)
Stage 2: H/8 × W/8 (패치 크기 2×2)
Stage 3: H/16 × W/16 (패치 크기 2×2)
Stage 4: H/32 × W/32 (패치 크기 2×2)

각 스테이지에서 Efficient Self-Attention을 사용하는데, 여기가 진짜 재밌다. Reduction ratio를 써서 K, V의 시퀀스 길이를 줄인다. 예를 들어 Q가 N개면, K와 V는 N/R개로 줄여버린다.

이게 왜 중요하냐면, 계산량이 O(N²)에서 O(N²/R)로 줄어든다. 실제로 R=8 정도만 써도 성능은 거의 유지하면서 속도는 확 빨라진다.

Mix-FFN: Positional Encoding 없애기

가장 인상 깊었던 부분은 Mix-FFN이다. 기존 FFN에 3×3 depth-wise convolution을 끼워넣었다.

# 대충 이런 느낌
x = Linear(x)
x = DepthwiseConv3x3(x)  # 여기가 핵심!
x = GELU(x)
x = Linear(x)

처음엔 "아니 트랜스포머에 왜 convolution을?" 했는데, 이게 positional encoding을 대체하는 역할을 한다. Convolution이 원래 local inductive bias를 가지고 있으니까, 따로 위치 정보를 주입하지 않아도 자연스럽게 공간 정보를 학습한다.

이 부분이 진짜 "아하!" 모먼트였다. Positional encoding을 없애려고 그렇게 고민했는데, 답은 의외로 간단했다.

All-MLP Decoder: 극도로 간단하게

디코더는 정말 충격적일 정도로 간단했다. 4개 스테이지에서 나온 feature를 받아서:

1. MLP로 채널 수 통일 (256차원)
2. Bilinear interpolation으로 크기 맞추기
3. 다 더하기
4. 마지막 MLP로 클래스 예측

끝이다. 진짜 이게 끝이다.

처음엔 "이게 될까?" 싶었는데, 실험 결과를 보니 복잡한 디코더들보다 오히려 더 잘 된다. 인코더가 충분히 강력하면 디코더는 단순해도 된다는 걸 제대로 보여준 셈이다.

실험 결과가 말해주는 것

ADE20K에서 SegFormer-B5가 84.0% mIoU를 달성했다. 당시 SOTA인 Swin Transformer보다 높은 성능이다. 그런데 더 중요한 건 효율성이다.

SegFormer-B0는 겨우 3.7M 파라미터로 37.4% mIoU를 달성했다. 비슷한 크기의 다른 모델들보다 훨씬 좋은 성능이다.

특히 인상적이었던 건 다양한 해상도에서의 robustness다. 512×512로 학습하고 768×768로 테스트해도 성능이 잘 유지된다. Positional encoding을 없앤 효과가 제대로 나타난 거다.

저자들이 한 ablation study도 재밌었다. Mix-FFN을 빼면 성능이 확 떨어진다. All-MLP 디코더를 복잡한 디코더로 바꿔도 성능은 비슷한데 속도만 느려진다.

이 연구가 주는 교훈

SegFormer를 공부하면서 느낀 건, 복잡한 게 항상 좋은 건 아니다는 거다.

트랜스포머가 나오고 나서 다들 "어떻게 하면 더 복잡하고 정교한 구조를 만들까" 고민했는데, SegFormer는 오히려 단순함으로 승부했다. Positional encoding 문제도 Mix-FFN이라는 간단한 해결책으로 풀어냈고, 디코더도 과감하게 단순화했다.

실제로 써보려면 몇 가지 고려해야 할 점들이 있다. 메모리는 확실히 덜 먹지만, hierarchical 구조 때문에 구현이 좀 까다로울 수 있다. 그리고 작은 물체 검출 성능이 어떤지는 더 확인해봐야 할 것 같다.

개인적으로는 Mix-FFN 아이디어를 다른 곳에도 적용해볼 수 있을 것 같다. Vision Transformer의 다른 변형들에도 이런 식으로 convolution을 살짝 섞어주면 어떨까?

그리고 All-MLP 디코더의 성공을 보면서, 다른 dense prediction 태스크들도 디코더를 더 단순화할 여지가 있지 않을까 생각했다. Object detection이나 instance segmentation에서도 비슷한 시도를 해볼 만하다.

다음에는 SegFormer의 후속 연구들을 좀 더 파봐야겠다. 특히 최근에 나온 efficient transformer 계열 논문들이 SegFormer의 아이디어를 어떻게 발전시켰는지 궁금하다.