이 논문은 의료 영상 분할(Medical Image Segmentation)을 위한 새로운 아키텍처인 TransUNet을 제안합니다. 의료 영상 분할은 질병 진단 및 치료 계획에 필수적인 전제 조건으로 여겨집니다.
기존 방법의 한계:
- CNN(Convolutional Neural Network) 기반 방법: FCN(Fully Convolutional Networks), 특히 U-Net과 같은 U자형 아키텍처는 의료 영상 분할 분야에서 지배적이었고 상당한 성공을 거두었습니다. 자기공명(MR) 영상의 심장 분할, 컴퓨터 단층 촬영(CT) 영상의 장기 분할, 대장 내시경 영상의 용종 분할 등 다양한 의료 애플리케이션에서 성공적으로 사용되었습니다.
- 하지만 CNN 기반 접근 방식은 컨볼루션 연산의 내재적 지역성(intrinsic locality) 때문에 장거리 의존성(long-range dependency)을 명시적으로 모델링하는 데 한계를 보입니다. 이로 인해 질감, 모양 및 크기 면에서 환자 간 변동성이 큰 대상 구조에 대해 약한 성능을 보일 수 있습니다. 이러한 한계를 극복하기 위해 CNN 특징 기반 자기 어텐션 메커니즘을 구축하는 연구가 있었습니다.
- 순수 Transformer 기반 방법: 순차 예측을 위해 설계된 Transformer는 내재적인 전역 자기 어텐션(innate global self-attention) 메커니즘을 가진 대안 아키텍처로 부상했습니다. Transformer는 CNN과 달리 컨볼루션 연산 없이 오직 어텐션 메커니즘에만 의존하며, 전역적 컨텍스트(global contexts) 모델링에 강력할 뿐만 아니라 대규모 사전 학습 하에서 다운스트림 작업에 대한 우수한 전이 가능성(transferability)을 보여주었습니다. Vision Transformer (ViT)는 전역 자기 어텐션을 전체 이미지에 직접 적용하여 ImageNet 분류에서 최첨단 성능을 달성했습니다.
- 하지만 순수 Transformer를 의료 영상 분할에 그대로 적용하는 것은 만족스러운 결과를 얻지 못함을 발견했습니다. 이는 Transformer가 입력을 1D 시퀀스로 처리하고 모든 단계에서 전역 컨텍스트 모델링에만 집중하기 때문에, 상세한 위치 정보(detailed localization information)가 부족한 저해상도 특징을 생성하기 때문입니다. 이러한 정보는 직접적인 업샘플링으로는 효과적으로 복구될 수 없어 거친 분할 결과를 초래합니다.
TransUNet 제안:
- 논문은 Transformer와 U-Net의 장점을 결합한 TransUNet을 제안합니다. TransUNet은 의료 영상 분할을 위한 강력한 대안 아키텍처로 제시됩니다.
- TransUNet은 순차 예측 관점에서 자기 어텐션 메커니즘을 구축하는 최초의 의료 영상 분할 프레임워크입니다. 또한 최초의 Transformer 기반 의료 영상 분할 프레임워크이며, 성공적인 ViT를 기반으로 합니다.
TransUNet 아키텍처:
TransUNet은 하이브리드 CNN-Transformer 인코더와 캐스케이드 업샘플러(Cascaded Upsampler)로 구성된 U자형 아키텍처입니다. (Figure 1(b) 참조)
- 하이브리드 CNN-Transformer 인코더 (Hybrid CNN-Transformer Encoder):
- 순수 Transformer를 인코더로 사용하는 대신, TransUNet은 CNN을 먼저 특징 추출기로 사용하여 입력 영상에서 특징 맵을 생성합니다.
- Transformer는 원본 이미지 패치가 아닌 CNN 특징 맵에서 추출된 1x1 패치에 대해 패치 임베딩을 적용하여 입력 시퀀스를 만듭니다.
- 이 설계는 다음과 같은 이유로 선택되었습니다: 1) 디코딩 경로에서 중간 고해상도 CNN 특징 맵을 활용할 수 있습니다. 2) 순수 Transformer 인코더보다 하이브리드 CNN-Transformer 인코더가 더 나은 성능을 보임을 발견했습니다.
- Transformer 인코더는 여러 개의 Multihead Self-Attention (MSA) 및 Multi-Layer Perceptron (MLP) 블록으로 구성됩니다. 입력 시퀀스는 패치 임베딩과 위치 임베딩(position embedding)을 더하여 생성됩니다. Layer Normalization이 적용됩니다. (Figure 1(a) 참조)
- 저자들은 ImageNet에서 사전 학습된 ResNet-50과 ViT를 결합한 "R50-ViT"를 하이브리드 인코더로 사용했습니다.
- 캐스케이드 업샘플러 (Cascaded Upsampler - CUP):
- Transformer에 의해 인코딩된 특징 시퀀스 zL ∈ RHW/P² × D를 최종 분할 마스크를 출력하기 위해 디코딩하는 역할을 합니다.
- 인코딩된 특징을 H/P × W/P 크기로 재구성한 후, CUP는 여러 단계의 업샘플링 블록을 통해 H/P × W/P 해상도에서 전체 해상도 H × W로 도달합니다. 각 블록은 2배 업샘플링 연산, 3x3 컨볼루션 레이어, ReLU 레이어로 구성됩니다. 기본 패치 크기 16x16의 경우, 전체 해상도에 도달하기 위해 네 개의 2배 업샘플링 블록이 연속적으로 사용됩니다.
- 순수 Transformer 인코더와 단순 업샘플링만 사용한 경우(ViT-None)는 고해상도 정보를 잃어버려 저해상도 특징이 상세 정보를 복구하지 못하는 문제가 있습니다. CUP는 이러한 정보 손실을 보상하기 위해 고안되었습니다.
- U자형 아키텍처와 스킵 커넥션 (U-shaped Architecture and Skip-connections):
- CUP와 하이브리드 인코더는 U자형 아키텍처를 형성하며, 스킵 커넥션(skip-connections)을 통해 서로 다른 해상도 수준에서 특징 결합(feature aggregation)을 가능하게 합니다.
- Transformer에 의해 인코딩된 자기 어텐션 특징은 업샘플링되어 인코딩 경로에서 건너뛴(skipped) 고해상도 CNN 특징 맵과 결합되어 정밀한 위치 정보(precise localization)를 가능하게 합니다.
- 이러한 디자인을 통해 TransUNet은 Transformer의 장점을 유지하면서도 의료 영상 분할에 중요한 미세한 공간 디테일(finer spatial details)을 향상시킵니다.
실험 및 결과:
- 데이터셋: Synapse 다중 장기 CT 데이터셋과 ACDC 심장 MR 데이터셋을 사용했습니다.
- 평가 지표: 평균 DSC (Dice Score) 및 평균 Hausdorff Distance (HD)를 사용했습니다.
- 주요 비교 결과 (Table 1, Table 5 참조):
- TransUNet은 Synapse CT 데이터셋에서 V-Net, DARR, R50 U-Net, R50 AttnUNet 등 기존 최첨단 방법보다 우수한 성능을 달성했습니다. 평균 DSC에서 기존 최고 성능인 R50-AttnUNet보다 1.91% 향상되었습니다.
- 순수 Transformer (ViT)에 단순 업샘플링을 사용한 경우(ViT-None)는 만족스럽지 못했고, CUP 디코더를 사용한 경우(ViT-CUP) 성능이 향상되었습니다. 이는 CUP 디자인이 단순 업샘플링보다 나은 디코딩 전략임을 시사합니다.
- ViT 인코더 대신 하이브리드 R50-ViT 인코더를 사용한 경우(R50-ViT-CUP) 성능이 추가로 향상되었습니다. 이는 하이브리드 인코더의 효과를 보여줍니다.
- 하이브리드 R50-ViT 인코더와 CUP 디코더를 사용한 R50-ViT-CUP는 V-Net 및 DARR보다 뛰어나지만, 순수 CNN 기반 R50-U-Net 및 R50-AttnUNet보다는 여전히 낮은 성능을 보였습니다.
- TransUNet은 R50-ViT-CUP에 U-Net 구조의 스킵 커넥션을 결합함으로써 R50-ViT-CUP 및 기존 최고 R50-AttnUNet보다 뛰어난 결과를 달성하여 최신 기술을 경신했습니다. 이는 TransUNet이 의료 영상 분할에 중요한 고수준 의미 특징(high-level semantic features)과 저수준 디테일(low-level details)을 모두 학습하는 강력한 능력을 보여줍니다.
- 분석 연구 (Ablation Studies):
- 스킵 커넥션 수: U-Net과 유사한 스킵 커넥션을 통합하면 저수준 공간 정보를 복구하여 미세한 분할 디테일을 향상시키는 데 도움이 됩니다. 스킵 커넥션 수를 늘릴수록 일반적으로 분할 성능이 향상됨을 확인했습니다. CUP의 세 가지 중간 업샘플링 단계(1/2, 1/4, 1/8 해상도) 모두에 스킵 커넥션을 삽입했을 때 최고의 평균 DSC와 HD를 달성했습니다. 특히 소형 장기(대동맥, 담낭, 신장, 췌장)의 성능 향상이 대형 장기(간, 비장, 위)보다 더 뚜렷했습니다. 스킵 커넥션에 가산(additive) Transformer를 적용하면 성능이 더욱 향상될 수 있음을 보였습니다. (Figure 2 참조)
- 입력 해상도: 기본 입력 해상도 224x224 외에 512x512에서도 실험했습니다. 512x512 입력은 Transformer의 유효 시퀀스 길이를 약 5배 증가시킵니다. 해상도를 224x224에서 512x512로 변경하면 평균 DSC가 6.88% 향상되었지만, 계산 비용이 훨씬 커집니다. (Table 2 참조)
- 패치 크기/시퀀스 길이: 패치 크기가 작을수록(즉, 유효 시퀀스 길이가 길수록) 일반적으로 더 높은 분할 성능이 얻어졌습니다. 패치 크기 16x16은 시퀀스 길이 196에 해당하고, 패치 크기 32x32는 시퀀스 길이 49에 해당합니다. Transformer가 더 긴 입력 시퀀스에 대해 요소 간에 더 복잡한 종속성을 인코딩하기 때문입니다. (Table 3 참조)
- 모델 스케일: "Base"와 "Large" 두 가지 모델 크기로 실험했습니다. "Large" 모델이 더 나은 성능을 보였습니다. (Table 4 참조)
- 시각화 결과 (Figure 3 참조): Synapse 데이터셋에 대한 정성적 비교 결과를 통해 다음을 확인했습니다: 1) 순수 CNN 기반 방법(U-Net, AttnUNet)은 장기를 과분할하거나 과소분할하는 경향이 더 컸습니다. 2) TransUNet은 다른 방법에 비해 거짓 양성(false positives) 예측이 적었습니다. 3) Transformer 기반 모델 중 R50-ViT-CUP의 예측은 경계와 모양이 TransUNet보다 거칠어 보였습니다. TransUNet은 고수준 전역 컨텍스트 정보와 저수준 디테일의 이점을 모두 누리기 때문에 더 미세한 분할과 상세 모양 정보 보존이 가능합니다. 이는 U-Net과 유사한 스킵 커넥션을 Transformer 설계에 통합하여 정밀한 위치 정보 학습을 가능하게 한다는 초기 아이디어를 다시 한번 입증합니다.
- 다른 데이터셋으로의 일반화: ACDC MR 데이터셋에서 자동 심장 분할 실험을 수행했습니다. TransUNet은 순수 CNN 기반 방법 및 다른 Transformer 기반 기준선보다 일관된 성능 향상을 보였으며, 이는 Synapse CT 데이터셋에서의 결과와 유사했습니다. 이는 TransUNet의 일반화 능력을 보여줍니다.
결론:
논문은 의료 영상 특징을 시퀀스로 처리하여 강력한 전역 컨텍스트를 인코딩할 뿐만 아니라, U자형 하이브리드 아키텍처 설계를 통해 저수준 CNN 특징을 잘 활용하는 TransUNet을 제안했습니다. TransUNet은 지배적인 FCN 기반 접근 방식에 대한 대안 프레임워크로서, CNN 기반 자기 어텐션 방법을 포함한 다양한 경쟁 방법보다 우수한 성능을 달성했습니다.