[논문 리뷰] LKM-UNet: Large Kernel Vision Mamba UNet for Medical Image Segmentation

Jinhong Wang, Jintai Chen외 2명|arXiv (Cornell University)|2024. 03. 12.

Medical Image Segmentation Techniques인용 수 4

한 줄 요약

이 논문은 2D 및 3D 의료 영상 분할을 위한 대용량 창 기반 Mamba U-Net인 LMa-UNet을 제안한다. 이 모델은 Mamba의 선형 복잡도를 활용하여 대용량 창 SSM을 통해 넓은 수신 영역를 구현한다. 피그먼트 수준 및 패치 수준의 SSM을 갖춘 계층적이고 이방향인 Mamba 블록(LM 블록)을 도입함으로써 우수한 국소 및 전역 특징 모델링 성능을 달성하였으며, 다중 기관 분할 벤치마크에서 CNN, Transformer 및 이전의 Mamba 기반 모델들을 능가하는 DSC 및 NSD 점수를 확보하였다.

ABSTRACT

In clinical practice, medical image segmentation provides useful information on the contours and dimensions of target organs or tissues, facilitating improved diagnosis, analysis, and treatment. In the past few years, convolutional neural networks (CNNs) and Transformers have dominated this area, but they still suffer from either limited receptive fields or costly long-range modeling. Mamba, a State Space Sequence Model (SSM), recently emerged as a promising paradigm for long-range dependency modeling with linear complexity. In this paper, we introduce a Large Kernel Vision Mamba U-shape Network, or LKM-UNet, for medical image segmentation. A distinguishing feature of our LKM-UNet is its utilization of large Mamba kernels, excelling in locally spatial modeling compared to small kernel-based CNNs and Transformers, while maintaining superior efficiency in global modeling compared to self-attention with quadratic complexity. Additionally, we design a novel hierarchical and bidirectional Mamba block to further enhance Mamba's global and neighborhood spatial modeling capability for vision inputs. Comprehensive experiments demonstrate the feasibility and the effectiveness of using large-size Mamba kernels to achieve large receptive fields. Codes are available at https://github.com/wjh892521292/LKM-UNet.

연구 동기 및 목표

• CNN과 Transformer의 의료 영상 분할에서의 제한점, 특히 제한된 수신 영역과 장거리 모델링에 대한 높은 계산 비용을 해결하기 위함.
• 구조적 상태공간모델(SSM)인 Mamba의 잠재력을 탐색하여 선형 복잡도로 대규모 수신 영역를 갖춘 공간 모델링을 가능하게 하기 위함.
• 시각 작업에서 Mamba의 단방향 성향과 위치 인식 부족 문제를 해결하기 위해 이방향이고 계층적인 SSM 아키텍처를 설계하기 위함.
• 대용량 창 기반 Mamba 모듈이 2D 및 3D 의료 영상에서 국소 및 전역 의존성을 효과적으로 모델링할 수 있음을 입증하기 위함.
• 다중 기관 분할 데이터셋에서 체계적인 아블레이션 및 벤치마크를 통해 제안된 LMa-UNet의 유효성을 검증하기 위함.

제안 방법

• 표준 컨볼루션 및 자기주의 모듈을 대용량 창 SSM으로 대체하여 넓은 수신 영역를 확보한 대용량 창 기반 Mamba U-Net(LMa-UNet)을 제안한다.
• 피그먼트 수준 SSM(PiM)과 패치 수준 SSM(PaM)으로 구성된 새로운 계층적이고 이방향인 Mamba 블록(LM 블록)을 도입한다. PiM은 국소 이웃 모델링을 위해, PaM은 전역 의존성 학습을 위해 사용된다.
• 토큰화 후 공간적 불연속성으로 인한 SSM의 기억 상실 문제를 완화하기 위해 PiM을 활용하여 국소 특징 표현을 향상시킨다.
• 선형 복잡도로 패치 간 장거리 의존성을 모델링하여 자기주의의 제곱형 비용을 피함으로써, PaM을 사용한다.
• 순서 모델링에서 입력 순서 편향을 줄이고 위치 인식 능력을 향상시키기 위해 이방향 Mamba(BiM)를 적용한다.
• PyTorch를 사용하여 nnU-Net 프레임워크에 모델을 구현하였으며, 각 스테이지별로 적응형 창 크기를 적용하고 표준 학습 초모수(Adam, 가중치 감쇠, 학습률)를 사용한다.

실험 결과

연구 질문

• RQ1 작은 커널 크기의 CNN과 창 크기 제약이 있는 Transformer에 비해, 대용량 창 크기를 갖춘 Mamba 기반 모델이 의료 영상 분할에서 더 나은 국소 및 전역 특징 모델링 성능을 달성할 수 있는가?
• RQ2 LM 블록 내 PiM과 PaM의 계층적 설계가 시각 작업에서 Mamba의 표현 능력을 얼마나 향상시키는가?
• RQ3 이방향 SSM이 내재한 위치 편향과 순서 민감도 문제를 이방향 Mamba가 어느 정도 완화하는가?
• RQ4 Mamba 기반 모델의 창 크기를 증가시키면 의료 영상 분할 성능에 측정 가능한 향상이 초래되는가?
• RQ5 PiM, PaM, BiM 구성 요소가 LMa-UNet의 전체 성능에 기여하는 상대적 기여도는 어느 정도인가?

주요 결과

• Abdomen CT 데이터셋에서 LMa-UNet은 DSC 86.82와 NSD 90.02를 기록하여, CNN 기반, Transformer 기반 및 이전의 Mamba 기반 모델들을 능가하였다.
• Abdomen MR 데이터셋에서 LMa-UNet은 DSC 77.35와 NSD 83.80를 기록하여, 평가된 모든 벤치마크에서 최신 기술 수준의 성능을 보였다.
• 창 크기를 [10,5,5,5,5,5,5]에서 [40,20,20,10,10,5,5]로 증가시켰을 때, DSC는 75.89에서 77.35로, NSD는 82.26에서 83.80으로 향상되어 넓은 수신 영역의 유용성을 확인하였다.
• 아블레이션 연구 결과, PaM보다 PiM이 더 큰 성능 향상 기여를 하였으며, 이는 국소 수신 영역 강화가 전역 모델링만으로는 더 중요하다는 것을 시사한다.
• PiM + PaM + BiM 모든 구성 요소를 포함한 완전한 LMa-UNet이 최고의 성능을 기록하였으며, DSC 77.35와 NSD 83.80를 달성하여 계층적이고 이방향 설계의 효과성을 입증하였다.
• Mamba의 선형 복잡도 덕분에 높은 효율성을 유지하였으며, 표준 CNN 및 Transformer가 처리하기 어려운 대용량 창 모델링이 가능하였다.