首页
学习
活动
专区
圈层
工具
发布
首页
学习
活动
专区
圈层
工具
MCP广场
社区首页 >专栏 >注意力驱动与卷积并行:KAN-Mamba FusionNet 优化医学图像分割技术 !

注意力驱动与卷积并行:KAN-Mamba FusionNet 优化医学图像分割技术 !

作者头像
未来先知
发布2024-12-19 19:00:36
发布2024-12-19 19:00:36
8100
举报
文章被收录于专栏:未来先知未来先知

医学图像分割在机器人手术、疾病诊断和治疗方案中具有重要意义。本研究提出了一种创新的方法,将Kolmogorov-Arnold Networks (KAN)与适应性Mamba层相结合,用于医学图像分割。 提出的KAN-Mamba FusionNet框架通过集成注意力驱动机制、卷积并行训练和自回归部署,提高图像分割效果,同时保留可解释性,与依赖Mamba进行疾病定位和准确诊断的最新技术形成对比。 作者在三个不同的医学图像分割数据集(BUSI、Kvasir-Seg和GlaS)上评估了提出的KAN-Mamba FusionNet模型。 结果表明,与最新方法相比,KAN-Mamba FusionNet在各项指标上始终具有更好的IoU和F1分数。 此外,作者通过消融研究对模型的行为进行了分析,研究了各种组件的影响,并评估了它们对所提模型整体性能的贡献。 研究发现,该方法对于依赖医学图像分割的可靠性和有效性,提供了一种独特的处理医疗保健中复杂视觉数据问题的方法。

1 Introduction

医学图像分割有助于准确地定位解剖结构特征,从而及时发现异常并实现患者治疗。最近对开发高级深度学习和计算机视觉方法的研究为解决疾病器官定位的挑战提供了一个途径。可靠的图像分割技术对于帮助医学专家做出明智的决策并提高诊断结果至关重要。这可能潜在地增加患者的寿命。

卷积神经网络(CNNs)广泛应用于 Token 个体像素以进行分割任务,以识别肿瘤位置、器官和相关解剖结构。类似于UNet模型的经典基于CNN的方法[27],是基于编码器和解码器结构的图像分割。UNet的变体,如UNet[27]、UNet++[36]和UNet3+[13],使用 Shortcut 将编码器中来自相关低级详细特征图的高 Level 语义特征图与解码器中来自相关低级详细特征图的低 Level 详细特征图相结合。由于这些模型使用层次技术从输入图像中提取特征,因此它们无法获取全局上下文依赖关系,这对于从医学图像中分割大小和形状各异的解剖结构至关重要。

不同于CNNs,ViT(视觉 Transformer )即视觉Transformer [7] 使用注意力机制来捕获全局上下文依赖关系,并可以接受不同的输入尺寸,使其更适合医学图像分割任务 [5]。然而,由于长期特征提取的高计算要求,Transformer 存在较大的模型大小、高内存和计算要求。

为了克服 Transformer 的高计算复杂性,提出了使用线性循环神经网络(RNNs)的状态空间模型(SSM)作为Mamba架构[9]。Mamba架构的不同变体,如U-Mamba [23],Vision Mamba [32]和SegMamba [33],通过使用选择性状态表示,动态地根据输入数据进行调整,从而提高医学图像分割的整体效率。这些模型能够实现时间复杂度为,这比 Transformer 低得多。

基于Mamba的模型使用类似于多层感知器(MLP)的卷积层。在过去的MLP层中,为了捕捉图像中的非线性复杂性,被KAN架构[22]所取代。例如,图像特征之间的边界可能不能捕捉到线性关系,这可能导致从医学图像中丢失一些重要细节,从而可能影响到正确的疾病诊断。为了解决这个问题,基于U-KAN的方法被提出用于医学图像分割[18],通过学习激活函数来捕捉非线性特征。然而,基于KAN架构的方法与Mamba相比具有一个缺点,即无法捕捉到图像中的复杂空间特征进行自注意力学习[34]。

在本文中,作者提出了一种名为KAN-Mamba FusionNet的医学图像分割模型,以解决现有技术的挑战。所提出的架构在Mamba层中引入了KAN,并实现了非线性变换的集成,以提高复杂医学图像数据的表示。在这里,作者利用KAN和基于Mamba架构的协同作用,克服了传统CNN和Transformer的局限性,尤其是在捕捉长程依赖关系和处理医学图像分割数据集方面的效率。

此外,在MLP、Mamba和KAN模型中使用的标准激活函数如ReLU、sigmoid和tanh具有不同的优缺点。例如,ReLU提供了效率和稀疏性,但可能会遇到“死ReLU”问题,而sigmoid和tanh适用于概率建模,但容易受到消失梯度的影响。为了解决这些问题,作者在提出的方法中提出了Bag-of-Activation(BoA)函数,通过可调整参数的加权求和方式将许多激活函数集成在一起。这种方法使得神经网络能够灵活地利用各种激活函数的优势,从而增强其表示复杂数据模式的能力,并提高整体效率。

作者贡献的摘要如下:

  1. 作者提出了一种新颖的神经网络架构KAN-Mamba FusionNet,该架构在U-KAN模型中引入了Mamba块,以引入自注意力层,从而实现对输入数据进行复杂上下文理解的局部化改进。
  2. 作者还创新性地在Mamba架构中添加了一个额外的KAN层,以进一步集成在输入中的任何非线性复杂性上的学习。
  3. 此外,作者在Mamba块中引入了支持Bag of Activation(BoA)功能,该功能动态地将多个激活函数组合在一起,为特征提供更强大的表示。

作者在医学图像数据集BUSI [1],GlaS [29] 和Kvasir-Seg [14]上评估了作者提出的架构,并在最后进行了消融研究。

2 Related Work

医学图像分割对于实现精确诊断和治疗规划至关重要。具有编码器-解码器结构的架构如U-Net在医学图像分割领域取得了显著的成功[27]。扩展如Unet++ [36]和CE-Net [10] 引入了复杂的模块,用于多尺度特征融合和上下文编码,以提高分割精度。除了卷积方法外,使用注意力机制捕获长程依赖关系的 Transformer 模型[31]如Vision Transformer [8],Medical Transformer [30]和TransUNet [5]已得到应用。3D分割 [15]和多维门控循环单元[2]在需要 Voxel 理解的任务中展现出巨大的潜力。最近,基于Mamba架构的体系结构出现,利用状态空间模型(SSMs)进行快速线性推理和交叉扫描模块[9, 20]。它们在生物医学图像分割数据集上的应用,突显了它们在医学成像中的日益重要性。

Kolmogorov-Arnold网络(KANs)为传统MLPs提供了医学分割的创新的替代方案。基于Kolmogorov-Arnold定理,KANs使用一元变换高效逼近多变量连续函数,从而实现对高维数据中的非线性捕捉[12]。KANs的特点在于其可解释性、鲁棒性和高效性,使其适用于需要精确度的应用场景,如医学影像。虽然其在医学分割中的应用还处于初级阶段,但像U-KAN这样的模型,通过使用可学习权重和非线性变换,在提取各种医学数据集的洞察力方面显示出潜力。

基于Mamba的模型旨在解决卷积神经网络(CNNs)在捕捉长程依赖关系方面的限制。视觉Mamba(Vim)和V Mamba通过将状态空间模型与注意力机制相结合,在视频理解和遥感任务中取得了改进的有效性[20]。在医学影像领域,这些模型在分割方面特别高效,利用生物医学图像分割数据集的迁移学习方法。对这些架构与CNNs和视觉 Transformer (ViTs)进行比较研究,突显了在数据需求、可解释性和性能方面的权衡。

生物医学图像分割算法已取得了多项算法改进。混合模型如基于HAAM的U-Net [4]和注意力引导的U-Net [26]利用注意力机制提高了特征提取的精度和准确度。其他设计如U2-MNet [19]和SU-Next [37]采用了诸如残差层析分离卷积和多层感知机框架等先进技术,以提高分割速度和准确度。像特征金字塔非局部网络 [28]和级联卷积神经网络 [3]在Breast UltraSound Images (BUSI)等数据集上改善了性能指标,实现了显著的dice分数。迁移学习和新型注意力机制正在推进图像分割的发展,KANs和Mambab基于的模型为传统的CNN结构提供了可行的替代方案。

在本文中,作者提出了一种新颖的架构KAN-Mamba FusionNet,该架构利用KAN捕捉的非线性复杂性,同时利用Mamba的能力从输入图像中捕获隐含的空间特征。

3 Proposed Architecture

提出的KAN-Mamba FusionNet架构将注意机制与健壮的KAN非线性模型[12]相结合。该思想是,用KAN块替换传统Mamba[9]中在SSM模型之前的卷积和批量归一化层(类似于MLP)。图1代表了作者提出的模型的总体架构流程。它利用了U-KAN架构[18]作为基础。图1中的整体架构图中的符号或颜色编码表示以下内容:

C1/C2/C3 分别表示卷积块()。

P1/P2: 代表 Patch 嵌入层。

M1:代表修改后的Mamba-KAN块。

L1/L2:表示 LayerNorm 。

D1/D2/D3/D4/D5: 代表深度卷积()模块。

K1/K2/K3/K4:代表KAN块。

O1:代表卷积层,输出最终结果。

所提出的架构流程受到过去在传统 MLP 层替换中使用的 KAN 块 [12] 的启发。为了捕捉非线性的复杂性,KAN 使用可学习权重来表示边缘。同样,过去曾使用 SSM 和注意层来选择性地处理输入并提取相关信息 [20]。结合这些,作者提出使用一个 Mamba 块与 KAN 块。此外,由于 Mamba 块包含卷积和批量归一化层(这与传统 MLP 的处理方法相似),作者提出了一种替换,即将这些层替换为单个 KAN 层。此外,作者提出了一种通过用一组激活函数替换单个激活函数来实现激活函数变化的想法,其中模型学习分配给每个函数的权重。接下来,作者将详细解释所提出的架构块的实现。

卷积模块:如图2所示,卷积模块C1/C2/C3由2D卷积层、批量归一化层、最大池化层和ReLU激活函数组成。

让作者将这种块变换表示为 ,其中:

在卷积块中, 表示输入的值。

如图3所示,作者通过插值操作替换了卷积层B中的池化操作,从而表达了深度卷积块:

Mamba块:Mamba块在第一个卷积块()之后添加,如图1所示。输入数据经过线性投影层时,保持原始输入不变,进行最终的线性组合。线性投影之后,应用卷积核层,在2D空间维度上进行卷积。进一步,通过批量归一化2D层和空间注意力层进行处理。从卷积到批量归一化的过渡组织成一个系列,形成一个三 mini-block 结构,空间注意力层将这些块分隔开。然后,通过激活函数处理输出,并进一步输入到状态空间模型(SSM)块。此外,在最终处理中还应用了空间注意力层。

注意层在卷积块注意力模块(Convolutional Block Attention Module)上可以表示为[34]:

在这里, 代表了一个在更大卷积函数 上作用的 sigmoid 函数,该函数的滤波器大小为 7 x 7,以捕获更多的空间特征。

Reshape 后,这个输出(形成 Shortcut )与输入(相同输入)的激活函数输出线性相加。

激活函数的集合:Mamba块中使用的单激活函数sulu被激活函数的集合所替代。模型学习到集合中激活函数的权重,以找到最有效的函数。

激活函数的集合可以表示为:

在作者的架构中,作者将在激活函数袋中包含以下激活函数 - ReLU、Tanh、Softplus、GELU 和 SiLU,保持相同的初始权重。其中,α_p 形成了与激活函数 ψ_p 相关的可学习权重。

KAN in Mamba:如图4所示,该块的主要结构在SSM之前使用的卷积层上有所不同。这些卷积层被KAN块[18]所取代。输入数据首先通过一个patch embedding层,然后与一个单一的KAN层进行处理。进一步,输出经过一个激活函数处理,然后输入SSM块。经过空间变换后,结果经过一个空间注意力层处理,然后经过一个线性投影层。投影层中进行的缩放操作有助于线性组合输入(通过 Shortcut )和同一输入上的激活层输出。

具有K层的MLP可以表示为以下形式[18]:

其中,,,表示变换矩阵,是激活函数,是输入数据。

类似地,作者可以写出KAN形式化问题[18]:

其中,Φ 使用可学习的激活函数 表示,使得

KAN层包括KAN线性变换[18],然后是一个深度卷积层(由Conv2D层、批量归一化和ReLU激活函数组成)。KAN层的基本单元可以表示为:

这样的转换过程重复三次,得到KAN层。用于表示包含KAN层、归一化层和恒等函数的KAN块,如图5所示。

因此,输出可以数学表示为:

在KAN模块中传递的输入为。

现在,作者将详细说明修改后的Mambo-KAN块的数学公式。

首先经过一个patch embedding层。

表示了在 Tokenized KAN 短语操作中与 Seq. U-KAN 模型 [18] 中提出的类似 Token 化步骤。然后将其输入到 KAN 块中,如下所示:

此前提到的操作代表KAN块变换。这一过程进一步通过一层的Bag of Activation (BoA)函数进行处理,以便

此外,输出经过状态空间模型(SSM)进行转换,得到以下输出:

此后,添加了一个注意层,结果如下:

其中 代表上述提到的注意力函数。

图4:修改后的Mambo-KAN块 图5:KAN块

图3:深度卷积块()最后,与 Shortcut 以及并行BoA函数的组合,给出最终输出为

此外,整体架构的损失使用二进制交叉熵(BCE)损失函数和 dice 损失函数的组合来计算,衡量了 N 个像素的 ground truth 标签偏差。

其中,和分别表示实际测量和估计的像素级标签。为了处理除零问题, dice 损失中包含一个小常数 。

4 Datasets

作者在三个不同的医学图像分割数据集上评估了作者的模型:乳腺癌超声图像(BUSI)[1],分段息肉图像[14]和乳腺肿块分割图像(GlaS)[29]。鉴于每个数据集的独特特性,这些评估为测试作者的方法的有效性提供了强有力的支持。

商业[1]: 数据集包含来自超声扫描的医学图像及其对应的分割 Mask ,用于识别与乳腺癌相关的肿瘤。整个数据集包含708张图像,其中210、437和133分别代表恶性、良性及正常乳腺癌病例的图像数量。作者使用了代表乳腺癌的图像。图像被统一缩放到256 x 256像素。

Kvasir-SeG数据集包含1000张胃肠道息肉图像(息肉是结直肠癌的前体)及其相应的分割 Mask ,这些图像由一名医生手动标注,随后由经验丰富的消化科医生进行验证。所有图像均统一缩放到256 x 256像素。

GlaS[29]: 该数据集包括腺体分割图像,并与西班牙巴塞罗那的医院诊所相关联。对于作者正在进行的研究,作者使用了数据集中165张图像,所有图像都重新缩放到256 x 256像素。

5 Experiments and Results

Implementation Details

对于这三个数据集,作者将学习率设置为1e-4,使用Adam优化器进行训练,并集成了一个余弦退火学习率调度器,最小学习率为1e-5。使用的损失函数是二进制交叉熵和 dice 损失的组合。每个数据集分别以 4:1 的比例进行训练和验证。训练过程持续了400个周期,输入数据应用了基本数据增强,如随机旋转和翻转。

Evaluation Details

为了评估模型的性能,作者在所有数据集上报告了验证IoU(预测与真实边界框的重叠度)和F1分数,并与最先进的方法进行了比较。IoU衡量了预测边界框和真实边界框之间的重叠度,而F1分数提供了精确率和召回率的平衡衡量指标。这些指标对于理解模型在不同分割任务上的准确性和可靠性至关重要。报告的结果是三个独立运行的平均值。作者还提出了计算成本指标,如GFLOPs和模型参数的总数,并将其与最先进的方法进行了比较。此外,作者还对作者的模型进行了消融研究,借助于其他有用的参数。

Performance Comparisons With State-of-Art Methods

表1展示了在三个数据集上的图像分割性能比较,作者与五种最先进的方法进行了比较,包括U-Net [27],U-NeXt [30],Rolling-UNet [21],U-Mamba [24]和Seg. U-KAN [18]。作者的实验结果表明,在所有数据集上,作者提出的KAN-Mamba FusionNet方法在与其他最先进方法的比较中,始终表现出良好的性能。

作者的模型在IoU和F1得分这两项指标上的表现,明显展示了其在处理各种分割挑战上的健壮性和灵活性。此外,F1得分提供了精确率和召回率的平衡衡量,在医学诊断中,该模型非常擅长识别实际阳性病例,最大限度地减少了假阴性(即未能检测到实际存在的病例)的数量,这样的检测失败可能导致严重的后果,对任何人来说都可能带来严重的影响。

此外,如表2所示,作者计算了GFLOPs和模型参数,以强调作者的模型可以在计算成本方面实现最小折衷,与以往的方法相比。

Ablation Studies

作者对提出的KAN-Mamba FusionNet进行了全面的评估,特别关注一些新增特征在模型中的影响。

Effect of Changes in Mamba Block: 表3展示了对Mamba块中修改进行消融研究的结果,以验证Mamba块中修改的影响。如表3所示,Mamba+MLP组合代表了经典的Mamba,之后是MLP层。然而,经典的模型表现不佳,而用KAN层替换MLP层则导致IoU、F1、准确性、AUC、精确率以及召回率等验证指标的提高,如表中的验证指标所示。此外,在经典的Mamba中加入一袋激活函数,然后用KAN层替换MLP层(如作者所提出的模型那样),可以进一步提高列出的性能参数。

Effect of Bag of Activation Functions: 如图4所示,当未在Mamba块中添加额外的激活函数层时,模型表现不佳,这与作者提出的架构一致。然而,当引入一个激活函数(例如,ReLU)时,模型的性能显著提高。此外,通过将多个激活函数(即“激活函数袋”)结合,实际上可以提高模型的学习能力,从而实现更好的结果。这些结果强调了添加“激活函数袋”的重要性,表明激活函数的多样性在提高模型性能中起着重要作用。

不同的激活函数捕获数据的不同方面,每个都为神经网络带来自己的优势和劣势。例如,ReLU(Rectified Linear Unit)[25]被定义为:

ReLU 引入了稀疏性,通过仅激活正输入,可以实现更快的收敛和降低计算复杂性。然而,ReLU 可能会遇到“死 ReLU”问题,即如果神经元持续输出零,则可能导致神经元变得无效并停止学习。另一方面,Sigmoid 函数 [11] 被定义为:

它擅长建模概率,在需要输出值在0和1之间的场景中非常有用。然而,它容易导致梯度消失,这可能会阻碍深度网络的训练。同样,双曲正切(tanh)[17]激活函数:

捕捉到零周围的对称关系,输出范围在-1到1之间,有助于数据中心化和提高收敛性。然而,它也存在消失梯度问题。

为了克服这些限制,作者提出了一种将多个激活函数合并为一个Bag of Activation(BoA)函数的方法,该方法定义在方程(4)中。这种形式使得网络能够探索更丰富的函数空间,从而增强其捕捉数据中多样化模式的能力。通过在训练过程中优化每个激活函数相关的权重,模型可以利用每个激活函数的优势,同时减轻它们的个别弱点。

例如,这种组合可以动态地平衡ReLU引入的稀疏性,同时利用Sigmoid的概率建模和tanh的零中心输出。这种自适应加权是通过基于梯度的优化实现的,确保网络为不同的任务或数据特性选择最合适的激活行为。

此外,网络的不同部分可能需要不同的激活行为,从而实现更专业和自适应的非线性变换。通过使网络能够自适应地结合多个激活函数,Bag of Activation Functions方法提供了一种强大和多功能的机制,以增强神经网络在各种机器学习任务中的性能。

6 Conclusion

在本文中,提出了KAN-Mamba FusionNet架构,用于医学图像分割。

该架构融合了 Kolmogorov-Arnold Networks (KAN)和Mamba网络,并引入了激活函数袋(bag of activation functions)以降低复杂性,捕捉输入图像中的非线性特征,并自适应地选择有效的激活函数以提高模型性能。

作者在BUSI、Kvasir-Seg和Glas数据集上进行的实验验证了我们的假设,与最先进的方法相比,展示了更高的IoU(交并比)和F1(Dice)分数。

此外,消融研究显示了每个架构组件的关键作用,特别是突出了激活函数袋如何对整体性能提升做出贡献。 总结来说,KAN-Mamba FusionNet提高了医学图像分割的准确性和鲁棒性,并且在计算效率方面表现优异。

激活函数袋的使用提供了一种灵活的方法来增强神经网络的效能,解决了医疗保健领域中的复杂视觉数据问题,并通过更准确的分割掩模可能改善患者预后。

通过自适应地整合各种激活函数,BoA提高了神经网络在多种机器学习应用中的多样性和鲁棒性。

参考文献

[0]. KAN-Mamba FusionNet: Redefining Medical Image Segmentation with Non-Linear Modeling.

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

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

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

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

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 1 Introduction
  • 2 Related Work
  • 3 Proposed Architecture
  • 4 Datasets
  • 5 Experiments and Results
    • Implementation Details
    • Evaluation Details
    • Performance Comparisons With State-of-Art Methods
    • Ablation Studies
  • 6 Conclusion
  • 参考文献
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档