[논문 리뷰] UNet-2022: Exploring Dynamics in Non-isomorphic Architecture

Jiansen Guo, Hong-Yu Zhou외 2명|arXiv (Cornell University)|2022. 10. 27.

Radiomics and Machine Learning in Medical Imaging인용 수 8

한 줄 요약

UNet-2022는 동적 공간 및 채널별 특징 가중치를 동시에 활용하기 위해 자기주의성과 깊이 분할 컨벌루션 모듈을 병렬화한 비동형 U-Net 아키텍처를 제안한다. 이로 인해 복부 다기관 및 심장 분할과 같은 다양한 의료 영상 분할 과제에서 최신 기준(SOTA) 성능을 달성하며, nnUNet보다 최대 4% 높은 Dice 스코어를 기록한다.

ABSTRACT

Recent medical image segmentation models are mostly hybrid, which integrate self-attention and convolution layers into the non-isomorphic architecture. However, one potential drawback of these approaches is that they failed to provide an intuitive explanation of why this hybrid combination manner is beneficial, making it difficult for subsequent work to make improvements on top of them. To address this issue, we first analyze the differences between the weight allocation mechanisms of the self-attention and convolution. Based on this analysis, we propose to construct a parallel non-isomorphic block that takes the advantages of self-attention and convolution with simple parallelization. We name the resulting U-shape segmentation model as UNet-2022. In experiments, UNet-2022 obviously outperforms its counterparts in a range segmentation tasks, including abdominal multi-organ segmentation, automatic cardiac diagnosis, neural structures segmentation, and skin lesion segmentation, sometimes surpassing the best performing baseline by 4%. Specifically, UNet-2022 surpasses nnUNet, the most recognized segmentation model at present, by large margins. These phenomena indicate the potential of UNet-2022 to become the model of choice for medical image segmentation.

연구 동기 및 목표

• 자기주의성과 컨벌루션의 상호보완적 강점을 의료 영상 분할에서 직관적으로 설명하기 위해.
• 기존 하이브리드 아키텍처에서 자기주의성과 컨벌루션을 조합할 때 명확한 근거 없이 해석이 어려운 문제를 해결하기 위해.
• 공간 및 채널 차원에서 동시에 동적 가중치 할당을 가능하게 하는 새로운 빌딩 블록을 설계하기 위해.
• 간단하면서도 효과적인 아키텍처 혁신을 통해 다양한 의료 영상 과제에서의 분할 성능 향상시키기 위해.
• 제안된 UNet-2022 프레임워크에서 ImageNet 기반 사전 훈련 및 추론 전략 튜닝의 효과를 검증하기 위해.

제안 방법

• 자기주의성과 깊이 분할 컨벌루션(DwConv) 모듈을 별도로 처리하는 병렬 비동형 블록(PI 블록)을 제안한다.
• 공간 및 채널별로 적응적인 가중치를 할당하는 동적 가중치 할당 메커니즘을 도입하여 두 모듈의 강점을 통합한다.
• 파rameter 수를 줄이고 과적합을 완화하면서도 성능을 유지하기 위해 깊이 분할 컨벌루션을 활용한다.
• 스킵 연결을 갖춘 U-Net 유사 인코더-디코더 구조를 사용하며, 표준 인코더 대신 제안된 PI 블록을 도입한다.
• 일반화 성능 향상을 위해 ImageNet 기반 사전 훈련을 적용한다. 특히 데이터가 적은 환경에서 유의미한 효과가 있다.
• 특징 집합을 향상시키기 위해 작은 스텝 크기(0.2× 컷 크기)로 슬라이딩 윈도우 테스트를 최적화하여 추론을 수행한다.

실험 결과

연구 질문

• RQ1 자기주의성과 컨벌루션을 조합한 하이브리드 아키텍처가 의료 영상 분할에서 순수 컨벌루션 모델이나 자기주의성 기반 모델보다 우수한 이유는 무엇인가?
• RQ2 자기주의성과 컨벌루션의 동적 가중치 할당 메커니즘을 공식적으로 비교하고 통합된 프레임워크에서 어떻게 활용할 수 있는가?
• RQ3 자기주의성과 DwConv를 병렬로 구현한 비동형 연산을 도입할 경우 다양한 의료 영상 과제에서 분할 성능에 어떤 영향을 미치는가?
• RQ4 ImageNet 사전 훈련이 제안된 UNet-2022 모델의 성능을 기존의 랜덤 초기화 훈련 대비 얼마나 향상시키는가?
• RQ5 슬라이딩 윈도우 추론에서 추론 스텝 크기가 최종 분할 품질에 어떤 영향을 미치는가?

주요 결과

• 제안된 PI 블록은 Synapse 데이터셋에서 평균 Dice 스코어 86.46과 HD95 11.34를 기록하며, Swin Transformer 블록(84.42 DSC, 20.74 HD95)과 ConvNeXt 블록(84.96 DSC, 16.60 HD95)을 모두 능가한다.
• PI 블록에서 자기주의성 모듈을 제거하면 평균 DSC는 85.20으로 감소하고 HD95는 14.96으로 증가하여 공간적 동적 특성의 중요성을 확인한다.
• ImageNet 기반 사전 훈련은 UNet-2022의 평균 DSC를 1.5%p 향상시키며, 랜덤 초기화 훈련 대비 HD95는 약 5mm 향상시킨다.
• 더 작은 추론 스텝 크기(0.2× 컷 크기)를 사용할 경우 0.5× 컷 크기 대비 HD95가 2.6mm 향상되어 더 밀도 높은 추론 샘플링의 이점이 입증된다.
• 복부 다기관 및 피부 병변 분할과 같은 다양한 과제에서 UNet-2022는 널리 사용되는 nnUNet 모델보다 최대 4% 높은 Dice 스코어를 기록한다.
• 제거 실험 결과, 자기주의성과 DwConv를 병렬로 통합하는 것이 단독으로 사용하거나 순차적으로 사용하는 것보다 더 효과적임을 확인한다.