---
title: 'LoG-VMamba: Efficient 2D/3D Segmentation'
url: https://www.emergentmind.com/topics/log-vmamba
type: topic
---

# LoG-VMamba: Efficient 2D/3D Segmentation

LoG-VMamba (Local-Global Vision Mamba) is a neural network module designed for computationally efficient and accurate 2D and 3D medical image segmentation. It extends the State Space Model (SSM)-based “Mamba” architecture with explicit local and global spatial context encoding and is demonstrated to substantially outperform both CNN- and Transformer-based baselines in medical imaging benchmarks. LoG-VMamba achieves this by enforcing spatial adjacency among local features, compressing global features, and maintaining linear complexity relative to image size, addressing the challenges of high-dimensional medical imaging data [2408.14415].

## 1. State-Space Models and Vision Mamba Foundations

State Space Models (SSMs) describe dynamical systems via latent state trajectories and have been adapted for deep learning. The continuous-time linear SSM is described as:

\[
\dot h(t) = A\,h(t) + B\,x(t), \qquad y(t) = C\,h(t),
\]
where \(x(t)\in\mathbb R^D\), \(h(t)\in\mathbb R^N\), parameters \(A, B, C\) are learnable.

Mamba leverages a selective SSM (S6/Mamba) discretization, enabling recurrent computation over sequences:

\[
\overline A = f_A(A,\Delta),\quad \overline B = f_B(B,\Delta)
\]
\[
h_t = \overline A\,h_{t-1} + \overline B\,x_t, \quad y_t = C\,h_t
\]

Mamba’s distinctiveness lies in making parameters \(B, C, \Delta\) input-dependent via small MLPs, achieving efficient hardware-aware kernels with \(\mathcal O(N)\) compute.

The original Vision Mamba (VSS) block processes visual data as:

- Input \(X\in\mathbb R^{H\times W\times C}\).
- LayerNorm, channel expansion, then splits to:
  - Depthwise convolution (DWC) followed by SiLU activation, then flattening and 1D SSM scan (optionally multiple directions).
  - Pointwise SiLU.
- Branches are multiplied, projected to \(C\), and added residually to \(X\).

While VSS grants global receptive fields at linear complexity, it struggles to jointly capture spatial locality and global context, particularly in high-dimensional medical images due to its sequential processing constraints [2408.14415].

## 2. Architectural Innovations in LoG-VMamba

LoG-VMamba introduces explicit mechanisms for capturing both local and global dependencies by replacing the vanilla token extractor with two parallel modules—a Local Token eXtractor (LTX) and Global Token eXtractor (GTX):

- **LTX**: Squeezes channel dimension of \(X\) by factor \(S\) with DWC. An \(R \times R\) “unfolding” window extracts all spatial neighbors for every position and flattens their features, guaranteeing each token encodes local context.

  \[
  X_s = \mathrm{DWC}_{S\!:\!1}(X) \in \mathbb R^{H\times W\times \frac C S}
  \]
  Each of the \(H\,W\) positions becomes a token of dimension \(C^L = \frac{C}{S} R^2\).

- **GTX**: Applies a dilated DWC with stride \(K \times K\) for spatial compression, followed by channel grouping and linear projection to \(C^L\)-dimensional tokens. This yields a compressed global summary.

  \[
  G = \mathrm{DWC}_{\text{dil}=K}\bigl(X_s\bigr) \in \mathbb R^{\frac H K\times\frac W K\times \frac C S}
  \]
  Total number of global tokens: \(\frac{C/S}{\gamma} \times \frac{H W}{K^2}\).

- **Token Concatenation**: Tokens from LTX and GTX are interleaved to produce \(X^{LG} \in \mathbb R^{(L+L_2)\times C^L}\). Interleaving, rather than appending, ensures both local and global contexts are available early in SSM processing.

- **LoG-VMamba Block**: The full processing pipeline is:

  1. LayerNorm of input.
  2. LTX and GTX extraction, interleaving tokens.
  3. Single-direction horizontal SSM scan:
     \[
     h_t = \bar A h_{t-1} + \bar B X^{LG}_t,\quad y_t = C h_t
     \]
  4. Spatial stacking, projection, gating, and addition to the residual branch.

This approach allows local and global dependencies to be encoded prior to sequential SSM computation, overcoming the locality/globality tension that limits prior SSM-based vision models.

## 3. Network Architectures for 2D and 3D Segmentation

LoG-VMamba is instantiated within U-shaped segmentation networks for both 2D and 3D data:

- **2D Model**: Based on a Swin-UMamba encoder with patch merging and ImageNet pretraining. The decoder consists of four upsampling stages, each concatenating skip connections and applying 1–2 LoG-VMamba blocks at fixed spatial resolution and channel count.

- **3D Model**: Based on U-Mamba-Enc, replacing all encoder Mamba blocks with LoG-VMamba. Input volumes vary (e.g., \(128^3\) for BraTS, \(16\times256\times224\) for ACDC). The decoder retains transposed-convolution upsampling, with skip-connections concatenated pre-LoG-VMamba.

A summary of feature-map sizes, per-stage pooling, kernel, and other design choices is provided in the original work. The architecture is explicitly designed to maintain linear complexity and facilitate efficient fusion of hierarchical features [2408.14415].

## 4. Computational Complexity and Memory Analysis

LoG-VMamba’s dominant computation is the single horizontal SSM scan over the concatenated token sequence:

\[
\mathcal O\left((L+L_2) C^L\right) \sim \mathcal O(N C), \quad N = H W
\]

- **Self-attention (ViT)**: \(\mathcal O(N^2 C)\) per layer and \(\mathcal O(N^2)\) memory.
- **CNNs**: \(K^2 N C^2\) for \(K \times K\) convolutions, memory \(\mathcal O(N C)\); global context grows only with depth.

LoG-VMamba’s efficient scan and controlled channel expansions/compressions reduce computational cost by 20–50% in GFLOPs relative to ViT-based decoders for the studied medical segmentation tasks, without sacrificing the capacity for global or local spatial modeling.

## 5. Experimental Protocols and Comparative Results

LoG-VMamba’s efficacy was evaluated on representative 2D/3D segmentation benchmarks:

- **2D Endoscopy (MICCAI EndoVis '17)**: Custom train/val splits (1440/360), RGB input \(640\times384\).
- **2D Cell (NeurIPS '22)**: 800 train / 200 val.
- **3D BraTS 2020**: 236 train, 59 val, 74 test.
- **3D ACDC**: 160/40/100 train/val/test.

Augmentations (random flip, elastic, color jitter, intensity scaling, noise), Adam optimizer, Dice+Cross-Entropy loss, and 5-fold validation were used. Metrics included Dice, IoU, Normalized Surface Dice (NSD, 2D), and Hausdorff 95% (HD95, 3D).

**Main Results:**

| Task         | Method             | Params | GFLOPs | Dice %         | IoU %         | NSD %         | HD95 (mm)       |
|--------------|--------------------|--------|--------|----------------|---------------|---------------|-----------------|
| 2D Endo      | Swin-UMamba†       | 27.5M  | 45.4G  | 71.23 ± 1.00   | 67.81 ± 0.99  | 72.77 ± 1.02  | -               |
|              | **Ours (LoG)**     | 30.3M  | 48.6G  | 75.17 ± 0.24   | 71.68 ± 0.23  | 76.83 ± 0.25  | -               |
| 2D Cell      | Swin-UMamba†       | -      | -      | 73.50 ± 0.86   | -             | 83.31 ± 0.66  | -               |
|              | **Ours (LoG)**     | -      | -      | 76.21 ± 0.10   | -             | 86.44 ± 0.09  | -               |
| 3D BraTS     | U-Mamba-Enc        | -      | -      | 87.01 ± 0.10   | -             | -             | 4.38 ± 0.09     |
|              | SegMamba           | -      | -      | 87.62 ± 0.16   | -             | -             | 4.73 ± 0.22     |
|              | **Ours (LoG)**     | -      | -      | 88.06 ± 0.08   | -             | -             | 3.97 ± 0.04     |
| 3D ACDC      | U-Mamba-Enc        | -      | -      | 91.65 ± 0.31   | -             | -             | 1.13 ± 0.02     |
|              | U-Mamba-Bot        | -      | -      | 91.94 ± 0.05   | -             | -             | 1.26 ± 0.15     |
|              | **Ours (LoG)**     | -      | -      | 92.18 ± 0.13   | -             | -             | 1.10 ± 0.00     |

LoG-VMamba consistently outperformed prior baselines across all metrics.

Ablation studies showed LTX and GTX each contributed to improvements in Dice and NSD/HD95, with their combination yielding the highest gains. Interleaved token placement was more effective than naive concatenation. Multi-direction SSM scans did not provide further benefit beyond single-horizontal direction [2408.14415].

## 6. Component Analysis and Design Choices

The table below delineates the contribution of each architectural component (results from Endoscopy/BraTS):

| Block      | Dice % (Endo/BraTS) | NSD % / HD95 (Endo/BraTS) |
|------------|--------------------|---------------------------|
| VSS        | 71.23 / 87.01      | 72.77 / 4.38              |
| +GTX only  | 72.64 / 87.99      | 74.24 / 4.05              |
| +LTX only  | 74.15 / 87.71      | 75.81 / 4.12              |
| **LoG**    | 75.17 / 88.06      | 76.83 / 3.97              |

Early global context (GTX) and explicit local adjacency (LTX) are each responsible for ~1–2% Dice gains; using both yields a further improvement. Interleaved token placement maximizes performance.

A single-direction SSM (\(M=1\)) suffices for effective context propagation; adding additional scanning directions did not increase accuracy. This suggests LoG-VMamba’s local/global token structure enables sufficient information flow within a linear scan [2408.14415].

## 7. Implementation Outline and Summary

The recipe for LoG-VMamba integration in segmentation training consists of:

```python
for epoch in 1…N_epochs:
  for (X_batch, Y_batch) in train_loader:
    # Forward
    feats = encoder(X_batch)   # Swin/HPC or 3D U-Mamba-Enc
    predictions = decoder_with_LoG_VMamba(feats)
    loss = DiceCE(predictions, Y_batch)
    # Backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Inference
def inference(volume):
  patches = sliding_window(volume)
  out_patches = [ forward(p) for p in patches ]
  return stitch(out_patches)
```
where each `LoG_VMamba` block encapsulates the LTX/GTX extraction, interleaved sequencing, SSM scan, gating, and residual addition. Networks employing LoG-VMamba demonstrate marked segmentation accuracy benefits with strict linear complexity scaling, making this approach particularly suitable for high-dimensional 2D/3D medical imaging tasks [2408.14415].

Source: https://www.emergentmind.com/topics/log-vmamba