[논문 리뷰] ConvFormer: Plug-and-Play CNN-Style Transformers for Improving Medical Image Segmentation

Xian Lin, Zengqiang Yan외 3명|arXiv (Cornell University)|2023. 09. 09.

Advanced Neural Network ApplicationsComputer Science인용 수 3

한 줄 요약

이 논문은 주어진 의료 영상 세그멘테이션 작업에서 주의 집중 붕괴를 완화하기 위해 풀링, CNN 스타일의 자기주도 어텐션(CSA), 그리고 컨볼루션형 피드포워드 네트워크를 사용하여 표준 비전 트랜스포머 구성 요소를 대체하는 플러그 앤 플레이형 CNN 스타일 트랜스포머 모듈인 ConvFormer을 제안한다. 2D 특징 맵에서 직접 작동하고 적응형, 커널 스케일러블 어텐션을 가능하게 함으로써 ConvFormer은 장거리 의존성 학습을 향상시키고, ACDC, ISIC, ICH 데이터셋에서 여러 최신 트랜스포머 기반 모델에 대해 일관된 성능 향상을 달성한다.

ABSTRACT

Transformers have been extensively studied in medical image segmentation to build pairwise long-range dependence. Yet, relatively limited well-annotated medical image data makes transformers struggle to extract diverse global features, resulting in attention collapse where attention maps become similar or even identical. Comparatively, convolutional neural networks (CNNs) have better convergence properties on small-scale training data but suffer from limited receptive fields. Existing works are dedicated to exploring the combinations of CNN and transformers while ignoring attention collapse, leaving the potential of transformers under-explored. In this paper, we propose to build CNN-style Transformers (ConvFormer) to promote better attention convergence and thus better segmentation performance. Specifically, ConvFormer consists of pooling, CNN-style self-attention (CSA), and convolutional feed-forward network (CFFN) corresponding to tokenization, self-attention, and feed-forward network in vanilla vision transformers. In contrast to positional embedding and tokenization, ConvFormer adopts 2D convolution and max-pooling for both position information preservation and feature size reduction. In this way, CSA takes 2D feature maps as inputs and establishes long-range dependency by constructing self-attention matrices as convolution kernels with adaptive sizes. Following CSA, 2D convolution is utilized for feature refinement through CFFN. Experimental results on multiple datasets demonstrate the effectiveness of ConvFormer working as a plug-and-play module for consistent performance improvement of transformer-based frameworks. Code is available at https://github.com/xianlin7/ConvFormer.

연구 동기 및 목표

• 소규모 의료 영상 데이터셋에 적용할 때 비전 트랜스포머에서 발생하는 주의 집중 붕괴 문제를 해결하기 위해.
• 트랜스포머 기반 의료 영상 세그멘테이션 모델에서 자기주도 어텐션 맵의 수렴성과 다양성을 향상시키기 위해.
• 기존 아키텍처를 대체하지 않고도 순수 트랜스포머 및 CNN-트랜스포머 하이브리드 아키텍처 모두를 향상시킬 수 있는 플러그 앤 플레이 모듈을 설계하기 위해.
• 자기주도 어텐션을 학습 가능한 컨볼루션 커널로 모델링하여 확장 가능하고 적응형 장거리 특징 학습을 가능하게 하기 위해.
• 다양한 데이터셋과 백본 아키텍처에서 일관된 성능 향상을 입증하기 위해.

제안 방법

• 공간 국소성을 유지하고 해상도를 감소시키기 위해 ViT의 1D 토큰화를 2D 컨볼루션과 최대 풀링으로 대체한다.
• 각 픽셀의 장거리 의존성을 위해 적응형, 스케일러블 컨볼루션 커널로 구성된 CNN 스타일 자기주도 어텐션(CSA)을 도입한다.
• CSA 이후 특징을 정제하기 위해 컨볼루션 피드포워드 네트워크(CFFN)를 사용하여 2D 공간 구조를 유지한다.
• CSA에서 수용 영역 크기를 제어하기 위해 학습 가능한 스케일링 인자 α를 사용하여 국소적 및 전역적 맥락 간의 동적 적응을 가능하게 한다.
• 표준 ViT에서 사용하는 1D 시퀀스 변환을 피하기 위해 입력 영상을 2D 그대로 처리한다.
• 플러그 앤 플레이 모듈로 설계되어 기존의 트랜스포머 기반 세그멘테이션 프레임워크에 쉽게 통합될 수 있다.

Figure 1: Visualization of attention maps from the selected layers of the first head in different transformer frameworks. The darker the color, the closer the dependency.

실험 결과

연구 질문

• RQ1 CNN 스타일 트랜스포머 설계가 저데이터 의료 영상 세그멘테이션에서 주의 집중 붕괴를 완화할 수 있는가?
• RQ2 1D 시퀀스 대신 2D 특징 맵에서 작동할 경우 주의도 다양성과 모델 수렴성이 향상되는가?
• RQ3 적응형, 커널 스케일러블 어텐션 기반 플러그 앤 플레이 모듈이 다양한 트랜스포머 기반 세그멘테이션 아키텍처에서 일관되게 성능 향상을 이룰 수 있는가?
• RQ4 CSA에서 수용 영역의 스케일러비리티가 의료 영상 기준 평가에서 세그멘테이션 정확도에 어떤 영향을 미치는가?
• RQ5 제안된 모듈이 의료 영상 세그멘테이션에서 다양한 데이터셋과 백본 모델에 일반화되는가?

주요 결과

• Pure-transformer 백본인 SETR에 통합했을 때, ACDC 데이터셋에서 평균 Dice 스코어가 3.86% 향상된다.
• ISIC 데이터셋에서 SETR와 함께 사용했을 경우 Dice 스코어가 1.38% 향상되었고, Patcher와 함께 사용했을 경우 1.07% 향상되었다.
• ICH 데이터셋에서 SETR과 함께 사용했을 경우 Dice 스코어가 1.39% 향상되었고, Patcher와 함께 사용했을 경우 1.15% 향상되어 다양한 데이터셋에서 일관된 성능 향상을 보였다.
• TransUNet 및 TransFuse와 같은 CNN-트랜스포머 하이브리드 모델에서도 CNN의 지배적인 영향에도 불구하고 ConvFormer은 안정적인 성능 향상을 제공한다.
• 제거 실험 결과, α = 0.4일 때 최고의 성능(ACDC에서 91.00% Dice)을 기록하여 최적의 수용 영역 스케일링이 이루어짐을 확인했다.
• 시각화 결과, 기준 모델 대비 ConvFormer이 더 다양하고 구조화된 주의도 맵을 생성함으로써 주의 집중 붕괴를 효과적으로 완화함을 확인했다.

Figure 2: Comparison between vanilla vision transformer and ConvFormer. CBR is short for the combination of convolution, batch normalization, and Relu. Multiple heads are omitted for simplicity.