[論文レビュー] 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 more|arXiv (Cornell University)|Jul 1, 2024

Brain Tumor Detection and Classification被引用数 5

ひとこと要約

論文は xLSTM-UNet を導入し、CNN・Transformer・Mamba ベースのセグメンテーションモデルを 2D および 3D 医用画像セグメンテーションタスクで上回る UNet 的アーキテクチャである。

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-, and Mamba-based baselines を diverse 2D/3D データセットで上回るセグメンテーション性能を示す。
• 再現性を促進する実装およびデータセットの詳細を提供。

提案手法

• UNet ライクなエンコーダ–デコーダ構造を採用。
• 複数のエンコーダ層に xLSTM ブロックを残差ブロックとインスタンス正規化とともに組み込む。
• 中間特徴をフラット化して正規化し、ViL/xLSTM ブロックへ入力して出力をデコーダ経路に結合する。
• 2 つのバリアントを探索する：ours_bot（ボトルネックのみ xLSTM）と ours_enc（全エンコーダブロックに xLSTM）。
• AdamW を用いて高性能 GPU 上で Dice とクロスエントロピー損失を用いたエンドツーエンドの訓練。

実験結果

リサーチクエスチョン

• 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 データセットでは、両方の xLSTM-UNet バリアントが最高の 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 を達成し、ベースラインを上回る。