---
title: 'Weak-Mamba-UNet: Weakly Supervised Segmentation'
url: https://www.emergentmind.com/topics/weak-mamba-unet
type: topic
---

# Weak-Mamba-UNet: Weakly Supervised Segmentation

Weak-Mamba-UNet is a weakly supervised medical image segmentation framework that fuses Convolutional Neural Networks (CNN), Vision Transformers (ViT), and Visual Mamba (VMamba) architectures within a collaborative, cross-supervisory learning loop. It is specifically designed for applications where annotations are sparse or imprecise, such as scribble-based labels, by leveraging the complementary strengths of encoder–decoder networks built from convolution, attention, and state-space modules [2402.10887].

## 1. Constituent Architectures

Weak-Mamba-UNet comprises three U-shaped encoder–decoder subnetworks, termed “views,” each adopting a distinct architectural paradigm yet sharing an identical high-level connection structure. The three subnetworks are:

1. **CNN-based UNet**: Employs four-level downsampling and upsampling pathways, with double 3×3 convolutions, ReLU activations, and Batch Normalization at each stage. Skip connections concatenate feature maps from encoder to decoder, optimizing for local spatial detail.

2. **Swin Transformer-based SwinUNet**: Utilizes patch embedding, shifted-window multi-head self-attention (SWA), and standard MLP blocks in a three-level encoder–decoder structure. Patch merging and skip connects preserve multi-scale context; the SWA layers efficiently harvest global contextual dependencies via windowed self-attention.

3. **VMamba-based Mamba-UNet**: Implements state-space modeling via Visual Mamba blocks for long-range spatial dependencies. The encoder and decoder both operate in three resolution stages. Each VMamba block is mathematically formulated as a continuous-discrete SSM:

   $$
   \frac{dh(t)}{dt} = A h(t) + B u(t),\quad y(t) = C h(t) + D u(t)
   $$

   Discretized (with $\Delta t=1$), yielding:

   $$
   h_k = \Phi h_{k-1} + \Gamma x_k,\quad z_k = C h_k + D x_k
   $$

   where $x_k$ is the input feature (flattened token), $h_k$ is the hidden state, and $(\Phi, \Gamma, C, D)$ are learned. Outputs are post-processed via feed-forward networks and normalization layers.

This tripartite architectural ensemble is essential to achieving complementary modeling of local detail (CNN), global context (ViT), and efficient long-range interactions (VMamba).

## 2. Collaborative and Cross-Supervisory Training Mechanism

The distinguishing hallmark of Weak-Mamba-UNet is its collaborative, cross-supervised weakly-supervised learning protocol. The system is trained with two forms of supervision:

- **Partial Cross-Entropy Loss (`pCE`)** on scribble-labeled pixels: The loss is applied only to pixels annotated in the scribbled mask set $\Omega_L$:
  $$
  \mathcal{L}_{\mathrm{pce}}^i = -\sum_{p\in\Omega_L}\sum_{k=1}^K y^{(p)}_{\mathrm{s},k}\log(y^{(p)}_{\mathrm{p},k})
  $$
  where $y_{\mathrm{s}}$ denotes the sparse annotation, and $y_{\mathrm{p}}$ is the predicted softmax probability for class $k$.

- **Pseudo-Label Dice Loss across the ensemble**: At each iteration, each subnetwork predicts a segmentation mask, and a soft pseudo-label is synthesized by convex combination:
  $$
  Y_{\mathrm{pseudo}} = \alpha Y_{\mathrm{cnn}} + \beta Y_{\mathrm{vit}} + \gamma Y_{\mathrm{mamba}}
  $$
  with $\alpha+\beta+\gamma=1$ and $\alpha, \beta, \gamma \sim \mathrm{Uniform}(0,1)$ per batch. For each network $i$, the Dice loss to this pseudo-label is:
  $$
  \mathcal{L}_{\mathrm{dice}}^i = 1 - \frac{2\sum_p\sum_k y^{(p)}_{\mathrm{p},k} Y^{(p)}_{\mathrm{pseudo},k}}{\sum_p\sum_k y^{(p)}_{\mathrm{p},k} + \sum_p\sum_k Y^{(p)}_{\mathrm{pseudo},k}}
  $$

The final objective summed over all models is:
$$
\mathcal{L}_{\mathrm{total}} = \sum_{i\in\{\mathrm{cnn,vit,mamba}\}}\Bigl(\mathcal{L}_{\mathrm{pce}}^i + \mathcal{L}_{\mathrm{dice}}^i\Bigr)
$$

This loop enables each subnetwork to iteratively refine itself and its peers via mutual pseudo-label supervision, which is crucial for learning from highly incomplete annotation.

## 3. Algorithmic Workflow and Inference

The Weak-Mamba-UNet workflow consists of the following steps:

1. **Forward Pass**: Each network processes input $X$ to yield $Y_{\mathrm{cnn}}, Y_{\mathrm{vit}}, Y_{\mathrm{mamba}}$.
2. **Pseudo-Label Construction**: New convex weights are drawn, and $Y_{\mathrm{pseudo}}$ is computed.
3. **Loss Computation**: Each network is penalized with $\mathcal{L}_{\mathrm{pce}}$ on labeled pixels and $\mathcal{L}_{\mathrm{dice}}$ with respect to $Y_{\mathrm{pseudo}}$ over the entire image.
4. **Parameter Update**: Stochastic gradient descent steps update each network’s parameters.
5. **Test-Time Output**: At inference, the final prediction is the argmax after averaging the logits from the three subnetworks:
   $$
   Y_{\mathrm{final}}(p) = \arg\max_k \left[\frac{1}{3}\left(f_{\mathrm{cnn}}(X)_k + f_{\mathrm{vit}}(X)_k + f_{\mathrm{mamba}}(X)_k\right)\right]
   $$
   Optionally, the VMamba-based model alone may be used for efficient deployment.

## 4. Scribble-Based Annotation Paradigm

Scribble annotation provides weak spatial supervision by marking 1–3-pixel-wide lines within each label region. The mask preprocessing strategy follows Valvano et al. (2021), converting dense ground truth into sparse, narrow strokes that annotate only ≈2% of the available pixels. These sparse labels serve as hard anchors for loss calculation, allowing the model to explore plausible segmentations in the vast unlabeled regions by relying on model-driven consistency and pseudo-label refinement.

## 5. Experimental Results and Quantitative Assessment

Empirical validation on the ACDC MRI cardiac segmentation dataset (100 patients, 4 anatomical classes, 224×224 resolution) demonstrates the efficacy of the Weak-Mamba-UNet strategy. Metrics include mean Dice, Accuracy, Precision, Sensitivity, Specificity (higher is better), and 95%-Hausdorff Distance (HD) and Average Surface Distance (ASD, both lower is better). Performance on weak supervision (scribbles) is as follows:

| Framework + Network             | Dice ↑  | Acc ↑   | Pre ↑   | Sen ↑   | Spe ↑   | HD ↓   | ASD ↓ |
|---------------------------------|---------|---------|---------|---------|---------|--------|-------|
| pCE + UNet                      | 0.7620  | 0.9807  | 0.6799  | 0.9174  | 0.9823  | 151.06 | 54.65 |
| Gated CRF + UNet                | 0.9046  | 0.9955  | 0.8890  | 0.9304  | 0.9922  | 7.43   | 2.08  |
| Gated CRF + SwinUNet            | 0.8995  | 0.9955  | 0.8920  | 0.9175  | 0.9904  | 6.66   | 1.62  |
| **Weak-Mamba-UNet (ours)**      | **0.9171** | **0.9963** | **0.9095** | **0.9309** | **0.9920** | **3.96** | **0.88**  |

Ablation studies indicate that model heterogeneity (CNN, SwinUNet, VMamba) is essential to reaching peak performance; homogeneous tri-ensembles (such as 3×SwinUNet) underperform considerably (Dice drops to 0.7446).

## 6. Significance, Insights, and Variations

The integration of state-space modeling via VMamba blocks provides efficient and expressive long-range context propagation at linear computational cost, a critical factor for scalable weak supervision. The cross-supervisory pseudo-label loop mitigates overfitting to sparse supervisory signals and enables robust mask refinement. This collaborative training schema enables the ensemble to outpace single-model weakly supervised approaches and homogeneous ensembles.

A plausible implication is that introducing additional architectural diversity (e.g., fusing other SSMs or explicit uncertainty modeling) could further improve label efficiency in low-annotation regimes.

## 7. Reproducibility and Implementation

Weak-Mamba-UNet is publicly available (https://github.com/ziyangwang007/Mamba-UNet), is implemented in PyTorch ≥1.8, and requires CUDA 11 support. The repository provides scripts for both training (configurable chosen backbones, annotation paths, hyperparameters) and inference (ensemble or single-subnetwork) [2402.10887].

## References

- "Weak-Mamba-UNet: Visual Mamba Makes CNN and ViT Work Better for Scribble-based Medical Image Segmentation" [2402.10887]

Source: https://www.emergentmind.com/topics/weak-mamba-unet