[논문 리뷰] D-Former: A U-shaped Dilated Transformer for 3D Medical Image Segmentation

Yixuan Wu, Kuanlun Liao외 5명|arXiv (Cornell University)|2022. 01. 03.

AI in cancer detection인용 수 9

한 줄 요약

D-Former는 3D 의료 영상 분할을 위한 U자형 확장 트랜스포머를 제안하며, 토큰 수를 늘리지 않고도 확장된 멀티헤드 자기주의를 사용하여 패치 간 장거리 상관관계를 포착함으로써 계산 비용을 감소시킨다. 동적 위치 인코딩 기법을 도입하여 공간 모델링을 향상시켜, 경쟁 기반의 CNN 및 트랜스포머 모델보다 낮은 FLOPs와 파라미터 수로 Synapse 및 ACDC 데이터셋에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Computer-aided medical image segmentation has been applied widely in diagnosis and treatment to obtain clinically useful information of shapes and volumes of target organs and tissues. In the past several years, convolutional neural network (CNN) based methods (e.g., U-Net) have dominated this area, but still suffered from inadequate long-range information capturing. Hence, recent work presented computer vision Transformer variants for medical image segmentation tasks and obtained promising performances. Such Transformers model long-range dependency by computing pair-wise patch relations. However, they incur prohibitive computational costs, especially on 3D medical images (e.g., CT and MRI). In this paper, we propose a new method called Dilated Transformer, which conducts self-attention for pair-wise patch relations captured alternately in local and global scopes. Inspired by dilated convolution kernels, we conduct the global self-attention in a dilated manner, enlarging receptive fields without increasing the patches involved and thus reducing computational costs. Based on this design of Dilated Transformer, we construct a U-shaped encoder-decoder hierarchical architecture called D-Former for 3D medical image segmentation. Experiments on the Synapse and ACDC datasets show that our D-Former model, trained from scratch, outperforms various competitive CNN-based or Transformer-based segmentation models at a low computational cost without time-consuming per-training process.

연구 동기 및 목표

• 3D 의료 영상 분할 트랜스포머에서 글로벌 자기주의의 높은 계산 비용을 해결하기 위해.
• 공간 해상도를 유지하면서 3D 의료 영상에서 장거리 특징 모델링을 향상시키고 파라미터 수를 줄이기 위해.
• 제한된 계산 자원이 있는 3D 의료 영상에 적합한 효율적이고 확장 가능한 트랜스포머 아키텍처를 설계하기 위해.
• 입력에 따라 변화하는 순서 기반 위치 관계를 반영할 수 있는 동적 위치 인코딩을 통합하여 공간 이해도를 향상시키기 위해.

제안 방법

• 확장된 방식으로 번갈아가며 국소적이고 글로벌 자기주의를 계산하는 확장된 자기주의 메커니즘을 제안하여, 패치 수를 늘리지 않고도 수용역을 확장한다.
• 국소 범위 모듈(LSM)과 글로벌 범위 모듈(GSM)을 번갈아가며 사용하는 U자형 인코더-디코더 아키텍처인 D-Former을 도입한다.
• 첫 번째 국소 자기주의 모듈 이전에 위치한 동적 위치 인코딩(DPE) 블록을 도입하여 유연하게 위치 인식 표현을 학습한다.
• 다중 스테이지의 계층적 설계를 적용하여, 각 스테이지에 번갈아가며 LSM과 GSM 블록을 배치하여 다중 척도의 문맥적 특징을 점진적으로 포착한다.
• 입력에 따라 변화하는 공간 관계에 적응하는 학습된 위치 임베딩을 적용하여 3D 공간 구조의 모델링을 향상시킨다.
• 사전 훈련 없이 모델을 처음부터 훈련시켜 낮은 계산 비용으로도 높은 성능을 달성한다.

실험 결과

연구 질문

• RQ1 확장된 자기주의 메커니즘이 3D 의료 영상 분할에서 계산 비용을 줄이면서도 장거리 상관관계 모델링을 유지할 수 있는가?
• RQ2 제안된 D-Former 아키텍처는 3D 의료 영상 작업에서 표준 트랜스포머 및 CNN과 비교해 분할 정확도와 FLOPs 측면에서 어떻게 성능을 내는가?
• RQ3 3D 공간 표현 향상을 위해 동적 위치 인코딩의 최적의 위치와 설계는 무엇인가?
• RQ4 글로벌 자기주의를 확장된 글로벌 범위 모듈로 대체함으로써 토큰 수를 늘리지 않으면서도 성능 향상을 이룰 수 있는가?
• RQ5 Synapse 및 ACDC와 같은 3D 의료 영상 데이터셋에서 D-Former 모델이 처음부터 훈련되었을 때 얼마나 효과적인가?

주요 결과

• D-Former는 Synapse 데이터셋에서 평균 DSC 88.83%를 기록하여 Swin-Unet 및 nnFormer를 포함한 최신 기술 수준의 모델들을 능가한다.
• 표준 트랜스포머 모델과 비교해 FLOPs와 파라미터 수를 줄였지만도 뛰어난 분할 정확도를 유지한다.
• 제안된 동적 위치 인코딩(DPE)은 절대 위치 인코딩 대비 평균 DSC를 4.05% 향상시키고, 상대 위치 인코딩 대비 2.48% 향상시켰다.
• DPE 블록을 첫 번째 국소 자기주의 모듈 이전에 배치할 경우 가장 높은 성능을 기록하여, 공간 모델링에서의 중요성을 확인했다.
• 확장된 자기주의를 사용하는 글로벌 범위 모듈(GSM)은 국소 범위 모듈만을 사용하는 경우보다 정확도를 0.66% 향상시켜 장거리 모델링에서의 효과를 입증했다.
• 제거 실험 결과, 확장된 자기주의 메커니즘과 DPE가 모두 낮은 계산 비용으로 최신 기술 수준의 성능을 달성하는 데 핵심 요소임을 확인했다.