SegFormer 완벽 가이드 — 위치 인코딩이 없는 트랜스포머가 세그멘테이션에 효과적인 이유

SegFormer 이전에도 Transformer를 의미론적 분할에 도입하려는 시도는 있었지만, 대부분 위치 인코딩과 복잡한 디코더에 의존했습니다. 따라서 두 가지 문제에 직면했습니다. 테스트 해상도가 훈련 해상도와 다를 경우 정확도가 떨어질 수 있었고, 계산 비용이 많이 들 수 있었습니다. SegFormer(Xie, Wang, Yu, Anandkumar, Alvarez, and Luo, NeurIPS 2021)는 위치 인코딩이 없는 계층적 인코더와 컨볼루션을 전혀 사용하지 않는 All-MLP 디코더라는 놀랍도록 간단한 설계로 이 두 가지 문제를 모두 해결합니다. 이 논문에서는 원 논문(arXiv:2105.15203)을 바탕으로 이러한 설계가 효과적인 이유를 단계별로 살펴봅니다.

본 글은 원 논문 "SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers"(Xie et al., NeurIPS 2021, arXiv:2105.15203)를 기반으로 작성되었습니다.

0. 학습 내용

• SegFormer가 기존 Transformer 기반 분할 모델의 문제점을 해결하고자 했던 이유
• 계층적 MiT(Mix Transformer) 인코더가 위치 인코딩 없이 작동하는 방식
• 컨볼루션이 전혀 없는 MLP만으로 구성된 디코더가 높은 정확도를 달성할 수 있는 이유
• Cityscapes 및 ADE20K 데이터셋에서 B0~B5에 대한 mIoU, 파라미터 개수 및 계산 비용 측정 결과
• 이미지 손상에 대한 강건성을 포함하여 실제 모델 선택 방법

1. SegFormer란 무엇인가?

SegFormer는 위치 인코딩이 없는 계층적 Transformer 인코더(MiT)와 컨볼루션을 전혀 사용하지 않는 경량 All-MLP 디코더를 결합한 의미론적 분할 모델입니다. 핵심 아이디어는 인코더와 디코더 간의 역할 분담을 재설계하는 것입니다. 인코더가 여러 해상도의 특징을 생성하면 디코더를 크게 단순화할 수 있습니다.

2. 왜 이러한 설계가 필요했을까요?

Vision Transformer(ViT)를 의미론적 분할의 인코더로 직접 사용하는 것은 두 가지 문제를 야기합니다. 첫째, ViT에서 생성된 특징은 단일 해상도만 가집니다. 객체는 하나의 이미지 내에서 다양한 크기를 가질 수 있으며, FCN 및 U-Net과 같은 기존 CNN 기반 분할 모델은 여러 해상도의 특징을 결합하여 이를 해결합니다. 단일 해상도만 생성하는 ViT는 이러한 다중 스케일 표현 능력을 상실합니다.

두 번째 문제는 위치 인코딩에 대한 의존성입니다. ViT는 이미지를 패치로 나누어 Transformer에 입력하지만, Transformer 자체는 이미지에서 각 패치의 위치에 대한 정보를 가지고 있지 않습니다. 따라서 고정되거나 학습된 위치 인코딩을 별도로 추가해야 합니다. 이 인코딩은 학습에 사용된 입력 해상도에 맞춰 구축되므로, 추론 시 해상도가 다른 이미지의 경우 위치 인코딩을 보간해야 하며, 이로 인해 정확도가 떨어질 수 있습니다. 분할 작업은 종종 객체 탐지보다 고해상도 입력을 요구하기 때문에 이러한 해상도 의존성은 특히 문제가 됩니다.

SegFormer는 계층적 인코더를 사용하고 위치 인코딩을 완전히 제거함으로써 이 두 가지 문제를 근본적으로 해결합니다.

3. 입력은 무엇인가요?

다른 분할 모델과 마찬가지로 입력은 RGB 이미지 x \in \mathbb{R}^{H \times W \times 3}입니다. SegFormer는 Cityscapes 데이터셋의 경우 고해상도 이미지(1024×2048)에서, ADE20K 데이터셋의 경우 약 512×512 해상도의 이미지에서 평가되었습니다. ViT는 16×16 크기의 패치로 다운샘플링하는 반면, SegFormer의 인코더는 4×4 크기의 패치부터 시작하여 특징 추출 초기 단계에서 더 많은 디테일을 보존합니다.

4. 예측 결과는 무엇인가요?

출력은 입력 해상도의 1/4 크기로 축소된 세그멘테이션 맵입니다. 각 픽셀의 클래스는 \hat{y} \in \mathbb{R}^{\frac{H}{4} \times \frac{W}{4} \times N_{cls}} (N_{cls}는 클래스 개수)로 표현됩니다. 이후 최근접 이웃 보간법 또는 유사한 방법을 사용하여 원래 해상도로 복원됩니다. 핵심 아이디어는 인코더가 충분히 풍부한 멀티스케일 특징을 생성하면 디코더는 간단한 선형 레이어 조합만으로 이러한 특징들을 통합하고 정확한 마스크를 생성할 수 있다는 것입니다.

5. 기본 아키텍처

SegFormer는 크게 두 가지 블록으로 구성됩니다. 하나는 4가지 해상도에서 특징을 추출하는 MiT(Mix Transformer) 인코더이고, 다른 하나는 이러한 특징들을 결합하여 세그멘테이션 마스크를 출력하는 All-MLP 디코더입니다.

그림 1 — MiT 인코더는 4가지 해상도(1/4, 1/8, 1/16, 1/32)의 특징을 생성하고, All-MLP 디코더는 선형 레이어만 사용하여 이러한 특징을 결합하여 분할 마스크를 생성합니다. 유일한 컨볼루션은 인코더의 Mix-FFN에 있습니다. 디코더는 말 그대로 전체가 선형 레이어로 구성되어 있습니다.

인코더가 4가지 해상도로 특징을 미리 구축하기 때문에 디코더는 특징을 정렬하고 결합하는 최소한의 처리만으로 정확한 마스크를 얻을 수 있습니다. 이러한 분업이 SegFormer의 전체 파라미터 수를 낮게 유지하는 주요 이유입니다.

6. 구성 요소의 기술적 세부 사항

겹치는 패치 병합 — 4단계 계층 구조

ViT처럼 이미지를 한 번만 패치로 나누는 대신, MiT 인코더는 4단계에 걸쳐 해상도를 점진적으로 낮춥니다. 출력 해상도는 입력의 1/4, 1/8, 1/16, 1/32로, ResNet과 같은 CNN 백본에서 사용하는 계층 구조를 재현합니다.

패치 분할 방식 또한 ViT의 겹치지 않는 패치와 다릅니다. 인접한 패치는 병합될 때 겹칩니다. 첫 번째 단계에서는 커널 크기 K=7, 스트라이드 S=4, 패딩 P=3를 사용하고, 이후 단계에서는 K=3, S=2, P=1을 사용합니다. 이러한 오버랩을 통해 패치 경계를 가로지르는 로컬 연속성, 즉 인접한 패치들이 공유하는 정보를 유지하면서 특징 맵의 크기를 줄입니다.

효율적인 셀프 어텐션

일반적인 멀티 헤드 셀프 어텐션은 길이가 N인 시퀀스에 대해 O(N^2)의 계산량을 필요로 합니다. 고해상도 분할의 경우 N가 매우 커지기 때문에 이는 심각한 병목 현상입니다. SegFormer는 시퀀스 축소(Sequence Reduction) 기능을 도입하여 어텐션을 계산하기 전에 키(Key) 및 값(Value) 시퀀스를 R의 축소 비율로 압축합니다.

K, V 시퀀스는 사용되기 전에 길이가 N에서 N/R로 줄어듭니다. 1~4단계의 길이 축소 비율 R는 [64, 16, 4, 1]로 설정됩니다. 얕은 단계는 고해상도와 긴 시퀀스를 사용하므로 적극적으로 압축하고, 깊은 단계는 해상도가 낮고 짧은 시퀀스를 사용하므로 전혀 압축하지 않습니다. 이는 계산 비용을 O(N^2)에서 O(N^2/R)로 줄여 고해상도 입력에서도 셀프 어텐션을 실용적으로 사용할 수 있게 합니다.

Mix-FFN — 위치 인코딩 제거 방법

SegFormer는 피드포워드 네트워크(FFN)에 3×3 컨볼루션을 혼합하여 Mix-FFN을 생성함으로써 위치 인코딩을 완전히 제거할 수 있습니다.

3×3 컨볼루션은 이미지 경계에 제로 패딩을 사용합니다. 이 패딩된 경계는 "이미지 내 현재 위치"에 대한 단서를 컨볼루션에 간접적으로 전달하여, 명시적인 위치 인코딩 없이 Mix-FFN을 통과하는 것만으로 위치 정보가 특징에 자연스럽게 녹아들게 합니다. 결과적으로, 학습 해상도와 다른 해상도에서 테스트할 때 위치 인코딩을 보간할 필요가 없어 정확도 저하를 방지할 수 있습니다.

모든 MLP 디코더 - 컨볼루션이 불필요한 이유

이 디코더는 선형 레이어(MLP)만을 사용하여 다음 네 단계로 구성됩니다.

1. 채널 통합: 각 단계의 특징 F_i을 독립적인 선형 레이어를 사용하여 동일한 채널 수 C로 매핑합니다.

2. 업샘플링 및 연결: 모든 특징을 1/4 해상도로 업샘플링하고 채널 차원을 따라 연결합니다.

3. 융합(Fuse): 연결된 4C 차원 특징을 선형 레이어를 사용하여 C 차원으로 매핑합니다.

4. 예측(Predict): 마지막 선형 레이어를 사용하여 각 클래스에 대해 하나의 채널을 가진 분할 마스크를 출력합니다.

원 논문은 유효 수용 영역(ERF) 분석을 통해 이 설계를 뒷받침합니다. DeepLabv3+와 같은 CNN 기반 디코더는 넓은 수용 영역을 얻고 높은 정확도를 달성하기 위해 확장 컨볼루션 및 멀티스케일 풀링과 같은 복잡한 메커니즘이 필요합니다. 이와 대조적으로 MiT 인코더는 추가적인 메커니즘 없이 셀프 어텐션을 통해 얕은 단계에서는 자연스럽게 로컬 ERF를, 깊은 단계에서는 전체 이미지를 커버하는 비로컬 ERF를 획득합니다. 인코더 자체가 로컬에서 글로벌 규모에 이르는 정보를 포함하는 특징을 이미 구축하고 있기 때문에 디코더는 복잡한 수용 영역 확장 메커니즘이 필요하지 않습니다. 이것이 바로 간단한 선형 레이어 조합으로 충분한 이론적 이유입니다.

학습 레시피

원 논문에서는 특별한 기법 없이 간단한 학습 설정을 사용합니다. MiT 인코더는 ImageNet-1K에서 사전 학습되었고, 디코더는 무작위 초기화를 통해 학습됩니다. 최적화에는 AdamW 알고리즘이 사용되었으며, 학습률은 초기값 0.00006에서 시작하여 폴리 스케줄링(계수 1.0의 멱법칙 감쇠)을 적용하여 점진적으로 감소합니다. 반복 횟수는 ADE20K와 Cityscapes 데이터셋에서는 16만 회, COCO-Stuff 데이터셋에서는 8만 회입니다. 학습 데이터 증강은 0.5~2.0 비율의 무작위 크기 조정, 무작위 좌우 반전, 무작위 자르기(ADE20K의 경우 512×512, Cityscapes의 경우 1024×1024, ADE20K의 B5 데이터셋에서는 640×640)로 제한됩니다. 논문에서는 OHEM(Online Hard Example Mining), 보조 손실 함수, 클래스 균형 손실 함수와 같은 널리 사용되는 기법을 사용하지 않았다고 명시적으로 밝히고 있습니다. 목표는 MiT 인코더와 All-MLP 디코더가 이러한 추가 기능 없이도 높은 mIoU를 달성할 수 있음을 보여주는 것입니다.

7. 모델 변형 비교

SegFormer는 MiT-B0부터 MiT-B5까지 6가지 크기로 제공됩니다. 다음은 원 논문의 표 2에 보고된 Cityscapes 및 ADE20K 측정 결과입니다.

출처: 원 논문의 표 2(Xie et al., NeurIPS 2021). Cityscapes mIoU는 멀티스케일 및 좌우 반전(MS) 테스트를 통해 측정됩니다. 파라미터 수는 전체 인코더와 디코더를 모두 포함하며, B5의 경우에도 디코더는 전체의 약 4%에 불과합니다. 참고로, B5는 싱글스케일(SS) 테스트에서 82.4 mIoU, MS 테스트에서 84.0 mIoU를 달성합니다.

이 논문은 당시 최첨단 기술과 이러한 결과를 비교하여 SegFormer-B4가 6410만 개의 파라미터로 ADE20K에서 50.3% mIoU를 달성하여 당시 최고 기술보다 2.2포인트 높은 성능을 보였으며, 모델 크기는 거의 5분의 1에 불과하다고 보고합니다. B0는 380만 개의 파라미터만으로도 76.2의 실용적인 Cityscapes mIoU를 유지하며, 이러한 수치는 에지 지향 설계의 타당성을 뒷받침합니다.

8. 한계점

SegFormer의 대부분의 약점은 일반적인 Transformer 기반 모델들이 공통적으로 가지고 있는 약점입니다. 첫째, 대형 모델(B3 이상)은 학습 및 추론 비용이 높습니다. Cityscapes 데이터셋의 FLOPs는 B0의 125.5G에서 B5의 1460.4G로 10배 이상 증가했습니다. 고해상도 분할의 경우 이러한 증가는 실시간 처리를 어렵게 만들 수 있습니다.

대규모 사전 학습 데이터에 대한 의존성 또한 중요한 문제입니다. MiT 인코더는 ImageNet-1K 데이터셋으로 사전 학습되었기 때문에, 미세 조정 데이터의 품질과 양이 의료 영상이나 위성 영상처럼 사전 학습 데이터와 크게 다른 영역으로의 전이 성능에 큰 영향을 미칩니다.

전선이나 표지판과 같은 매우 작은 객체나 얇은 구조물 역시 SegFormer에게는 근본적으로 어려운 문제입니다. 겹치는 패치 병합(Overlapped Patch Merging)은 첫 번째 단계에서도 4×4로 다운샘플링하기 때문에, 폭이 몇 픽셀에 불과한 구조물은 더 깊은 단계에 도달하기 전에 정보 손실이 발생할 수 있습니다.

동시에, 원 논문에서 보고된 Cityscapes-C 데이터셋에 대한 견고성 평가는 분명한 강점입니다. 가우시안 노이즈가 추가된 손상된 이미지의 경우, SegFormer-B5는 Xception-71 백본을 사용하는 DeepLabv3+ 대비 최대 588%의 상대적 mIoU 향상을 기록했습니다. 눈과 같은 노이즈가 추가된 손상 이미지의 경우, 최대 상대적 향상은 295%입니다. 이러한 측정 결과는 일반적인 CNN 기반 모델보다 이미지 손상에 대한 내성이 더 뛰어나다는 것을 보여줍니다. 이러한 견고성은 셀프 어텐션이 로컬 노이즈에 덜 영향을 받기 때문인 것으로 생각됩니다.

9. 실제 모델 선택 방법

엣지 디바이스 및 실시간 처리에는 SegFormer-B0 또는 B1이 현실적인 선택입니다. B0은 380만 개의 파라미터와 Cityscapes 데이터셋에서 7620만 mIoU의 성능을 보여주며 정확도와 컴팩트함의 균형을 잘 맞추고 있어 임베디드 디바이스 및 드론에 배포하기에 적합합니다.

자율 주행에서의 주행 가능 영역 인식과 같이 고해상도 이미지와 일관된 정확도가 요구되는 애플리케이션에는 B2 또는 B3가 적절한 절충안이 될 수 있습니다. B4와 B5는 가장 높은 정확도를 제공하지만, Cityscapes FLOPs가 1200G를 초과하므로 차량용 GPU에서 실시간 추론을 수행하려면 일반적으로 양자화 또는 TensorRT 변환과 같은 최적화가 필요합니다.

의료 또는 위성 영상의 오프라인 정밀 분석의 경우 실시간 성능보다 정확도가 우선시됩니다. B4 또는 B5가 적합하며, 필요에 따라 도메인별 데이터에 맞춰 미세 조정할 수 있습니다.

농업용 로봇이나 실외 감시 카메라와 같이 잦은 악천후나 조명 변화가 있는 실외 환경에서는 Cityscapes-C에서 측정된 견고성이 실질적인 이점입니다. 유사한 정확도를 가진 CNN 기반 모델과 비교했을 때, SegFormer는 이미지 손상 환경에서 정확도 손실이 적을 것으로 예상됩니다.

객체 탐지 및 의미론적 분할의 기초를 배우려면 "객체 탐지 및 의미론적 분할 소개"를 참조하십시오. 최근 시맨틱 분할 관련 동향은 “시맨틱 분할 기술 동향”을 참조하십시오.

10. 세 줄 요약

• SegFormer는 위치 인코딩 없이 다중 스케일 특징을 생성하는 계층적 인코더(MiT)와 컨볼루션을 사용하지 않는 All-MLP 디코더를 결합합니다.

• Mix-FFN의 3×3 컨볼루션은 위치 정보를 암묵적으로 유출하여 위치 인코딩의 필요성을 없애고, Transformer의 넓은 유효 수용 영역 덕분에 간단한 MLP 디코더를 사용할 수 있습니다.

• 이 제품군은 B0(380만 파라미터, 도시 풍경 데이터셋에서 76.2 mIoU)부터 B5(8470만 파라미터, 84.0 mIoU)까지 다양한 파라미터를 지원하며, 측정 결과 이미지 손상에 대한 강건성도 입증되었습니다.

참고 자료

• SegFormer: Transformer를 이용한 의미론적 분할을 위한 간단하고 효율적인 설계 (NeurIPS 2021, arXiv:2105.15203)
• arXiv의 SegFormer 논문 (전문)
• 공식 SegFormer 구현 (NVlabs/SegFormer, GitHub)
• NeurIPS 2021 학술대회 논문집: SegFormer