PaperlayerView exact revision on arXiv

Source-linked AI summary

Semi-Mamba-UNet: Pixel-Level Contrastive and Pixel-Level Cross-Supervised Visual Mamba-based UNet for Semi-Supervised Medical Image Segmentation

Chao Ma, Ziyang Wang

arXiv:2402.07245v3eess.IVcs.CV

TL;DR

Medical image segmentation faces limited access to expert annotations, while Mamba-based segmentation under limited annotation settings remains an area for exploration. Semi-Mamba-UNet combines a Mamba-based UNet with a CNN-based UNet in semi-supervised pixel-level cross-supervision and contrastive learning. The authors report state-of-the-art performance on two public benchmarks and a tendency to predict test images with higher IoU than baseline methods.

  • Problem

    Training advanced segmentation networks with limited annotations remains an area for further exploration, while annotations are costly and difficult to obtain.

    Source: §1 · p. 3, §2.2 · p. 5

  • Method

    Semi-Mamba-UNet combines Mamba- and CNN-based segmentation networks in semi-supervised pixel-level cross-supervision and uses projectors for pixel-level contrastive learning.

    Source: §1 · p. 4, §3 · p. 7

  • Results

    The authors report state-of-the-art performance on two public benchmark datasets and that Semi-Mamba-UNet was more likely than baseline methods to predict test images with high IoU scores.

    Source: §1 · p. 4, §4 · p. 17

  • Takeaways & Limitations

    The evaluations demonstrate the effectiveness of Semi-Mamba-UNet for the studied semi-supervised medical image segmentation settings.

    Source: §5 · p. 19

    The contrastive-learning procedure treats labelled data as unlabelled data to expand the dataset during consistency regularisation, differing from conventional SSL strategies.

    Source: §3.3 · p. 9

Abstract

from arXiv · show

Medical image segmentation is essential in diagnostics, treatment planning, and healthcare, with deep learning offering promising advancements. Notably, the convolutional neural network (CNN) excels in capturing local image features, whereas the Vision Transformer (ViT) adeptly models long-range dependencies through multi-head self-attention mechanisms. Despite their strengths, both the CNN and ViT face challenges in efficiently processing long-range dependencies in medical images, often requiring substantial computational resources. This issue, combined with the high cost and limited availability of expert annotations, poses significant obstacles to achieving precise segmentation. To address these challenges, this study introduces Semi-Mamba-UNet, which integrates a purely visual Mamba-based U-shaped encoder-decoder architecture with a conventional CNN-based UNet into a semi-supervised learning (SSL) framework. This innovative SSL approach leverages both networks to generate pseudo-labels and cross-supervise one another at the pixel level simultaneously, drawing inspiration from consistency regularisation techniques. Furthermore, we introduce a self-supervised pixel-level contrastive learning strategy that employs a pair of projectors to enhance the feature learning capabilities further, especially on unlabelled data. Semi-Mamba-UNet was comprehensively evaluated on two publicly available segmentation dataset and compared with seven other SSL frameworks with both CNN- or ViT-based UNet as the backbone network, highlighting the superior performance of the proposed method. The source code of Semi-Mamba-Unet, all baseline SSL frameworks, the CNN- and ViT-based networks, and the two corresponding datasets are made publicly accessible.

1. Introduction

The introduction motivates Semi-Mamba-UNet as a Mamba-based approach to medical image segmentation with limited annotations. It combines pixel-level contrastive learning and cross-supervision within semi-supervised learning, and is evaluated against CNN- and ViT-based baselines.

1. Introduction

  • Semi-Mamba-UNet explores Visual Mamba in a U-shaped segmentation network and, to the authors’ knowledge, is the first study to apply Mamba to medical image segmentation with limited annotations.Mamba is described as capturing global context with efficient computational costs.Source: §1 · pp. 3–4
  • The framework combines semi-supervised learning with pixel-level cross-supervision, using one network’s pseudo-labels to train the other, and pixel-level contrastive learning with a pair of projectors.The contrastive strategy uses labelled and unlabelled data.Source: §1 · p. 4
  • The evaluation compares the Mamba-based approach with CNN-based UNet and ViT-based Swin-UNet across semi-supervised learning frameworks on two public benchmark datasets.The authors report state-of-the-art performance on the two benchmarks.Source: §1 · p. 4
  • The paper makes the Semi-Mamba-UNet source code and baseline methods publicly available.The introduction also describes public availability of the CNN- and ViT-based networks.Source: §1 · p. 4

2. Related Work

Related work traces Mamba-based segmentation networks and reviews semi-supervised and self-supervised approaches for learning with limited annotations.

2.1. Mamba in Medical Image Segmentation

  • Mamba-based medical segmentation includes UNet-style and non-UNet 2D methods as well as 3D methods, while training these networks with limited annotations remains an area for further exploration.The subsection situates Mamba among earlier CNN, Transformer, and ViT approaches to medical image segmentation.Source: §2.1 · p. 5

2.2. Medical Image Segmentation with Limited Annotations

  • Limited-annotation segmentation work uses consistency regularisation, student–teacher pseudo-labeling, and adversarial learning, while contrastive learning is explored to improve feature learning.The paper proposes combining SSL and self-supervised learning for Mamba-based networks.Source: §2.2 · pp. 5–6

3. Methodology

Semi-Mamba-UNet combines Mamba-based and conventional UNet segmentation networks, training them with labeled supervision, unlabeled cross-supervision, and pixel-level contrastive learning.

3. Methodology

  • The framework combines labeled supervision with unlabeled-data pseudo-label cross-supervision and self-supervised contrastive learning across Mamba-based UNet and conventional UNet.The total objective includes supervision, semi-supervision, and contrastive losses.Source: §3 · pp. 6–7

3.1. Mamba-UNet

  • Mamba-UNet uses two Visual Mamba blocks in a U-shaped encoder–decoder; Visual Mamba adds cross-scan and convolutional operations.UNet uses two-layer CNN blocks, while Swin-UNet uses two Swin Transformer blocks.Source: §3.1 · pp. 7–8, Figure 3 · p. 8

3.2. Pixel-Level Cross-Supervised Learning

  • Mamba-UNet and UNet exchange predictions as pseudo-labels on unlabeled images, directly cross-supervising one another.Each network’s prediction can serve as a pseudo-label to retrain the other network.Source: §3 · pp. 6–7, §3.2 · p. 8

3.3. Pixel-Level Contrastive Learning

  • Paired projectors extract pixel-level features; adaptive average pooling filters background pixels, while channel-wise L2 regularization sparsifies features and targets perturbation resistance.The method is motivated by cardiac data with many background pixels and focuses contrastive learning on pixel variability across samples.Source: §3.3 · p. 9
  • Inter-network feature consistency uses mean squared error, and labeled images are also treated as unlabeled during consistency regularization.The authors describe this as expanding the dataset for that consistency-training process.Source: §3.3 · p. 9

4. Experiments and Results

Semi-Mamba-UNet was evaluated against SSL baselines on cardiac and prostate MRI with limited labelled cases. Its test-set IoU distributions favored higher scores, and backbone ablations reported significant differences between network settings.

Datasets

  • Evaluation used ACDC cardiac MRI (100 cases, four classes) and PROMISE12 prostate MRI (50 cases, binary segmentation), with three or seven cardiac and eight or 12 prostate training cases labelled.Images were resized to 224 × 224; the same randomly selected labelled and unlabelled cases were used for Semi-Mamba-UNet and all baselines.Source: §4 · p. 10

Implementation Details

  • Training ran for 30,000 iterations with batch size 16 and stochastic gradient descent; reported runtimes were about 5–8 hours for ACDC and 2–5 hours for PROMISE12.The setup used an RTX 3090 GPU and evaluated validation performance every 200 iterations.Source: §4 · p. 10

Baseline Segmentation Networks

  • Experiments compared SSL frameworks using UNet and Swin-UNet backbones, with Semi-Mamba-UNet evaluated alongside them.The baseline SSL methods included Mean Teacher, DAN, ICT, ADVENT, UAMT, and deep co-training.Source: §4 · pp. 10, 15

Evaluation Metrics

  • Evaluation covered Dice, accuracy, precision, sensitivity, specificity, 95% Hausdorff distance, and average surface distance; higher similarity scores and lower difference scores are preferable.For four-class segmentation, the reported difference measures are averaged across classes.Source: §4 · pp. 15, 17

Evaluation on Each Image of Test Set

  • Semi-Mamba-UNet was more likely than baseline methods to predict test images with high IoU scores, based on per-image distributions across the two datasets.The cardiac histograms cover three and seven labelled cases; the section also reports histograms for both datasets.Source: §4 · p. 17, Figure 8 · p. 19

Ablation Study

  • Backbone ablations found P-values below 0.05 between Semi-Mamba-UNet and other methods, supporting the use of different segmentation networks as an SSL network perturbation.The authors interpret the differing networks as multi-view learning, with local-information and global-feature learning strengths that can benefit one another.Source: §4 · p. 18, Table 6 · p. 18, Table 7 · p. 18

5. Conclusion

The paper integrates Visual Mamba into UNet for semi-supervised medical image segmentation and combines pixel-level cross-supervision with contrastive learning. It reports extensive experimental evaluations and identifies volumetric segmentation and limited-supervision refinement as future directions.

5. Conclusion

  • The proposed SSL strategy combines pixel-level cross-supervision and contrastive learning to use Visual Mamba for medical image segmentation.The authors report extensive experimental evaluations demonstrating the effectiveness of Semi-Mamba-UNet.Source: §5 · p. 19
  • Future work will extend the approach to volumetric segmentation and further refine it for limited supervised learning scenarios.The authors plan to continue leveraging Visual Mamba in these directions.Source: §5 · p. 19

Related papers

Loading 2402.07245v3…