[논문 리뷰] CiT-Net: Convolutional Neural Networks Hand in Hand with Vision Transformers for Medical Image Segmentation

Tao Lei, Rui Sun외 4명|arXiv (Cornell University)|2023. 06. 06.

Advanced Neural Network Applications참고 문헌 38인용 수 5

한 줄 요약

이 논문은 의료 영상 분할을 위한 하이브리드 CNN-Transformer 아키텍처인 CiT-Net을 제안한다. 이 모델은 CNN 브랜치에 동적 탈형 변형 컨볼루션을 통합하고, Transformer 브랜치에 이동 창 기반 적응형 보완 주의 모듈을 적용한다. 사전 훈련 없이도 기존 최고 성능(SOTA) 방법들보다 더 낮은 파라미터 수와 FLOPs를 기록하면서도, ISIC2018 및 LiTS-Liver 데이터셋에서 Swin-Unet과 TransUNet 등의 모델을 능가하는 최고 성능을 달성한다.

ABSTRACT

The hybrid architecture of convolutional neural networks (CNNs) and Transformer are very popular for medical image segmentation. However, it suffers from two challenges. First, although a CNNs branch can capture the local image features using vanilla convolution, it cannot achieve adaptive feature learning. Second, although a Transformer branch can capture the global features, it ignores the channel and cross-dimensional self-attention, resulting in a low segmentation accuracy on complex-content images. To address these challenges, we propose a novel hybrid architecture of convolutional neural networks hand in hand with vision Transformers (CiT-Net) for medical image segmentation. Our network has two advantages. First, we design a dynamic deformable convolution and apply it to the CNNs branch, which overcomes the weak feature extraction ability due to fixed-size convolution kernels and the stiff design of sharing kernel parameters among different inputs. Second, we design a shifted-window adaptive complementary attention module and a compact convolutional projection. We apply them to the Transformer branch to learn the cross-dimensional long-term dependency for medical images. Experimental results show that our CiT-Net provides better medical image segmentation results than popular SOTA methods. Besides, our CiT-Net requires lower parameters and less computational costs and does not rely on pre-training. The code is publicly available at https://github.com/SR0920/CiT-Net.

연구 동기 및 목표

• 표준 CNN의 고정된 수신장과 파라미터 공유 문제로 인한 의료 영상 분할의 한계를 해결하기 위해.
• 표준 Transformer에서 채널 및 교차 차원 주의의 부족으로 인해 복잡한 의료 영상에서 성능 저하가 발생하는 문제를 해결하기 위해.
• MLP를 경량 퍼셉트론 모듈로 대체하여 모델 복잡도를 줄이고 사전 훈련에 대한 의존도를 제거하기 위해.
• 기존 하이브리드 CNN-Transformer 모델보다 더 낮은 계산 비용과 더 적은 파라미터로 높은 분할 정확도를 달성하기 위해.
• 국소적이고 전역적 특징을 모두 유지하는 강력하고 효율적이며 일반화 능력이 뛰어난 의료 영상 분할 아키텍처를 개발하기 위해.

제안 방법

• 고정 커널 제약 조건을 피하기 위해 커널 오프셋을 학습함으로써 적응형 공간 특징 추출이 가능한 동적 탈형 컨볼루션(DDConv)을 CNN 브랜치에 도입한다.
• 공간적 특징과 채널 특징 간의 교차 차원적 의존성을 적응형 주의를 통해 모델링하기 위해 Transformer 브랜치에 이동 창 기반 적응형 보완 주의 모듈((S)W-ACAM)을 제안한다.
• Transformer의 표준 MLP를 대체하기 위해 경량 퍼셉트론 모듈(LPM)을 사용하여 모델 파라미터와 계산 비용을 크게 감소시킨다.
• CNN과 Transformer 브랜치가 병렬로 특징을 처리하도록 이중 브랜치 아키텍처를 설계하여 국소적 및 전역적 표현을 효과적으로 융합한다.
• Transformer 브랜치에서 장거리 의존성 모델링을 가능하게 하면서도 효율성을 유지하기 위해 압축된 컨볼루션 투영을 사용한다.
• 사전 훈련 없이 종단 간(end-to-end)으로 모델을 훈련시어 과적합 위험을 줄이고 제한된 의료 데이터셋에서의 일반화 능력을 향상시킨다.

실험 결과

연구 질문

• RQ1 표준 CNN에 비해 동적 탈형 컨볼루션은 의료 영상 분할에서 국소적 특징 표현을 향상시키는가?
• RQ2 Transformer 브랜치에 교차 차원 주의를 통합하면 복잡한 의료 영상에서 분할 정확도가 향상되는가?
• RQ3 경량 퍼셉트론 모듈은 성능 저하 없이 표준 MLP를 효과적으로 대체하여 모델 크기와 계산 비용을 줄일 수 있는가?
• RQ4 사전 훈련 없이도 제안된 하이브리드 아키텍처는 분할 정확도, 파라미터 수, FLOPs 측면에서 SOTA 방법들과 비교해 어떻게 성능을 내는가?
• RQ5 적응형 주의와 동적 컨볼루션의 통합은 의료 영상에서 국소적 및 전역적 특징의 유지에 얼마나 기여하는가?

주요 결과

• CiT-Net-B는 ISIC2018 데이터셋에서 90.88%의 Dice 스코어를 기록하여, Swin-Unet, TransUNet, CvT보다 각각 1.20%, 1.03%, 1.01% 높은 성능을 보였으며, 사전 훈련 없이도 성능을 확보했다.
• CiT-Net-T는 단지 11.58M 파라미터와 4.53 GFLOPs로 ISIC2018에서 최고의 Dice 스코어 90.72%를 기록하여 가장 효율적인 SOTA 모델이 되었다.
• LiTS-Liver 데이터셋에서 CiT-Net-B는 DI, VOE, RVD, ASD, RMSD 등 모든 다섯 가지 지표에서 뛰어난 성능을 보이며 강력한 일반화 능력과 내구성을 입증했다.
• 절단 실험 결과, DDConv와 (S)W-ACAM을 함께 사용할 경우 기준 모델 대비 1.2%의 Dice 향상이 이루어졌으며, 최적의 성능을 달성했다.
• LPM는 최적 설정에서 모델 파라미터를 9.67M로 줄여 계산 비용을 크게 낮추면서도 높은 정확도를 유지했다.
• ISIC2018에서 CiT-Net은 파라미터 수와 FLOPs가 더 적은 BAT, CvT, CrossForm보다 각각 1.02%, 3.00%, 3.79% 높은 Dice 스코어를 기록했다.