xLSTM-UNet can be an Effective Backbone for 2D & 3D Biomedical Image Segmentation Better than its Mamba Counterparts

Authors: st tianrun chen
lanyun  zhu
tao  xu
yan  wang
ying  zang
Venue: IEEE_org_EMBS_BHI_2024_Conference
Published: 2024-07-06 14:03:12

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 has recently been proposed as the successor of Long Short-Term Memory (LSTM) networks and has 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, we provide the first integration of xLSTM with image segmentation backbone -- namely xLSTM-U, which extend the success of xLSTM in the biomedical image segmentation domain. By integrating the local feature extraction strengths of convolutional layers with the long-range dependency-capturing abilities of xLSTM, the proposed 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 paper highlights the potential of xLSTM-based architectures in advancing biomedical image analysis in both 2D and 3D. We believe this new finding will be of interest to the research community and may inspire future studies. The code, models, and datasets are publicly available at https://github.com/tianrun-chen/xLSTM-UNet-PyTorch/tree/main.

Main paper: 
 i. introduction: 
 Biomedical image segmentation is a critical task in medical imaging, enabling precise delineation of anatomical structures and anomalies essential for diagnosis, treatment planning, and research [1] , [2] . In recent years, deep learning methods have achieved remarkable success in tumor segmentation [3] and organ segmentation in 3D Computed Tomography (CT) scans [4] , as well as in cell segmentation in microscopy images [5] , [6] . These advancements underscore the transformative impact of deep learning on the landscape of biomedical image segmentation, paving the way for more accurate and efficient diagnostic and treatment planning tools. Traditionally, Convolutional Neural Networks (CNNs) have been the backbone of this domain in deep learning-enabled methods, leveraging their powerful local feature extraction capabilities [7] - [9] . More recently, Vision Transformers (ViTs) have gained popularity by offering a robust alternative, capable of capturing global context through self-attention mechanisms [10] - [12] . Despite their successes, both CNNs and ViTs face inherent limitations. CNNs struggle with long-range dependencies due to their localized receptive fields, while ViTs encounter substantial computational overhead [13] , [14] , especially with highresolution images or high-dimensional imaging modalities like 3D images or hyper-spectral imaging like stimulated Raman scattering (SRS) imaging [15] - [17] or mid-infrared (IR) spectroscopic imaging [18] .To address these challenges, recent work has proposed to integrate computation modules that have long-range dependencies and also exhibit linear computational and memory complexity w.r.t. sequence length. Among these computation modules, State Space Models (SSMs) [19] - [21] , like Mamba [22] , has demonstrated its huge success. SSMs excel in handling long-range dependencies and have been successfully integrated into conventional UNet architectures. Variants like UMamba [23] , VM-Unet [24] - [26] , Mamba-Unet [27] , Swin-UMamba [28] , and SegMamba [29] , have demonstrated their considerable success.Meanwhile, Extended Long-Short Term Memory (xLSTM) has recently emerged as a powerful successor to Long Short-Term Memory (LSTM) networks, challenging Transformers in sequence modeling [30] . Like SSMs, xLSTM can handle long-range dependencies and maintain linear computational and memory complexity. However, xLSTM has demonstrated superior performance in neural language processing (NLP) and image classification (in its Vision-LSTM (ViL) implementation [31] ). This success naturally raises the question: Can xLSTM, or ViL, also excel in image segmentation, specifically in the field of medical image segmentation?The answer is Yes! In this technical report, we introduce xLSTM-UNet, the first xLSTM-enabled U-Net image segmentation network that can perform both 2D and 3D medical image segmentation tasks and achieves state-of-the-art (SOTA) results. We conducted comprehensive experiments in various 2D and 3D medical segmentation scenarios, including organs in abdominal MRI, instruments in endoscopy, cells in microscopy, and cancer segmentation in 3D brain MRI volumes. The results show that xLSTM-UNet outperforms existing CNN-based and Transformer-based segmentation methods, as well as its Mamba-based counterparts. These findings highlight the potential of xLSTM-based architecture to set new benchmarks in the field of medical image segmentation, offering improved accuracy and efficiency across a wide range of applications. To further advance research in this area, we will release the model and code, enabling future explorations in various fields such as automated pathology detection, camouflaged image segmentation, precision agriculture, environmental monitoring, satellite imagery analysis, and industrial inspection, and so on.Medical image segmentation is a crucial application of computer vision, involving the delineation of pathological and anatomical structures from volumetric data for diagnosis, treatment planning, and patient monitoring. This task is complex due to the multidimensional nature of the data, requiring highly accurate segmentation outputs.Deep learning has significantly advanced medical image segmentation, particularly with convolutional neural networks (CNNs) [32] . Early works [8] , [33] used CNNs as a backbone for segmentation tasks. A key breakthrough came with U-Net [7] , which introduced a symmetric encoder-decoder structure that captured both fine-grained spatial information and abstract features. Extensions like Unet++ [9] improved performance through nested skip connections and deep supervision. However, CNNs often struggle with long-range dependencies.Transformers [34] addressed this limitation by modeling global dependencies with self-attention mechanisms, as seen in hybrid models like TransUNet [10] , which combine CNNs for local features and Transformers for global context.While Transformers improve segmentation, they are computationally expensive. To mitigate this, State Space Models (SSMs) [19] - [21] , such as Mamba [22] , offer efficient longrange dependency handling with linear computational complexity, as demonstrated in Unet-based models [23] - [29] .Extended Long-Short Term Memory (xLSTM) networks [30] present another alternative, efficiently modeling longrange dependencies with reduced computational demands. Vision-xLSTM [31] integrates global context with computational efficiency, making it a promising approach for developing precise and efficient medical image segmentation models.1) Vision-xLSTM: Vision-xLSTM (ViL) is an adaptation of xLSTM specifically for computer vision tasks. The configuration contains interchanging mLSTM segments where even-indexed segments handle flattened features starting from the top left moving to the bottom right, while odd-indexed segments proceed from the bottom right to the top left. This two-way operation allows ViL segments to grasp strong worldwide dependencies within the input.The flattened features, after being normalized, are projected into an embedding domain, doubling their size. These enlarged embeddings are split into two paths: x mlstm ∈ R N ×2Z and y ∈ R N ×2Z . x mlstm undergoes a 1D causal convolution process with the SiLU activation function applied to it. The intermediate outcome (X ∈ R N ×2Z ) is subsequently mapped to query, key, and value vectors, which are akin to the vectors used in the Transformer model. These vectors are then forwarded to the mLSTM cell. The mLSTM sublayer consists of n parallel attention heads, each equipped with a matrix memory unit.At first, input and forget gate pre-activations, ĩ ∈ R N ×d h and f ∈ R N ×d h respectively, are calculated by linearly projecting the concatenated Q, K, and V matrices as expressed below mathematically:EQUATIONEQUATIONHere, B ∈ R N ×d h is the bias matrix. The formation of query, value, and key is mathematically defined as: EQUATION)Where Q belongs to R N ×d , K resides in R N ×d , V is in R N ×d , these are query, key, and value matrices. The W Q resides in R 2Z×d , W K is within R 2Z×d , and W V is found in R 2Z×d , they are the adjustable weight matrices for generating the query, key, and value per head. Herein, d symbolizes the intended dimension for queries, keys, and values.The pre-activation outputs further derive the gate decay matrix D ∈ R N ×d h following this equation:EQUATIONHerein, log σ denotes the log-sigmoid activation function. The decay matrix D is stabilized to ensure subsequent exponentiation of D results in stable output. The final gate decay matrix output D ∈ R N ×d h is represented as:EQUATIONHere, exp represents the exponentiation operation. The Q and K vectors perform a dot product to derive the attention scores S ∈ R N ×N , analogous to the self-attention mechanism in Transformers. A causal mask M ∈ R N ×N ensures that attention is solely from the previous patch to the current one. It is mathematically represented as:EQUATIONThe attention score matrix undergoes element-wise multiplication with D to form the combination matrixEQUATIONFinally, the cell state ht is updated as:ht = C ⊗ V (8)Then, the embedding y from the second path is multiplied with ht , followed by a down-projection operation. This final output o is projected into a Z-dimensional embedding space using the projection matrix W D ∈ R 2Z×Z . Mathematically, this can be expressed as:EQUATION2) The Architecture of xLSTM-UNet: Fig. 1 showcases the xLSTM-UNet network architecture. xLSTM-UNet follows a conventional UNet-like structure. The input information first passes a convolution layer for initial down-sampling. Then several subsequent layers that are constructed using the aforementioned xLSTM building blocks to capture both local features and long-range dependencies form the main part of the encoder. Note that the xLSTM-UNet is designed with the goal of harnessing the best aspects of both U-Net and xLSTM for improved global comprehension in medical image understanding. Therefore, instead of only applying xLSTM in the compressed latent space after the down-sampling has finished, we hereby use the xLSTM in multiple layers in the encoder, in which each layer contains two successive Residual blocks with one plain convolution and an Instance Normalization (IN) and followed by a xLSTM block as in [31] . Specifically, the image features passing the residual blocks has a shape of (B, C, H, W, D), which is first flattened and transposed to (B, H × W × D, C), followed by a layer normalization, and then feed to the ViL block. Such the practice of involving xLSTM in multiple layers helps the feature extractions in multiple resolutions/perception fields, and this information extracted by xLSTM blocks is reshaped to (B, C, H, W, D) and concatenated to the layers in the decoding steps to facilitate the segmentation mask generation.After encoding, the decoder, comprising Residual blocks and transposed convolutions, concentrates on the meticulous recovery of detailed local information. Additionally, we inherit the skip connection from the U-Net architecture to interconnect the hierarchical features from the encoder to the decoder. The final decoder feature is fed into a 1 x 1 convolutional layer, coupled with a Softmax layer, to generate the ultimate segmentation probability map. Furthermore, following [23] in 2D segmentation, we also implemented a variant where the U-xLSTM block is exclusively utilized in the bottleneck, denoted as 'ours bot', while 'ours enc,' denotes the network that applies xLSTM block in all encoder blocks in all 2D segmentation tasks. In 3D segmentation, the xLSTM blocks is added in the bottleneck.To validate the effectiveness of our xLSTM-UNet, we utilized several representative medical image segmentation datasets, covering organ, instrument, and cell segmentation across different resolutions and modalities. This comprehensive evaluation demonstrates the applicability and superiority of xLSTM-UNet across diverse medical imaging scenarios.Abdomen MRI: We used the Abdomen MRI dataset from the MICCAI 2022 AMOS Challenge [4] for abdominal organ segmentation. The dataset, annotated by radiologists using MedSAM [2] and ITK-SNAP [35] , includes 60 MRI scans for training and 50 for testing, with 13 organ classes. For 2D tasks, images were cropped to 320x320 pixels, and for 3D tasks, patches were 48x160x224 pixels, following U-Mamba settings [23] .Endoscopy images: The endoscopy image dataset was sourced from the MICCAI 2017 EndoVis Challenge [36] , focusing on the segmentation of seven surgical instruments from endoscopic images. We used the official split of 1800 training and 1200 testing frames, with images resized to 384x640 pixels for nnU-Net.Microscopy images: The microscopy image dataset was obtained from the NeurIPS 2022 Cell Segmentation Challenge [6] , which focuses on cell segmentation in various microscopy images. We trained on 1000 images and tested on 101, converting the task to semantic segmentation as in [23] , with images resized to 512x512 pixels for nnU-Net.BraTS2023: The BraTS2023 dataset [37] - [39] comprises 1,251 3D brain MRI volumes. Each volume features four imaging modalities (T1, T1Gd, T2, and T2-FLAIR) and three segmentation targets: Whole Tumor (WT), Enhancing Tumor (ET), and Tumor Core (TC). For training, we use a random crop size of 128×128×128 to process the 3D data.The network is implemented based on UMamba [23] . The loss function used is the sum of Dice loss and cross-entropy loss. We employ the AdamW optimizer with a weight decay of 0.05. The learning rates were set to 0.005 for the Abdomen MRI dataset, 0.01 for the Endoscopy dataset, 0.007 for training xLSTM-UNet Bot on the Microscopy dataset, 0.0015 for training xLSTM-UNet Enc on the Microscopy dataset, and 0.01 for training on the BraTS2023 dataset. The batch sizes were set as follows: 2 for the 3D Abdomen MRI dataset, 30 for the 2D Abdomen MRI dataset, 2 for the Endoscopy dataset, 12 for the Microscopy dataset, and 4 for the BraTS2023 dataset. All networks were trained from scratch for 1000 epochs on a single NVIDIA A100 GPU. For more implementation details, please refer to our codebase.In 2D medical segmentation, to ensure a fair comparison, we follow the evaluation protocol in UMamba [23] . We selected two CNN-based segmentation networks (nnU-Net [4] and SegResNet [40] ) and two Transformer-based networks (UNETR [12] and SwinUNETR [11] ), as well as the UMamba itself, which has two variations: U-Mamba Bot and U-Mamba Enc. Similar to our configuration, U-Mamba Bot is applied only at the bottleneck, while U-Mamba Enc is used in each encoder. We used the Dice Similarity Coefficient (DSC) and Normalized Surface Distance (NSD) as evaluation metrics for the semantic segmentation tasks on the Abdomen MRI and Endoscopy datasets [41] . For cell segmentation on the Microscopy dataset, we employed the F1 score.In 3D medical segmentation, for the 3D Abdomen MRI dataset, the baseline methods and tasks remain consistent with those used in 2D medical segmentation. For the BraTS2023 dataset, to ensure a fair comparison, we follow the evaluation protocol outlined in SegMamba [29] . We use the same baseline methods, including three CNN-based methods (SegresNet [40] , UX-Net [42] , MedNeXt [43] ), three transformer-based methods (UNETR [12] , SwinUNETR [11] , and SwinUNETR V2 [44] ), and the Mamba-based method SegMamba itself. Following previous evaluation protocols, Dice and HD95 were used as evaluation metrics.Table I presents the segmentation performance of various methods on the Abdomen MRI 2D, Endoscopy, and Microscopy datasets for 2D image segmentation task. Our proposed xLSTM-UNet outperforms all baseline methods and achieves state-of-the-art (SOTA).Notably, both variations of xLSTM-UNet show superior performance across all datasets. Specifically, xLSTM-UNet Enc demonstrates the highest performance with a DSC of 0.7747 and an NSD of 0.8374 on the Abdomen MRI 2D dataset, outperforming the previous state-of-the-art (SOTA) model, U-Mamba by a significant margin. Additionally, xLSTM-UNet Bot achieves DSC and NSD scores of 0.7636 and 0.8322, respectively, surpassing the similarly structured U-Mamba Bot. Similarly, on the Endoscopy dataset, both xLSTM-UNet Bot and xLSTM-UNet Enc achieve the best DSC and NSD scores of 0.6843 and 0.7001, respectively. For the Microscopy dataset, xLSTM-UNet Enc and xLSTM-UNet Bot achieve F1 scores of 0.6036 and 0.5818, respectively, both surpassing the previous SOTA results, indicating their robustness in cell segmentation tasks.The visualized segmentation examples of 2D medical images further illustrate the effectiveness of xLSTM-UNet. As shown in Figure 2 , xLSTM-UNet is more robust to heterogeneous appearances and exhibits fewer segmentation outliers compared to other models. This visual evidence underscores the quantitative results, highlighting the superior performance and reliability of xLSTM-UNet in diverse medical image segmentation tasks.3D medical image segmentation is generally more challenging compared to its 2D counterpart, as it involves processing a larger amount of information. The increased dimensionality leads to a dramatic surge in computational complexity, with the resolution increase causing a cubic rise in the number of computations. Accurate spatial relationship modeling is also essential for achieving satisfactory segmentation results. These factors make the xLSTM-based building blocks, with their computational efficiency, well-suited for this task.We conducted evaluations on the 3D segmentation dataset in BraTS2023 and Abdomen MRI 3D. Table II shows the performance comparison on the BraTS2023 dataset, including metrics for whole tumor (WT), tumor core (TC), and enhancing tumor (ET) regions. Our proposed method demonstrates superior performance across all evaluated metrics, including Dice and HD95, compared to other baseline methods such as SegresNet, UX-Net, MedNeXt, UNETR, SwinUNETR, SwinUNETR-V2, and SegMamba. Specifically, our method achieves the highest average Dice score of 91.80, highlighting its effectiveness in accurately segmenting brain tumor regions. Table III shows the performance comparison of the Organs in the Abdomen MRI 3D dataset. Our proposed method, xLSTM-UNet Bot, achieves the highest Dice score of 0.8483 and the best NSD score of 0.9153, surpassing other methods such as nnU-Net, SegResNet, UNETR, SwinUNETR, and U-Mamba Bot. This demonstrates the robustness and accuracy of our approach in segmenting abdominal organs in MRI images.The ability of xLSTM to effectively model semantic information in complex spatial domains has been a key factor in our success. The superior experimental results also underscore the suitability of xLSTM-based building blocks for tackling the challenge of semantic segmentation in complex imaging applications.To further validate the improvements introduced by our newly proposed modules, we conducted a series of ablation studies. The process began by using a baseline model where all blocks in the U-Net architecture were fully convolutional layers. Starting from this baseline, we incrementally replaced the convolutional blocks with xLSTM modules. The results of these ablation studies are presented in Table IV , illustrating the effectiveness of each module on overall model performance.In our ablation experiments, replacing even a single convolutional layer with an xLSTM module resulted in significant improvements, particularly on the Dataset704 Endovis17 dataset, where DCS and NSD metrics increased by 18.9% and 19.1%, respectively. As more xLSTM layers were added, accuracy consistently improved across all datasets. These results demonstrate that integrating xLSTM blocks into the U-Net architecture effectively enhances segmentation accuracy in medical imaging.The efficiency analysis of different methods, as shown in Table V , highlights the performance and computational complexity of models. The scatter plot (Figure 3 ) visualizes the relationship between FLOPs and F1 scores for various models for the task of cell segmentation in microscopy image, with the circle size representing the number of parameters. Our models, Ours bot and Ours enc, demonstrate a significant FLOPs, which is significantly more efficient than nnFormer and other high-parameter models like U-Mamba Bot and U-Mamba Enc. This analysis underscores the superior balance achieved by our models between computational efficiency and segmentation accuracy, making them particularly suitable for practical applications where both performance and resource consumption are critical. Our models' lower parameter counts and FLOPs, combined with high F1 scores, clearly demonstrate their superiority over existing methods, affirming the effectiveness of our approach in optimizing model architecture for both efficiency and accuracy. We also evaluated the inference speed of the models as part of the ablation studies, focusing on the 2D segmentation task from the Dataset702 AbdomenMR dataset. The results, shown in the last column of Table IV, provide the average inference time per nii.gz file, highlighting the trade-off between perfor-mance and computation time. While the baseline FCN model had the fastest inference speed, the xLSTM-based models progressively increased inference time with each additional layer. However, the improvements in segmentation accuracy justify the added computation time for tasks requiring high precision.This study demonstrates that xLSTM, a model with linear computational complexity, can be an effective component in image segmentation networks. Our experimental results clearly show that xLSTM-UNet outperforms Mamba-based counterparts, underscoring the promising future of xLSTM. Given the recent huge interest in Mamba in academia, we believe that it is important to also recognize and investigate the potential of xLSTM, which has shown remarkable efficacy in this domain.Meanwhile, medical image segmentation is inherently challenging. General image segmentation foundation models, such as Segment Anything, often fail when applied to medical images [45] - [47] . Furthermore, medical imaging datasets are typically small. In this study, the datasets used were limited in size, which restricts our ability to explore the effects of different network scales and dataset sizes on segmentation outcomes. Investigating whether xLSTM-driven image algorithms adhere to scaling laws remains an interesting question for future research.Currently, xLSTM lacks dedicated optimization for hardware such as NVIDIA GPUs, which presents an opportunity for the community to contribute. Collaborative efforts are essential to optimize xLSTM for various vision tasks, leveraging its full potential. By releasing our code, we aim to encourage and facilitate further research and development, enabling the community to build on our initial findings and drive progress in this area. This research represents an initial exploration into the application of xLSTM in medical image segmentation. There are many peaks to climb and numerous scenarios to test in this field. We hope that our comprehensive experiments and tests will demonstrate the significant potential of xLSTM in practical applications, encouraging scholars to continue exploring this promising model. With further development and optimization, we envision xLSTM achieving success comparable to that of Mamba and even Transformers, becoming a cornerstone in image segmentation and beyond.In this paper, we introduce xLSTM-UNet, the first U-Net architecture enhanced with Extended Long-short-memory (xLSTM) / ViL for both 2D and 3D medical image segmentation tasks. Through extensive experiments across a variety of medical imaging scenarios-including abdominal MRI, endoscopy, microscopy, and brain MRI-we have demonstrated that xLSTM-UNet significantly outperforms existing CNN-based and Transformer-based methods, as well as its Mamba-based counterparts. These findings underscore the effectiveness of xLSTM in handling complex segmentation tasks, particularly in the challenging domain of 3D medical image segmentation.Our results show that the xLSTM-based architecture can achieve state-of-the-art (SOTA) performance with reasonable memory and computation cost, offering enhanced accuracy and efficiency. This marks a significant advancement and future potential of xLSTM or similar building blocks, not only in the field of medical image segmentation, but also with potential applications extending beyond healthcare. We believe that this new findings will be interesting to the community and can inspire future researches. 
 
 Appendix:  



Reviews: review: title: title: xLSTM-UNet can be an Effective Backbone for 2D & 3D Biomedical Image Segmentation Better than its Mamba Counterparts
paper_summary: paper_summary: The paper introduces the xLST-Unet architecture for biomedical image segmentation. The method is successfully applied to three different biomedical datasets, and it has been found to give better results.
strength_weakness: strength_weakness: - There are no ablation studies.
- No comparisons of the number of parameters per method.
- No speed comparisons.
questions: questions: Please address the weaknesses by answering:
- Any ablation studies?
- Compare the number of parameters in each method.
- Add speed comparisons.
title: title: Review for submission 427
paper_summary: paper_summary: The paper proposed xLSTM-UNet, which integrated xLSTM with UNet for medical image segmentation. Experiments on both 2D and 3D dataset shows improvement.
strength_weakness: strength_weakness: The model is a combinable of designed CNN decoder with Vision-LSTM encoder. There lacks of ablation study for decoder design. 

Minor comments: it would be better to include some arrows to emphasize the difference in Fig.2.
questions: questions: See weekness.
title: title: An image segmentation model that combines xLSTM with U-Net that outperforms existing methods in biomedical image segmentation tasks.
paper_summary: paper_summary: The paper introduces xLSTM-UNet, a novel image segmentation architecture that integrates Vision-LSTM (xLSTM) with the traditional U-Net structure. This model aims to address the limitations of Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) in handling long-range dependencies and computational overhead in biomedical image segmentation. xLSTM-UNet combines the local feature extraction capabilities of convolutional layers with the long-range dependency management of xLSTM. The proposed architecture outperforms existing CNN-based, Transformer-based, and Mamba-based models in various biomedical segmentation tasks, including abdominal MRI, endoscopy, and microscopy images. The study highlights xLSTM-UNet’s potential in advancing 2D and 3D medical image analysis and suggests its broader applicability across multiple domains. The paper’s experimental results and comprehensive validations demonstrate the model's superior performance and efficiency.
strength_weakness: strength_weakness: It lacks detailed discussion on potential limitations or drawbacks of the proposed model. While the paper demonstrates superior performance of xLSTM-UNet, it does not extensively address the increased computational complexity and resource demands that may accompany the integration of xLSTM into the U-Net architecture.
questions: questions: No question