[논문 리뷰] Narrowing the semantic gaps in U-Net with learnable skip connections: The case of medical image segmentation

Haonan Wang, Peng Cao외 3명|arXiv (Cornell University)|2023. 12. 23.

COVID-19 diagnosis using AI인용 수 4

한 줄 요약

이 논문은 U-Net 기반의 의료 영상 분할 프레임워크인 UDTransNet을 제안한다. 이는 표준 스킵 커넥션 대신 학습 가능한 어텐션 메커니즘을 도입하여 인코더와 디코더의 특징 간 의미적 격차를 좁힌다. 다중 척도 특징 융합을 위한 이중 어텐션 트랜스포머(DAT)와 단계 간 일관성을 향상시키기 위한 디코더 유도 재보정 어텐션(DRA) 모듈을 도입함으로써, 더 적은 파라미터로도 다양한 의료 영상 벤치마크에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Most state-of-the-art methods for medical image segmentation adopt the encoder-decoder architecture. However, this U-shaped framework still has limitations in capturing the non-local multi-scale information with a simple skip connection. To solve the problem, we firstly explore the potential weakness of skip connections in U-Net on multiple segmentation tasks, and find that i) not all skip connections are useful, each skip connection has different contribution; ii) the optimal combinations of skip connections are different, relying on the specific datasets. Based on our findings, we propose a new segmentation framework, named UDTransNet, to solve three semantic gaps in U-Net. Specifically, we propose a Dual Attention Transformer (DAT) module for capturing the channel- and spatial-wise relationships to better fuse the encoder features, and a Decoder-guided Recalibration Attention (DRA) module for effectively connecting the DAT tokens and the decoder features to eliminate the inconsistency. Hence, both modules establish a learnable connection to solve the semantic gaps between the encoder and the decoder, which leads to a high-performance segmentation model for medical images. Comprehensive experimental results indicate that our UDTransNet produces higher evaluation scores and finer segmentation results with relatively fewer parameters over the state-of-the-art segmentation methods on different public datasets. Code: https://github.com/McGregorWwww/UDTransNet.

연구 동기 및 목표

• 의료 영상 분할에서 U-Net의 고정된 스킵 커넥션의 한계를 조사하기 위해.
• 모든 스킵 커넥션이 동일하게 유용한 것은 아니며, 최적의 조합은 데이터셋에 따라 다를 수 있음을 규명하기 위해.
• 다중 척도 인코더 특징 간 및 인코더-디코더 단계 간 의미적 격차를 해결하기 위해.
• U-Net 아키텍처 전반에 걸쳐 특징 호환성을 향상시키는 학습 가능한 어텐션 기반 메커니즘을 설계하기 위해.
• 낮은 파라미터 수로 최신 기술 수준의 분할 성능을 달성하기 위해.

제안 방법

• 다중 척도 인코더 특징 간 채널 및 공간적 의존성을 모델링하기 위해 채널 기반 융합 어텐션(CFA)과 공간 기반 선택 어텐션(SSA)을 조합한 이중 어텐션 트랜스포머(DAT) 모듈을 제안한다.
• CFA에서는 자기 어텐션을 사용해 다양한 인코더 단계 간 채널을 정렬하고, SSA에서는 교차 어텐션을 사용해 스케일 간 중요한 공간 영역을 강조한다.
• 디코더 특징을 사용해 DAT 출력을 재보정함으로써 의미적 일관성을 감소시키는 디코더 유도 재보정 어텐션(DRA) 모듈을 도입한다.
• U-Net 기반 아키텍처의 스킵 커넥션에 DAT와 DRA를 통합하여 어텐션 기반 특징 융합을 가능하게 하며, 엔드 투 엔드 학습을 지원한다.
• 메인 인코더나 디코더에 어텐션을 적용하는 대신, CNN 기반 인코더와 스킵 커넥션에 트랜스포머 기반 어텐션 모듈을 융합한 하이브리드 아키텍처를 사용한다.
• 어 attention 헤드 수와 트랜스포머 레이어 수 등의 하이퍼파라미터를 최적화하여, 4개의 헤드와 4개의 레이어에서 최적의 설정을 도출한다.

실험 결과

연구 질문

• RQ1 U-Net의 모든 스킵 커넥션은 다양한 의료 영상 분할 작업에서 동일하게 효과적인가?
• RQ2 다중 척도 인코더 특징 간 및 인코더-디코더 단계 간 의미적 격차는 분할 성능에 어떤 영향을 미치는가?
• RQ3 스킵 커넥션에 학습 가능한 어텐션 메커니즘이 도입되면 특징 정렬이 향상되고 분할 정확도가 향상되는가?
• RQ4 스킵 커넥션에 어텐션을 통합하는 것이 인코더나 디코더에 별도로 어텐션을 적용하는 것보다 더 효과적인가?
• RQ5 제안된 어텐션 모듈에서 어 헤드 수와 레이어 수의 최적 설정은 무엇인가?

주요 결과

• UDTransNet은 세 가지 공개 데이터셋에서 최신 기술 수준의 Dice 스코어를 달성하였다: GlaS에서 91.03±0.56, Synapse에서 89.28±0.58로 기존 방법들을 초월하였다.
• ResNet-34를 동일한 백본으로 사용할 경우, UDTransNet은 GlaS에서 91.03±0.56의 Dice 스코어를 기록하여 MCTrans(88.44±0.74)와 TransUNet(87.41±1.15)를 크게 앞섰다.
• Swin Transformer를 백본으로 사용할 경우에도 UDTransNet은 GlaS에서 88.03±1.27의 Dice 스코어를 기록하여 동일한 백본을 사용한 Swin-UNet(89.04±0.88)보다 뛰어났다.
• 모델은 하이퍼파라미터 변화에 대해 강건하며, 최적의 성능을 4개의 어 헤드와 4개의 레이어에서 기록하여 이러한 설정에 대해 낮은 민감도를 보였다.
• 제거 실험을 통해 DAT와 DRA 모듈이 모두 필수적임을 확인하였으며, 특히 DRA가 인코더와 디코더 간 특징 일관성을 크게 향상시켰다.
• 최신 기술 수준의 모델들보다 더 적은 파라미터로도 높은 성능을 달성하여 효율성과 효과성을 동시에 입증하였다.