From Convolution to Transformer: A Comparative Study of U-Net Variants for Brain Tumor and Retinal Vessel Segmentation

기술 요약: 컨볼루션에서 트랜스포머까지

문제 정의
의료 영상 분할(Segmentation)은 컴퓨터 보조 진단, 치료 계획 수립 및 질병 모니터링에 있어 매우 중요합니다. U-Net 아키텍처는 인코더-디코더 구조와 스킵 연결(skip connections) 덕분에 생물 의학 분할 분야의 표준이 되었으나, 기존의 컨볼루션 기반 변형 모델들은 한계에 직면해 있습니다. 구체적으로, 표준 U-Net은 국소적인 컨볼루션 연산에 의존하기 때문에 장거리 의존성(long-range dependencies)과 전역적 문맥 정보(global contextual information)를 포착하는 능력이 제한됩니다. 이러한 한계는 불규칙한 형태, 이질적인 외관, 모호한 경계를 보이는 뇌종양 분할과 같은 복잡한 작업에서 특히 치명적입니다. 또한, 다양한 U-Net 변형 모델(잔차, 어텐션 기반 또는 트랜스포머 기반)의 성능은 영상 양상(3D MRI 대 2D 안저 영상)과 대상 구조(종양 하위 영역 대 가는 망막 혈관)에 따라 크게 달라질 수 있습니다.

방법론
본 연구는 두 가지 서로 다른 벤치마크 데이터셋을 대상으로 다섯 가지 대표적인 U-Net 기반 아키텍처를 비교 평가합니다:

1. BraTS 2023: 뇌종양 분할을 위한 다중 모달 볼륨형 MRI 데이터셋 (전체 종양, 강화 종양, 종양 핵 하위 영역 포함).
2. DRIVE: 가늘고 분지되는 형태와 낮은 대비를 특징으로 하는 2D 안저 영상 데이터셋 (망막 혈관 분할).

평가된 모델들은 다음과 같은 설계 철학의 스펙트럼을 나타냅니다:

• 3D U-Net: 슬라이스 간 해부학적 연속성을 포착하기 위해 3D 컨볼루션을 사용하여 기존 2D 아키텍처를 볼륨형 데이터로 확장함.
• Residual U-Net: 잔차 블록(residual blocks)을 도입하여 그래디언트 전파와 특징 학습의 안정성을 개선하고, 복잡하고 이질적인 구조 추출을 도움.
• Attention U-Net: 스킵 연결에 어텐션 게이트(attention gates)를 도입하여 무관한 배경 정보를 억제하고 작업 관련 영역을 강조함.
• UNETR: 컨볼루션 인코더를 비전 트랜스포머(ViT)로 대체하여 이미지 패치에 대한 셀프 어텐션 메커니즘을 통해 전역적 의존성을 모델링함.
• Swin UNETR: 계층적 표현 학습과 이동 윈도우 셀프 어텐션(shifted window self-attention)을 활용하여 국소적 세부 사항과 전역적 문맥을 모두 효율적으로 포착하도록 U-Net 프레임워크에 Swin Transformer 블록을 통합함.

모든 모델은 일관된 전처리(정규화, 패딩/크로핑) 및 데이터 증강 전략을 사용하여 훈련되었습니다. 성능은 Dice 점수, 정밀도(Precision), 재현율(Recall), F1 점수 및 훈련 손실 수렴도를 포함한 정량적 지표를 통해 평가되었습니다.

주요 기여
본 논문은 세 가지 주요 기여를 제공합니다:

1. 비교 평가: 볼륨형 뇌 MRI 및 2D 망막 영상 작업 모두에서 다섯 가지 U-Net 변형(컨볼루션, 잔차, 어텐션 기반 및 트랜스포머 기반)을 체계적으로 비교함.
2. 메커니즘 분석: 잔차 학습, 어텐션 게이트, 트랜스포머 셀프 어텐션과 같은 특정 아키텍처 메커니즘이 서로 다른 의료 영상 시나리오에서 분할 성능에 미치는 영향을 분석함.
3. 실무적 가이드: 복잡한 볼륨형 분할에 적합한 아키텍처와 미세 구조 추출에 효과적인 아키텍처를 구분함으로써 모델 선택에 대한 실무적 통찰을 제공함.

실험 결과

• BraTS 2023 (뇌종양): Swin UNETR이 평균 Dice 점수 0.8965(WT: 0.8966, ET: 0.9088, TC: 0.8841)로 가장 우수한 전체 성능을 달 기록했습니다. UNETR이 평균 Dice 0.8540으로 그 뒤를 이었습니다. Residual U-Net(0.7341), 3D U-Net(0.6890), Attention U-Net(0.6522)은 더 낮은 성능을 보였습니다. 결과는 트랜스포머 기반의 전역 문맥 모델링이 볼륨형 종양 분할에 매우 유익하다는 것을 나타냅니다.
• DRIVE (망막 혈관): Swin UNETR이 다시 한번 가장 높은 Dice 점수(0.8078)를 기록했으며, Residual U-Net(0.8069)과 UNETR(0.8015)이 근소한 차이로 뒤를 이었습니다. Attention U-Net은 0.7855를 기록했고, 3D U-Net은 가장 낮은 훈련 손실(0.1030)에도 불구하고 가장 낮은 Dice 점수(0.5967)를 기록했습니다. 이러한 불일치는 아키텍처가 데이터 구조와 맞지 않을 경우(예: 2D 미세 구조에 3D 모델 사용), 낮은 훈련 손실이 반드시 더 나은 분할 품질을 보장하지는 않는다는 점을 강조합니다.

의의 및 주장
본 논문은 Swin UNETR이 계층적 인코더와 이동 윈도우 어텐션을 통해 국소적 공간 세부 사항 보존과 장거리 문맥 모델링 사이의 균형을 성공적으로 맞춤으로써 가장 효과적인 아키텍처라고 주장합니다. 연구는 다음과 같이 제안합니다:

• 트랜스포머 기반 변형 모델(특히 Swin UNETR)은 뇌종양 경계 설정과 같이 전역적 문맥 이해가 필요한 복잡한 분할 작업에 우수합니다.
• **잔차 학습(Residual learning)**은 Residual U-Net이 트랜스포머 모델과 경쟁력 있는 성능을 보였던 것처럼, 망막 혈관 분할과 같이 미세한 구조적 세부 사항 보존이 필요한 작업에서 여전히 매우 효과적입니다.
• 모델 선택은 작업에 따라 달라져야 합니다. 볼륨형 및 복잡한 작업에는 트랜스포머 기반 접근 방식이 권장되는 반면, 미세 구조 추출에는 잔차 또는 하이브리드 접근 방식이 충분할 수 있습니다.

저자들은 향-후 연구에서 다양한 임상 데이터셋에 대한 견고성과 일반화 능력을 더욱 향상시키기 위해 트랜스포머 기반의 문맥 모델링과 컨볼루션 또는 잔차 모듈을 결합한 하이브리드 아키텍처를 탐구해야 한다고 결론지었습니다.