基于MAMBA的跨模式文本信息弱监督医学图像分割

论文题目MAMBA-Based Weakly Supervised Medical Image Segmentation with Cross-Modal Textual Information摘要在医学图像分割中获取像素级的标注数据代价很高。虽然半监督和弱监督方法减少了标注依赖但它们仍然需要一些像素级的标注。相比之下利用对应于医学图像的文本描述作为分割的监督信息更有前景。文字描述更容易获得因为用户只需要提供病变的位置和外观细节。提出了一种基于Mamba的文本-图像融合分割算法TIFCMamba。该框架并行处理图像和文本以建立跨模式对应通过对比学习来对齐片段编码的特征。该架构采用了基于Mamba的图像编码器与传统的Transformer模型相比它降低了计算复杂性。我们提出Mamba Fusion(MF)模块通过二维融合(BiDF)将文本和图像特征结合在一起在保持计算效率的同时实现了模式内细化和模式间交互。在息肉和皮肤病变数据集上的实验证明了与完全监督方法和最新的弱监督方法相比具有竞争力的性能。1、介绍医学图像分割已成为临床诊断的重要工具。随着深度学习的发展医学图像分割通常依赖于完全监督的范例这需要像素级的标注。然而医学图像的标注需要专业知识而且高质量的标注成本很高限制了分割模型的发展[171823]。近年来人们提出了半监督和弱监督方法来降低标注代价。与完全监督的方法相比这些方法仍然需要一些像素级的标签来进行训练但减少了对全面注释的依赖[1014]。然而医学文本注释更容易获得因为用户只需要提供病变位置和描述。这使得医学文本成为医学图像分割中监督的一个很有前途的解决方案[2835]。基于变换的注意机制通过改进图像和文本信息的融合显著地促进了多模式医学图像分割[924]。然而全注意力模型的二次计算和内存需求带来了挑战特别是在处理大图像和冗长的文本描述时[1526]。状态空间模型(SSM)如Mamba[13]及其变种Vmanba[20]由于其线性复杂性和全局接受域而显示出良好的前景。然而SSM在多模式融合方面的研究还不够深入主要集中在多模式图像融合。ReMamba[33]指出传统的令牌连接方法对于Mamba是无效的因为其线性结构限制了令牌交互导致融合不足和性能降低。在文本监督范式中只有与图像对应的语义文本被用作监督信息而没有任何像素级掩码注释并且模型的训练由文本-图像的语义或特征对齐来驱动[1625]。图文对齐的方法在许多作品中被广泛采用。具体地说图像-文本对齐方法通常使用图像编码器和文本编码器将两者对齐到联合嵌入空间中。通过这种方式零样本传递技术可用于允许两个编码器生成分段输出而无需专门的注释[29]。这种方法在培训和测试之间造成了不一致。在训练过程中图文对齐是基于整个图像的语义特征而在测试过程中目标是将文本语义与特定的图像区域对齐。由于模型在训练期间可能不学习局部文本语义和图像区域之间的关系[1936]因此这种未对准可能导致次优性能。为了应对这些挑战我们提出了一种基于Mamba的文本监督切分框架TIFCMamba。该方法的主要贡献如下1)引入了TIFCMamba利用医学文本监督和多模式对比学习来降低标注代价同时避免了基于Transformer的模型的高计算代价。2)提出了一种具有二维融合机制的Mamba融合块增强了文本特征的交互解决了Mamba中令牌融合的局限性。3)提出了一种图文互对齐机制用于训练和测试过程中图文段之间的精确对齐。2、方法2.1、Mamba的初步研究来自控制系统理论的结构化状态空间模型[1213]通过隐藏状态h(T)∈RN将输入序列x(T)∈R变换为输出y(T)∈R该隐藏状态h(T)RN由其中A∈RN×N是状态矩阵B∈RN是输入矩阵C∈RN是输出矩阵D∈R是跳跃连接(此后省略即D0)。由于连续时间系统不能直接用于数字计算我们使用时间步长∆进行离散其中AEXP(∆A)B≈B和CC。在S4[13]模型中参数(A、B、C、∆)通过梯度下降学习但保持独立于输入限制了上下文建模。为了解决这个问题Mamba引入了S6机制它使B、C和∆依赖于输入。具体地对于输入序列x∈RB×L×C(具有批次大小B、序列长度L和特征维度C)参数计算如下这里线性(·)表示线性变换软加(·)确保非负性因为∆∆˜是一种可学习的偏差。这种依赖输入的设计增强了模型的适应性和捕获上下文信息的能力。2.2、总体框架文本监督分割的核心挑战是建立图像和文本之间的语义对应关系。如图1所示我们的TIFCMamba框架操作于医学图像-文本数据集D{(XI 1XT 1)(XI 2XT 2)···(XI iXT i)···(XI nXT n)}其中图像缺乏像素级标签并且仅用语义描述进行标注。我们使用对比学习来联合训练图像分割器ΦI和文本分割器ΦT以跨通道对准分割区域。对于每个图像-文本对(XIXT i)关键字选择器[5]从XT i(例如红色息肉左侧)提取关键字WT。图像分割器ΦI使用XI I和WT来生成图像掩码MI而文本分割器ΦT处理XT I和WT以产生文本掩码MT。通过用MI裁剪XI并随机填充背景来获得掩模图像ˆI并且类似地使用MT构造完整的文本ˆT。最后Clip[21]‘S图像编码者EI和文本编码者Et分别从ˆI和ˆT中提取特征对比学习对齐它们的表示。2.3、图文分割基于Mamba的图像分割器如图1所示我们的图像分割器以Mamba融合块为中心它集成了视觉和文本形态。来自每个块的中间特征被跳过连接到解码器以产生最终分割掩码。融合块包括两个模块。首先视觉状态空间(VSS)[20]模块将图像特征视为令牌序列并采用2D选择性扫描(SS2D)[20]机制该机制在四个方向上扫描特征以捕获远程空间相关性。其次二维融合(BiDF)模块寻求融合文本特征和图像特征。它计算图像特征和文本特征两种模式之间的互相关性并随后将该信息传播到每个图像块的特征。BiDF模块如图2(A)所示BiDF模块分两个阶段运行。在第一阶段BIDF模块首先将文本特征FT扩展为ˆFT∈Rh×w×CT然后将图像特征Fi与扩展后的文本特征ˆFT进行融合得到适合的∈Rh×w×CIT。这一过程允许每个图像补丁包含文本信息。最后沿着通道维度将这三个特征串联起来。这一过程形式化如下其中Wi∈Rc1×C0和WT∈Rct×C0为可学习参数。FIW·(ˆFTWT)T∈Rh×w通过1×1卷积变换得到合适的∈Rh×w×CIT。然后级联特征FCAT分两个阶段进行融合空间融合(使用VSS模块的2D选择性扫描)然后是信道融合(利用1D扫描的状态空间模型)从而产生最终的融合特征FFUSE。如图2(C)所示在第二阶段为了增强跨通道交互我们设计了两种融合机制一种沿着通道维度另一种沿着空间维度。对于信道维度我们采用VMamba的一维选择性扫描而对于空间维度我们采用2D选择性扫描。这一过程形式化如下文本分割器文本分割器ΦT其处理输入文本xti和一组名词{nj}j j1以生成名词特定的词掩码。具有两个可学习的多头注意层的片段文本编码器[21]提取单词特征XT˜Et(Xti)∈RL×C其中L是记号的数量C是特征维度。对于给定的名词Nj(嵌入Nj∈Rc)特定于单词的逻辑被计算为ℓjw·(XT.nj)b其中具有可学习参数w和b并且每个令牌计算一个点积。假设每个单词要么属于J个名词相关片段中的一个要么不属于因此J个片段上的Softmax加上额外的“None”类别产生词掩码MT[MT i]L i1由然后如果单词i对于J个片段之一达到最高概率则通过设置pi1来生成伪标签向量p∈{01}L否则设置pi0文本分割损失Ltxt被定义为MT和p之间的交叉熵指导ΦT正确地分割文本。为了对齐图像和文本模式我们的框架利用图像区域和相应文本段之间的对比学习。具体地片段编码器从图像区域ˆi提取区域嵌入EiEi(ˆi)并从文本段ˆT提取单词嵌入EtEt(ˆT)。对于一批B个三元组(每个由图像、其配对的文本和所选名词组成)我们计算相似度矩阵S∈RB×B其中每个元素Sij是Ei和Et j之间的余弦相似度。然后应用对称的InfoNCE[21]损失使用可学习的温度参数τ请注意即使多次选择相同的名词相应的区域和文本段也保持不同以确保有效的对齐。图像和文本完成器在进行片断图文编码时为了避免图文空白区域对编码特征对齐造成不利影响我们分别引入了图文完成器。我们通过将图像Iˆ和文本ˆT分别与图像-文本目标区域以外的区域中的像素或单词进行随机互补来避免不利影响ˆiXI·MIFilli(1−MI)ˆTXT·MTFillT(1−MT)。双模对比性对准损失对于图像分割器我们使用了TCL[7]的损失函数Limg它只依赖于图像-文本对进行训练。总损失是图像分割损失Limg、文本分割损失Ltxt和对比损失Lign的加权和其中我们将损耗系数设置为λimg1.0、λtxt1.0、λAlign0.5。3、实验3.1、实验设置数据集我们在息肉医学图像数据集上进行了实验并添加了文本线索以提高分割性能。对于息肉分割使用了三个公开可用的结肠镜数据集CVC-ClinicDB[2]、CVC-ColonDB[27]、ETIS-LaribPolypDB[22]和ISIC2017[3]。我们使用GPT-4[1]来生成数据集中图像的描述并调整了一些文本描述包括位置、外观、边界框和病变所占图像的比例并对一些文本描述进行了调整。数据集的文本描述可以在我们的代码中找到。实施细节使用并集的平均交集(MIUU)和平均骰子系数来评估性能。所有数据集按73的比例划分为训练集和测试集每幅图像的大小调整为480×480。我们设计了基于VMamba的TIFCMamba结构的三种图像分割器变体。训练过程包括两个阶段首先使用ADAM优化器在CC3M[8]数据集上进行初始预训练(每个图文对随机选择两个名词)然后在冻结文本分支参数的情况下对息肉和皮肤数据集进行微调。实验在4个NVIDIA3090GPU上进行批大小为4学习速率为1×10−5。3.2、结果与分析在表1中我们比较了现有的前五个弱监督模型和两个全监督模型。其中WeakPolyp[30]使用边界框作为监督而SimTxtSeg[32]、TCL[7]、Code[31]和XCoOp[4]使用文本作为监督。与标注代价较高的全监督方法相比我们的方法获得了与之相当的分割性能。与其他最先进的弱监督分割模型相比我们的模型在ClinicDB、殖民地数据库、LaribPolypDB和ISIC2017数据集上的mDice和Miou分别提高了1.26%和1.77%1.16%和0.89%0.61%和1.36%以及0.67%和0.25%。与五种文本监督模型相比对分割结果的定性分析如图3所示。3.3、消融实验文字描述的影响。我们比较了使用不同的文本描述对我们模型的分割性能的影响。如图3所示不同的文本描述对分割结果有很大的影响特别是当文本在方向上描述前景的位置不准确时此外当文本描述过于冗余时也会影响模型的分割性能。融合模式的影响。比较了图像-文本特征融合时我们采用的跨模式特征融合方法对分割性能的影响。如表2所示单独使用空间融合和信道融合时两者都对模型性能有正向影响但信道融合对分割性能的改善贡献更大。模型效率比较。我们在息肉数据集上与VMamba、SwinUnet的模型效率进行了比较。由于我们的模型使用了一个额外的文本分割器所以我们只在测试阶段将图像分割器的效率与上述两个模型进行了比较。如图4所示我们的TIFCMamba-T实现了mDice与Mac和型号尺寸之间的最佳平衡。