[논문 리뷰] H-vmunet: High-order Vision Mamba UNet for Medical Image Segmentation

Renkai Wu, Yinghao Liu외 2명|arXiv (Cornell University)|2024. 03. 20.

Image Processing Techniques and Applications인용 수 4

한 줄 요약

이 논문은 고차원 2차원 선택적 스캔(H-SS2D)을 도입하여 부적절한 정보를 줄이고 전역 수용 영역을 유지함으로써 의료 영상 분할 성능을 햖थ하는 H-vmunet을 제안한다. UNet 아키텍처 내에 새로운 고차원 시각 상태공간(H-VSS) 모듈에 H-SS2D를 통합함으로써 H-vmunet는 VM-UNet 대비 67.28% 적은 파라미터를 사용하며, ISIC2017, Spleen, CVC-ClinicDB 데이터셋에서 여러 베이스라인을 능가한다.

ABSTRACT

In the field of medical image segmentation, variant models based on Convolutional Neural Networks (CNNs) and Visual Transformers (ViTs) as the base modules have been very widely developed and applied. However, CNNs are often limited in their ability to deal with long sequences of information, while the low sensitivity of ViTs to local feature information and the problem of secondary computational complexity limit their development. Recently, the emergence of state-space models (SSMs), especially 2D-selective-scan (SS2D), has had an impact on the longtime dominance of traditional CNNs and ViTs as the foundational modules of visual neural networks. In this paper, we extend the adaptability of SS2D by proposing a High-order Vision Mamba UNet (H-vmunet) for medical image segmentation. Among them, the proposed High-order 2D-selective-scan (H-SS2D) progressively reduces the introduction of redundant information during SS2D operations through higher-order interactions. In addition, the proposed Local-SS2D module improves the learning ability of local features of SS2D at each order of interaction. We conducted comparison and ablation experiments on three publicly available medical image datasets (ISIC2017, Spleen, and CVC-ClinicDB), and the results all demonstrate the strong competitiveness of H-vmunet in medical image segmentation tasks. The code is available from https://github.com/wurenkai/H-vmunet .

연구 동기 및 목표

• CNN의 장거리 의존성 파악 능력 부족과 ViT의 국소적 특징 민감도 및 계산 비용 문제를 해결하기 위해.
• 상태공간 모델(SSM)의 표준 2차원 선택적 스캔(SS2D) 연산에서 발생하는 부정확한 중복성과 국소적 특징 학습의 열악함을 극복하기 위해.
• Mamba 기반 모듈과 함께 고차원 공간 상호작용을 통합함으로써 UNet 프레임워크를 의료 영상 분할에 최적화하기 위해.
• 강력한 전역 맥락을 유지하면서 노이즈와 중복성을 최소화하는 파라미터 효율적이고 고성능의 분할 모델을 개발하기 위해.

제안 방법

• 고차원 2차원 선택적 스캔(H-SS2D)을 제안하여, 특징 처리 과정에서 점진적으로 중복 정보를 줄이는 단계별 프레이저드 SS2D 메커니즘을 구현한다.
• 각 순서의 상호작용에서 국소적 특징 학습을 향상시키기 위해 Local-SS2D 모듈을 도입한다.
• H-SS2D 기반의 고차원 시각 상태공간(H-VSS) 모듈을 설계하여 기존의 VSS를 대체함으로써 계층적이고 순서 인식 특징 추출을 가능하게 한다.
• 스킵 연결을 갖춘 인코더-디코더 아키텍처에 H-VSS를 통합하여 고차원 시각 Mamba UNet(H-vmunet)을 구성한다.
• 성능과 계산 비용의 균형을 맞추기 위해 인코더 및 디코더 모듈에서 [2,3,4,5]의 혼합 순서 전략을 사용한다.
• 세 가지 공개 데이터셋에서의 추론 및 비교를 통해 H-SS2D 및 H-VSS 구성 요소의 유효성을 검증한다.

실험 결과

연구 질문

• RQ1 SS2D 연산에서 고차원 공간 상호작용이 중복 정보를 줄이면서도 전역 수용 영역을 유지할 수 있는가?
• RQ2 H-SS2D를 시각 Mamba 모듈에 통합함으로써 국소적 특징 학습과 분할 정확도가 향상되는가?
• RQ3 의료 영상 분할 벤치마크에서 H-vmunet은 VM-UNet 및 MHorUNet에 비해 성능과 파라미터 효율성 측면에서 어떻게 비교되는가?
• RQ4 분할 정확도와 계산 효율성 측면에서 H-SS2D에 대한 최적의 순서 구성은 무엇인가?

주요 결과

• H-vmunet는 ISIC2017, Spleen, CVC-ClinicDB 세 가지 의료 영상 데이터셋에서 모두 최신 기준 성능을 달성하며, MHorUNet 및 VM-UNet을 모두 능가한다.
• 기본 모델인 VM-UNet 대비 파라미터를 67.28% 감소시켜 파라미터 효율성 측면에서 뚜렷한 향상을 이룬다.
• 절단 분석 결과 H-SS2D 및 Local-SS2D 모듈이 필수적임을 확인하였으며, 전체 H-SS2D 설정에서 가장 높은 Dice 스코어를 기록했다.
• 혼합 순서 구성 [2,3,4,5]는 단일 순서 또는 균일한 순서 설정보다 더 우수한 성능을 내어 계층적 상호작용의 이점이 있음을 시사한다.
• 시각화 결과 H-vmunet는 VM-UNet에 비해 더 선명한 병변 경계와 더 나은 소형 병적 특징 민감도를 보였다.
• H-SS2D 메커니즘은 중복 정보를 효과적으로 최소화하면서도 전역 맥락을 유지하여 설계 원칙의 타당성을 입증했다.