[논문 리뷰] MaxViT-UNet: Multi-Axis Attention for Medical Image Segmentation

Abdul Rehman, Asifullah Khan|arXiv (Cornell University)|2023. 05. 15.

Brain Tumor Detection and Classification인용 수 8

한 줄 요약

이 논문은 MaxViT-UNet을 제시한다, 의료 영상 분할을 위한 하이브리드 CNN-Transformer 인코더-디코더로, 다축 자기 주의를 갖춘 하이브리드 디코더를 사용하여 적은 메모리와 계산으로 핵 분할을 개선한다.

ABSTRACT

Since their emergence, Convolutional Neural Networks (CNNs) have made significant strides in medical image analysis. However, the local nature of the convolution operator may pose a limitation for capturing global and long-range interactions in CNNs. Recently, Transformers have gained popularity in the computer vision community and also in medical image segmentation due to their ability to process global features effectively. The scalability issues of the self-attention mechanism and lack of the CNN-like inductive bias may have limited their adoption. Therefore, hybrid Vision transformers (CNN-Transformer), exploiting the advantages of both Convolution and Self-attention Mechanisms, have gained importance. In this work, we present MaxViT-UNet, a new Encoder-Decoder based UNet type hybrid vision transformer (CNN-Transformer) for medical image segmentation. The proposed Hybrid Decoder is designed to harness the power of both the convolution and self-attention mechanisms at each decoding stage with a nominal memory and computational burden. The inclusion of multi-axis self-attention, within each decoder stage, significantly enhances the discriminating capacity between the object and background regions, thereby helping in improving the segmentation efficiency. In the Hybrid Decoder, a new block is also proposed. The fusion process commences by integrating the upsampled lower-level decoder features, obtained through transpose convolution, with the skip-connection features derived from the hybrid encoder. Subsequently, the fused features undergo refinement through the utilization of a multi-axis attention mechanism. The proposed decoder block is repeated multiple times to segment the nuclei regions progressively. Experimental results on MoNuSeg18 and MoNuSAC20 datasets demonstrate the effectiveness of the proposed technique.

연구 동기 및 목표

• 의료 영상 분할에서 지역적 및 전역 맥락을 모두 포착하기 위해 하이브리드 CNN-Transformer 아키텍처의 필요성을 제기한다.
• 다축 주의를 사용하여 업샘플된 디코더 피처와 건너뛰기 연결을 융합하는 하이브리드 디코더를 갖춘 UNet 유사 인코더-디코더인 MaxViT-UNet을 제안한다.
• 배경으로부터 핵을 구분하는 판별력을 높이면서 메모리 및 계산 부담을 줄인다.

제안 방법

• 의료 영상 분할을 위한 UNet 유형의 하이브리드 CNN-Transformer 아키텍처인 MaxViT-UNet를 도입한다.
• 전치 합성곱(transpose convolution)을 통해 업샘플된 하위 레벨 디코더 피처를 하이브리드 인코더에서 얻은 skip-connection 피처들과 결합하는 하이브리드 디코더를 설계한다.
• 특징 정제를 위해 각 디코딩 단계에서 다축 주의 메커니즘을 적용하여 특징을 정제한다.
• 제안된 디코더 블록을 여러 차례 반복하여 핵 영역을 점진적으로 분할한다.
• MoNuSeg18 및 MoNuSAC20 데이터셋에서 효과를 입증하기 위해 평가한다.

실험 결과

연구 질문

• RQ1 다축 주의가 적용된 CNN-Transformer 하이브리드 아키텍처가 표준 UNet이나 순수 Transformer 접근법에 비해 핵 분할을 개선할 수 있는가?
• RQ2 제안된 하이브리드 디코더가 주의 기법으로 인코더-디코더 피처를 효과적으로 융합하여 분할 경계를 향상시키는가?
• RQ3 의료 영상에서 객체와 배경을 구분하는 데 다축 주의의 영향은 무엇인가?
• RQ4 MoNuSeg18 및 MoNuSAC20 데이터셋에서 분할 품질과 효율성 측면에서 이 방법의 성능은 어떠한가?

주요 결과

• MaxViT-UNet 아키텍처는 MoNuSeg18 및 MoNuSAC20 데이터셋에서 효과를 보인다.
• 디코더 내 다축 주의를 도입하면 핵과 배경 영역 간의 구별력이 향상된다.
• 하이브리드 디코더는 정제하기 전에 업샘플된 디코더 피처를 인코더의 skip-connection 피처와 융합하여 점진적인 핵 분할을 가능하게 한다.
• 이 방법은 분할 성능 향상을 달성하면서 메모리 및 계산 부담을 최소화하는 것을 목표로 한다.