[OUTTA Alpha팀 논문 리뷰] Part 8-4. Swin-UNet: Unet-like Pure Transformer for Medical Image Segmentation

논문 링크: https://arxiv.org/pdf/2105.05537

 

---

1. Introduction

• 의료 영상 분할은 진단·수술 등 임상 응용에 핵심이 되며, 오랜 기간 U-Net(CNN 기반)이 주류를 이루어 옴.
• 하지만 CNN의 국소적 특성(convolution operation이 인접 영역에 집중)으로 인해, 글로벌 문맥(long-range dependency) 학습에 한계가 존재.
• Transformer가 NLP 분야에서 성공한 이후, Vision 분야에도 적용(ViT, DeiT, Swin Transformer 등)되어 글로벌 상호작용을 학습하는 데 강점이 있음을 보임.
• 본 논문에서는 Swin Transformer를 활용하여 U-Net 유사 구조를 전부 Transformer 기반으로 설계한 Swin-Unet을 제안 → 의료 영상 분할에서의 가능성 및 성능 검증.

---

2. Related Work

1. CNN 기반 의료 영상 분할
   • 전통적으로 U-Net(2015)이 큰 성과를 거두었고, 이어서 Res-UNet, Dense-UNet, U-Net++, UNet3+ 등이 등장.
   • 3D 의료 영상 분야도 3D-UNet, V-Net 등 CNN 기반 기법이 다수.
   • 그러나 컨볼루션의 국소성으로 인해 전역 정보 학습에 제약.
2. Vision Transformers
   • NLP에서 Transformer가 대성공(→ ViT, DeiT, Swin Transformer 등 Vision 분야로 확장).
   • Swin Transformer는 윈도우 기반 다단계 구조로, 분류·검출·세그멘테이션 등 다양한 Vision 과제에서 탁월한 성능.
3. Self-Attention/Transformer + CNN 혼합 시도
   • 일부 연구[2,27,1 등]는 Transformer와 CNN을 병합해 성능을 높이는 전략.
   • 그러나 본 논문은 CNN 없이, 순수 Transformer 구조만으로 U-Net을 구성하는 접근을 제안.

---

3. Method

3.1 Architecture Overview

• U-Net 형태: Encoder, Bottleneck, Decoder, 그리고 여러 스킵 연결로 구성.
• Encoder:
  • 입력 이미지를 패치 크기(4×4)로 나눈 뒤(중첩되지 않음), 각 패치(4×4×3=48차원)를 토큰으로 변환.
  • 선형 레이어(“linear embedding”)를 통해 채널 크기 $C$로 투영.
  • 여러 개 Swin Transformer 블록과 패치 병합(Patch Merging) 레이어를 반복하면서 점진적 다운샘플링 수행.
• Bottleneck:
  • 인코더의 최종 출력(가장 작은 해상도)에 대해 Swin Transformer 블록 2개만 적용해 깊은 특징(Deep Feature)을 학습.
  • 이 단계에서 해상도와 채널 차원은 변하지 않음.
• Decoder:
  • 인코더와 대칭적으로 Swin Transformer 블록 + 패치 확장(Patch Expanding) 레이어를 사용해 업샘플링.
  • 스킵 연결을 통해 인코더 각 스케일의 특징을 결합, 공간 정보 손실 보완.
  • 마지막에는 4배 업샘플링을 수행해 원본 해상도로 복원하고, 선형 투영으로 픽셀 단위 분할 예측을 출력.

3.2 Architecture

1. Swin Transformer Block

   • 

   • Swin Transformer는 기존 MSA와 달리, 윈도우(창) 기반으로 어텐션을 계산.
   • 두 개의 연속된 Swin Transformer 블록이 각각 W-MSA(Window-based MSA)와 SW-MSA(Shifted Window-based MSA)를 수행:
     1. W-MSA 블록 → 지역적(윈도우 단위) 어텐션
     2. SW-MSA 블록 → 시프트 윈도우로 전역 정보 보강
   • 각 블록: LN(LayerNorm) → W-MSA 또는 SW-MSA → Residual Connection → LN(LayerNorm) →  MLP (2-layer, GELU) → Residual Connection구조.

   • 

   • 

     • Multi-head Attention의 내부 계산은 기존 Transformer와 유사
     • relative position bias로 공간적 정보 유지.
2. Encoder
   1. 입력:
      • 원본 이미지를 (4×4) 패치로 분할 → 토큰화(해상도 #\frac{H}{4}\times\frac{W}{4}#, 채널 $C$).
   2. Swin Transformer 블록 2개 연속 적용 → 형상 차원/해상도 그대로 유지, 특징 학습.
   3. 패치 병합(Patch Merging) 레이어
      • 2×2 블록의 토큰을 연결해 해상도를 절반, 채널 수를 2배로 조정(2배 다운샘플링).
      • 이를 3번 반복하여 총 4단계(인코더)로 계층적 특징 추출.
3. Bottleneck
   • 가장 작은 해상도(예: $\frac{H}{32}\times\frac{W}{32}$​)에서 Swin Transformer 블록 2개 추가.
   • 해상도나 채널 차원은 변하지 않고, 깊은 표현 학습만 담당.
4. Decoder
   1. Patch Expanding 레이어
      • 인코더와 반대 방향:
        • 채널 수를 선형 변환으로 늘린 후, 2×2 패치를 재배열하여 해상도 2배로 업샘플링.
        • 채널은 반으로 줄여서 인코더에 대응되는 스케일을 맞춤.
   2. Swin Transformer block:
      • 업샘플링된 특징에 대해서도 Swin Transformer를 적용해 특징을 정제.
   3. Skip Connection:
      • 인코더에서 추출된 다중 스케일 특징을 디코더 단계별로 합쳐, 공간 정보 손실을 보완.
   4. 최종적으로 4배 Upsampling까지 수행한 뒤,
      • 선형 투영 레이어로 픽셀 단위(원본 해상도) 분할 결과 출력.

---

4. Experiments

4.1 Datasets

1. Synapse Multi-organ CT
   • 30명 환자, 3779장의 축상 CT 슬라이스.
   • 장기(간, 췌장, 비장, 신장 등) 8개 클래스에 대한 분할.
   • 18명(훈련), 12명(테스트) 분할 → DSC, HD로 평가.
2. ACDC Cardiac MRI
   • 좌심실(LV), 우심실(RV), 심근(MYO) 분할.
   • 70(훈련), 10(검증), 20(테스트) 환자 → 주로 DSC로 평가.

4.2 Implementation details

• 입력 크기: 224×224 (패치 크기 4×4)
• 사전 학습: ImageNet에서 Swin Transformer pretrained 가중치 사용.
• Optimizer: SGD(momentum=0.9), weight decay=1e-4, batch size=24.
• 증강: 뒤집기, 회전 등.

4.3 Experiment results on Synapse dataset

• Swin-Unet가 DSC=79.13%, HD=21.55 등 기존 CNN(U-Net, Res-UNet 등) 대비 경계(HD) 정확도가 특히 개선.
• TransUnet(Transformer+CNN 혼합)보다도 HD 지표에서 ~10% 향상.
• 시각적으로도 과잉·부족 세분화 현상을 줄이는 데 성공 → 전역 상호작용 학습 덕분.

4.4 Experiment results on ACDC dataset

• DSC=90.00%로 강력한 일반화 능력 확인 (MRI 모달리티에서도 우수).
• CNN 기반과 비교 시 여전히 전역적 특성을 더 잘 포착.

4.5 Ablation study

1. Up-sampling 기법 비교 (Patch Expanding vs. ConvTranspose vs. Bilinear):
   • Patch Expanding 방식이 DSC↑, HD↓ 측면에서 가장 우수.
2. Skip Connection 개수:
   • (0,1,2,3 개) 비교 → 스킵이 많을수록 성능↑, 최종 3개 사용.
3. 입력 해상도:
   • 224→384로 증가 시 일부 성능 개선 있으나 연산량 크게 증가 → 224×224 유지.
4. 모델 크기(Tiny, Small 등):
   • 단순히 깊이를 늘린다고 성능이 비례 향상되지 않음 → Tiny 설정이 효율적.

---

5. Conclusion

• 본 논문에서는 의료 영상 분할에 순수 Transformer만을 사용한 U자형 구조(Swin-Unet)를 제안.
• 윈도우 기반 Swin Transformer 블록 + Patch Expanding + Skip Connection을 활용하여 전역 문맥과 다중 스케일 특징을 함께 학습.
• 실험(다기관 CT, 심장 MRI)에서 CNN 대비 경계 정확도, 일반화 능력 우수성을 입증.
• 한계 및 향후 과제:
  • ImageNet 사전 학습에 의존 → 의료 이미지 도메인 자체 사전 학습 연구 필요.
  • 2D slice만 다룸 → 3D volume 데이터 확장 연구가 필요.

---

요약

Swin-Unet은 순수 Transformer 기반 U-Net 구조로, CNN 없이도 의료 영상 분할에서 전역 상호작용 학습을 성공적으로 수행함을 보임. 인코더-디코더 모든 계층에 Swin Transformer 블록을 적용하고, 업샘플링에는 Patch Expanding Layer를 도입해 CNN식의 Conv/Interpolation을 대체함.

실험 결과, 다중 기관 분할(CT)·심장 분할(MRI) 작업에서 높은 정확도와 우수한 경계 복원 능력을 달성. 이는 Transformer가 의료 영상 분석에서도 효과적임을 시사하는 중요한 시도라 할 수 있음.