[논문 리뷰] AFTer-UNet: Axial Fusion Transformer UNet for Medical Image Segmentation

Xiangyi Yan, Hao Tang외 4명|arXiv (Cornell University)|2021. 10. 20.

Advanced Neural Network Applications인용 수 10

한 줄 요약

AFTer-UNet는 축 방향 융합 트랜스포머를 U-Net 아키텍처에 통합하여 3D 의료 영상 분할을 위한 엔드 투 엔드 프레임워크를 제안한다. 이는 동일한 영역 내 및 영역 간 장거리 의존성을 공동으로 모델링한다. 축 방향과 공간 방향에 대해 별도로 자기주의를 계산한 후 융합함으로써 계산 비용과 GPU 메모리 사용량을 줄이며, 4.34%의 평균 Dice 점수 향상(비교 기준 U-Net 대비)과 2.56%의 향상(비교 기준 CoTr 대비)을 달성하여 최신 기술 수준의 성능을 확보한다.

ABSTRACT

Recent advances in transformer-based models have drawn attention to exploring these techniques in medical image segmentation, especially in conjunction with the U-Net model (or its variants), which has shown great success in medical image segmentation, under both 2D and 3D settings. Current 2D based methods either directly replace convolutional layers with pure transformers or consider a transformer as an additional intermediate encoder between the encoder and decoder of U-Net. However, these approaches only consider the attention encoding within one single slice and do not utilize the axial-axis information naturally provided by a 3D volume. In the 3D setting, convolution on volumetric data and transformers both consume large GPU memory. One has to either downsample the image or use cropped local patches to reduce GPU memory usage, which limits its performance. In this paper, we propose Axial Fusion Transformer UNet (AFTer-UNet), which takes both advantages of convolutional layers' capability of extracting detailed features and transformers' strength on long sequence modeling. It considers both intra-slice and inter-slice long-range cues to guide the segmentation. Meanwhile, it has fewer parameters and takes less GPU memory to train than the previous transformer-based models. Extensive experiments on three multi-organ segmentation datasets demonstrate that our method outperforms current state-of-the-art methods.

연구 동기 및 목표

• 기존의 2D 및 3D 트랜스포머 기반 모델이 3D 볼륨에서 축(영역 간) 맥락 정보를 효과적으로 활용하지 못하는 한계를 해결하기 위해.
• 체적 데이터에서 3D 자기주의의 높은 GPU 메모리 소비를 줄이면서도 장거리 모델링 능력을 유지하기 위해.
• 특히 길쭉하거나 해부학적으로 변동성이 큰 장기의 분할 정확도를 향상시키기 위해 다수의 영역 간 맥락 신호를 융합하기 위해.
• 이전의 트랜스포머 기반 모델보다 파라미터 수가 적고 메모리 사용량이 낮은 효율적인 모델을 유지하기 위해.
• 임상 적용을 고려해 일반적인 GPU와 호환되는 확장 가능한 엔드 투 엔드 프레임워크를 개발하기 위해.

제안 방법

• 축 방향과 공간 방향에 대해 별도로 자기주의를 계산한 후 출력을 융합하는 축 융합 메커니즘을 도입한다.
• U-Net의 인코더와 디코더 사이에 축 융합 트랜스포머를 통합하여 U-Net의 U자형 스킵 연결을 유지한다.
• 특징 추출을 위해 2D 컨볼루션 네트워크 인코더와 디코더를 사용하고, 축 융합 트랜스포머를 통해 영역 간 장거리 의존성을 캡처한다.
• 축 방향과 각 2D 영역 내에서 다중 헤드 자기주의를 적용하고 상대적 위치 인코딩을 사용하여 장거리 맥락을 모델링한다.
• 영역 간 주의의 조밀함을 제어하기 위해 샘플링 주파수 $N_f$를 도입하여 가까운 인접 영역에 더 중점을 두어 효율성과 성능을 향상시킨다.
• 축 주의와 공간 주의를 분리함으로써 계산 복잡도를 $O(HW D^2)$에서 $O(HW D + H^2W^2)$로 감소시켜 단일 11GB GPU에서의 학습을 가능하게 한다.

실험 결과

연구 질문

• RQ1 영역 내 및 영역 간 주의의 축 융합이 순수한 2D 또는 3D 트랜스포머 모델을 초월하여 3D 의료 영상 분할 성능을 향상시킬 수 있는가?
• RQ2 축 주의와 공간 주의를 분리함으로써 GPU 메모리 사용량을 줄일 수 있고, 장거리 모델링 능력은 유지되는가?
• RQ3 영역 간 축 주의에 포함된 인접 영역 수가 길쭉한 장기의 분할 성능에 미치는 영향은 어떠한가?
• RQ4 축 융합 트랜스포머 레이어의 수를 늘릴수록 분할 정확도는 얼마나 향상되는가?
• RQ5 축 방향의 샘플링 주파수가 모델 성능과 효율성에 어떤 영향을 미치는가?

주요 결과

• Thorax-85 데이터셋에서 AFTer-UNet는 U-Net 대비 2.13%의 평균 Dice 점수 향상과 CoTr 대비 0.69%의 향상을 기록한다.
• 도전적인 길쭉한 장기인食도에서 AFTer-UNet는 U-Net 대비 2.91% 향상되고 CoTr 대비 1.45% 향상된다.
• 다중 장기 데이터셋에서 AFTer-UNet는 간에 대해 U-Net 대비 3.17% 향상, 왼쪽 신장에 대해 7.30%, 오른쪽 신장에 대해 7.11%, 췌장에 대해 6.64% 향상된다.
• CoTr 대비 AFTer-UNet는 간에 대해 3.85% 높은 Dice 점수, 왼쪽 신장에 대해 5.26%, 오른쪽 신장에 대해 6.61%, 췌장에 대해 5.68% 향상된다.
• AFTer-UNet는 GPU 메모리 사용량을 크게 줄여 기존의 3D 트랜스포머 모델이 다중 GPU가 필요로 했던 것과는 달리 단일 RTX-2080Ti(11GB)에서 학습이 가능하다.
• 41.5M개의 파라미터를 유지하면서도 TransUNet(43.5M) 및 CoTr(41.9M)와 유사한 모델 크기를 확보하여 파라미터 증가 없이도 효율성을 입증한다.