KAN-Mamba FusionNet: Redefining Medical Image Segmentation with Non-Linear Modeling通过非线性建模重新定义医学图像分割

这篇文章提出了一种名为KAN-Mamba FusionNet的新型神经网络架构，用于医学图像分割。该架构结合了Kolmogorov-Arnold网络（KAN）和Mamba层，旨在解决现有医学图像分割方法在处理复杂视觉数据时的局限性。以下是文章的主要内容总结：

1. 背景与动机：

   • 医学图像分割在疾病诊断、治疗计划和机器人手术中至关重要。

   • 现有的方法（如CNN和Transformer）在处理长距离依赖关系和非线性特征时存在局限性。

   • KAN能够捕捉非线性特征，而Mamba模型通过状态空间模型（SSM）提高了计算效率。

2. 提出的方法：

   • KAN-Mamba FusionNet：结合了KAN的非线性建模能力和Mamba的注意力机制，提升了医学图像分割的准确性和效率。

   • 激活函数袋（BoA）：引入了多种激活函数的加权组合，增强了模型的表达能力。

   • 架构设计：在Mamba块中引入了KAN层，替换了传统的卷积和批量归一化层，以更好地捕捉图像中的非线性特征。

3. 实验与结果：

   • 在三个医学图像分割数据集（BUSI、Kvasir-Seg和GlaS）上进行了评估。

   • 实验结果表明，KAN-Mamba FusionNet在IoU（交并比）和F1分数上优于现有的最先进方法。

   • 消融研究验证了KAN层和BoA对模型性能的贡献。

4. 主要贡献：

   • 提出了KAN-Mamba FusionNet架构，结合了KAN和Mamba的优势。

   • 引入了激活函数袋（BoA），增强了模型的非线性表达能力。

   • 在多个医学图像分割数据集上验证了模型的有效性，展示了其在处理复杂视觉数据方面的优势。

5. 结论：

   • KAN-Mamba FusionNet在医学图像分割中表现出色，具有较高的准确性和计算效率。

   • 该模型通过自适应地结合多种激活函数，提升了神经网络的鲁棒性和多功能性，有望在医疗保健领域带来更准确的诊断和治疗方案。

这篇文章提出了一种创新的医学图像分割方法，通过结合KAN和Mamba的优势，显著提升了分割性能，并通过实验验证了其有效性。这里是自己的论文阅读记录，感兴趣的话可以参考一下，如果需要阅读原文的话可以看这里，如下所示：

摘要

医学图像分割在机器人手术、疾病诊断和治疗计划中至关重要。本研究提出了一种创新的方法，将Kolmogorov-Arnold网络（KAN）与改进的Mamba层相结合，用于医学图像分割。提出的KAN-Mamba FusionNet框架通过将注意力驱动机制与卷积并行训练和自回归部署相结合，改进了图像分割，同时保持了可解释性。与现有技术相比，现有技术仅依赖Mamba进行病灶定位和准确诊断。我们在三个不同的医学图像分割数据集（BUSI、Kvasir-Seg和GlaS）上评估了提出的KAN-Mamba FusionNet模型。结果表明，与现有技术相比，KAN-Mamba FusionNet在IoU和F1分数上表现更好。此外，我们通过消融研究提供了对模型行为的见解，检查了各个组件的影响并评估了它们对模型整体性能的贡献。研究结果说明了该方法在可靠医学图像分割中的强大性和有效性，为解决医疗保健中的复杂视觉数据问题提供了一种独特的方法。

1 引言

医学图像分割有助于准确定位解剖特征，从而及时发现异常并进行患者治疗。最近的研究致力于开发先进的深度学习和计算机视觉方法，为解决病灶器官定位中的挑战提供了途径[6, 16, 27, 35]。可靠的图像分割技术对于帮助医学专家做出明智决策和改善诊断结果至关重要，这可能会增加患者的预期寿命。

卷积神经网络（CNN）通常用于标记单个像素以进行分割任务，以识别肿瘤位置、器官和相关解剖特征。经典的基于CNN的方法（如UNet模型[27]）基于编码器和解码器架构进行图像分割。UNet的变体（如UNet[27]、UNet++[36]和UNet3+[13]）使用跳跃连接将编码器中的低层详细特征图与解码器中的高层语义特征图相结合。由于这些模型使用分层技术从输入图像中提取特征，它们无法获取全局上下文依赖关系，而这些依赖关系对于从医学图像中分割不同大小的解剖部位至关重要。

与CNN不同，视觉变换器（ViT）[7]使用注意力机制来捕捉全局上下文依赖关系，并且可以接受不同大小的输入，使其更适合医学图像分割任务[5]。然而，变换器由于长时间的特征提取而面临模型规模大、内存和计算需求高的问题。

为了克服变换器计算复杂度高的问题，提出了使用线性RNN（递归神经网络）的状态空间模型（SSM）作为Mamba架构[9]。Mamba架构的不同变体（如U-Mamba[23]、Vision Mamba[32]和SegMamba[33]）通过采用基于输入数据动态调整的选择性状态表示，提高了医学图像分割的整体效率。这些模型能够实现O(nlog(n))的时间复杂度，远低于变换器。

Mamba模型使用类似于多层感知器（MLP）的卷积层。过去，MLP层已被KAN架构[22]取代，以捕捉图像中的非线性复杂性。例如，图像特征之间的边界可能不呈现线性关系，导致医学图像中丢失一些重要细节，从而可能导致错误的疾病诊断。为了解决这个问题，提出了基于U-KAN的方法用于医学图像分割[18]，通过学习激活函数来捕捉非线性特征。然而，KAN架构中的方法在捕捉图像中复杂空间特征的自我注意力学习方面存在劣势[34]。

在本文中，我们提出了KAN-Mamba FusionNet模型用于医学图像分割，以解决现有技术面临的挑战。所提出的架构在Mamba层中集成了KAN，并引入了非线性变换以改进复杂医学成像数据的表示。我们利用KAN和Mamba架构的协同作用，克服了传统CNN和变换器的局限性，特别是在捕捉长距离依赖关系和处理图像分割数据集方面。

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

以下是我们的贡献总结：

• 我们提出了一种新颖的神经网络架构KAN-Mamba FusionNet，将Mamba块引入U-KAN模型中，以引入自注意力层，从而从输入数据中进行复杂的上下文理解，以改进定位。

• 我们进一步通过在Mamba架构中添加另一个KAN层来进一步集成对输入中任何非线性复杂性的学习。

• 此外，我们在Mamba块中引入了“激活函数袋”（BoA）的支持，动态结合多个激活函数以提供更鲁棒的特征表示。

• 我们在医学图像数据集（BUSI[1]、GlaS[29]和Kvasir-Seg[14]）上评估了所提出的架构，并在最后提供了消融研究。

2 相关工作

医学图像分割对于实现准确的诊断和治疗计划至关重要。像U-Net这样的架构以其编码器-解码器结构在该领域取得了显著成功[27]。扩展版本如Unet++[36]和CE-Net[10]结合了复杂的模块用于多尺度特征融合和上下文编码，以提高分割精度。除了基于卷积的方法外，变换器模型[31]如Vision Transformer[8]、Medical Transformer[30]和TransUNet[5]也被应用，使用注意力机制捕捉长距离依赖关系。像3D分割[15]和多维门控循环单元[2]这样的方法在需要体积理解的任务中显示出潜力。最近，基于Mamba的架构出现，利用状态空间模型（SSM）进行快速线性时间推理和交叉扫描模块[9, 20]。它们在生物医学图像分割数据集上的应用突出了它们在医学成像中的重要性。

Kolmogorov-Arnold网络（KAN）为医学分割提供了传统MLP的创新替代方案。基于Kolmogorov-Arnold定理，KAN通过单变量变换有效地近似多元连续函数，从而能够捕捉高维数据中的非线性[12]。KAN以其可解释性、鲁棒性和效率著称，适用于需要精确性的应用，如医学成像。虽然它们在医学分割中的应用尚处于初期阶段，但像U-KAN这样采用可学习权重和非线性变换的模型已经显示出从各种医学数据集中提取洞察力的潜力。

基于Mamba的模型旨在解决卷积神经网络（CNN）在捕捉长距离依赖关系方面的局限性。Vision Mamba（Vim）和VMamba通过将状态空间模型与注意力机制相结合，在视频理解和遥感任务中展示了改进的效果[20]。在医学成像中，这些模型在分割方面特别高效，利用迁移学习在生物医学图像分割数据集上进行训练。这些架构与CNN和视觉变换器（ViT）的比较研究突出了在数据需求、可解释性和性能方面的权衡。

生物医学图像分割已经看到了几种算法的改进。混合模型如基于混合自适应注意力模块（HAAM）的U-Net[4]和注意力引导的U-Net[26]通过使用注意力机制增强特征提取，实现了高精度和准确性。替代设计如U2-MNet[19]和SU-Next[37]采用残差初始深度可分离卷积和多层感知器框架来提高分割速度和准确性。像特征金字塔非局部网络[28]和级联CNN[3]这样的模型在乳腺超声图像（BUSI）等数据集上提高了性能指标，实现了显著的Dice分数。迁移学习和新颖的注意力机制等技术正在推进图像分割，KAN和基于Mamba的模型为传统CNN结构提供了可行的替代方案。

在本文中，我们提出了一种新颖的架构KAN-Mamba FusionNet，它利用KAN捕捉的非线性复杂性，并利用Mamba的能力从输入图像中捕捉隐含的空间特征。

3 提出的架构

提出的KAN-Mamba FusionNet架构将注意力驱动机制与鲁棒的KAN非线性模型[12]相结合。其思想是将传统Mamba[9]中SSM模型之前使用的卷积和批量归一化层（类似于MLP）替换为KAN块。图1展示了我们提出的模型的整体架构流程。它使用U-KAN架构[18]作为基础。整体架构图（图1）中的符号或颜色编码表示以下内容：

• C1/C2/C3：表示卷积块（ConvB）。

• P1/P2：表示补丁嵌入层。

• M1：表示改进的Mamba-KAN块。

• L1/L2：表示层归一化。

• D1/D2/D3/D4/D5：表示深度卷积（DwConvB）块。

• K1/K2/K3/K4：表示KAN块。

• O1：表示Conv2D层，生成最终输出。

提出的架构流程受到过去在替换传统MLP层时使用的KAN块的启发[12]。为了捕捉非线性复杂性，KAN在边缘上使用可学习的权重。同样，过去也有使用SSM和注意力层来选择性处理输入并提取相关信息[20]。结合这些，我们提出在Mamba块中使用KAN块。此外，由于Mamba块由卷积和批量归一化层组成（类似于传统MLP的处理方法），我们提出将这些层替换为单个KAN层。此外，我们提出了一种激活函数的变体，通过将单个激活函数替换为“激活函数袋”，其中模型学习分配给每个函数的权重。所提出的架构块的实现细节如下：

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

我们将此块变换表示为ConvB(⋅)，使得：

Mamba块：Mamba块添加在第一个卷积块ConvBC1之后，如图1所示。输入数据通过线性投影层，同时保持原始输入不变以进行最终的线性组合。在线性投影之后，应用卷积核层，在2D空间维度上进行卷积。然后通过批量归一化2D层和空间注意力层进行处理。从卷积到批量归一化的步骤被组织成一系列，形成一个三迷你块结构，空间注意力层分隔这些块。输出然后通过激活函数，进一步馈入状态空间模型（SSM）块。此外，应用空间注意力层进行最终处理。

卷积块注意力模块上的注意力层可以表示为[34]：

其中，σ表示作用于较大卷积函数的sigmoid函数，滤波器大小为7 x 7，以捕捉更多的空间特征。

在重塑后，此输出与输入（形成跳跃连接）和激活函数的输出线性组合，其中相同的输入被馈入。

激活函数袋：Mamba块中使用的单个激活函数sulu被替换为“激活函数袋”。模型学习分配给“激活函数袋”的权重，以找到最有效的函数。

“激活函数袋”可以表示为：

Mamba中的KAN：如图4所示，此块的架构主要在于Mamba中在SSM之前使用的卷积层。这些卷积层被替换为KAN块[18]。输入数据首先通过补丁嵌入层，然后通过单个KAN层进行处理。然后，输出通过激活函数，进一步馈入SSM块。在空间变换之后，结果通过空间注意力层处理，然后通过线性投影层。投影层中的调整有助于将输出与输入（来自跳跃连接）和激活层的输出线性组合。

具有K层的MLP可以表示为[18]：

这种变换重复三次以获得KAN层。KANB(⋅)用于表示由KAN层、归一化层和恒等函数组成的KAN块，如图5所示。

因此，输出可以数学表示为：

之后，添加一个注意力层，结果如下：

图4：改进的Mamba-KAN块

图3：深度卷积块（DwConvB）

最后，与跳跃连接和并行BoA函数层的组合给出最终输出为：

4 数据集

我们在三个不同的医学图像分割数据集上评估了我们的模型：乳腺超声图像（BUSI）[1]、分割息肉图像[14]和腺体分割图像（GlaS）[29]。鉴于每个数据集的独特特征，这些评估为测试我们方法的有效性提供了强有力的支持。

BUSI[1]：该数据集包含来自超声扫描的医学图像及其相应的分割掩码，用于识别与乳腺癌相关的肿瘤。整个数据集包含708张图像，其中210、437和133张分别代表恶性、良性和正常乳腺癌病例的图像。我们使用了代表乳腺癌的图像。图像被统一调整为256 x 256像素。

Kvasir-Seg[14]：Kvasir-SEG数据集包含1000张胃肠道息肉图像（息肉是结直肠癌的前兆）及其相应的分割掩码，这些掩码由医生手动注释并由经验丰富的胃肠病学家验证。所有图像被统一调整为256 x 256像素。

GlaS[29]：该数据集包含腺体分割图像，并与西班牙巴塞罗那的医院诊所相关联。在我们的研究中，我们使用了数据集中的165张图像，所有图像被调整为256 x 256像素。

5 实验与结果

实现细节

对于所有三个数据集，我们将学习率设置为1e-4，并使用Adam优化器进行训练，结合最小学习率为1e-5的余弦退火学习率调度器。使用的损失函数是二元交叉熵和Dice损失的组合。每个数据集按4:1的比例分为训练集和验证集。训练进行了400个epoch，并对输入应用了随机旋转和翻转等基本数据增强。

评估细节

为了评估模型的性能，我们报告了所有数据集上的验证IoU（交并比）和F1分数，并与现有技术方法进行了比较。IoU测量预测和真实分割之间的重叠，而F1分数提供了精度和召回率的平衡度量。这些指标对于理解模型在各种分割任务中的准确性和可靠性至关重要。报告的结果是三次独立运行的平均值。我们还提供了计算成本指标（如GFLOPs和模型参数总数），并与现有技术方法进行了比较。此外，我们使用其他有用参数对我们的模型进行了消融研究。

与现有技术方法的性能比较

表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和模型参数，以强调我们的模型可以在计算成本上实现最小的权衡，与以前的方法相比。

消融研究

我们对提出的KAN-Mamba FusionNet在各种设置下进行了全面评估，特别是为了理解我们模型中一些添加功能的效果。

Mamba块变化的影响：表3显示了消融研究的结果，以验证Mamba块中修改的效果。如表3所示，Mamba + MLP组合表示经典的Mamba后跟MLP层。然而，经典模型表现不佳，而将MLP层替换为KAN层则导致IoU、F1、准确性、AUC、精度和召回率的提高，如表中的验证指标所示。此外，在经典Mamba中结合“激活函数袋”，然后将MLP层替换为KAN层（如我们模型中提出的），进一步提高了列出的性能参数。

激活函数袋的影响：如表4所示，当没有额外的激活函数层添加到Mamba块中时，模型最初表现不佳，如我们架构中提出的那样。然而，当引入单个激活函数（如ReLU）时，模型的性能显著提高。此外，通过结合多个激活函数的组合，即“激活函数袋”，模型的学习能力进一步增强，导致更好的结果。这些结果强调了添加“激活函数袋”的重要性，表明激活函数的多样性在增强模型性能方面起着重要作用。

不同的激活函数捕捉数据的不同方面，每个激活函数都为神经网络带来了各自的优势和劣势。例如，修正线性单元（ReLU）[25]定义为：

捕捉围绕零的对称关系，提供-1到1范围内的输出，这有助于数据居中并提高收敛性。然而，它也容易受到梯度消失问题的影响。

为了解决这些局限性，我们提出了将多个激活函数组合成“激活函数袋”（BoA），如公式(4)中所定义。这种公式允许网络探索更丰富的函数空间，增强其建模数据中多样化模式的能力。通过在训练期间优化与每个激活函数相关的权重，模型可以利用每个激活函数的优势，同时减轻其各自的劣势。

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

此外，网络的不同部分可能受益于不同的激活行为，从而实现更专业和自适应的非线性变换。通过使网络能够自适应地组合多个激活函数，“激活函数袋”方法为增强神经网络在多样化机器学习任务中的性能提供了一种鲁棒且多功能的机制。

表1：各种分割方法在BUSI、Kvasir-Seg和GlaS数据集上的性能比较。

表2：各种方法的模型复杂度比较。

6 结论

在本文中，我们提出了KAN-Mamba FusionNet架构用于医学图像分割。它集成了Kolmogorov-Arnold网络（KAN）与Mamba，并通过“激活函数袋”减少复杂性，捕捉输入图像中的非线性，并自适应选择有效的激活函数以提升模型性能。我们在BUSI、Kvasir-Seg和GlaS数据集上的实验验证了我们的假设，展示了与现有技术方法相比更好的IoU（交并比）和F1（Dice）分数。此外，消融研究显示了每个架构组件的关键作用，特别是“激活函数袋”如何对整体性能提升做出贡献。总之，KAN-Mamba FusionNet提高了医学图像分割的准确性和鲁棒性，并且在计算上是高效的。“激活函数袋”的使用提供了一种灵活的方法来增强神经网络的效能，解决医疗保健中的复杂视觉数据问题，并可能通过更准确的分割掩码改善患者结果。通过自适应地集成各种激活函数，BoA提高了神经网络在多样化机器学习应用中的多功能性和鲁棒性。