[논문 리뷰] UNesT: Local Spatial Representation Learning with Hierarchical Transformer for Efficient Medical Segmentation

Xin Yu, Qi Yang외 13명|arXiv (Cornell University)|2022. 09. 28.

Advanced Neural Network Applications인용 수 4

한 줄 요약

UNesT는 블록 단위 특징 집약을 통해 국소적 공간 표현을 향상시키는 계층적 3D 트랜스포머 기반 의료 영상 분할 모델을 제안한다. 이는 전체 뇌(133개 클래스) 및 신장 하위구조 분할 과제에서 최신 기술 수준의 성능을 달성하며, 27개 네트워크 앙상블보다도 뛰어난 성능을 보인다. Colin 데이터셋에서 평균 DSC 0.8505, CANDI에서 0.7906을 기록했으며, 소규모 데이터셋에서도 높은 데이터 효율성과 강인성을 보였다.

ABSTRACT

Transformer-based models, capable of learning better global dependencies, have recently demonstrated exceptional representation learning capabilities in computer vision and medical image analysis. Transformer reformats the image into separate patches and realizes global communication via the self-attention mechanism. However, positional information between patches is hard to preserve in such 1D sequences, and loss of it can lead to sub-optimal performance when dealing with large amounts of heterogeneous tissues of various sizes in 3D medical image segmentation. Additionally, current methods are not robust and efficient for heavy-duty medical segmentation tasks such as predicting a large number of tissue classes or modeling globally inter-connected tissue structures. To address such challenges and inspired by the nested hierarchical structures in vision transformer, we proposed a novel 3D medical image segmentation method (UNesT), employing a simplified and faster-converging transformer encoder design that achieves local communication among spatially adjacent patch sequences by aggregating them hierarchically. We extensively validate our method on multiple challenging datasets, consisting of multiple modalities, anatomies, and a wide range of tissue classes, including 133 structures in the brain, 14 organs in the abdomen, 4 hierarchical components in the kidneys, inter-connected kidney tumors and brain tumors. We show that UNesT consistently achieves state-of-the-art performance and evaluate its generalizability and data efficiency. Particularly, the model achieves whole brain segmentation task complete ROI with 133 tissue classes in a single network, outperforming prior state-of-the-art method SLANT27 ensembled with 27 networks.

연구 동기 및 목표

• 시각 트랜스포머가 국소적 공간 정보를 유지하는 데서 약점이 있으며, 3D 의료 영상에서 대규모이고 이질적인 조직 구조를 처리하는 데 어려움을 겪는 문제를 해결하기 위해.
• 의료 영상에서 흔한 저데이터 환경에서 데이터 효율성과 모델 강인성을 향상시키기 위해.
• 뇌 및 신장 하위구조와 같은 수백 개의 조직 클래스와 계층적으로 연결된 구조를 정확하게 분할할 수 있도록 하기 위해.
• 앙상블 모델을 능가하는 성능과 효율성을 갖춘 단일 모델 솔루션을 개발하기 위해.
• 검사 및 일반화를 위한 검증을 위해 방사선의사가 주관한 내부 제작된 신장 하위구조 CT 데이터셋을 구축하기 위해.

제안 방법

• UNesT는 공간적으로 인접한 패치를 블록으로 묶는 계층적 3D 블록 집약 모듈을 사용하여 국소적 소통을 가능하게 하면서도 계층적 수준에서 전역 자기주의를 유지한다.
• 모델은 U-Net 유사 아키텍처를 사용하며, 중첩된 트랜스포머 인코더와 컨볼루션 디코더가 다중 해상도에서 스킵 연결을 통해 연결되어 있다.
• 국소적 구조에 대한 인덕티브 바이어스를 향상시키면서도 전역적 맥락을 유지하기 위해 트랜스포머 인코더를 계층적으로 스택한다.
• 계산 비용을 줄이면서도 고해상도 3D 볼륨에서 성능을 유지하기 위해 단순화되고 빠르게 수렴하는 트랜스포머 설계를 통합한다.
• 인코더와 디코더 간의 스킵 연결을 통해 다중 해상도 특징 융합이 가능해져, 작은 또는 복잡한 구조의 국소화 정확도가 향상된다.
• 다양한 모odalities(예: T1-가중 MRI 및 CT) 데이터셋에서 엔드 투 엔드로 훈련되며, 광범위한 데이터 증강 및 정규화 기법이 적용된다.

실험 결과

연구 질문

• RQ1 계층적 트랜스포머 아키텍처는 3D 의료 영상에서 전역 맥락을 유지하면서도 국소적 공간 표현을 효과적으로 학습할 수 있는가?
• RQ2 기존의 모델, 특히 앙상블 방법을 포함하여, UNesT는 133개 조직 클래스를 가진 전체 뇌 분할 과제에서 최고 성능을 달성할 수 있는가?
• RQ3 희귀하거나 복잡한 하위구조(예: 신장 구성요소)에 대해 저데이터 환경에서 UNesT의 성능은 어떠한가?
• RQ4 모델은 다양한 해부학적 영역, 모달리티, 조직 복잡도 수준 간에 얼마나 잘 일반화되는가?
• RQ5 제안된 방법은 광범위한 수동 애너테이션 없이도 방사선의사 수준의 재현 가능성을 신장 하위구조 분할에서 달성할 수 있는가?

주요 결과

• UNesT는 Colin 데이터셋을 사용한 전체 뇌 분할 과제에서 평균 DSC 0.8505를 기록했으며, 이는 이전 최고 기록인 0.7444보다 뚜렷한 향상이다.
• CANDI 데이터셋에서 UNesT는 평균 DSC를 0.6968에서 0.7025로 향상시켜 도전적인 공개 벤치마크에서 성능 향상을 입증했다.
• 27개 네트워크 앙상블인 SLANT27보다도 단일 모델로 더 높은 정확도를 달성하여 추론 비용과 복잡도를 감소시켰다.
• 신장 하위구조 분할 과제에서 UNesT는 DSC 0.7906을 기록하여 복잡하고 계층적으로 연결된 신장 구성요소에서 강력한 성능을 보였다.
• 모델는 방사선의사와 비교해 체적 측정 평균 차이가 0.03에 불과할 정도로 높은 재현 가능성을 보였으며, 이는 상호 평가자 간 일致성(0.29)을 초월했다.
• 이상치 분석에서 90%의 세그멘테이션이 신뢰할 수 있었고, 10%는 과소 또는 과대 세그멘테이션을 보였으며, 이는 시각적 검토로 감지 가능했고 후처리를 통해 보정 가능할 수 있었다.