Hybrid Transformer-Mamba Architecture for Weakly Supervised Volumetric Medical Segmentation

2025-12-11

AI Summary

In weakly supervised 3D medical image segmentation, conventional 2D encoders neglect volumetric structural cues, leading to inaccurate lesion localization. Method: We propose TranSamba, a hybrid architecture featuring the first Transformer–Mamba collaboration: a Vision Transformer (ViT) models local–global spatial relationships within slices, while a cross-plane Mamba module enables linear-complexity inter-slice contextual modeling. Crucially, Mamba’s state-space dynamics directly optimize Class Activation Map (CAM) quality, enhancing localization fidelity. The framework supports end-to-end weakly supervised training, with computational complexity scaling linearly in depth and constant batch-wise memory footprint. Results: TranSamba achieves state-of-the-art performance on three multi-modal 3D medical datasets, outperforming existing methods by a significant margin in segmentation accuracy under weak supervision.

Technology Category

• Computer Vision: 3D Computer Vision
• Machine Learning: Multimodal Learning
• Constraint Satisfaction and Optimization: Satisfiability Modulo Theories

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Weakly supervised semantic segmentation offers a label-efficient solution to train segmentation models for volumetric medical imaging. However, existing approaches often rely on 2D encoders that neglect the inherent volumetric nature of the data. We propose TranSamba, a hybrid Transformer-Mamba architecture designed to capture 3D context for weakly supervised volumetric medical segmentation. TranSamba augments a standard Vision Transformer backbone with Cross-Plane Mamba blocks, which leverage the linear complexity of state space models for efficient information exchange across neighboring slices. The information exchange enhances the pairwise self-attention within slices computed by the Transformer blocks, directly contributing to the attention maps for object localization. TranSamba achieves effective volumetric modeling with time complexity that scales linearly with the input volume depth and maintains constant memory usage for batch processing. Extensive experiments on three datasets demonstrate that TranSamba establishes new state-of-the-art performance, consistently outperforming existing methods across diverse modalities and pathologies. Our source code and trained models are openly accessible at: https://github.com/YihengLyu/TranSamba.

Problem

• Develops hybrid Transformer-Mamba architecture for 3D medical segmentation
• Addresses 2D encoder limitations in weakly supervised volumetric imaging
• Enables efficient 3D context modeling with linear computational complexity

Innovation

• Hybrid Transformer-Mamba architecture for 3D context
• Cross-Plane Mamba blocks enable efficient cross-slice information exchange
• Linear complexity and constant memory usage for volumetric modeling