https://arxiv.org/abs/2409.03062
MobileUNETR: A Lightweight End-To-End Hybrid Vision Transformer For Efficient Medical Image Segmentation
Skin cancer segmentation poses a significant challenge in medical image analysis. Numerous existing solutions, predominantly CNN-based, face issues related to a lack of global contextual understanding. Alternatively, some approaches resort to large-scale T
arxiv.org
안녕하세요, 이번 글에서는 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 또는 트랜스포머 기반 접근법의 단점을 보완하여 로컬 및 글로벌 특징을 효과적으로 학습함.
MobileUNETR의 특징은 계층적 하이브리드 인코더-디코더 아키텍처로, 거친 특징과 세밀한 특징을 추출하고 결합하는 엔드투엔드 프레임워크입니다. 저수준 특징 추출과 공간적 다운샘플링을 위한 경량 컨볼루션 스템, 각 단계에서 지역 및 전역 특징을 효율적으로 추출하는 하이브리드 인코더 블록, 전역 특징 추출 및 장거리 맥락 이해를 위한 MobileViT 블록, 그리고 디코더 단계 동안 모델이 동적으로 특징을 적응할 수 있게 하면서 지역/전역 특징을 효율적으로 업스케일하고 결합하는 새로운 디코더 블록으로 구성됩니다
2-1. Encoder
인코더의 특징 추출 모듈은 크게 두 단계로 나눌 수 있음; CNN 기반 지역 특징 추출 및 다운 샘플링(by. 경량 컨볼루션 스템), 하이브리드 Transformer/CNN 기반 지역 및 전역 표현 학습(by. MobileViT 블록)
- 경량 컨볼루션 스템
- 초기 단계에서 CNN활용하여 로 특징 추출과 공간적 다운샘플링을 담당. Transformer 입력 크기 줄여 계산 비용 절감.
- 구성 요소:
- 3X3 Conv, BN, SiLU 활성화 함수
- MobileNet v2 블록 -> 효과적인 특징 추출
- 하이브리드 인코더 블록
- CNN과 트랜스포머의 장점을 결합하여 로컬 및 글로 특징을 동시에 학습.
- 구성 요소: MobileNet v2 블록 + MobileViT 블록
- Point-wise 1×1 Conv (stride=1, BN, SiLU): 특징 차원 변환 및 선형 변환 수행.
- Depth-wise 3×3 Conv (stride=2, BN, SiLU): 공간적 특징을 학습하며 계산 비용 절감.
- Piece-wise 1×1 Conv (stride=1, BN, SiLU): 특징 맵을 다시 조정.
- MobileViT 블록: 트랜스포머 기반 글로벌 특징 학습 수행.
- MobileViT 블록
- 트랜스포머를 컨볼루션과 결합하여 장거리 의존성 및 글로벌 컨텍스트 학습.
- 과정:
- CNN을 사용해 로컬 특징을 추출.
- 특징 맵을 패치로 변환하여 트랜스포머 입력으로 변형.
- 트랜스포머 블록을 통해 패치 간 관계 학습 및 글로벌 특징 강화.
- CNN을 사용해 트랜스포머 출력을 원래 공간 구조로 재구성.
- 트랜스포머로 학습한 글로벌 특징과 CNN으로 학습한 로컬 특징을 결합.
2-2. Decoder
- 하이브리드 디코더 블록
- CNN 기반 디코더의 한계: 특징을 계층적으로 업샘플링하며 skip connections과 결합하지만, 글로벌 컨텍스트와 정렬되지 않는 경우 발생.
- 과정:
- Transpose Conv: 이전 단계 특징 맵을 업샘플링 -> 공간 해상도 증가
- Concat: 업샘플링된 특징과 인코더에서 전달된 skip connection 특징을 결합. 로컬 특징과 글로벌 컨텍스트 정보를 효과적으로 통합.
- MobileNet v2 블록: 기존 CNN 디코더의 특징 정제 과정과 유사한 역할. 로 특징 정제
- MobileViT 블록: 글로벌 컨텍스트 반영
3. Experiment
- 데이터셋
- ISIC 2016: 피부 병변 분할을 위한 첫 번째 표준 데이터셋 (900개 훈련, 300개 테스트 이미지)
- ISIC 2017: 데이터셋 크기 확장 (2500개 훈련, 600개 테스트 이미지)
- ISIC 2018: 가장 포괄적인 피부 병변 분할 데이터셋 (2694개 훈련, 1000개 테스트 이미지)
- 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 적용 모델과 비교하여도 성능이 유사하거나 더 우수.
- 경량 아키텍처가 대형 모델과 경쟁 가능함을 입증.
- 과도한 매개변수화 없이도 최적의 성능을 달성 가능함을 강조.