[논문 리뷰] SegTransVAE: Hybrid CNN -- Transformer with Regularization for medical image segmentation

Quan-Dung Pham, Hai Nguyen-Truong외 2명|arXiv (Cornell University)|2022. 01. 21.

AI in cancer detection참고 문헌 12인용 수 5

한 줄 요약

SegTransVAE는 의료 영상 분할을 위한 하이브리드 3D CNN-Transformer 아키텍처를 제안하며, 변분 자동차오디터(VAE) 브랜치를 통해 국소적 특징 추출(3D CNN), 전역적 맥락 모델링(Transformer), 자기지도 학습을 통한 재구성(VAE)을 결합하여 일반화 성능을 향상시킨다. BraTS 및 KiTS19 데이터셋에서 3D U-Net, UNETR, SegresnetVAE를 능가하는 Dice Score와 95%-Hausdorff 거리 성능을 기록하면서도 경량 CNN과 유사한 추론 속도를 유지한다.

ABSTRACT

Current research on deep learning for medical image segmentation exposes their limitations in learning either global semantic information or local contextual information. To tackle these issues, a novel network named SegTransVAE is proposed in this paper. SegTransVAE is built upon encoder-decoder architecture, exploiting transformer with the variational autoencoder (VAE) branch to the network to reconstruct the input images jointly with segmentation. To the best of our knowledge, this is the first method combining the success of CNN, transformer, and VAE. Evaluation on various recently introduced datasets shows that SegTransVAE outperforms previous methods in Dice Score and $95\%$-Haudorff Distance while having comparable inference time to a simple CNN-based architecture network. The source code is available at: https://github.com/itruonghai/SegTransVAE.

연구 동기 및 목표

• 의료 영상 분할에서 CNN의 장거리 공간적 상관관계 파악 능력 부족과 Transformer의 인도크티브 바이어스 부족, 저자료 환경에서의 일반화 문제를 해결하기 위해.
• VAE 브랜치를 통한 분할 및 자기재구성 동시 학습을 통해 과적합을 줄이고 특징 학습을 향상시켜 분할 정확도를 향상시키기 위해.
• 3D CNN(국소 불변성), Transformer(전역 주의), VAE(정규화)의 장점을 통합한 유일한 인코더-디코더 프레임워크를 구성하기 위해.
• 어려운 의료 영상 기준점에서 최신 기술 수준(SOTA)의 성능을 달성하면서도 기존 모델과 경쟁 가능한 추론 속도를 확보하기 위해.

제안 방법

• 인코더는 인스턴스 정규화와 Leaky ReLU를 적용한 수정된 ResNet 블록을 스택하여 입력 볼륨으로부터 계층적 3D 국소 특징을 추출한다.
• 3D 특징 맵은 선형 투영을 통해 학습 가능한 가중치를 가진 d차원 임베딩 공간으로 변환되며, 공간적 구조를 유지하기 위해 학습 가능한 위치 임베딩이 추가된다.
• 멀티헤드 어텐션과 피드포워드 네트워크 레이어의 스택으로 구성된 트랜스포머 인코더는 모든 패치 간 자기어텐션을 통해 전역적 특징 모델링을 가능하게 한다.
• 특징 매핑 모듈은 트랜스포머 출력을 다시 3D 특징 맵으로 재구성하고, 컨볼루션 블록을 통해 채널 차원을 감소시켜 디코더 입력으로 제공한다.
• 최종 트랜스포머 레이어에 VAE 브랜치를 부착하여 입력 영상을 재구성함으로써 정규화를 제공하고, 제한된 데이터에서의 일반화 성능을 향상시킨다.
• 디코더는 인코더의 스킵 커넥션을 활용한 3D 트랜스포지션 컨볼루션을 사용하여 고해상도 분할 마스크를 재구성한다.

실험 결과

연구 질문

• RQ1 VAE 기반 자기지도 학습 재구성 기능을 갖춘 하이브리드 CNN-Transformer 아키텍처가 의료 영상 분할 작업에서 성능 향상에 기여하는가?
• RQ2 VAE 정규화 통합이 과적합을 줄이고, 특히 소규모 의료 데이터셋에서 일반화 성능을 향상시키는가?
• RQ3 국소적(CNN)과 전역적(Transformer) 특징 학습의 조합이 순수 CNN 또는 순수 Transformer 모델보다 분할 정확도에서 뛰어나지 않는가?
• RQ4 기존 최신 기술 수준의 모델과 비교해 본다면, 제안된 아키텍처는 추론 속도와 파rameter 효율성 측면에서 어떤가?

주요 결과

• BraTS 2021 데이터셋에서 SegTransVAE는 강화된 종양(ET)의 Dice Score가 85.48%, 전체 종양(WT)이 90.52%, 종양 코어(TC)가 92.60%를 기록하여 3D U-Net, UNETR, SegresnetVAE를 모두 능가했다.
• SegTransVAE는 ET의 95%-Hausdorff 거리가 2.89 mm, WT가 3.57 mm, TC가 5.84 mm로 기준 모델 대비 뚜렷한 향상을 보였다.
• KiTS19 데이터셋에서 SegTransVAE는 신장 분할의 Dice Score가 95.28%, 종양 분할이 66.31%를 기록하여 3D U-Net, UNETR, SegresnetVAE 모두를 능가했다.
• 신장 분할의 95%-Hausdorff 거리는 3.28 mm로, 3D U-Net(6.32 mm) 및 UNETR(8.84 mm)보다 유의미하게 낮았다.
• SegTransVAE는 4470만 개의 파라미터를 가지며 추론 시간이 0.45초로, 3D U-Net(560만 파라미터, 0.45초)과 유사하고 SegresnetVAE(750만 파라미터, 0.55초)보다 빠르게 작동했다.
• 3D U-Net보다 더 많은 파라미터를 가졌음에도 불구하고, SegTransVAE는 3D U-Net(1000+ GFLOPs)보다 낮은 GFLOPs(607.5)를 기록하여 더 높은 계산 효율성을 보였다.