들어가며: 의료 이미지 분할과 기존의 한계
의료 이미지 분할(Medical image segmentation)은 질병 진단 및 치료 계획을 수립하는 최첨단 의료 시스템에서 매우 중요한 역할을 합니다. 그동안 이 분야에서는 대칭적인 인코더-디코더 구조와 스킵 연결(skip-connection)을 특징으로 하는 U-Net 구조가 엄청난 성공을 거두며 표준으로 자리 잡았습니다.
하지만 기존 CNN(합성곱 신경망) 기반의 U-Net은 치명적인 한계가 있었습니다. 바로 합성곱 연산 고유의 '국소성(locality)' 때문에 이미지 내의 멀리 떨어진 영역 간의 관계(long-range dependencies)를 효과적으로 모델링하기 어렵다는 점입니다. 환자마다 장기의 모양, 크기, 텍스처 변동성이 큰 의료 데이터에서는 이러한 단점이 더욱 두드러집니다.
3D TransUNet의 등장: CNN과 Transformer의 융합
이러한 한계를 극복하기 위해 연구진은 글로벌 문맥(global context)을 파악하는 데 탁월한 Transformer의 Self-attention 메커니즘에 주목했습니다.
이번에 소개할 3D TransUNet은 기존 2D TransUNet 구조를 3D 네트워크로 확장한 모델입니다. 성능이 입증된 nnU-Net 아키텍처를 기반으로 하며, 인코더와 디코더 설계 모두에서 Transformer의 잠재력을 최대한 끌어냈습니다. 핵심은 CNN이 가진 고해상도 공간 디테일 처리 능력과 Transformer의 글로벌 문맥 파악 능력을 하이브리드 형태로 결합한 것입니다.
핵심 혁신 포인트 두 가지 💡
1. 글로벌 문맥을 이해하는 Transformer 인코더 (Transformer Encoder)
- CNN 피처 맵(Feature map)에서 추출한 이미지 패치들을 토큰화(tokenize)하여 글로벌 컨텍스트를 추출합니다.
- 복잡한 복부 장기들 사이의 상호 관계를 파악하는 것이 필수적인 다중 장기 분할(Multi-organ segmentation) 작업에서 특히 뛰어난 성능을 발휘합니다.
2. 후보 영역을 정교하게 다듬는 Transformer 디코더 (Transformer Decoder)
- 기존의 전통적인 픽셀 단위 분할을 '마스크 분류(mask classification)' 문제로 재정의하여, 예측 후보들을 학습 가능한 쿼리(learnable queries)로 설정합니다.
- 이 쿼리들은 후보 제안(candidate proposals)과 U-Net 피처 간의 교차 주의(cross-attention)를 통해 지속적으로 정제됩니다.
- 특히 종양 분할(tumor segmentation)과 같이 작고 까다로운 타겟을 처리할 때 훨씬 더 유용함이 입증되었습니다[cite: 18, 58].
- 예측의 전경(foreground)에 집중하여 이전 단계의 노이즈를 줄이고 점진적으로 분할 정확도를 높이는 Coarse-to-fine attention refinement 기법을 도입했습니다.
압도적인 성능과 유연성 🏆
흥미로운 점은 수행하는 의료 태스크에 따라 최적의 아키텍처 디자인이 다르다는 것입니다.
연구진의 실험에 따르면, 다중 장기 분할에서는 인코더 중심의 설계가 유리하고, 미세한 병변이나 종양을 찾을 때는 디코더 중심의 설계가 더 적합합니다.
3D TransUNet은 다중 장기 분할, 췌장 종양 분할, 간 혈관 분할 등 다양한 의료 애플리케이션에서 경쟁 모델들을 압도하는 성능을 보여주었습니다. 특히 뇌종양 분할 대회인 BraTS2021 챌린지에서는 기존 최고 솔루션(nnU-Net-Large)의 성능을 뛰어넘는 성과를 기록했습니다.
마치며
3D TransUNet은 단순히 최신 기술을 이어 붙인 것을 넘어, 태스크의 특성(큰 장기 vs 작은 종양)에 맞게 Transformer를 전략적으로 배치하는 통찰력을 보여준 연구입니다. 의료 AI 분야에서 아키텍처 설계의 새로운 기준을 제시한 이 모델이 앞으로 실제 임상 환경에서 어떻게 활약할지 기대가 됩니다.
