MambaMIL: Enhancing Long Sequence Modeling with Sequence Reordering in Computational Pathology（MICCA

MambaMIL: Enhancing Long Sequence Modeling with Sequence Reordering in Computational Pathology

MambaMIL：在计算病理学中通过序列重排序增强长序列建模

研究团队：香港科技大学

代码地址：https://github.com/isyangshu/MambaMIL

一、任务背景

• 多示例学习( Multiple Instance Learning，MIL )已成为计算病理学中提取整张切片图像( Whole Slide Images，WSIs )中判别性特征表示的主流范式。尽管取得了显著的进展，

• MIL最广泛使用的范式是使用预训练的模型将实例转换为低维特征，然后将这些特征聚合为包层表示用于后续分析。在该范式下，MIL将WSI分析概念化为一个长序列建模问题，旨在建模实例之间的相关性以及整个包内的整体上下文信息，以捕获判别性信息。

• 总体而言，现有方法在全面挖掘长序列内部的上下文信息方面存在局限性，阻碍了其性能的提升。

• 最近，结构化状态空间序列( Structured State Space Sequence，S4 ) 作为一种高效的架构被引入，以解决长序列建模的瓶颈问题。此外，选择性扫描空间状态序列模型，即Mamba，通过采用依赖输入的选择机制和硬件感知算法，在离散数据建模中提出了S4，使Mamba能够在不牺牲全局感受野的情况下实现线性复杂度。

• 然而，对于固有的非序列视觉数据，直接将Mamba应用于一个斑块化和扁平化的图像，必然会导致感受野的限制。这种局限性源于Mamba仅允许每个斑块与先前扫描位置之间的相互作用，从而无法估计与未扫描斑块之间的关系。

• 与典型的视觉模态不同，WSIs包含零散且稀少的正块，表现出较弱的空间相关性，这使得它们非常适合利用Mamba强大的序列建模能力。最近，S4MIL 将S4模型作为实例序列的多实例学习器引入到WSI分析中，证明了SSM在捕获长程依赖方面的有效性。值得注意的是，它直接采用S4模型，而没有充分考虑WSIs的独特特征，从而导致次优的结果。

二、本文贡献

( 1 )本文将Mamba框架引入到MIL中以解决长序列建模和过拟合问题，这标志着Mamba在计算病理学中的首次应用。

( 2 )提出序列重排序Mamba ( SR-Mamba )感知实例的顺序和分布，擅长捕获分散的正例实例之间的长距离依赖关系。作为MambaMIL的核心部件，SR - Mamba在顺序排序和转置排序中都可以学习实例之间的相关性，显著增强了原始Mamba捕获更多判别性特征的能力。

( 3 )为了验证MambaMIL的有效性，本文在9个数据集上进行了全面的实验，包括两个具有挑战性的任务的整体比较和消融研究。实验结果表明，MambaMIL可以取得优于当前最好水平的性能。

三、研究内容

1.预备知识（Mamba相关）

2.MambaMIL概述

• 为了有效地捕获长实例序列中的综合上下文信息，本文引入了一种新的方法，MambaMIL，通过将Mamba框架集成到MIL中，如图1所示。通过继承Mamba的属性，MambaMIL使得每个实例可以通过压缩的隐藏状态与先前扫描的任何实例进行交互，这有利于对长序列进行有效建模，同时降低了计算复杂度。

• 具体来说，给定一个WSI，将组织区域划分为L个块序列{ p1，p2，..，pL }，然后通过特征提取器将所有块映射为实例特征$X∈R^{L × D}$，其中D是特征维数。随后，将输入X通过线性投影进行降维。然后将输出输入到一系列堆叠的SR - Mamba模块中，这些模块负责对长序列进行建模。最后，利用聚合模块得到下游任务的包层表示。

3.Sequence Reordering Mamba ( SR-Mamba )。

为了处理受限的感受野，本文设计了感知实例顺序和分布的序列重排序Mamba ( SR-Mamba )，它利用了嵌入在实例中的内在有价值的信息。如图1所示，考虑到零散且稀少的阳性斑块，本文在vanilla Mamba上建立基于SSM的并行分支，以增强长序列建模。SR - Mamba模型对两个具有不同序列顺序的长序列进行建模，每个序列都与一个独特的压缩隐藏状态相关联，便于学习更具判别性的特征。

• 具体地，给定实例特征$X∈R^{L × D}$ ，首先将实例序列划分为大小为R的非重叠片段，并从整个序列中获得N = L / R片段。对于长度不被R整除的序列，将其补零，以便后续重新排序。

• 然后将X输入到两个独立的分支中。对于第一个分支，保留了X的原始排序，并将其反馈给后续的因果卷积层和状态空间模型( State Space Model，SSM )进行序列建模。整个过程可以表述为：

  

  • 然后X也被用来为从SSM得到的Y生成门控值。

• 对于第二个分支，提出了一个Sequence reordering操作作为SR - Mamba的核心部件。

• 具体来说，输入的样本特征被重塑为一个2 - D特征映射，$X∈R^{L × D}→X_{2d}∈R^{R × N × D}$ 。然后，沿着X_2d的第二个维度从每个不重叠的片段中依次采样样本，这可以被视为特征重嵌入re-embedding。

  • 通过执行此操作，生成了具有新排序的实例特征Xr，它可以利用Mamba固有的位置敏感特性来嵌入更具判别力的特征。

  • 整个Sequence reordering操作如图2所示。然后利用后续的因果卷积层和状态空间模型对Xr进行建模，

  • 

• 对于增强后的X′r，通过划分和置换操作将序列重新排列成原始排序，并通过Z将实例特征进行门操作：

  

  • 式中：ψ表示序列恢复操作。在对具有不同排序的长序列建模后，可以得到两个具有判别性的实例特征X′′和X′′r，并将它们聚合得到X输出。最后将聚合操作设计为这两个特征的逐元素相加：

  • 

与原始的Mamba不同，本文保持了顺序的排序和分布，同时从全局角度生成新的实例排序，用于特征重嵌入。SR - Mamba是在vanilla Mamba的基础上进行裁剪的，它能够鲁棒地理解和感知从WSI中分离出来的冗长的实例序列。基于堆叠的SR - Mamba模块，MambaMIL能够以线性复杂度建模长程依赖关系，从而实现有效的模型泛化。

四、实验

1.数据集和评估指标

为了验证本文提出的MambaMIL的有效性，在9个公开的具有挑战性的数据集上对两个具有代表性的下游任务进行了广泛的实验。为了研究泛化性和鲁棒性，使用了两组不同的特征，分别来自于在ImageNet上预训练的ResNet - 50 和在200k病理图像-文本对上预训练的PLIP。

• 生存预测。作者在来自TCGA的7个公开的具有挑战性的癌症数据集( BLCA、BRCA、COADREAD、KIRC、KIRP、LUAD、STAD)上进行了全面的实验，其中包含了带有生存结果注释的WSI。为了减少数据分裂对模型评估的影响，采用5折交叉验证的方法，将数据按照4：1的比例划分为训练和验证子集。使用交叉验证的一致性指数( C-Index )及其标准差( std )来评估提出的MambaMIL的有效性。

• 癌症分型。在两个公开的具有挑战性的数据集BRACS 和NSCLC上进行了对比实验。为了保证对比实验的稳健评估，采用10倍蒙特卡洛交叉验证，其中按照8：1：1的比例将数据划分为训练集、验证集和测试集。此外，为了与现有方法进行公平的比较，还在BRACS数据集的官方拆分上进行了实验，标记为表2中的⋆。根据设定的标准，采用曲线下面积( AUC )和准确率( ACC )指标以及它们的标准差( std )进行评估，这提供了对类不平衡不敏感的可靠评估。

2.实施细则

• 给出了MambaMIL在9个数据集上的实验结果，并与以下方法进行了比较：( 1 )传统的池化方法，包括Mean Pooling和Max Pooling；( 2 ) ABMIL 和三个不同的变体，包括CLAM - MB，DSMIL 和DTFDMIL；( 3 )基于Transformer的TransMIL；( 4 )基于Ssm的S4Mil。

• 遵循通常的设置，采用与CLAM 中相同的数据预处理，并为这些方法设置2 × 10 - 4的学习率，以确保最优的结果并实现公平的比较。相反，为了减少SR - Mamba模块中原子操作在反向传播过程中引入的随机性,对不同的数据集采用不同的学习率进行训练。详细的超参数可以在附录中找到。特殊调整旨在减小梯度差异对收敛性的影响，从而保证稳定性和可重复性。

3.对比结果

生存预测。如表1所示，在7个TCGA癌症数据集上进行了两种不同特征设置的对比实验。结果表明，MambaMIL的总体性能最好.与最先进的方法相比的基准。在这两个特征集下，MambaMIL在所有7个数据集上的平均性能分别比次优方法提高了2.6 %和2.7 %。

癌症分型。表2给出了在两个数据集上的实验结果，包括二分类任务和多分类任务。与当前最先进的方法相比，本文提出的MambaMIL表现出优异的性能，在BRACS数据集上达到了80.4 %的AUC，在NSCLC数据集上达到了95.9 %的AUC。值得注意的是，MambaMIL使用了与ABMIL相同的聚合模块，但显著优于ABMIL，在BRACS和NSCLC数据集上的AUC分别显著提高了3.9 %和2.1 %。

4.消融实验

• 为了评估SR - Mamba的有效性，作者进行了大量的实验来比较Mamba块的不同变体：vanilla Mamba，双向Mamba ( Bi-Mamba ) 和本文提出的SR - Mamba，在生存预测数据集上的性能。为了公平地比较每个特定的数据集，使用相同的设置来训练这些变体。如表3所示，SR - Mamba的性能超过了Mamba和Bi - Mamba，说明了序列重排序的有效性。

• 同时，过拟合也给MIL方法应用于WSI分析带来了巨大挑战，特别是对于Trans MIL等基于变压器的方法。如图3所示，在训练过程中，TransMIL在验证集上表现出明显的过拟合迹象，表现为验证损失显著增加，ACC和AUC指标均有所下降。相比之下，MambaMIL在整个评估期间表现出稳定的性能，表明其具有较强的缓解过拟合的能力。这种能力来源于提取的更具有判别性的表示从各种序列排序，类似于数据增强的效果，这显著增强了本文提出的模型的鲁棒性。