点击下方卡片，关注 「AI视界引擎」 公众号
U-Net已成为各种视觉应用（如图像分割和扩散概率模型）的基石。尽管通过融入 Transformer 或MLPs引入了许多创新设计和改进，但这些网络仍然局限于线性地建模模式，并且在可解释性方面存在不足。
为了解决这些挑战，作者的直觉受到了Kolmogorov-Arnold网络（KANs）在准确性和可解释性方面令人印象深刻结果的启发，它们通过源自Kolmogorov-Anold表示定理的可学习非线性激活函数堆叠来 Reshape 神经网络学习。
具体来说，在本文中，作者探索了KANs在改进视觉任务 Backbone 网络方面的未开发潜力。作者研究、修改并重新设计已建立的U-Net Pipeline ，通过在标记化的中间表示上整合专用的KAN层，称之为U-KAN。严格的医学图像分割基准测试验证了U-KAN在准确性甚至计算成本方面的优越性。
作者进一步探讨了U-KAN作为扩散模型中的U-Net噪声预测器的潜力，展示了它在生成面向任务的模型架构方面的适用性。
这些努力揭示了有价值的见解，并照亮了前景：通过U-KAN，您可以构建强大的 Backbone 网络，用于医学图像的分割和生成。
1 Introduction
在过去的十年中，许多研究致力于开发高效且健壮的医学成像分割方法，这主要是由于计算机辅助诊断和图像引导手术系统的需求所推动。其中，U-Net [71]是一项里程碑式的工作，它最初证明了带跳跃连接的编码器-解码器卷积网络在医学图像分割中的有效性。近年来，UNet已成为几乎所有领先医学图像分割方法的基础，并在许多图像转换任务中显示出良好的效果。此外，最近的扩散模型也利用了U-Net，通过训练它迭代预测在每一步去噪中要移除的噪声。
所提出的U-KAN得益于KAN网络在非线性建模能力和可解释性方面的吸引人属性，使其在普遍的U-Net架构中显著区别于其他结构。在严格的医学分割基准上的实证评估，无论是定量还是定性，都强调了U-KAN的卓越性能，即使没有更少的计算成本，也超越了已建立的U-Net主干，提高了准确性。作者的研究进一步探讨了U-KAN作为扩散模型中U-Net噪声预测器的潜力，证实了其在生成面向任务的模型架构中的相关性。简而言之，U-KAN标志着将受数学理论启发的算子融入高效视觉 Pipeline 设计的一步，预示着它在广泛视觉应用中的前景。
作者的贡献可以总结如下：
- 作者首次尝试将新兴的KAN的优势融入到成熟的U-Net Pipeline 中，使其更准确、高效和可解释。
- 作者提出了一个标记化的KAN模块，有效地引导KAN操作符与现有的基于卷积的设计兼容。
- 作者在广泛的医学分割基准上实证验证了U-KAN，取得了令人印象深刻的准确性和效率。
- 将U-KAN应用于现有的扩散模型作为改进的噪声预测器，展示了其在支撑生成任务和更广泛的视觉设置中的潜力。
2 相关工作
U-Net Backbone for Medical Image Segmentation
医学图像分割是一项具有挑战性的任务，深度学习方法已被广泛应用于此领域，并在近年来取得了突破性的进展[40; 49; 62; 71; 77]。U-Net [71] 是一种流行的医学图像分割网络结构。
它的编码器-解码器架构有效地捕捉图像特征。CE-Net [20] 进一步整合了上下文信息编码模块，增强了模型的感受野和语义表示能力。Unet++ [94] 提出了一种嵌套U-Net结构，通过融合多尺度特征来提高分割精度。除了基于卷积的方法外，基于Transformer的模型也引起了关注。Vision Transformer [14] 展示了Transformer在图像识别任务中的有效性。Medical Transformer [80] 和 TransUNet [9] 进一步将Transformer应用于医学图像分割，取得了令人满意的表现。
此外，像注意力机制[76]和多尺度特征融合[31]这样的技术在医学图像分割任务中得到了广泛应用。3D分割模型，如多维门控循环单元[2]和高效多尺度3D CNN[34]，也取得了令人称赞的结果。总之，医学图像分割是一个活跃的研究领域，深度学习方法已经取得了显著的进展。
最近，Mamba [18] 通过将选择机制和硬件感知算法整合到之前的工作中[19; 21; 57]，实现了线性时间推理和高效的训练过程，取得了划时代的里程碑。基于Mamba的成功，对于视觉应用，Vision Mamba [53] 和 VMamba [95] 分别使用双向Vim块和跨扫描模块来获取数据依赖的全局视觉上下文。
同时，U-Mamba [56] 和其他工作[72; 88]在医学图像分割中展现了卓越的性能。由于Kolmogorov-Arnold网络（KAN）[55]已成为MLP的一个有前景的替代方案，并展示了其精度、效率和可解释性，作者相信现在是探索其在视觉 Backbone 网络中更广泛应用的时候了。
U-Net Diffusion Backbone for Image Generation
扩散概率模型，作为生成模型的前沿类别，已成为研究领域的焦点，特别是在与计算机视觉相关的任务中[26; 69; 70]。与其他生成模型类别，如变分自编码器（VAE）[36]，生成对抗网络（GANs），以及向量量化方法[15; 82]不同，扩散模型引入了一种新颖的生成范式。这些模型使用一个固定的马尔可夫链来映射潜在空间，促进复杂的映射，捕捉数据集中固有的复杂结构。
最近，它们在生成方面的出色表现，从高级细节到生成样本的多样性，推动了在图像合成，图像编辑，图像到图像翻译，以及视频生成等计算机视觉应用中的突破性进展。
扩散模型由一个“扩散过程”和一个“去噪过程”组成。在扩散过程中，逐渐将高斯噪声添加到输入数据中，最终将其腐蚀为近似纯高斯噪声。
在“去噪过程”中，通过学习的一系列逆向扩散操作，从其噪声状态恢复原始输入数据。通常，卷积U-Net[71]，作为事实上的主干网络选择，被训练以在每次去噪步骤中迭代预测要移除的噪声。
与之前专注于利用预训练的扩散U-Net进行下游应用的工作不同，最近的工作致力于探索扩散U-Net的内在特征和结构性质。Free-U研究有策略地重新评估U-Net的跳跃连接和主干特征图对U-Net架构两个组成部分的优势的贡献。RINs[33]为DDPMs引入了一种基于注意力的新颖高效架构。DiT[66]提出了将纯Transformer与扩散结合的方案，展示了其可扩展性。在本文中，作者展示了一个整合了U-Net和KAN的主干方案在生成方面的潜力，推动了生成主干边界的拓展和选择。
Kolmogorov-Arnold Networks (KANs)
科尔莫哥洛夫-阿诺德定理[37]假定任何连续函数都可以表示为有限变量的连续一元函数的组合，这为构建通用神经网络模型提供了理论基础。Hornik等人[28]进一步证实了这一点，他们证明了前馈神经网络具有普遍逼近能力，为深度学习的发展铺平了道路。借鉴科尔莫哥洛夫-阿诺德定理，学者们提出了一种新颖的神经网络架构，称为科尔莫哥洛夫-阿诺德网络（KANs）[29]。
KANs由一系列级联的科尔莫哥洛夫-阿诺德层组成，每一层包含一组可学习的一维激活函数。这种网络结构在逼近高维复杂函数方面被证明是有效的，并在各种应用中展现出强大的性能。
KANs的特点是具有强的理论可解释性和可解释性。黄等人[30]分析了KANs的优化特征和收敛性，验证了它们优秀的逼近能力和泛化性能。梁等人[47]进一步引入了深度KAN模型，并将其应用于图像分类等任务。邢等人[87]将KANs应用于时间序列预测和控制问题。尽管取得了这些进展，但在将具有坚实基础理论的KAN新型神经网络模型广泛融入到通用视觉网络的实际应用中仍显不足。
相比之下，本文进行了初步探索，试图设计一种融合KAN的通用视觉网络架构，并在广泛的分割和生成任务上进行验证。
3 Method
图1展示了所提出的U-KAN的总体架构，遵循由卷积阶段和标记化的科尔莫哥洛夫-阿诺德网络（Tok-KAN）阶段组成的两阶段编码器-解码器架构。输入图像通过编码器，其中最初三个块使用卷积操作，接着是两个标记化的MLP块。解码器由两个标记化的KAN块组成，之后是三个卷积块。每个编码器块将特征分辨率减半，而每个解码器块将其加倍。此外，编码器和解码器之间集成了跳跃连接。卷积阶段和Tok-KAN阶段中每个块的通道数分别由超参数至和至确定。
KAN as Efficient Embedder
本研究旨在将知识感知网络（KANs）融入UNet框架，或者独立利用KANs进行医学图像分割任务。这种方法的基础是KANs已证明的高效率和可解释性，如[55]所述。包含层的多层感知器（MLP）可以描述为转换矩阵和激活函数的交互作用。这可以用数学表达式表示为：
其中它通过多层非线性变换的序列来模拟复杂的函数映射。尽管其潜力巨大，但这种结构内部的固有模糊性严重阻碍了模型的解释性，从而对直观理解底层决策机制造成了相当大的挑战。
为了缓解MLPs中固有的参数效率低下和解释性有限的问题，Liu et al.[55] 提出了科莫戈罗夫-阿诺德网络（KAN），灵感来源于科莫戈罗夫-阿诺德表示定理[38]。与MLP类似，层的KAN可以表征为多个KAN层的嵌套：
其中表示整个KAN网络的第层。每个KAN层，具有维输入和维输出，包含个可学习的激活函数：
KAN网络从第层到第层的计算结果可以用矩阵形式表示：
总之，KANs通过在边上使用可学习的激活函数，以及将参数化激活函数作为权重，从而无需线性权重矩阵，与传统的MLPs区分开来。这种设计使得KANs能够以较小的模型尺寸实现相当或更优的性能。此外，其结构在不牺牲性能的情况下增强了模型的解释性，使其适用于各种应用。
U-KAN Architecture
3.2.1 Convolution Phrase
每个卷积块由三个组件组成：一个卷积层（Conv）、一个批量归一化层（BN）和一个ReLU激活函数。作者使用3x3的核大小，1的步长和1的填充量。编码器中的卷积块集成了一个2x2池窗口大小的最大池化层，而解码器中的卷积块包含一个双线性插值层以放大特征图。形式上，给定一个图像 ，每个卷积块的输出可以表述为：
其中 表示第 层的输出特征图。考虑到卷积短语中有 个块，最后的输出导出为 。
3.2.2 Tokenized KAN Phrase
分词在分词的KAN块中，作者首先通过将卷积短语输出的特征 Reshape 为一个扁平化的2D Patch 序列来进行分词[10; 14]，其中每个 Patch 的大小为，是特征 Patch 的数量。作者首先使用可训练的线性投影将向量化的 Patch 映射到一个潜在的维嵌入空间，如下所示：
图1：U-KAN Pipeline 概述。在卷积短语中的几个卷积块进行特征提取之后，中间的地图被分词并由分词的KAN短语中的堆叠的Tok-KAN块处理。时间嵌入仅在应用于扩散U-KAN时注入到KAN块中。
线性投影是通过一个核大小为3的卷积层实现的，正如[86]所示，一个卷积层足以编码位置信息，并且它的性能实际上优于标准的位置编码技术。像ViT中那样的位置编码技术，当测试和训练的分辨率不同时，需要进行插值，通常会导致性能降低。
通过KAN层嵌入给定获得的标记，作者将它们送入一系列KAN层（）。在每个KAN层之后，特征通过一个有效的深度卷积层（DwConv）[8]和批归一化层（BN）以及ReLU激活。在这里作者使用残差连接，并将原始标记作为残差添加。然后作者应用层归一化（LN）[4]并将输出特征传递到下一个块。正式地，第个分词的KAN块的输出可以详细表述为：
其中是第层的输出特征图。假设在分词的KAN短语中有个块，最终的输出为。在作者的实现中，作者设置和。
3.2.3 U-KAN Decoder
作者采用了通常使用的U形架构与密集的跳跃连接来构建U-KAN。U-Net及其变体在医学图像分割任务中已经显示出显著的效率[39; 89; 91]。这种架构利用跳跃连接恢复低级细节，并采用编码器-解码器结构来提取高级信息。
给定来自KAN短语中第层的跳跃连接特征以及来自最后一个上采样块的特性，第个上采样块的输出特征为：
其中表示特征拼接操作。同样，给定来自卷积短语中第层的跳跃连接特征以及来自最后一个上采样块的特征，第个上采样块的输出特征为：
在语义分割任务的背景下，最终的分割图可以从第0层的输出特征图得出，其中是语义类别的数量，表示 GT 分割。因此，分割损失可以是：
其中表示逐像素的交叉熵损失。
Extending U-KAN to Diffusion Models
上述讨论集中在通过U-KAN生成给定输入图像 的分割 Mask 。在本节中，作者将U-KAN进一步扩展到扩散版本，称为扩散U-KAN，它释放了KANs的生成能力。遵循去噪扩散概率模型（DDPM）[26]，扩散U-KAN能够从随机高斯噪声 逐步去除噪声，生成图像。这个过程可以通过预测给定噪声输入的噪声来实现：，其中 是被高斯噪声 脏化的图像 ， 是控制噪声强度的时步，而 。
为此，作者基于分割U-KAN进行了两项修改，以将其提升到扩散版本。首先，与仅在不同的隐藏层之间传播特征不同，作者在每个块中注入可学习的时间嵌入，使网络具有时间感知能力（见图1中的虚线“时间嵌入”），并移除DwConv和残差连接，因此将方程7改为以下格式，以适应生成任务的目标：
其中 是线性投影， 表示给定时间步 [26] 的时间嵌入。其次，作者修改了预测目标，以实现基于扩散的图像生成。扩散U-KAN不是预测给定图像的分割 Mask ，而是预测给定噪声图像 和随机时间步 的噪声 ，通过以下MSE损失进行优化：
通过上述损失函数优化后，使用DDPM采样算法[26]生成图像，该算法利用训练有素的扩散U-KAN进行去噪。
4 Experiments
Datasets
作者对提出的方法在三个独特且异构的数据集上进行了彻底的评估，这些数据集各自展现出独特的特性、不同的数据大小和各异的图像分辨率。这些数据集通常用于图像分割和生成等任务，为检验作者方法的有效性和适应性提供了坚实的测试平台。
Busi数据集[1]由超声图像组成，展示了正常、良性以及恶性的乳腺癌案例及其对应的分割图。在作者的研究中，作者使用了647张超声图像，代表了良性与恶性的乳腺肿瘤。所有这些图像都被一致地调整为的尺寸。该数据集提供了辅助检测和区分各种乳腺肿瘤类型的全面图像集合，为医疗专业行人和研究行人提供了宝贵的见解。
GlaS数据集[80]由612个标准定义(SD)帧组成，来自31个序列。每个帧的分辨率为，从23名患者那里收集而来，与位于西班牙巴塞罗那的Hospital Clinic相关联。该数据集中的序列是使用Olympus Q160AL和Q165L设备以及Extra II视频处理器记录的。作者的研究特别使用了165张来自GlaS数据集的图像，所有这些图像都被调整到的尺寸。
CVC-ClinicDB数据集[5]（简称“CVC”）是用于结肠镜视频中息肉诊断的公开可用资源。它总共包含612张图像，每张图像的分辨率为，仔细地从31个不同的结肠镜序列中提取。这些帧提供了各种息肉实例，对于息肉检测算法的开发和评估特别有用。为了确保作者研究中使用的不同数据集之间的一致性，CVC-ClinicDB数据集中的所有图像都被统一调整为。
Implementation Details
作者使用NVIDIA RTX 4090 GPU上的Pytorch实现了U-KAN。对于BUSI、GlaS和CVC数据集，作者将批处理大小设置为8，学习率设为1e-4。作者使用Adam优化器训练模型，并采用余弦退火学习率调度器，最小学习率为1e-5。损失函数是二进制交叉熵（BCE）和Dice损失的组合。作者将每个数据集随机划分为的训练集和的验证集。这些数据集上的所有结果都是三次随机运行的平均值。只应用了包括随机旋转和翻转在内的基本数据增强。总共训练了400个周期。作者使用各种指标，如IoU和F1分数，从定性和定量两个方面比较输出的分割图像。作者还报告了与计算成本相关的指标，如Gflops和参数数量（Params）。
扩散U-KAN图像被裁剪并调整为进行无条件生成。作者在相同的训练设置下对所有方法进行基准测试：1e-4的学习率，1000个周期，Adam优化器，以及余弦退火学习率调度器。为了评估每种方法的生成能力，作者使用随机高斯噪声作为输入生成了2048个图像样本。然后作者使用各种指标，如Frechet初始距离（FID）[65]和初始得分（IS）[75]，定性和定量地比较生成的图像。这些指标提供了对生成图像的多样性和质量方面的洞察。
Comparison with State-of-the-arts on Image Segmentation
表1展示了所提出的U-KAN与所有基准数据集上所有比较方法的对比结果。作者将U-KAN与近年来受到青睐的医学图像分割框架进行了比较，与卷积 Baseline 模型如U-Net[71]，U-Net++[94]进行了基准测试。作者还评估了与基于注意力的模型，包括Att-UNet [63] 和最先进的效率变形器变种，U-Mamba [56]的性能对比。此外，由于KAN作为MLP的有前途的替代品，作者进一步与基于MLP的高级分割网络进行了比较，包括U-Next [81] 和 Rolling-UNet [54]。在性能指标方面，作者使用交并比（IoU）和F1分数这两个标准指标来评估图像分割任务。结果显示，在所有数据集上，作者的U-KAN超越了所有其他方法的性能。
除了准确性的优势之外，本文还展示了当作为网络 Baseline 时作者方法的效率。如表2所示，作者报告了模型在不同数据集上的参数量（M）和Gflops以及分割准确性。结果表明，作者的方法不仅在分割准确性上超过了大多数分割方法，而且在效率上也显示出显著优势或与之相当的水平，除了U-Next之外。总的来说，在分割准确性和效率之间的权衡上，作者的方法表现出最佳的性能。
作者进一步在图2中展示了所有数据集上的综合定性比较。首先，从结果中可以明显看出，基于纯CNN的方法，如U-Net和U-Net++，更容易对器官进行过分割或欠分割，这表明这些模型在编码全局上下文和区分语义方面的局限性。相比之下，作者提出的U-KAN与其他方法相比，假阳性更少，表明其在抑制噪声预测方面的优越性。当与基于Transformer的模型和基于MLP的高效架构并列时，U-KAN的预测通常在边界和形状方面表现出更精细的细节。这些观察突显了U-KAN在保持复杂形状信息的同时进行精细分割的能力。这进一步证实了作者的初步直觉，强调了引入KAN层所带来的优势。
与图像生成领域的最先进技术比较
作者探讨了作者提出的U-KAN作为生成任务主干网的潜力。作者将U-KAN与基于传统U-Net的各种扩散变体模型进行了比较，以评估这种架构对不同生成任务的效能。结果在表3中呈现，作者报告了三个数据集上的FID [65]（弗雷歇初始距离）和IS [75]（初始分数）指标。弗雷歇初始距离是两个分布之间的距离度量，在本例中，是生成图像分布与真实图像分布之间的距离。FID越低，生成图像越接近真实图像。另一方面，初始分数通过评估这些图像能否被分类到特定类别来衡量生成图像的质量。IS越高，生成图像被正确分类的可能性越大。作者实验的结果清楚地表明，与其他领域内的最先进模型相比，作者的方法在生成性能上表现出色。这表明U-KAN的架构特别适合于生成任务，为生成高质量图像提供了一种有效且高效的方法。
图3展示了作者生成结果的一些可视化。观察发现，作者的方法可以在多个不同数据集上生成真实且多样的内容，这证明了其在生成高质量图像方面的灵活性和有效性。这进一步支持了U-KAN在生成任务上有显著优势的观点，使其成为未来在这一领域研究和发展的有力候选者。### 消融研究
为了彻底评估所提出的TransUNet框架并在不同设置下验证性能，进行了以下多种消融研究。
KAN层的数量如前所述，在U-KAN中包含KAN层已被证明是有益的，它通过显式融入高效率的嵌入，促进了更精细分割细节的建模。这项消融研究的目的是评估融入不同数量的KAN层的影响。作者按照表4所示，将KAN层的数量从一层修改为五层。观察到具有三个KAN层的配置提供了最优秀的性能。这些结果证实了作者的初步假设，即在U-KAN中策略性地融入适量的KAN层可以有效捕捉到复杂的与分割相关的细微差别。
为了进一步证实KAN层在提高模型性能中的作用，作者进行了一系列消融实验，结果如表5所示。在这些实验中，作者用传统的多层感知机（MLPs）替换了引入的KAN层，以观察这样的修改是否会导致性能下降。这种方法使作者更直观地理解了KAN层在提高模型的可解释性和整体性能方面的重要性。最初，作者修改了一个已经包含KAN层的模型，用一个或几个标准的MLPs替换了KAN层。随后，使用相同的 数据集和训练参数，作者重新训练了修改后的模型，并记录了它在各种任务上的表现。结果显示，当用MLPs替换KAN层时，在多个任务上的性能都有明显下降，特别是在需要强大特征提取和表示能力的复杂任务中。这些发现强调了KAN层在增强模型表达能力和提升整体性能中的关键作用。
最后，作者对U-KAN的不同模型大小进行了消融研究。具体来说，作者考察了U-KAN的不同配置，称之为_Small_和_Large_模型。这些变体之间的主要区别在于它们的通道设置，即在第一到第三KAN层的通道数（-），详细情况见表6。_Small_模型的通道设置为64-96-128，而_Large_模型的通道数设置为256-320-512。相比之下，作者默认模型的通道数设置为128-160-256。作者观察到，较大的模型与性能提升相关联，这与整合了KAN的模型的缩放法则特性相一致。最终，为了在性能和计算成本之间取得平衡，作者选择在实验中采用默认的基础模型。
5 Conclusion
本文介绍了U-KAN，并展示了Kolmogorov-Arnold网络（KANs）在提升如U-Net这样的 Backbone 网络用于各种视觉应用中的显著潜力。
通过将KAN层集成到U-Net架构中，可以构建一个在视觉任务方面具有令人印象深刻的准确性、效率和可解释性的强大网络。作者在几个医学图像分割任务下对作者的方法进行了实证评估。
此外，U-KAN的适应性和有效性也凸显了它作为扩散模型中噪声预测的U-Net优越替代品的潜力。
这些发现强调了探索像KANs这样的非传统网络结构对于推进更广泛的视觉应用的重要性。
参考
[1].U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation.
点击上方卡片，关注 「AI视界引擎」 公众号