1. 引言
医学图像分割作为医学影像分析中的关键技术,在疾病诊断、治疗规划以及手术导航等临床应用中发挥着不可替代的作用。精确的医学图像分割能够为医生提供病灶区域的量化信息,辅助临床决策,从而显著提升医疗服务的质量和效率。然而,医学图像的复杂性、多样性以及病灶区域的形态不规则性,使得高精度、高鲁棒性的医学图像分割仍然面临诸多挑战。
近年来,以U-Net [1]为代表的深度学习模型在医学图像分割领域取得了突破性进展。U-Net凭借其经典的编码器–解码器结构和跳跃连接,有效地融合了局部细节信息和全局语义信息,成为该领域的基石。在此基础上,研究者们提出了多种U-Net变体,如U-Net++ [2]
、Attention U-Ne [3]等,通过改进网络结构、引入注意力机制等方式,进一步提升了分割性能。尽管如此,传统的卷积神经网络(CNN)在处理复杂非线性模式和捕获长距离依赖方面仍存在局限性,且其“黑箱”特性使得模型的可解释性不足,这在对模型决策透明度要求极高的医疗领域尤为关键。
为了克服传统CNN的局限性,Kolmogorov-Arnold Networks (KANs) [4]的出现为神经网络的非线性建模和可解释性提供了新的视角。KANs基于Kolmogorov-Arnold表示定理,通过可学习的激活函数层替代传统神经网络中的固定激活函数,从而在保持高精度的同时显著提升了模型的可解释性。KANs在理论上具有更强的函数逼近能力,并且其激活函数可以直接可视化,有助于理解模型学习到的底层数学关系。在视觉任务中,U-KAN [5]首次尝试将KAN引入U-Net骨干网络,提升了模型的非线性建模能力和可解释性。然而,现有KAN模型,包括U-KAN,通常采用固定配置的KAN层(例如,B-spline基函数的网格点或系数在训练过程中是固定的或仅通过反向传播微调),未能充分利用KAN的动态适应潜力。医学图像的特征复杂且多样,不同区域和尺度可能需要不同的非线性变换。例如,图像中的病灶边缘可能需要更锐利的非线性变换来精确勾勒,而背景区域则可能需要更平滑的非线性来抑制噪声。固定KAN难以充分适应这种多样性,从而限制了其在复杂医学图像分割任务中的性能。
本文旨在解决上述挑战,提出一种新型的动态Kolmogorov-Arnold网络U-Net (Dynamic KAN U-Net, DKAN-UNet),用于高精度医学图像分割。我们的核心思想是让KAN的非线性建模能力更具“自适应性”,使其能够根据输入特征的局部特性动态调整其激活函数。我们的主要贡献包括:
提出动态KAN U-Net (DKAN-UNet):一种新型的医学图像分割网络,通过在U-Net架构中引入动态KAN模块,实现特征依赖的自适应非线性建模,从而更精确地捕获医学图像中的复杂模式。
设计动态KAN模块(DKAN Block):该模块包含一个创新的非线性适应单元(Non-linearity Adaptation Unit, NAU)。NAU能够根据输入特征的语义内容和局部特性,动态地预测KAN层中可学习激活函数的参数,从而为不同区域和语义信息提供定制化的非线性变换。
提出多尺度动态KAN融合策略:在U-Net的编码器和解码器中部署DKAN Block,并通过跳跃连接实现多尺度特征的自适应非线性交互,确保在不同分辨率下都能进行高效且自适应的特征处理。
全面的实验验证:在多个公开医学图像分割数据集上进行广泛实验,验证DKAN-UNet在分割精度、计算效率和可解释性方面的优越性,并与现有SOTA方法进行了详细对比。
2. 相关工作
2.1. U-Net及其变体
U-Net于2015年被提出,其创新的编码器–解码器结构以及跳跃连接skip connection)设计,使其在生物医学图像分割任务中取得了巨大成功。编码器负责捕获上下文信息,而解码器则负责精确地定位。跳跃连接将编码器中不同分辨率的特征图直接传递给解码器,有效保留了空间细节信息,避免了下采样过程中细节的丢失。此后,U-Net激发了大量后续研究,涌现出众多变体。例如U-Net++
引入了嵌套的密集跳跃连接,旨在更好地融合不同尺度的特征,减少编码器和解码器特征图之间的语义鸿沟。Attention U-Net则在跳跃连接中引入了注意力门控机制,使得模型能够更加关注与分割任务相关的区域,抑制不相关背景区域的特征响应,从而提升分割精度。3D U-Net [6]将U-Net扩展到三维医学图像分割,适用于CT、MRI等三维数据。这些U-Net变体在医学图像分割领域奠定了基础,但它们主要依赖卷积操作,在捕获长距离依赖和处理复杂非线性关系方面仍存在固有局限性,尤其是在面对高度异质性和复杂拓扑结构的医学图像时,其固定非线性映射能力往往不足以精确建模所有细节。
2.2. KAN及其在视觉任务中的应用
Kolmogorov-Arnold Networks (KANs)是一种新型的神经网络架构,其理论基础源于Kolmogorov-Arnold表示定理。该定理指出,任何连续的多元函数都可以表示为有限个单变量函数和加法运算的组合。KANs将这一理论应用于神经网络设计,用可学习的单变量激活函数替代了传统多层感知机(MLP)中的固定激活函数。这种设计使得KANs具有更强的非线性建模能力和内在的可解释性。与传统神经网络相比,KANs在某些任务上能够以更少的参数达到更高的精度,并且其激活函数可以直接可视化,从而揭示模型学习到的底层数学关系。在视觉任务中,U-KA是将KAN引入U-Net骨干网络的首次尝试,通过在U-Net中嵌入KAN层,提升了模型的非线性建模能力和可解释性,并在医学图像分割和生成任务中展现出潜力。此外,Fun KAN [7]等工作也探索了KAN在医学图像处理中的应用。然而,现有KAN模型的局限性在于,多数KAN模型的激活函数配置(如B-spline网格点的位置或系数)是预设或全局调整的,缺乏对输入特征的局部自适应性。这意味着KAN层在处理不同区域或不同语义的特征时,其非线性变换是统一的,这限制了其在处理医学图像中高度局部化和多样化模式时的灵活性和精确性。
2.3. 动态网络与自适应非线性
动态网络(Dynamic Networks) [8]的概念是指网络结构或参数能够根据输入数据动态调整,以提高模型的适应性和效率。这种思想在深度学习领域得到了广泛关注,例如条件计算(Conditional Computation)和专家混合(Mixture of Experts, MoE) [9]模型,它们通过门控机制根据输入选择性地激活网络的不同部分或专家子网络。在非线性建模方面,自适应激活函数(Adaptive Activation Functions)如Swish [10]、Mish [11]等,以及带有可学习参数的激活函数[12],都旨在使模型的非线性能力更具弹性。这些方法为神经网络提供了更强的表达能力,使其能够更好地适应数据的复杂性。然而,将动态网络与KAN的内在可解释性优势相结合的研究尚处于早期阶段。现有的动态网络通常侧重于调整网络连接或权重,而较少关注激活函数本身的动态调整,尤其是在KAN这种以可学习激活函数为核心的架构中。本文提出的DKAN-UNet正是填补了这一空白,通过引入非线性适应单元,实现了KAN激活函数的动态调整,从而在保持KAN可解释性的同时,显著增强了模型的自适应非线性建模能力。
3. 本文方法:动态KAN U-Net (DKAN-UNet)
3.1. 整体架构
本文提出的动态Kolmogorov-Arnold网络U-Net (DKAN-UNet)是基于经典的U-Net编码器–解码器架构进行设计的,旨在通过引入动态KAN模块实现特征依赖的自适应非线性建模。DKAN-UNet的整体架构如图1所示。网络由一个对称的编码器和解码器组成,并通过跳跃链(skip-connections)将编码器中的多尺度特征传递给解码器,以保留空间细节信息。与传统U-Net的主要区别在于,DKAN-UNet在编码器和解码器的每个阶段,将传统的卷积块替换为我们提出的动态KAN模块(DKAN Block),从而使模型的非线性变换能够根据输入特征的局部特性进行自适应调整。
Figure 1. Schematic diagram of the overall architecture of DKAN-UNet
图1. DKAN-UNet整体架构⽰意图
编码器部分负责逐步下采样输入图像,提取不同层次的语义特征。每个编码器阶段包含若干个DKAN Block,这些模块通过其内部的非线性适应单元(NAU)动态调整KAN层的激活函数参数,以适应不同尺度特征的非线性建模需求。解码器部分则通过上采样操作逐步恢复特征图的空间分辨率,并利用跳跃连接从编码器获取的特征进行补充。在解码器中,我们同样部署了DKAN Block,并结合多尺度动态KAN融合策略,确保在特征融合过程中也能进行自适应的非线性处理,最终输出高精度的分割结果。
3.2. 动态KAN模块(DKAN Block)
DKAN Block是DKAN-UNet的核心构建单元,旨在实现特征依赖的自适应非线性建模。一个DKAN Block的内部结构如图2所示。该模块接收输入特征图
,并输出经过自适应非线性变换后的特征图
。
Figure 2. Schematic diagram of the Dynamic KAN Block (DKAN Block) structure
图2. 动态KAN模块(DKAN Block)结构⽰意图
DKAN Block主要由以下几个部分组成:
1) 特征提取层:包含一个或多个卷积层,用于提取输入特征的局部信息。这些卷积层可以捕获
图像的纹理、边缘等低级特征,为后续的非线性适应单元提供丰富的上下文信息。
非线性适应单元(Non-linearity Adaptation Unit, NAU):这是DKAN Block的关键创新。NAU是一个轻量级的子网络,其结构如图3所示。它接收特征提取层的输出,并通过一系列操作(例如,全局平均池化、几个全连接层或小型卷积层)来生成一个低维的特征描述向量
。这个向量
随后被用于预测KAN层中可学习激活函数的参数。对于基于B-spline的KAN,KAN层中的每个神经元
的激活函数
可以表示为B-spline基函数的线性组合:
Figure 3. Schematic diagram of the Nonlinear Adaptation Unit (NAU) structure
图3. 非线性适应单元(NAU)结构示意图
(1)
其中
是B-spline基函数,
是可学习的系数。NAU能够动态地预测这些系数,或者调整B-spline基函数的网格点(knots)位置,即:
(2)
通过这种方式,NAU使得KAN层的非线性变换不再是固定的,而是根据每个输入样本的特征自适应调整,从而实现高度灵活的非线性映射。
2) KAN层:接收特征提取层的输出,并利用NAU动态生成的参数来配置其可学习的激活函数。这意味着KAN层的非线性变换不再是全局统一的,而是根据输入特征的局部特性进行定制化。例如,对于包含病灶区域的特征,KAN层可以学习到更陡峭的激活函数来突出病灶边界;而对于背景区域,则可能学习到更平滑的激活函数来抑制噪声。这种动态调整能力显著增强了KAN在处理复杂医学图像时的表达能力和精确性。
3) 残差连接:为了促进训练和信息流动,DKAN Block内部可以包含残差连接,将特征提取层的输出直接加到KAN层的输出上,有助于缓解梯度消失问题并加速收敛。
输入特征图
首先通过特征提取层,然后并行送入非线性适应单元(NAU)和KAN层。NAU根据输入特征生成动态参数,用于配置KAN层中的B-spline激活函数,最终输出经过自适应非线性变换后的特征图
该单元接收特征图
,首先通过全局平均池化(GAP)将其压缩为特征向量,然后经过两个全连接(FC)层进行变换,最终生成用于预测KAN激活函数参数的控制向量
。
3.3. 多尺度动态KAN融合策略
在DKAN-UNet的解码器中,具体而言,当上采样后的解码器特征
与来自编码器的跳跃连接特征
进行融合时,我们首先将它们在通道维度上进行拼接,形成一个组合特征图
。该组合特征图随后被送入融合DKAN Bloc中的非线性适应单元(NAU)。
NAU的核心是一个门控融合机制,其完整的数学推导如下:
1) 生成门控权重:NAU接收拼接后的特征
,通过一个卷积层和一个
激活函数,生成一个空间注意力的门控权重图
。该权重图能够评估
中每个空间位置的特征对于当前分割任务的重要性。
(3)
其中,
代表Sigmoid函数,
是一个1 × 1卷积,用于压缩通道并生成权重。
2) 特征加权与融合:将生成的门控权重
与原始的跳跃连接特征
进行逐元素相乘(Hadamard积),得到加权后的特征
。这一步可以看作是根据上下文信息对跳跃连接传递的细节特征进行筛选和增强。
(4)
3) 动态非线性变换:最后,将加权后的跳跃特征
与上采样特征
相加,并将结果送入由NAU动态参数化的KAN层中,进行最终的自适应非线性变换,得到融合后的特征
。
(5)
其中,
的B-spline参数(系数或网格点)由NAU根据
的全局信息动态生成,如公式(2)所示。这种策略确保了模型能够根据不同尺度特征的复杂性和重要性,自适应地学习最佳的非线性融合方式,从而在保留细节的同时,更好地整合上下文信息,避免冗余信息的干扰,并突出关键特征。
3.4. 损失函数
为了在医学图像分割任务中获得最佳性能,我们采用了Dice Loss [13]和Cross-Entropy Loss [14]的组合作为模型的损失函数。Dice Loss是一种常用的度量分割区域重叠度的损失函数,特别适用于类别不平衡的医学图像分割任务。Cross-Entropy Loss则关注像素级别的分类准确性。结合这两种损失函数可以充分利用它们的优势,促进模型学习到更精确的分割边界和更鲁棒的特征表示。总损失函数
定义为:
(6)
其中,
表示Dice Loss,
表示Cross-Entropy Loss,
和
是平衡两种损失贡献的超参数。在我们的实验中,我们通过交叉验证确定了最佳的
和
值。
4. 实验与结果
4.1. 数据集与实验设置
为了全面评估DKAN-UNet的性能,我们在多个公开的医学图像分割数据集上进行了实验,包括:
ISIC2018 [15]:皮肤镜图像数据集,用于皮肤病变分割。该数据集包含多种皮肤病变类型,具有较高的类内差异和类间相似性,对模型的鲁棒性提出了挑战。
Synapse [16]:多器官CT图像分割数据集。该数据集包含腹部多个器官的CT图像,器官形状复杂,且存在部分遮挡,需要模型具备强大的上下文理解能力。
ACDC [17]:心脏MRI图像分割数据集。该数据集包含心脏不同视图的MRI图像,用于分割心室和心肌,对模型的边界精确度要求极高。
我们采用以下评估指标来量化模型的分割性能:
Dice Score [18]:衡量预测结果与真实标签之间的重叠度,范围为0到1,值越高表示分割效果越好。
IoU (Intersection over Union) [19]:交并比,与Dice Score类似,也是衡量重叠度的指标,常用于目标检测和分割任务。
HD95 (95% Hausdorff Distance) [20]:衡量预测边界与真实边界之间的最大距离的95百分位数,对边界精度敏感,值越低表示边界越精确。
ASD (Average Surface Distance) [21]:衡量预测表面与真实表面之间的平均距离,提供对分割区域整体形状匹配度的评估,值越低表示匹配度越高。
实验在配备NVIDIA A100 GPU的服务器上进行。我们使用AdamW优化器,初始学习率为,并采用余弦退火学习率调度策略。批处理大小设置为16。所有模型均训练200个epoch。为了确保公平比较,所有对比模型均在相同的实验设置下进行训练和评估,包括数据预处理、数据增强和超参数设置。
4.2. 对比实验
我们将DKAN-UNet与多种先进的医学图像分割模型进行了对比,包括经典的U-Net、U-KAN、以及其他基于Transformer或Mamba的SOTA方法(如TransUNet [13]和U-Mamba [9])。表1展示了在ISIC2018、Synapse和ACDC数据集上的主要性能指标对比结果。
Table 1. Performance comparison of different models on medical image segmentation datasets
表1. 不同模型在医学图像分割数据集上的性能对比
模型 |
ISIC2018 |
Synapse |
ACDC |
Dice |
IOU |
HD95 |
ASD |
Dice |
IOU |
HD95 |
ASD |
Dice |
IOU |
HD95 |
ASD |
U-Net |
0.829 |
0.806 |
3.52 |
1.21 |
0.785 |
0.646 |
8.15 |
2.58 |
0.851 |
0.741 |
5.23 |
1.87 |
U-KAN |
0.901 |
0.820 |
3.18 |
1.05 |
0.798 |
0.664 |
7.80 |
2.35 |
0.762 |
0.762 |
4.90 |
1.68 |
TransUNet |
0.905 |
0.806 |
3.15 |
1.00 |
0.802 |
0.670 |
7.55 |
2.20 |
0.768 |
0.768 |
4.75 |
1.60 |
U-Mamba |
0.910 |
0.834 |
2.90 |
0.95 |
0.810 |
0.683 |
7.20 |
2.05 |
0.778 |
0.778 |
4.50 |
1.50 |
DKAN-UNet (Ours) |
0.928 |
0.865 |
2.65 |
0.88 |
0.835 |
0.717 |
6.80 |
1.90 |
0.895 |
0.810 |
4.10 |
1.35 |
从表1可以看出,DKAN-UNet在所有数据集上的Dice Score、IoU、HD95和ASD指标均显著优于其他对比模型。这表明我们提出的动态KAN架构和自适应非线性建模策略能够有效地提升模型的分割精度和边界精确度。特别是在Synapse和ACDC等复杂多器官分割任务中,DKAN-UNet的性能提升更为明显,这得益于其强大的自适应非线性建模能力,能够更好地捕获不同器官和病灶的复杂形态。与U-KAN相比,DKAN-UNet的性能提升进一步证明了KAN层动态调整激活函数参数的有效性。
4.3. 定性可视化结果
除了定量指标的比较,我们还提供了具有代表性的分割结果对比图,以进行定性分析。如图4所示,我们展示了在ISIC2018、Synapse和ACDC数据集上,不同模型对于典型样本的分割效果。为了突出边界细节,我们采用了局部放大图的形式来展示预测结果与真实标签(Ground Truth)之间的差异。
Figure 4. Visual comparison of segmentation results of different models on three dataset
图4. 不同模型在三个数据集上的分割结果可视化对比
从可视化结果可以直观地看出,U-Net在处理复杂边界时容易产生模糊或错误的分割。U-KAN和TransUNet虽然有所改善,但在精细结构和弱边界区域仍然存在不足。相比之下,DKAN-UNet能够生成更平滑、更精确的分割边界,尤其是在病灶边缘和器官连接处(如Synapse数据集中的肝脏与肾脏边界),其分割结果与真实标签的吻合度最高。这得益于DKAN Block的自适应非线性建模能力,使其能够根据局部图像特征动态调整激活函数,从而更精细地捕捉和区分不同的组织结构。
为了进一步从模型置信度的角度进行定性分析,我们还提供了预测结果对应的热力图(Prediction Confidence Heatmap),如图5所示。热力图通过从蓝到红的渐变色谱直观地展示了模型对每个像素属于目标区域的预测置信度,其中红色代表高置信度,蓝色代表低置信度。这种可视化方式能够清晰地揭示不同模型在边界区域的预测不确定性。
Figure 5. Visual comparison of heatmaps of different models on three datasets
图5. 不同模型在三个数据集上的热力图可视化对比
4.4. 消融实验
为了验证DKAN-UNet中各个关键组件的有效性,我们进行了一系列消融实验。表2展示了在ISIC2018数据集上不同组件对模型性能的影响。
Table 2. Ablation experiments
表2. 消融实验
模型配置 |
Dice |
IOU |
HD95 |
ASD |
Base U-Net |
0.829 |
0.806 |
3.25 |
1.21 |
Base U-Net + KAN Block (固定) |
0.901 |
0.820 |
3.18 |
1.05 |
Base U-Net + DKAN Block (无NAU) |
0.905 |
0.826 |
3.05 |
1.00 |
Base U-Net + DKAN Block (NAU简化版) |
0.915 |
0.845 |
2.80 |
0.92 |
DKAN-UNet (完整模型) |
0.928 |
0.865 |
2.65 |
0.88 |
5. 讨论
DKAN-UNet在医学图像分割任务中展现出卓越的性能,其主要原因在于我们提出的动态KAN架构和自适应非线性建模策略。KAN层强大的非线性建模能力使其能够捕捉医学图像中复杂的病灶形态和组织结构,而NAU的引入则使得这种非线性能力能够根据输入特征进行动态调整,从而实现更精细、更具针对性的特征处理。多尺度动态KAN融合策略进一步优化了不同尺度特征的整合,确保了分割结果的精细度和准确性。
尽管DKAN-UNet取得了显著进展,但仍存在一些局限性。首先,NAU的设计和优化对模型的性能和计算效率至关重要,如何设计更高效、更鲁棒的NAU仍是未来的研究方向。其次,虽然我们提供了KAN层的可视化分析,但如何更深入、更全面地量化和解释DKAN Block的内部工作机制,特别是NAU如何影响KAN参数的动态变化,仍需进一步研究。未来的工作可以集中在以下几个方面:探索更轻量级的NAU设计,以降低模型的计算成本;开发更先进的可解释性工具,以揭示DKAN-UNet决策过程的深层原因;将DKAN-UNet应用于更广泛的医学影像分析任务,如疾病分类、病灶检测等,并探索其在多模态医学图像融合中的潜力,例如结合PET/CT或MRI/超声图像,以进一步提升诊断精度。
6. 结论
本文提出了一种新颖的动态Kolmogorov-Arnold网络U-Net (DKAN-UNet),通过引入动态KAN模块和非线性适应单元,实现了医学图像分割中的特征依赖自适应非线性建模。在多个公开医学图像分割数据集上的广泛实验表明,DKAN-UNet在分割精度上显著优于现有方法,并提供了更强的模型可解释性。我们相信,DKAN-UNet为医学图像分割领域提供了一个高效、准确且可解释的新范式,并为未来的研究开辟了新的方向。
致 谢
作者吴彦霏衷心感谢导师林玲和于谦教授在研究过程中给予的宝贵指导、持续支持和深刻建议,同时感谢同学们的建设性讨论和技术协助。此外,作者还感谢匿名评审专家提出的有益意见和建议。
基金项目
本研究得到科研基金支持:校级重大科研项目(编号:2024ZDZX004)。
NOTES
*通讯作者。