Loading [MathJax]/jax/output/CommonHTML/config.js
前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
圈层
工具
发布
首页
学习
活动
专区
圈层
工具
MCP广场
社区首页 >专栏 >综合 CNN 和 Transformer, Microscopic-Mamba 仅用 4M 参数揭示微观图像的秘密 !

综合 CNN 和 Transformer, Microscopic-Mamba 仅用 4M 参数揭示微观图像的秘密 !

作者头像
未来先知
发布于 2024-10-22 08:03:01
发布于 2024-10-22 08:03:01
3660
举报
文章被收录于专栏:未来先知未来先知

在医学显微图像分类(MIC)领域,CNN和Transformer模型被广泛研究。然而,CNN在建模远程依赖关系方面存在困难,限制了它充分利用图像中的语义信息的能力。相反,Transformer受到平方计算复杂性的阻碍。 为了解决这些挑战,作者提出了一种基于Mamba架构的模型:Microscopic-Mamba。 具体而言,作者设计了一种部分选择前馈网络(PSFFN),用以取代视觉状态空间模块(VSSM)的最后线性层,从而增强Mamba在局部特征提取方面的能力。 此外,作者引入了一种调制交互特征聚合(MIFA)模块,以有效地调制和动态地聚合全局和局部特征。作者还引入了一种并行VSSM机制,以提高跨通道信息交互的同时减少参数数量。 大量实验证明,作者的方法在五个公共数据集上实现了最先进的性能。 代码可在https://github.com/zs134/Microscopic-Mamba中找到。

I Introduction

显微镜成像技术在医学领域中至关重要,已成为现代医学研究和临床诊断的不可或缺的工具[1]。通过对显微镜图像进行分类,医学研究行人可以在组织、细胞和分子水平观察结构和动态变化,进一步理解疾病机制和病理过程[2]。然而,基于卷积神经网络(CNN)和视觉 Transformer (ViTs)的典型显微镜成像分类(MIC)任务存在局限性,这严重影响了它们捕捉长期信息的能力。虽然有些研究提出了一些方法来降低ViTs的计算复杂性,但牺牲了其在自注意力中捕捉全局信息的能力。因此,迫切需要一种新的架构,可以同时有效地捕捉全局和局部特征,同时保持线性复杂性。

状态空间模型(SSM)最近得到了广泛关注。基于经典的SSM研究,现代SSM如Mamba不仅建立了长期依赖关系,还具有与输入大小相关的线性复杂性,使其在轻量级MIC应用中与CNN和ViTs竞争[16]。

研究行人已将Mamba应用于医学成像领域,如Segamba[17]创建了SSM-CNN混合模型以有效地模拟像素级的长期依赖关系。T-Mamba[18]将共享位置编码和基于频率的特征集成到视觉Mamba中,以解决空间位置保留和频率域特征增强的局限性。Contrastive Masked Vim Autoencoder 是一种新的3D多模态数据表示学习方法,通过对比学习提高阿尔茨海默病的诊断效果。

然而,当直接应用于MIC任务时,这种有前景的架构的准确性无法匹配其效率。与具有明显目标特征的其他视觉任务不同,MIC需要捕捉全局上下文和局部详细特征。因此,本文提出了一种增强的Mamba架构用于MIC,以提高其在MIC任务中的应用。

总之,作者的贡献如下:

(1)作者提出了一种轻量级的混合架构Microscopic-Mamba,它结合了Mamba在学习和捕捉全局特征方面的优势以及CNN在提取局部特征方面的力量。从作者的知识中可以得知,这是Mamba在MIC任务中的首次应用,为未来高效的MIC研究提供了新的 Baseline 和参考。

(2)作者提出了一种新颖的双分支混合-卷积-SSM块,包括了一个卷积分支和一个SSM分支,其中以自适应残差视觉状态空间(ARVSS)块为核心。作者引入了一种部分选择的前馈网络(PSFFN)作为视觉状态空间模块(VSSM)中线性层的高级替换,提高了Mamba的局部信息感知能力。为了进一步降低模型参数,作者引入了一个并行VSSM机制。此外,作者提出了一种模量交互特征聚合(MIFA)模块,有效地聚合全局和局部信息。

(3)作者在五个开源公共显微镜图像数据集上进行了广泛的实验,结果表明Microscopic-Mamba在保持较低参数和计算复杂性的同时,超过了当前最先进的方法。

II Method

Overall Architecture

图1(a)显示了微镜-Mamba的整体架构。具体来说,它由四个阶段组成,这些阶段由堆叠的混合卷积-正则化交叉模型( Hybrid-Conv-SSM Blocks)构成。此外,每个阶段在其之前都有一个嵌入或融合层,用于实现空间下采样和通道扩展。作者将全局平均池层应用到最终输出,并将其发送到线性分类头。

混合卷积-正则化交叉模型(Hybrid-Conv-SSM Block) 混合卷积-正则化交叉模型是微镜-Mamba的核心组件。如图1(b)所示,它具有双分支结构,包括SSM分支和卷积分支。它通过通道分裂将特征图分成两组,然后分别由卷积分支和SSM分支处理。卷积分支由深度可分离卷积(DWConv)和点对点卷积(PWConv)组成,用于提取局部信息。SSM分支由并行有效视觉Mamba(PEVM)块组成,用于捕获全局上下文信息。最后,作者提出的调制互作用特征聚合(MIFA)模块将这两种类型的信息聚合,完成了全局-局部特征融合。

Parallel Efficient Vision Mamba Block

如图1(c)所示,并行高效的视觉Mamba(PEVM)模块包括一个 layer normalization(LN)层和四个自适应残差视觉状态空间(ARVSS)模块。以前的研究表明,并行网络可以帮助减少模型参数而不损失准确性。具体地说,输入 经过 LN 层后,特征图沿着通道分成四组:,等。这些组由 ARVSS 模块并行处理,得到 , 等。最后,将这些特征图在通道上连接起来。然后,通过 Channel Shuffle 增强模块间的信息交换,得到输出 。

自适应残差视觉状态空间块(ARVSS)。图1(e)描述了 ARVSS 模块。经过层规范化后的输入 被分成两个支路。这个过程可以表示为:

其中 表示层规范化过程。在第一个支路中,输入 经过一个线性层和一个 SiLU 激活函数,产生输出图 。在第二个支路中,输入 经过一个线性层、DWConv 和一个 SiLU 激活函数。然后,通过 2D 选择性扫描(SS2D)模块进行进一步的特征提取。最后,特征图进行规范化,得到输出 。然后,两个支路的特征图通过逐元素乘积进行融合,得到特征图 。最后,特征图经过部分选择性的前馈网络(PSFFN)进行高效的选定感受野提取,并将残差与原始输入连接起来,得到 ARVSS 模块的最终输出:。该过程定义如下:

  1. 6.

其中,表示线性处理,表示2D选择性扫描,表示深度可分卷积,是SiLU激活函数,表示逐元素乘法,表示逐元素加法,表示通过部分选择性前馈网络处理。

部分选择性前馈网络。为了进一步增强Mamba提取局部特征的能力,作者提出了一种新的部分选择性前馈网络(PSFFN),用PSFFN替换ARVSS块中的线性层。如图1(f)所示,作者主要使用PWConv和高效部分深度可分卷积(PDWConv)来完成前馈特征提取。具体来说,假设输入特征为,它首先通过PWConv降低channel维度,然后经过SiLU激活函数,得到。接着,作者将输入特征沿着channel分成两组,一组占据1/4的原始channel,另一组占据3/4,得到输出和。然后,PDWConv和SiLU进一步从中提取特征,得到输出。最后,PSFFN使用channel concatenation恢复channel维度的尺寸,同时使用channel shuffling在channel维度上对特征图进行Shuffle,以避免不同通道之间的信息损失。这个过程定义如下:

在论文中,作者引入了PSFFN,使得 ARVSS 块可以根据需要自适应地调整由 SS2D 产生的全局特征,从而降低参数数量并提升模型对细节特征的感知能力。

其中, 表示点卷积, 表示通道剪切, 表示高效的部分深度可分卷积, 表示通道切乱, 表示通道拼接。

Modulation Interaction Feature Aggregation Module

自视觉状态空间块专注于捕捉全局特征,因此作者引入了并行卷积分支和PSFFN以包含局部特征。然而,这些还不能有效地融合全局和局部特征。因此,作者引入了调制交互特征聚合(MIFA)模块,该模块位于两个分支之间,用于调制和调整特征和。具体来说,通过计算全局特征图和局部特征图和,两个输入分支被重新加权。计算过程可以表述为:

其中,表示GELU激活函数,表示ReLU激活函数,表示Sigmoid激活函数。(11)

接下来,通过矩阵点积,全局和局部特征图分别应用于各自的分支。在此之后,采用通道 ConCat 和ECA注意力[33]来完成特征图聚合。这个过程可以表述为:

(13)

其中,表示MIFA模块的最终输出,表示ECA注意力。

III Experiments

Datasets and Implementation Details

为了验证模型的性能,作者选择了五个公共医学图像数据集:视网膜色素上皮细胞(RPE)的数据集、简化的病理图像分析数据集(MHIST)、疟疾细胞图像数据集(SARS)、组织MNIST和MedFM结肠[34,35,36,37,38]。RPE数据集包含1862张视网膜色素上皮细胞的图像,分为四个类别。SARS数据集包括27558张疟疾细胞的图像,分为两个类别。MHIST包含3152张结肠息肉的二分类图像。MedMF结肠数据集对结肠镜病理切片中的肿瘤组织进行了分类,包括来自396名结肠癌患者的1009个病理组织贴片。组织MNIST包含了236386张人类肾皮质细胞的图像,分为八个组。所有数据集都按照6:2:2的比例分别划分成了训练、验证和测试集。

作者在PyTorch 2.0.0上实现了Microscopic-Mamba,并在NVIDIA RTX 3090(24GB)上训练,循环200次,每次批处理16。作者使用了Adam优化器,初始学习率为0.0001,权重衰减为1e-4,以及交叉熵损失来优化模型参数。此外,作者引入了余弦退火学习率衰减策略和10个epoch的预热训练阶段。

Comparison with SOTA Models

表1显示了作者的方法与目前最先进的方法的比较。这些方法包括基于CNN的,基于Transformer的,混合CNN-Transformer的以及基于Mamba的方法。从表1中可以看出,微观-Mamba在五个数据集的参数(Params),通用平均关联计算(GMACs),总体准确率(OA)和曲线下面积(AUC)方面取得了最佳结果。

Ablation Study

表2显示,作者对RFE和MHIST数据集进行了详细的消融实验。结果表明,MIFA和PSFFN模块可以提高模型性能,同时保持较低的参数数量。当同时使用MIFA和PSFFN时,获得最佳结果。此外,如表3所示,引入的并行VSSM机制不仅降低了模型复杂性,而且通过引入多尺度交互提高了性能。

IV Conclusion

在本文中,作者提出了一种基于Mamba的新网络:Microscopic-Mamba。它降低了在实际医疗场景中的计算约束,具有更低的计算复杂性。Microscopic-Mamba采用双分支结构,由SSM和CNN组成。

作者提出了调制交互特征聚合(MIFA)模块,以高效地调节和聚合两个分支的全局-局部特征。

此外,作者将Mamba模块中的最终线性层替换为部分选择的前馈网络(PSFFN),以进一步提高Mamba的局部感知能力。

同时,作者引入了一种并行VSSM机制,在降低计算复杂性的同时增强了不同通道之间的信息交互。

实验结果显示,作者的提出的方法在其他最新方法的五个公共数据集上取得了最佳性能。作者认为,作者的专用的Mamba架构设计可以为MIC提供新的视角。

参考

[1].Microscopic-Mamba: Revealing the Secrets of Microscopic Images with Just 4M Parameters.

本文参与 腾讯云自媒体同步曝光计划,分享自微信公众号。
原始发表:2024-10-21,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 未来先知 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
暂无评论
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档