[SNU LIVER FORMULA 논문 리뷰] SELFMedMAE: SELF PRE-TRAINING WITH MASKED AUTOENCODERS FOR MEDICAL IMAGE CLASSIFICATION AND SEGMENTATION

논문 리뷰: 2203.05573

깃허브 주소: SelfMedMAE/lib/models/mae3d.py at main · cvlab-stonybrook/SelfMedMAE

SelfMedMAE/lib/models/mae3d.py at main · cvlab-stonybrook/SelfMedMAE

---

1. Introduction

• 의료 영상 분석의 특징:
  • 해부학적 구조 간의 기능적, 기계적 상호작용이 중요함.
  • 예: X-ray에서 폐 병변 판단은 심장/종격동 등 주변 구조와의 관계가 필요함.
  • 예: 뇌종양은 주변 부종, 구조 변화 등 맥락 정보와 연관됨.
• Masked Autoencoder (MAE):
  • 최근 자연영상에서 ViT 기반 self-supervised pre-training에 효과적임이 입증됨.
  • 입력의 일부를 마스킹 → 전체 복원 학습 → 맥락 정보 통합 능력 향상.
• 제안 방식: Self Pre-training
  • 타 데이터셋 없이, target task의 학습 데이터 자체에서 pre-training 수행
  • ImageNet 수준의 대규모 의료영상 데이터셋이 없다는 현실적 문제 해결 가능
  • domain mismatch 방지

---

2. Methodology

2.1. Vision Transformer (ViT) 구조

• Patch Embedding:
  • 2D 또는 3D 입력(x ∈ ℝᴴ⨯ᵂ⨯ᴰ⨯ᶜ)을 non-overlapping patch로 변환 후 linear projection
  • 예: 3D 볼륨은 (16×16×16) 패치 단위로 나눔
• Position Embedding:
  • Pre-training: sine-cosine embedding 사용
  • Fine-tuning: learnable position embedding을 sine-cosine 값으로 초기화
• Transformer Block:
  • Multi-head Self-Attention + MLP block 구조

2.2. MAE Self Pre-training 구성

① Encoder

• 마스킹되지 않은 patch만 입력으로 사용
• positional embedding 추가
• 전체 복원 학습을 위해, partial input으로부터 전체 표현을 추론해야 함

② Decoder

• encoder 출력 + mask token으로 구성된 전체 시퀀스 입력
• positional embedding 포함
• decoder는 pre-training에만 사용, downstream task에는 사용되지 않음

③ Loss Function

• masked patch만 복원 (visible patch에는 손실 계산하지 않음)
• MSE Loss 사용
• 정규화된 픽셀/voxel 값을 예측 대상으로 사용하면 성능 향상

2.3. Downstream Task Architecture

• 분류 (ChestX-ray14):
  • ViT의 class token 출력 뒤에 linear classifier 추가
  • multi-label binary cross entropy loss 사용
• 분할 (BTCV, BraTS):
  • UNETR 구조 채택: ViT encoder + CNN decoder
  • encoder의 다양한 resolution 출력을 decoder에 skip-connection
  • decoder는 spatial dimension 복원, upsampling, concatenation

---

3. Experiments and Results

3.1. 실험 설정 및 전처리

• CXR14: 112,120 chest X-ray, multi-label 분류
  • 전처리: histogram equalization, random crop/flip (224×224)
• BTCV: 30 CT abdomen scan, 8개 장기 segmentation
  • 전처리: voxel intensity [-175, 250] → [0, 1], crop (96×96×96)
• BraTS (MSD): 484 multi-modal MRI, 3종 tumor segmentation
  • 전처리: non-zero region normalization, crop (128×128×128)

학습 설정

• optimizer: AdamW
• patch size: 16×16 (2D) / 16×16×16 (3D)
• pre-training:
  • CXR14: 800 epochs (batch size 256)
  • BTCV: 10,000 epochs (batch size 6)
  • BraTS: 500 epochs (batch size 6)
• fine-tuning:
  • learning rate: task별 다름 (1e-3, 8e-4, 4e-4)
  • cosine decay, layer-wise lr decay (ratio: 0.75), DropPath (10%)

3.2. 주요 실험 결과

MAE Reconstruction 성능

• CXR, BTCV, BraTS에서의 마스킹(75%) 및 복원 결과
• visible patch는 손실 계산되지 않으므로 블러 현상 존재
• 복원 품질보다 맥락적 representation 학습이 목적

Lung Disease Classification (CXR14)

→ 기존 ImageNet transfer learning보다 MAE Self가 0.8% 향상

Abdomen Multi-organ Segmentation (BTCV)

→ UNETR baseline 대비 4.7% 향상, ImageNet보다도 우수

Brain Tumor Segmentation (BraTS)

Ablation: 마스킹 비율 / 에폭 수 영향

• BraTS에서 best 결과: 마스킹 비율 12.5%
• 500→2000 epoch까지 pre-training 길게 할수록 성능 증가
• 그러나 overfitting 가능성도 존재 → 적절한 조절 필요

Qualitative Result

• BTCV: false positive 제거 및 더 명확한 장기 분할
• BraTS: 뇌종양 내 necrotic core를 더 잘 포착

---

4. Conclusion

• MAE self pre-training은 의료 영상 분석에서 강력한 성능 향상을 유도
• 특히 데이터가 적은 환경에서 ImageNet 기반 transfer learning보다 우수
• 3D 의료 영상 (CT/MRI)에도 적용 가능성 입증
• 향후 예후 예측 및 생존 분석과 같은 임상 응용으로 확장 가능