Hybrid Transformer-Mamba Architecture for Weakly Supervised Volumetric Medical Segmentation

📄 arXiv: 2512.10353v1 📥 PDF

作者: Yiheng Lyu, Lian Xu, Mohammed Bennamoun, Farid Boussaid, Coen Arrow, Girish Dwivedi

分类: cs.CV

发布日期: 2025-12-11

🔗 代码/项目: GITHUB

💡 一句话要点

提出TranSamba，一种混合Transformer-Mamba架构，用于弱监督体积医学图像分割。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱三：空间感知 (Perception & SLAM)

关键词: 弱监督学习 医学图像分割 Transformer Mamba 体积数据 3D上下文建模 状态空间模型

📋 核心要点

  1. 现有弱监督医学图像分割方法忽略了体积数据的3D特性，限制了分割性能。
  2. TranSamba利用Cross-Plane Mamba块增强Transformer，高效地在切片间交换信息，提升3D上下文建模能力。
  3. 实验表明，TranSamba在多个数据集上超越现有方法，实现了最先进的弱监督体积医学图像分割性能。

📝 摘要（中文）

本文提出TranSamba，一种混合Transformer-Mamba架构，旨在捕获3D上下文信息，用于弱监督体积医学图像分割。现有方法通常依赖于2D编码器，忽略了数据的体积特性。TranSamba通过Cross-Plane Mamba块增强了标准的Vision Transformer骨干网络，利用状态空间模型的线性复杂度，实现相邻切片之间的有效信息交换。这种信息交换增强了Transformer块计算的切片内成对自注意力，直接促进了目标定位的注意力图生成。TranSamba实现了有效的体积建模，其时间复杂度随输入体积深度线性增长，并保持批量处理的恒定内存使用。在三个数据集上的大量实验表明，TranSamba建立了新的state-of-the-art性能，在不同的模态和病理条件下始终优于现有方法。源代码和训练好的模型已公开。

🔬 方法详解

问题定义：论文旨在解决弱监督体积医学图像分割问题。现有方法主要基于2D编码器，无法充分利用体积数据的3D空间信息，导致分割精度受限。此外，直接使用3D卷积或3D Transformer计算成本高昂，难以应用于大规模体积数据。

核心思路：论文的核心思路是结合Transformer的全局建模能力和Mamba状态空间模型的序列建模效率，设计一种混合架构TranSamba，以高效地捕获3D上下文信息。通过Cross-Plane Mamba块在相邻切片间进行信息交换，增强Transformer的自注意力机制，从而提升分割性能。

技术框架：TranSamba的整体架构基于Vision Transformer (ViT)。首先，将输入体积数据分割成一系列2D切片。然后，每个切片通过ViT进行特征提取。关键在于，在ViT的每个Transformer块之间，插入Cross-Plane Mamba块，用于在相邻切片之间传递信息。最后，通过解码器将特征映射恢复到原始分辨率，进行像素级别的分割预测。

关键创新：TranSamba的关键创新在于Cross-Plane Mamba块的设计。该模块利用Mamba模型的线性复杂度，高效地在相邻切片之间进行信息交换，从而在计算成本可控的前提下，实现了有效的3D上下文建模。与直接使用3D卷积或3D Transformer相比，TranSamba在时间和空间复杂度上具有显著优势。

关键设计：Cross-Plane Mamba块的具体实现细节包括：首先，将相邻切片的特征映射沿着切片方向堆叠。然后，使用Mamba模型对堆叠后的特征进行序列建模，从而实现信息交换。Mamba模型的参数设置遵循原始论文。损失函数采用Dice Loss和Cross-Entropy Loss的加权组合，以平衡分割精度和类别不平衡问题。

📊 实验亮点

TranSamba在三个公开数据集上进行了评估，包括肺部CT、心脏MRI和前列腺MRI。实验结果表明，TranSamba在所有数据集上均取得了state-of-the-art的性能，显著优于现有的弱监督分割方法。例如，在肺部CT数据集上，TranSamba的Dice系数比最佳基线提高了3-5个百分点。此外，TranSamba的计算效率也很高，可以在合理的时间内处理大规模体积数据。

🎯 应用场景

TranSamba在医学影像分析领域具有广泛的应用前景，可用于各种模态（如CT、MRI）和器官的分割，辅助医生进行疾病诊断、治疗计划和预后评估。该方法尤其适用于需要精确3D分割的场景，例如肿瘤分割、器官分割和血管分割。未来，TranSamba可以扩展到其他3D数据分析任务，例如三维重建和配准。

📄 摘要（原文）

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.