[논문 리뷰] xLSTM-UNet can be an Effective 2D & 3D Medical Image Segmentation Backbone with Vision-LSTM (ViL) better than its Mamba Counterpart

Tianrun Chen, Chaotao Ding외 6명|arXiv (Cornell University)|2024. 07. 01.

Brain Tumor Detection and Classification인용 수 5

한 줄 요약

요점: 이 논문은 xLSTM-UNet를 소개합니다. 이는 xLSTM/ViL을 백본으로 사용하는 UNet 유사 아키텍처로, 2D 및 3D 의학 영상 분할 작업에서 CNN-, Transformer-, 및 Mamba 기반 분할 모델보다 우수한 성능을 보입니다.

ABSTRACT

Convolutional Neural Networks (CNNs) and Vision Transformers (ViT) have been pivotal in biomedical image segmentation, yet their ability to manage long-range dependencies remains constrained by inherent locality and computational overhead. To overcome these challenges, in this technical report, we first propose xLSTM-UNet, a UNet structured deep learning neural network that leverages Vision-LSTM (xLSTM) as its backbone for medical image segmentation. xLSTM is a recently proposed as the successor of Long Short-Term Memory (LSTM) networks and have demonstrated superior performance compared to Transformers and State Space Models (SSMs) like Mamba in Neural Language Processing (NLP) and image classification (as demonstrated in Vision-LSTM, or ViL implementation). Here, xLSTM-UNet we designed extend the success in biomedical image segmentation domain. By integrating the local feature extraction strengths of convolutional layers with the long-range dependency capturing abilities of xLSTM, xLSTM-UNet offers a robust solution for comprehensive image analysis. We validate the efficacy of xLSTM-UNet through experiments. Our findings demonstrate that xLSTM-UNet consistently surpasses the performance of leading CNN-based, Transformer-based, and Mamba-based segmentation networks in multiple datasets in biomedical segmentation including organs in abdomen MRI, instruments in endoscopic images, and cells in microscopic images. With comprehensive experiments performed, this technical report highlights the potential of xLSTM-based architectures in advancing biomedical image analysis in both 2D and 3D. The code, models, and datasets are publicly available at http://tianrun-chen.github.io/xLSTM-UNet/

연구 동기 및 목표

• 의료 영상 분할에서 장거리 의존성에 xLSTM/ViL을 활용하려는 동기 부여.
• ViL/xLSTM 블록을 인코더 층에 주입하는 UNet 유사 아키텍처(xLSTM-UNet) 설계.
• CNN-, Transformer-, 및 Mamba 기반 기준선 대비 다양한 2D/3D 데이터셋에서 향상된 분할 성능 시연.
• 재현성을 촉진하기 위한 구현 및 데이터셋 상세 정보 제공.

제안 방법

• UNet과 유사한 인코더–디코더 구조를 채택합니다.
• 여러 인코더 레이어에 xLSTM 블록을 잔차 블록과 인스턴스 정규화와 함께 도입합니다.
• 중간 특징을 평탄화하고 정규화한 뒤 ViL/xLSTM 블록에 공급하고 출력은 디코더 경로에 연결합니다.
• 두 가지 변형을 탐구합니다: ours_bot(병목에서만 xLSTM) 및 ours_enc(모든 인코더 블록에 xLSTM).
• Dice와 교차 엔트로피 손실을 사용해 엔드투엔드로 학습하며, 고성능 GPU에서 AdamW를 사용합니다.”] ,
• research_questions=[

실험 결과

연구 질문

• RQ1 Can xLSTM-UNet surpass CNN-, Transformer-, and Mamba-based segmentation networks in 2D and 3D medical image segmentation?
• RQ2 Do encoder-wide xLSTM insertions (ours_enc) provide more benefit than bottleneck-only usage (ours_bot)?
• RQ3 Is xLSTM-UNet robust across diverse modalities including abdomen MRI, endoscopy, microscopy, and brain MRI?
• RQ4 How does xLSTM-UNet scale between 2D and 3D segmentation tasks?

주요 결과

• xLSTM-UNet는 Abdomen MRI 2D, Endoscopy, 및 Microscopy 데이터셋에서 최첨단 성능을 달성하며, ours_enc가 Abdomen MRI 2D에서 DSC 0.7747 및 NSD 0.8374에 도달합니다.
• 두 가지 xLSTM 변형(ours_bot 및 ours_enc)이 2D 작업에서 U-Mamba 변형 및 다른 기준선보다 우수합니다.
• Endoscopy 데이터셷에서 두 변형 모두 최상의 DSC 및 NSD 점수를 달성합니다(각각 0.6843 및 0.7001).
• Microscopy 데이터셋에서 xLSTM-UNet 변형은 F1 점수 0.6036(ours_enc) 및 0.5818(ours_bot)을 달성하여 이전 SOTA를 능가합니다.
• 3D BraTS2023에서 xLSTM-UNet 변형이 다른 방법들보다 평균 Dice가 가장 높게 나타납니다(91.80).
• Abdomen MRI 3D에서 xLSTM-UNet_bot이 DSC 0.8483 및 NSD 0.9153로 기준선을 능가합니다。）