[2025-1] 임수연 - MobileUNETR

https://arxiv.org/abs/2409.03062

MobileUNETR: A Lightweight End-To-End Hybrid Vision Transformer For Efficient Medical Image Segmentation

---

안녕하세요, 이번 글에서는 MobileUNETR: A Lightweight End-To-End Hybrid Vision Transformer For Efficient Medical Image Segmentation (2024) 논문 리뷰를 해보도록 하겠습니다.

 

1. Introduction

연구 배경 및 동기

피부암 예시 사진

• 피부암 Segmentation의 필요성:
  • 피부암은 전 세계적으로 가장 흔하고 빠르게 증가하는 암 중 하나로, 조기 발견이 중요함.
  • 병변의 모양, 크기, 색상 등이 환자마다 다르게 나타나므로 정확한 분할이 어려움.
  • 피부과 전문의의 진단이 조명, 개인적 경험, 주관적 판단에 따라 달라질 수 있어 객관적인 평가 도구가 필요함.

• 기존 접근법과 한계:
  • 초기 CAD(Computer-Aided Diagnostic) 시스템: 색상 공간 변환, 주성분 분석, 수작업 특징 엔지니어링 활용 → 복잡한 피부 병변을 정확히 분할하는 데 한계.
  • CNN 기반 접근법:
    • CNN은 이미지의 local 특징을 잘 포착하지만, 장거리(global) 관계 이해가 부족함.
    • 컨텍스트 이해를 향상하기 위해 더 깊은 네트워크를 설계했지만, 모델 크기가 커지고 계산 비용이 증가하는 문제 발생.
  • Transformer 기반 접근법:
    • Vision Transformer(ViT)는 장거리 의존성을 학습하여 글로벌 컨텍스트를 이해하는 데 강점이 있음.
    • 입력 이미지마다 생성되는 큰 시퀀스 길이로 인해 계산적으로 복잡하고 무거움 -> 의료 영상 분할과 같은 리소스 제한적인 환경에서 큰 문제
• MobileUNETR의 필요성
  • CNN의 로컬 특징 추출 능력과 트랜스포머의 글로벌 컨텍스트 이해 능력을 결합하여 계산 비용을 줄이면서도 높은 성능을 유지하는 모델을 제안.
  • 기존 모델 대비 파라미터 수 10배, 계산 복잡성 23배 감소하면서도 SOTA 모델을 뛰어넘는 성능을 보임.

 

 

 

2. MobileUNETER Architecture

• (a) MobileUNETR는 계층적 하이브리드 CNN-트랜스포머 인코더-디코더 구조로 설계됨.
• 기존 CNN 또는 트랜스포머 기반 접근법의 단점을 보완하여 로컬 및 글로벌 특징을 효과적으로 학습함.

MobileUNETER Architecture: (b) 경량 컨볼루션 스템 (c) 하이브리드 인코더 블록 (d) MobileViT 블록 (e) 하이브리드 디코더 블록

 

MobileUNETR의 특징은 계층적 하이브리드 인코더-디코더 아키텍처로, 거친 특징과 세밀한 특징을 추출하고 결합하는 엔드투엔드 프레임워크입니다. 저수준 특징 추출과 공간적 다운샘플링을 위한 경량 컨볼루션 스템, 각 단계에서 지역 및 전역 특징을 효율적으로 추출하는 하이브리드 인코더 블록, 전역 특징 추출 및 장거리 맥락 이해를 위한 MobileViT 블록, 그리고 디코더 단계 동안 모델이 동적으로 특징을 적응할 수 있게 하면서 지역/전역 특징을 효율적으로 업스케일하고 결합하는 새로운 디코더 블록으로 구성됩니다

 

2-1. Encoder

인코더의 특징 추출 모듈은 크게 두 단계로 나눌 수 있음; CNN 기반 지역 특징 추출 및 다운 샘플링(by. 경량 컨볼루션 스템), 하이브리드 Transformer/CNN 기반 지역 및 전역 표현 학습(by. MobileViT 블록)

• 경량 컨볼루션 스템
  • 초기 단계에서 CNN활용하여 로 특징 추출과 공간적 다운샘플링을 담당. Transformer 입력 크기 줄여 계산 비용 절감.
  • 구성 요소:
    • 3X3 Conv, BN, SiLU 활성화 함수
    • MobileNet v2 블록 -> 효과적인 특징 추출
• 하이브리드 인코더 블록
  • CNN과 트랜스포머의 장점을 결합하여 로컬 및 글로 특징을 동시에 학습.
  • 구성 요소: MobileNet v2 블록 + MobileViT 블록
    1. Point-wise 1×1 Conv (stride=1, BN, SiLU): 특징 차원 변환 및 선형 변환 수행.
    2. Depth-wise 3×3 Conv (stride=2, BN, SiLU): 공간적 특징을 학습하며 계산 비용 절감.
    3. Piece-wise 1×1 Conv (stride=1, BN, SiLU): 특징 맵을 다시 조정.
    4. MobileViT 블록: 트랜스포머 기반 글로벌 특징 학습 수행.
• MobileViT 블록
  • 트랜스포머를 컨볼루션과 결합하여 장거리 의존성 및 글로벌 컨텍스트 학습.
  • 과정:
    1. CNN을 사용해 로컬 특징을 추출.
    2. 특징 맵을 패치로 변환하여 트랜스포머 입력으로 변형.
    3. 트랜스포머 블록을 통해 패치 간 관계 학습 및 글로벌 특징 강화.
    4. CNN을 사용해 트랜스포머 출력을 원래 공간 구조로 재구성.
    5. 트랜스포머로 학습한 글로벌 특징과 CNN으로 학습한 로컬 특징을 결합.

2-2. Decoder

• 하이브리드 디코더 블록
  • CNN 기반 디코더의 한계: 특징을 계층적으로 업샘플링하며 skip connections과 결합하지만, 글로벌 컨텍스트와 정렬되지 않는 경우 발생.
  • 과정:
    1.  Transpose Conv:  이전 단계 특징 맵을 업샘플링 -> 공간 해상도 증가
    2. Concat: 업샘플링된 특징과 인코더에서 전달된 skip connection 특징을 결합. 로컬 특징과 글로벌 컨텍스트 정보를 효과적으로 통합. 
    3. MobileNet v2 블록: 기존 CNN 디코더의 특징 정제 과정과 유사한 역할. 로 특징 정제
    4. MobileViT 블록: 글로벌 컨텍스트 반영

 

 

 

3. Experiment

• 데이터셋
  • ISIC 2016: 피부 병변 분할을 위한 첫 번째 표준 데이터셋 (900개 훈련, 300개 테스트 이미지)
  • ISIC 2017: 데이터셋 크기 확장 (2500개 훈련, 600개 테스트 이미지)
  • ISIC 2018: 가장 포괄적인 피부 병변 분할 데이터셋 (2694개 훈련, 1000개 테스트 이미지)
  • PH2: 비교적 소규모 데이터셋으로, 모델의 일반화 능력을 테스트하는 데 사용됨.

ISIC 2017 데이터셋, PH2 데이터셋 실험 결과

• 성능 비교
  • ISIC 2016: 기존 모델 대비 IoU 2.17%, Dice 1.21% 향상
  • ISIC 2017: IoU 2.47%, Dice 1.84% 향상
  • ISIC 2018: IoU 2.54%, Dice 1.71% 향상
  • PH2: IoU 2.68%, Dice 1.3% 향상

 

 

 

• AdvancedTrainingTechniques와의 비교 결과
  • 파라미터 효율적 미세 조정(PEFT) 기법과 비교 평가.
  • PEFT 적용 모델과 비교하여도 성능이 유사하거나 더 우수.
  • 경량 아키텍처가 대형 모델과 경쟁 가능함을 입증.
  • 과도한 매개변수화 없이도 최적의 성능을 달성 가능함을 강조.