UU-Mamba: Uncertainty-aware U-Mamba for Cardiovascular Segmentation

The development of deep learning techniques, particularly Convolutional Neural Networks (CNNs) [7], has significantly advanced cardiovascular segmentation [28, 29], driven by comprehensive datasets like ACDC [2], ImageCAS [9], and Aorta [10, 11]. These datasets address a range of cardiovascular structures, including heart chambers, coronary arteries, and aortic branches, and present challenges related to anatomical variability and disease-specific characteristics. Multi-modality cardiac imaging segmentation, which utilizes imaging modalities like Positron Emission Tomography (PET), Single Photon Emission Computed Tomography (SPECT), MRI, and CT, aims to precisely segment anatomical structures and pathological regions. However, inherent challenges such as phase alignment, resolution, and image quality imbalances complicate the process. Traditional methods, including registration-based segmentation with multi-atlas approaches, and fusion-based segmentation techniques, address these issues by combining information across modalities [30, 31, 32, 33], but these methods are computationally expensive and often require large datasets [34, 35]. Hybrid approaches, combining both traditional and deep learning methods, are emerging to improve the robustness and clinical applicability of cardiovascular segmentation [36].

Recent advances in deep learning have improved segmentation accuracy by enabling complex representations of cardiovascular structures. For instance, CNNs like U-Net [37, 38] and its variants have been effective for cardiovascular segmentation tasks, particularly on the ACDC dataset [2]. However, these methods are prone to overfitting and issues such as class imbalance, which complicates their application across diverse datasets like ImageCAS [9] and Aorta [10, 11]. Zeng et al. [9] tackled coronary artery segmentation challenges in ImageCAS, utilizing multi-scale feature extraction to capture finer details, but the method still struggles with arteries exhibiting atypical morphology, requiring precise hyperparameter tuning.

The Aorta dataset [10, 11] focuses on segmenting the aorta and its branches, where varying aortic diameters and pathologies like aneurysms present additional complexities. Imran et al. [10] proposed the CIS-UNet, a hybrid model incorporating Context-Aware Shifted Window Self-Attention mechanisms to address spatial variability in aortic structures, though its performance remains sensitive to imaging quality and resolution. To further improve segmentation, some studies have integrated CNNs with attention mechanisms and multi-scale processing. For example, Hu et al. [39] adapted the Segment Anything Model to medical images, incorporating multi-scale processing and CNN heads for enhanced segmentation across various datasets.

A key development in segmentation accuracy came from Isensee et al. [37], who combined U-Net and V-Net architectures in the nnUNet framework [37]. While this method has proven effective, it heavily depends on extensive annotated datasets, limiting its applicability in scenarios with scarce labeled data—a frequent issue in cardiovascular imaging. To address this, transfer learning has been employed, as demonstrated by Chen et al. [33], who pre-trained models on the ACDC dataset and fine-tuned them on smaller, less-annotated datasets. Despite improving performance, transfer learning remains susceptible to domain shift issues, particularly when applied to datasets like ImageCAS and Aorta with differing data distributions.

Standard segmentation methods often rely on basic loss functions like Cross-Entropy loss, which struggle to effectively manage class imbalance or capture the finer details necessary for accurate segmentation. Recent research highlights the need to optimize for flatter minima in the loss landscape to enhance model generalization. Caldarola et al. [40] demonstrated that such optimization improves model robustness and generalization, especially when dealing with noisy or ambiguous data, which is essential for reliable cardiovascular segmentation across diverse clinical scenarios.

The Mamba architecture [41] represents a substantial advancement in medical image segmentation by integrating the capabilities of Vision Transformers (ViTs) [42] and Convolutional Neural Networks (CNNs) [7]. It is essential to achieve precision in medical imaging duties by integrating global contextual information and managing long sequences [43]. U-Mamba [12] expands the conventional U-Net framework [37, 38] by integrating attention mechanisms and multi-scale processing, thereby improving the accuracy and robustness of segmentation. This is achieved by building upon this foundation. This method enables the model to concentrate on pertinent details within intricate anatomical structures and efficiently process information across a range of scales, from a broad contextual understanding to precise low-level details. Furthermore, U-Mamba incorporates deep supervision, which expedites training and enhances convergence, rendering it both efficient and dependable for clinical applications that require rapid and precise image processing. The Mamba architecture’s adaptability is further underscored by specialized variants such as Weak-Mamba-UNet [44], which are able to handle intricate scenarios with improved performance and excel in scribble-based segmentation tasks. In conclusion, models that are based on U-Mamba exhibit superior segmentation performance in a variety of medical applications, such as histopathological imaging and cardiac MRI.

The Mamba architecture’s computational efficacy is one of its major advantages, as it is the result of the strategic integration of CNNs [7] and ViTs [42]. This hybrid design capitalizes on the advantages of both architectures: ViTs’ global attention mechanisms facilitate the efficient management of intricate spatial relationships and long-range dependencies, while CNNs are adept at extracting local features with minimal computational overhead. Mamba reduces the computational burden that is typically associated with pure transformer-based models, which are resource-intensive due to their quadratic complexity in relation to the duration of the input sequence, by combining these two approaches. The hierarchical design of Mamba further improves its computational efficacy. The model is capable of processing images at various scales, enabling it to capture critical features at lower resolutions and subsequently refine these details at higher resolutions. By concentrating processing capacity in the most critical areas, this multi-scale approach minimizes the total number of computations, as opposed to implementing a uniform computation across all pixels.

Furthermore, Mamba employs efficient attention mechanisms and sparse operations to minimize the number of operations necessary for each layer. This optimization is especially advantageous in the segmentation of medical images, as the computational costs associated with high-resolution images can be prohibitive rapidly. The computational efficiency of U-Mamba is also enhanced by the deep supervision strategy, which facilitates quicker convergence during the training process. This implies that the computational resources required are further reduced by the necessity of fewer training epochs to attain high performance. Overall, the Mamba architecture optimizes the strengths of both CNNs [7] and ViTs [42] while minimizing computational demands, thereby achieving a balance between precision and efficiency. Thus, it is particularly well-suited for clinical applications that prioritize both accuracy and rapidity.

Sharpness-Aware Minimization (SAM) [13] has gained substantial attention for its capacity to enhance the generalization of deep learning models by optimizing both the training loss and the sharpness of the loss landscape. This method has been particularly advantageous in the field of medical image segmentation, where the ability to accurately detect boundaries and generalize them across a variety of datasets is essential.

SAM has exhibited significant enhancements in model robustness and accuracy within the context of medical image segmentation. For instance, Mariam et al. [45] implemented SAM in the RF-UNet model to segment retinal vessels. In addition to enhanced metrics such as accuracy, sensitivity, and specificity, their experiments on the DRIVE dataset demonstrated a substantial decrease in both training and validation losses. This emphasizes SAM’s capacity to improve generalization and mitigate overfitting in retinal segmentation tasks.

Additional research has investigated sophisticated variations of SAM for the purpose of medical segmentation. In the context of breast ultrasound image segmentation, Hassan et al. [46] assessed a variety of sharpness-based optimizers, such as SAM. Their results suggested that SAM consistently enhanced generalization across various models, surpassing other sharpness-based optimizers such as Adaptive SAM.

Random Sharpness-Aware Minimization (RSAM) by Liu et al. [47] is another innovation in SAM that incorporates randomness to enhance the efficiency and stability of SAM optimization. RSAM’s potential for improved generalization across domains suggests that it has the potential to be effective in medical imaging, despite the fact that it has not yet been explicitly applied to medical segmentation. Li et al. [48] also proposed Friendly SAM (FSAM), which further optimizes SAM for improved performance in diverse and complex data environments. This direction could be highly relevant to medical segmentation.

Collectively, these studies demonstrate that SAM [13] and its variants are essential for the advancement of medical image segmentation, particularly in tasks that necessitate high precision and robustness, such as retinal vessel extraction, breast ultrasound segmentation, and other medical imaging challenges. The integration of SAM into medical segmentation models is expected to result in even greater improvements in generalization and performance as SAM continues to develop. The notion that SAM is the most dependable sharpness-based method in medical image analysis was further substantiated by this study.

The U-Mamba network is designed to improve the accuracy of medical image segmentation and improve global context comprehension by combining the assets of the Mamba block [41] and U-Net [37, 38]. The Mamba block, which is specifically engineered for Selective Structured State Space Sequence Models (S6), is particularly well-suited for medical imaging duties due to its exceptional ability to manage long-range dependencies and sequential information.

State Space Models (SSM) [49] describe systems in terms of their internal states and observations over time, thereby facilitating effective sequence modeling through these underlying states. The fundamental form is denoted as follows: $\mathbf{x}_{t}$ is the input state vector, $\mathbf{u}_{t}$ is the control input, $\mathbf{w}_{t}$ is the process noise, $\mathbf{A}$ is the state transition matrix, and $\mathbf{B}$ is the control input matrix.

For observation $\mathbf{y}_{t}$, calculated using the observation matrix $\mathbf{C}$, feedthrough matrix $\mathbf{D}$, and observation noise $\mathbf{v}_{t}$, the formula is:

The S6 architecture advances traditional state space models by integrating selective attention mechanisms and structured parameterization. The selective attention mechanism can be represented as:

where $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$ are the query, key, and value matrices that are derived from the state vector $\mathbf{x}_{t}$, and $d_{k}$ is the dimension of the key vectors. This mechanism enables the model to effectively capture intricate dependencies by concentrating on pertinent components of the input sequence.

The integration of S6 into the Mamba block is especially crucial for sequential medical image processing tasks, such as cardiac MRI segmentation, which require the capture of temporal dynamics and structure [41]. The method, on the other hand, is exclusively concerned with per-image segmentation, which involves the application of the state transition and observation matrices ($\mathbf{A}$, $\mathbf{C}$, etc.) to individual images. Each image is treated independently.

U-Mamba capitalizes on Mamba’s linear scaling advantage to improve CNNs’ capacity to simulate long-range dependencies, all while circumventing the high computational costs associated with self-attention mechanisms employed in Transformers [50] such as ViT [42] and SwinTransformer [51]. The U-Mamba block, which is comprised of two sequential residual blocks followed by a Mamba block, is depicted in Figure 2.

Additionally, each block includes Leaky ReLU activation, Instance Normalization, and convolutional layers. Mamba blocks with two parallel branches: one with an SSM layer and one without, flatten, transpose, normalize, and process image features. The Hadamard product is then employed to merge these features, which are subsequently projected back to their original shape and transposed.

An encoder with these blocks is included in the complete U-Mamba network architecture to capture both local features and long-range dependencies, while a decoder composed of residual blocks and transposed convolutions is used to recover detailed local information and resolution. Skip connections are used to connect hierarchical features from the encoder to the decoder. The final decoder output is processed through a 1$\times$1$\times$1 convolutional layer and a Softmax layer to generate the final segmentation probability map.

Introducing uncertainty into loss functions entails allocating weights to distinct components of the loss according to the estimated uncertainty for each data point [52, 53]. This method allows the model to concentrate on learning from more dependable instances while simultaneously reducing the impact of potentially erroneous or ambiguous data. Kendall and Gal introduced the concept of adjusting loss functions by utilizing homoscedastic and heteroscedastic uncertainty [54]. Heteroscedastic uncertainty fluctuates between instances, while homoscedastic uncertainty remains constant across all data points. The model can improve its resilience and precision by focusing on confident predictions and reducing the impact of equivocal ones by adapting its learning process to capitalize on these uncertainties. This optimization enhances overall performance and improves the training process across diverse datasets [55, 56, 57].

To further boost segmentation accuracy, we employ an uncertainty-aware loss function that combines region-based, distribution-based, and pixel-based losses, capitalizing on their complementary strengths:

1. 1.

   Dice loss [15]: This region-based metric emphasizes the overlap between predicted and ground truth areas, ensuring accurate preservation of shape and boundary details in segmented regions.

2. 2.

   Cross-Entropy (CE) loss [16]: This distribution-based loss ensures precise categorization of individual pixels, thereby improving classification accuracy.

3. 3.

   Focal loss [17]: This pixel-level loss addresses class imbalance by assigning greater importance to challenging instances, enhancing the model’s ability to manage complex scenarios [58, 59, 60].

Let $p_{i}$ denote the predicted probability and $g_{i}$ the ground truth label, with the predicted segmentation and the corresponding ground truth mask. The Dice Similarity Coefficient (DSC) is a metric that quantifies the degree of overlap between the predicted segmentation and the ground truth. It is defined as follows:

DSC values range from 0 to 1, with 1 signifying complete overlap between the prediction and the ground truth and 0 indicating no overlap.

The Dice loss is defined as: in order to integrate this metric into a loss function for training segmentation models.

The Dice loss is designed to minimize the discrepancy between the predicted segmentation and the ground truth by optimizing the DSC. The model is trained to generate segmentations that exhibit a greater overlap with the ground truth by minimizing the Dice loss, thereby enhancing the accuracy of the segmentation.

The standard entropy formula is employed to determine the Cross-Entropy (CE) loss:

To address class imbalance, we utilize the Focal loss, which focuses more on difficult-to-classify samples:

where $\gamma$ is a focusing parameter default to 2.

The uncertainty-aware loss $\mathcal{L}_{UA}$ is defined by combining these loss components within an uncertainty-aware framework:

in which $M$ is the number of individual loss components, $\mathcal{L}_{m}$ represents each loss component (such as Dice, CE, and Focal loss), and $\sigma_{m}$ are learnable parameters that modify the contribution of each loss component based on the estimated uncertainty. To reduce the aggregate loss, these parameters are optimized during the training process.

By integrating uncertainty into the loss calculation, the model is able to dynamically modify the weights of each separate loss component. While mitigating the effects of class imbalance, this method strikes a balance between global and local accuracy. As an outcome, the model becomes more resilient to ambiguous or chaotic data, resulting in an overall improvement in segmentation performance.

To improve the U-Mamba model’s generalizaiton in segmenting cardiovascular images, such as those in the ACDC [2], ImageCAS [9], and Aorta dataset [10, 11], our methodology employs Sharpness-Aware Minimization (SAM) optimization [13]. The model’s generalizability is enhanced by the flattening of the loss landscape, implemented by SAM optimization. Despite the fact that conventional optimization techniques are designed to identify the lowest points in the loss landscape, these points are frequently precipitous, which results in inadequate generalization to new data. SAM, on the other hand, identifies gentler minima—regions in the parameter space where the model’s performance remains consistent and is less susceptible to perturbations.

SAM is employed due to its effective reduction of overfitting, a common issue in medical image segmentation. Performance on unseen data may be impaired by the narrow valleys in the loss landscape that are a common consequence of conventional optimization methods. SAM, in contrast, concentrates on the identification of flattened minima, which are linked to enhanced generalization. When dealing with complex and diverse datasets, this is especially beneficial, as the variability in cardiac MRI images can exacerbate overfitting if not properly managed.

Optimization of SAM is accomplished through a two-step iterative procedure. Parameters are initially adjusted to optimize loss for each mini-batch. After this, the model parameters are adjusted to reduce the maximum loss. This perturbation is designed to identify model parameters that are located in flatter regions of the loss landscape, which are typically associated with improved generalization and increased robustness to minor changes in the input data.

The model parameters shall be denoted by $\theta$, the loss function by $\mathcal{L}$, and the training dataset by $\mathcal{D}$. To investigate the loss landscape within a neighborhood around $\theta$ defined by the norm constraint $\|\epsilon\|_{2}\leq\rho$, the perturbation $\epsilon$ is introduced, with $\rho$ determining the size of this neighborhood.

Mathematically, the SAM optimization is expressed as:

The perturbation $\epsilon$ within the $\|\epsilon\|_{2}\leq\rho$ constraint is determined in the initial phase to maximize the loss. This method identifies the worst-case direction in the local neighborhood of $\theta$. Furthermore, this guarantees that the model parameters are directed toward regions of the loss landscape that are not precipitous. In order to enhance the parameters’ resilience to perturbations, the model parameters $\theta$ are modified in the second phase to reduce the loss at the worst-case perturbed location.

By applying these two stages iteratively, SAM directs the optimizer toward parameter configurations that are resilient to perturbations, thereby improving generalization and overall performance. The model is able to identify flattened minima in the loss landscape as a result of this approach, which results in improved generalization [13].

We conduct a comparison of UU-Mamba with five of the state-of-the-art segmentation models—TransUNet [61], Swin-Unet [62], nnUNet [37], nnFormer [63], and U-Mamba [12]—on the ACDC dataset [2]. Transformer-based networks are TransUNet and Swin-Unet, while nnUNet and nnFormer employ CNN-based architectures. U-Mamba is a hybrid architecture that combines components from both Transformer-based and CNN-based networks.

Using the Dice Similarity Coefficient (DSC) as the evaluation metric, we conducted a quantitative evaluation of UU-Mamba against these five 3D heart segmentation models on the ACDC dataset [2]. The segmentation results for the compared algorithms are depicted in Figure 3 on a few images from the ACDC dataset. Table 1 provides the average DSC scores across all regions, as well as the DSC scores for each model in three cardiac regions: the right ventricle (RV), myocardium (Myo), and left ventricle (LV). The scores that demonstrate the greatest performance are indicated in italics.

Our UU-Mamba model surpasses all other methods, attaining the highest overall performance with an average DSC of 92.79%. The DSC scores for each region are as follows: 92.41% for RV, 90.90% for Myo, and 95.04% for LV for each region. Our model’s flexibility and efficacy are demonstrated by its exceptional ability to accurately segment the right ventricle and myocardium. Despite a minor decrease in DSC for the left ventricle in comparison to other models, this is counterbalanced by the highest overall average DSC and the superior scores in other regions, as illustrated in Table 1.

In contrast, TransUNet attains an average DSC of 89.71%, with a relatively lower score for Myo. The average DSC of Swin-Unet is 90.00%, with the maximum DSC for LV and lower performance for RV. The nnUNet model achieves an average DSC of 91.61%, with significant enhancements in Myo segmentation. Providing robust performance in all regions, particularly the myocardium, the nnFormer model obtains an average DSC of 92.06%. U-Mamba’s average DSC of 92.22% indicates substantial improvements in the RV and Myo regions.

The superior performance of our UU-Mamba model in comparison to the other existing models is emphasized by this quantitative evaluation. Our method exhibits the potential to improve the accuracy of 3D heart segmentation in medical imaging, as it has the highest average DSC and notably strong segmentation in the right ventricle and myocardium. The effectiveness of our approach is underscored by the enhancements it achieves over models such as U-Mamba, nnFormer, and nnUNet. This approach employs sophisticated loss functions and optimization techniques to capitalize on both global and local features.

We perform an ablation study to investigate the impact of integrating Sharpness-Aware Minimization (SAM) optimization [13] into our UU-Mamba model, alongside traditional and uncertainty-aware loss functions. This study spans three datasets: ACDC [2], Aorta [10, 11], and ImageCAS [9], and employs 3D loss surface visualizations [27] using ParaView [64, 65, 66] to demonstrate the smoother loss landscapes indicative of model robustness and improved generalization.

We perform experiments to evaluate each method on the Mean Squared Error (MSE) of the DSC scores to assess their robustness quantitatively. The MSE is calculated as shown in Eq. (10). Results are shown in the Tables 2, 3, and 4. These results show that the uncertainty-aware loss reduces the MSE compared to the standard CE loss, reflecting its ability to better address the variability and uncertainty in the data. The inclusion of SAM optimization significantly decreases the MSE, achieving the lowest error value. This reduction in MSE highlights SAM’s role in minimizing errors and producing more accurate segmentation maps.

Figures 3, 4, and 5 show the MSE between the output segmentation and the ground truth for each method, providing a visual comparison of the segmentation quality. These visualizations complement the quantitative results by illustrating the error distribution and highlighting areas where the SAM optimization and uncertainty-aware loss contribute to more accurate and consistent segmentation outcomes.