[논문 리뷰] SegStitch: Multidimensional Transformer for Robust and Efficient Medical Imaging Segmentation

Shengbo Tan, Zeyu Zhang외 6명|arXiv (Cornell University)|2024. 08. 01.

Brain Tumor Detection and ClassificationNeuroscience인용 수 3

한 줄 요약

SegStitch는 축 방향 패치 기반 자기주의와 노이즈 제거 일반 미분방정식(OED) 블록을 조합하여 장거리 특징 모델링과 강건성을 향상시키는 새로운 3D 의료 영상 분할 프레임워크를 제안한다. UNETR 대비 파라미터 수를 36.7% 감소시키고 FLOPs를 10.7% 감소시키며 BTCV에서 mDSC가 11.48% 향상되고 ACDC에서 6.71% 향상되어 최신 기술 수준을 확립한다.

ABSTRACT

Medical imaging segmentation plays a significant role in the automatic recognition and analysis of lesions. State-of-the-art methods, particularly those utilizing transformers, have been prominently adopted in 3D semantic segmentation due to their superior performance in scalability and generalizability. However, plain vision transformers encounter challenges due to their neglect of local features and their high computational complexity. To address these challenges, we introduce three key contributions: Firstly, we proposed SegStitch, an innovative architecture that integrates transformers with denoising ODE blocks. Instead of taking whole 3D volumes as inputs, we adapt axial patches and customize patch-wise queries to ensure semantic consistency. Additionally, we conducted extensive experiments on the BTCV and ACDC datasets, achieving improvements up to 11.48% and 6.71% respectively in mDSC, compared to state-of-the-art methods. Lastly, our proposed method demonstrates outstanding efficiency, reducing the number of parameters by 36.7% and the number of FLOPS by 10.7% compared to UNETR. This advancement holds promising potential for adapting our method to real-world clinical practice. The code will be available at https://github.com/goblin327/SegStitch

연구 동기 및 목표

• 표준 비전 트랜스포머가 3D 의료 영상 분할에서 높은 계산 비용과 열악한 국소 특징 모델링 성능을 보이는 데 기인한 한계를 해결하기 위해.
• 하이브리드 주의 메커니즘을 통해 3D 장기 및 병변 분할에서 장거리 종속성 학습과 의미 일관성을 향상시키기 위해.
• 지속적인 특징 보정을 위한 신경 미분방정식(ODE) 블록 통합을 통해 모델의 강건성과 일반화 능력을 향상시키기 위해.
• 분할 정확도를 훼손하지 않고 모델 복잡성과 추론 비용을 감소시켜 임상 적용을 가능하게 하기 위해.

제안 방법

• SegStitch는 전체 볼륨 주의를 피하기 위해 3D 볼륨 전반에 걸쳐 공간적 및 의미 일관성을 유지하기 위해 축 방향 패치 기반 쿼리를 사용한다.
• 거시적 및 미시적 자기주의를 조합한 이중 경로 주의 메커니즘을 도입하여 전반적인 맥락과 국소 세부 정보를 더 잘 포착한다.
• 추론 중 지속적이고 동적인 특징맵 갱신을 가능하게 하기 위해 신경 기억 ODE(nmODE) 기반의 노이즈 제거 ODE 블록을 통합한다.
• 특징을 개선하기 위해 학습된 벡터장(예: sin², sin, tanh)을 사용하는 ODE 블록은 입력의 흐름 변화나 데이터 노이즈에 대한 강건성을 향상시킨다.
• 패치 간의 의미 일관성을 향상시켜 분할 일관성을 개선하기 위해 패치 전역에 걸쳐 공유 쿼리 메커니즘을 도입한다.
• 교차 엔트로피 손실과 Dice 손실을 사용하여 3D 의료 영상 데이터셋에서 엔드 투 엔드로 학습하며, ODE 변형에 대한 분석 실험을 수행한다.

실험 결과

연구 질문

• RQ1 공유 쿼리 기반의 축 방향 패치 주의는 3D 의료 영상 분할에서 의미 일관성을 향상시키고 계산 비용을 감소시킬 수 있는가?
• RQ2 노이즈가 많거나 왜곡된 의료 스캔이 존재할 경우 ODE 블록 통합이 모델의 강건성과 특징 보정 능력을 어떻게 향상시키는가?
• RQ3 제안된 SegStitch 프레임워크는 UNETR와 같은 최신 기술 수준의 트랜스포머보다 분할 정확도와 효율성에서 얼마나 뛰어나게 성능을 발휘하는가?
• RQ4 sin², 시그모이드, tanh 등 다양한 ODE 형태 중에서 어느 것이 학습 안정성, 일반화 능력, 분할 성능 간의 최적 균형을 이룬다?

주요 결과

• SegStitch는 BTCV 데이터셋에서 최신 기술 수준의 방법 대비 평균 덱스 유사도 계수(mDSC)를 11.48% 향상시켰다.
• ACDC 데이터셋에서 SegStitch는 기존 최신 기술 접근법 대비 mDSC를 6.71% 향상시켜 도전적인 심장 분할 작업에서 뛰어난 성능을 입증했다.
• UNETR 대비 모델 파라미터 수를 36.7% 감소시키고 FLOPs를 10.7% 감소시켜 뚜렷한 효율성 향상을 보였다.
• ACDC 데이터셋에서 sin² 활성화 함수를 사용한 ODE 블록이 검증 손실(-0.8895)과 mDSC(93.32%)에서 최고 성능을 기록했으며, 시그모이드 및 오토에코더 분석 대비 뛰어난 성능을 보였다.
• 정성적 결과에서는 긴장한 기관인 오른쪽 심실과 신장의 윤곽 예측이 뛰어나, 장거리 종속성 모델링 능력 향상이 확인되었다.
• 주의 맵 시각화 결과는 모델이 주목하는 것이 긴장한 구조물임을 확인했으며, 효과적인 전역 특징 학습이 이루어졌음을 시사한다.