💻 computer science

Hybrid Transformer-Mamba for Weakly Supervised Volumetric Medical Segmentation

The paper introduces TranSamba, a hybrid Transformer-Mamba architecture that leverages efficient cross-plane modeling to capture 3D context from plane-level labels, achieving state-of-the-art performance in weakly supervised volumetric medical segmentation.

Original authors: Yiheng Lyu, Lian Xu, Coen Arrow, Mohammed Bennamoun, Farid Boussaid, Girish Dwivedi

Published 2026-06-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yiheng Lyu, Lian Xu, Coen Arrow, Mohammed Bennamoun, Farid Boussaid, Girish Dwivedi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a computer to find a specific object inside a 3D medical scan, like a tumor in a brain or a cavity in a heart. Usually, to teach a computer this, you need a human to painstakingly draw a perfect outline around the object on every single slice of the scan. This is like asking a teacher to grade every single page of a student's notebook individually. It takes forever and is very expensive.

The Problem: The "2D Blindness"
To save time, researchers use "weak supervision." Instead of drawing outlines, they just tell the computer, "Yes, there is a tumor in this whole stack of slices," or "No, there isn't." It's like the teacher just saying, "You did a good job on this whole chapter," without pointing out exactly where the mistakes are.

The problem is that most computers are trained to look at medical scans like a stack of 2D photographs. They look at one slice, make a guess, then look at the next slice, make another guess, and never really talk to each other. They miss the fact that the human body is a 3D object. A tumor isn't just a flat circle on one page; it's a 3D blob that stretches across many pages.

The Solution: TranSamba (The "Team of Specialists")
The authors of this paper built a new AI model called TranSamba. Think of it as a hybrid team of two very different specialists working together to solve the puzzle:

The "Local Detective" (The Transformer): This part is great at looking at a single slice of the scan and saying, "Hey, I see a shape here that looks like the target." It's very good at focusing on the details within one specific picture.
The "Context Connector" (The Mamba): This is the new star. Imagine the Mamba as a super-fast messenger who runs between the pages of the book. Instead of just looking at one page, it quickly whispers to the detective on Page 5, "Hey, the thing on Page 4 is connected to the thing on Page 6." It helps the model understand how the object flows from one slice to the next without getting bogged down by the math.

How They Work Together
In the old way, trying to connect all the pages at once was like trying to have a conversation where everyone shouts at everyone else at the same time. It gets chaotic and slow (computationally expensive).

TranSamba changes the game. It uses the "Context Connector" (Mamba) to efficiently pass information between neighboring slices in a straight line. This is like a relay race where the baton is passed smoothly from one runner to the next. This helps the "Local Detective" (Transformer) make much better guesses because it now knows what's happening in the slices right next to it.

Why It's a Big Deal

Speed and Efficiency: Because the Mamba part is so efficient, the model doesn't get slower even if the scan has hundreds of slices. It stays fast and doesn't need a massive amount of computer memory to do it.
Better Results: The authors tested this on three different types of medical scans (brain tumors, kidney tumors, and heart cavities). In every case, TranSamba found the objects better than any other method that uses these "weak" labels. It was like the team that finally figured out how to read the whole book instead of just guessing based on the cover.

The Catch
The paper notes that while this model is great at finding where the object is, it still struggles a bit with very tiny objects (like a speck of dust compared to a boulder). Also, because it's trained on "weak" labels, it's currently a research tool and might need some extra steps before a doctor could use it directly in a hospital.

In a Nutshell
TranSamba is a new AI architecture that teaches computers to understand 3D medical scans by combining a sharp-eyed local observer with a fast, efficient messenger that connects the dots between slices. This allows the computer to learn from simple, cheap labels and still find medical issues with high accuracy, without needing a supercomputer to do the math.

Technical Summary: Hybrid Transformer-Mamba for Weakly Supervised Volumetric Medical Segmentation

Problem Statement
Volumetric medical image segmentation (e.g., CT, MRI) typically requires labor-intensive voxel-level annotations. Weakly supervised segmentation offers a solution by training models using only slice-level (plane-level) labels, which lack localization information. While Class Activation Mapping (CAM) can extract coarse localization from classification models, existing methods face significant challenges due to the volumetric nature of medical data. Most current approaches rely on 2D encoders that process individual slices independently, neglecting the 3D context. Furthermore, attempting to capture volumetric context via standard cross-plane self-attention (SA) is computationally prohibitive, as it scales quadratically with the number of planes ( $N$ ).

Methodology: TranSamba
The authors propose TranSamba, a hybrid architecture combining Vision Transformers (ViT) and the Mamba (Structured State Space Model) architecture to address these limitations. The core design philosophy is to leverage the efficiency of Mamba for cross-plane interactions while retaining the efficacy of Transformer self-attention for in-plane object localization.

Architecture Overview: The encoder consists of $L$ hybrid layers. Each layer contains a Cross-Plane Mamba (CPM) block followed by a Transformer block.
Cross-Plane Modeling (CPM): To capture 3D context efficiently, patch tokens from $N$ neighboring planes are stacked and reshaped into a single sequence of length $M \times N$ (where $M$ is the number of patches per plane). The CPM block uses a Mamba module to model interactions between these tokens. Crucially, class tokens are separated and excluded from the Mamba block to preserve global plane representation. This allows the model to learn inter-plane dependencies in linear time.
In-Plane Modeling (Transformer): The output of the CPM block is reshaped back into $N$ parallel sequences. These are processed by standard Transformer blocks using Class-to-Patch (C2P) self-attention. The enriched token representations from the cross-plane step enhance the in-plane attention maps, leading to better object localization.
Training and Inference: The model is trained using plane-level labels. The loss function combines cross-entropy terms derived from a global average pooling (GAP) of the class token and a global weighted ranking pooling (GWRP) of the patch tokens. During inference, attention maps are generated by aggregating the C2P attention weights from all Transformer blocks.

Key Contributions

Hybrid Architecture: The introduction of TranSamba, which integrates CPM blocks with a ViT backbone to enable efficient cross-plane modeling without sacrificing the strengths of in-plane self-attention.
Complexity Efficiency: The architecture achieves linear time complexity with respect to the number of planes ( $N$ ) and maintains constant space complexity for batch processing, in contrast to the quadratic scaling of standard cross-plane self-attention.
State-of-the-Art Performance: The method demonstrates that cross-plane modeling is a generalizable strategy for weakly supervised volumetric segmentation, achieving superior results across diverse modalities and pathologies without relying on complex multi-scale aggregation or self-supervised regularization strategies often used in prior work.

Experimental Results
TranSamba was evaluated on three distinct datasets:

BraTS: Brain tumor segmentation (T2 FLAIR).
KiTS: Kidney tumor segmentation.
LASC: Left atrium cavity segmentation.
Performance: TranSamba achieved State-of-the-Art (SOTA) results on all three datasets, outperforming the second-best methods by mean margins of 5.3% in Dice Similarity Coefficient (DSC) and 2.4% in Intersection-over-Union (IoU).
Ablation Studies:
- Cross-Plane vs. In-Plane: Replacing the CPM block with in-plane Mamba blocks (V2) yielded significantly lower performance, confirming that the improvement stems specifically from cross-plane modeling rather than just the use of Mamba.
- Modeling Scope: Modeling across 16 planes yielded the best results, suggesting that local cross-plane interactions are more beneficial than long-range modeling for continuous structures.
- Layer Design: Placing the CPM block before the Transformer block (CrossIn design) proved optimal, ensuring that attention extraction benefits from the enriched cross-plane context.
Object Size Analysis: Performance showed a positive correlation with object size. While TranSamba excelled on large and huge objects, it faced challenges with small objects (a known issue in medical segmentation), particularly in the KiTS dataset where objects were often small.

Significance and Claims
The paper claims that TranSamba effectively bridges the gap between the need for 3D context and the computational constraints of weakly supervised learning. By decoupling cross-plane and in-plane modeling, the authors demonstrate that efficient linear-time cross-plane modeling is a critical, generalizable component for improving object localization in volumetric medical segmentation. The work suggests that future clinical applications may require post-processing or multi-step approaches, but the core architectural innovation provides a robust foundation for weakly supervised tasks across diverse medical imaging scenarios.

Technical Summary: Hybrid Transformer-Mamba for Weakly Supervised Volumetric Medical Segmentation

More like this