论文精读-U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation

最新推荐文章于 2024-11-08 19:55:32 发布

江江酱₍ᐢ..ᐢ₎♡

最新推荐文章于 2024-11-08 19:55:32 发布

阅读量2.3k

点赞数 18

文章标签：深度学习人工智能计算机视觉

本文链接：https://blog.csdn.net/Jawy_/article/details/142096548

版权

论文链接:

https://arxiv.org/abs/2406.02918

论文代码：

https://yes-u-kan.github.io/

一、参考文献

[1] Li C, Liu X, Li W, et al. U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation[J]. arXiv preprint arXiv:2406.02918, 2024.

[2] Liu Z, Wang Y, Vaidya S, et al. Kan: Kolmogorov-arnold networks[J]. arXiv preprint arXiv:2404.19756, 2024.

二、研究背景

1、在过去的十年里，医学影像分割方法得到了广泛的发展。U-Net作为这一领域的标志性工作，通过其编码器-解码器结构和跳跃连接（skip connections），在医学图像分割任务中取得了显著的效果。同时，扩散模型近年来也利用U-Net作为骨干网络来预测去噪步骤中的噪声，推动了生成模型的进步。然而，尽管已有的基于U-Net的改进方法，如U-Net++、3D U-Net和V-Net等，提升了医学图像分割的精度和性能，但这些网络仍然面临非线性建模模式的限制和可解释性不足的问题；

2、为了应对这些挑战，研究者受到Kolmogorov-Arnold网络（KANs）在精度和可解释性方面出色表现的启发，提出通过集成KAN的非线性学习功能，来改进U-Net的视觉任务骨干结构。通过整合KAN层，希望能有效捕捉复杂的非线性模式，并提高网络的可解释性。这种设计旨在克服传统U-Net结构在捕捉空间依赖关系和复杂模式时的局限性，同时增强模型在医学影像中的临床决策支持能力。

三、研究方法

（一）Unet

本文以经典医学图像分割网络Unet网络为架构，并加入KAN模块一优化分割效果。Unet网络为U型对称的encoder-decoder结构，并添加跳跃连接将编码器中的高分辨率特征与解码器中的特征进行融合。

（二）MLP VS KAN

MLP先线性组合再非线性激活函数；KAN先非线性激活再组合，通过使用可学习的函数替代固定的激活函数，从根本上消除了对线性权重矩阵的依赖。

（三）MLP

根据上图可以看出，MLP存在以下几点明显的缺点：

MLP存在的缺点：

1、梯度消失和梯度爆炸：当使用传统的激活函数（如Sigmoid和Tanh）时，MLP在进行反向传播计算梯度时确实容易遇到梯度消失和梯度爆炸的问题，会出现激活函数的导数连乘积，当它非常小或者非常大时，网络很深，连续乘积会使得梯度趋于0（梯度消失）或者变得异常大；

2、参数效率低：MLP通常使用全连接层，这意味着每层的每个神经元都与前一层的所有神经元相连接，导致了参数数量迅速增加，尤其是对于输入维度很高的数据。这不仅增加了计算负担，也增加了模型过拟合的风险。这就是大模型的困境，参数量过多的被浪费，效率低下；

3、处理高维数据的能力有限：MLP没有利用数据的内在结构（如图像中的局部空间相关性或者文本数据的序列信息）。例如，在图像处理中，MLP无法有效地利用像素之间的局部空间关联，这使得其在图像识别等任务上地性能就不如卷积神经网络（CNN）；

4、长期依赖问题：虽然MLP理论上可以逼近任何函数，但是在实际应用中，它们很难捕捉到输入序列中地长期依赖关系（长时间跨度地相关信息）。这一点在处理时间序列或者自然语言处理任务时尤为明显，而RNN和Transformer在这些任务中表现更好。

（四）Kolmogorov–Arnold Networks(KAN)

1、网络理论基础

KAN 的理论基础，源自两位苏联数学家Vladimir Arnold和Andrey Kolmogorov的研究成果。他们的理论研究基于多变量连续函数的概念，根据这一理论，任何多变量连续函数 $f$ 都可以表示为有限个单变量连续函数的组合。

$f(\mathbf{x})=\sum_{q=1}^{2n+1}\Phi_q\left(\sum_{p=1}^n\phi_{q,p}(x_p)\right)$

其中， $\phi_{q,p}:[0,1]\to\mathbb{R},\Phi_q:\mathbb{R}\to\mathbb{R}$

2、KAN网络结构

图中展示的结构，使用了两种尺度或者分辨率的组合：粗粒度网和细粒度网，在保持任保持计算效率的同时，更加精确地捕捉和适应函数的变化。很多个类似四分之三周期的正弦函数组合起来就能拟合任意形状的函数。换句话说，两层的网络使用B-spline这一激活函数两次求和就能够得到对应的输出。

进一步，我们需要搭建更大的网络来对应不同的情况，因此需要将KAN网络结构加深：类似于MLP简单地堆叠更多KAN层。

接下来介绍一下KAN层：

KAN层的定义：输入维度为为 $n_{\mathrm{in}}$ ，输出维度为 $n_{\mathrm{out}}$ 的KAN层可以定义为一个由一维函数组成的矩阵：

$\mathbf{\Phi}=\{\phi_{q,p}\},\quad p=1,2,\cdots,n_{\mathrm{in}},\quad q=1,2\cdots,n_{\mathrm{out}}$

其中，函数 $\phi_{q,p}$ 具有可训练的参数。

因此，关键是需要找到输入与输出之间的转移矩阵。

进一步，得到转移矩阵：

$\mathbf{x}_{l+1}=\underbrace{\begin{pmatrix}\phi_{l,1,1}(\cdot)&\phi_{l,1,2}(\cdot)&\cdots&\phi_{l,1,n_{l}}(\cdot)\\\phi_{l,2,1}(\cdot)&\phi_{l,2,2}(\cdot)&\cdots&\phi_{l,2,n_{l}}(\cdot)\\\vdots&\vdots&&\vdots\\\phi_{l,n_{l+1},1}(\cdot)&\phi_{l,n_{l+1},2}(\cdot)&\cdots&\phi_{l,n_{l+1},n_{l}}(\cdot)\end{pmatrix}}_{\Phi_{l}}\mathbf{x}_{l}$

其中， $\mathbf{x}_{l}$ 为输入； $\mathbf{x}_{l+1}$ 为输出； $\Phi_{l}$ 是与第 l 个 KAN 层对应的函数矩阵。

同时，针对网络层加深的情况进一步放开了2n + 1的限制，隐层节点可以自行设置。原始两层KAN网络的结构是[n, 2n + 1, 1]，如下图所示：

一般的一般的 KAN 网络由 L 层组成：给定一个输入向量 $\mathbf{x}_0\in\mathbb{R}^{n_0}$ ，KAN 的输出为：

$\mathrm{KAN}(\mathbf{x})=(\mathbf{\Phi}_{L-1}^\circ\mathbf{\Phi}_{L-2}^\circ\cdots^\circ\mathbf{\Phi}_1^\circ\mathbf{\Phi}_0)\mathbf{x}$

进一步对比MLP：

$\mathrm{MLP}(\mathbf{x})=(\mathbf{W}_{L-1}\circ\sigma^\circ\mathbf{W}_{L-2}\circ\sigma^\circ\cdots\circ\mathbf{W}_1\circ\sigma^\circ\mathbf{W}_0)\mathbf{x}$

其中 MLP 中， $\mathbf{W}$ 表示线性权重参数， $\sigma$ 表示非线性激活函数。

可以看出，MLP中的权重参数和非线性激活函数被单变量函数参数所取代。

3、激活函数

进一步介绍KAN网络中的残差激活函数。

$\phi(x)=w_bb(x)+w_s\mathrm{spline}(x)$

其中， $b(x)$ 为基函数，类似于残差连接； ${spline}(x)$ 为样条函数； $w$ 用于控制激活函数整体的大小。

$b(x)=\text{silu}(x)=\frac{x}{(1+e^{-x})}$

$\mathrm{spline}(x)=\sum_ic_iB_i(x)$

其中， $c_i$ 为可学习的参数； $B_i(x)$ 为样条基函数，样条基函数是一组分段的多项式函数，通常被用来表示更复杂的曲线或函数。

4、KAN的优点

KAN的优点：

1、虽然 KAN 基于Kolmogorov-Arnold 定理，但它也从 MLP 中汲取灵感，结合了各自的优势，避免了各自的缺点。KAN 在外部结构上借鉴了 MLP，而在内部则采用了样条函数。因此，KAN 不仅能学习特征（得益于与 MLP 的外部相似性），还能通过优化这些学到的特征达到更高的精度（得益于与样条函数的内部相似性）；

2、同时，借助 spline 样条函数的表达能力，无需很多节点就能实现比较强的表达能力，因此总的来说，可以比 MLP 节省不少参数量。

（五）U-KAN

1、网络结构

下图为U-KAN的整体架构，采用了一个双阶段的编码器-解码器架构，包括卷积阶段和Token化的Kolmogorov-Arnold网络（Tok-KAN）阶段。

2、卷积阶段

每个卷积块由Conv+BN+ReLU组成，采用了 3x3 的卷积核大小，步长为 1，填充为 1。编码器中的卷积模块整合了一个 2x2 大小的最大池化层。给定一个图像 $X_0=I\in\mathbb{R}^{H_0\times W_0\times C_0}$ ，每个卷积模块的输出可以表示为：

$X_\ell=\operatorname{Pool}(\operatorname{Conv}(X_{\ell-1}))$

其中， $X_{\ell}\in\mathbb{R}^{H_{\ell}\times W_{\ell}\times C_{\ell}}$ 表示第 $\ell$ 层的输出特征图；假设在卷积阶段有 $L$ 个模块，最终输出为 $X_L$ 。

3、Tokenized KAN 阶段

（1）Tokenization：

这一阶段的Tokenization相当于Vision Transformer (ViT)中的Patch Embedding+Linear Projection+Positional Encoding，具体操作如下：

在 Token 化的 KAN 模块中，首先对卷积阶段的输出特征 $X_L$ ，进行 token 化，将其重塑为一系列展平的二维patch序列 $\{X_{L}^{i}\in\mathbb{R}^{P^{2}\cdot C_{L}}\mid i=1,..,N\}$ ，其中每个 patch 的大小为 $N={\frac{H_{L}\times W_{L}}{P^{2}}}$ ，并且 $P\times P$ 是特征patch的数量。然后，我们使用一个可训练的线性投影矩阵 $E\in\mathbb{R}^{(P^{2}\cdot C_{L})\times D}$ 将向量化的patch映射到一个 D 维的潜在嵌入空间中，表示为：

$Z_0=[X_L^1\mathbf{E};X_L^2\mathbf{E};\cdots;X_L^N\mathbf{E}]$

其中， $Z_0$ 是经过线性投影后得到的初始嵌入向量。

线性投影通过卷积实现，卷积层足以获取编码位置信息，并且实际上比标准的位置编码技术性能更好。

（2）Embedding by KAN Layer

给定获取到的 token，将它们传递到一系列 KAN 层中（N = 3）。在每个 KAN 层之后，特征通过一个高效的深度卷积层（DwConv），然后通过批归一化层（BN）和 ReLU 激活层处理。在这里使用残差连接，将原始的 token 作为残差加入。接着，我们应用层归一化（LN），并将输出特征传递到下一个模块。形式上，第 k 个 Token 化的 KAN 模块的输出可以表示为：

$Z_k=\mathrm{LN}(Z_{k-1}+\mathrm{DwConv}(\mathrm{KAN}(Z_{k-1})))$

其中， $Z_{_k}\in\mathbb{R}^{H_k\times W_k\times D_k}$ 是第k层的输出特征图。

（3）U-KAN Decoder

本文采用常用的U形结构，并结合跳跃连接来构建U-KAN。

给定来自KAN阶段第k层的跳跃连接特征 $Z_k$ 和来自最后一个上采样模块的特征 $Z_{k+1}^{\prime}$ ，第k个上采样模块的输出特征 $Z_k^{\prime}$ 为：

$Z_k^{\prime}=\mathrm{Cat}\left(Z_{k+1}^{\prime},(Z_k)\right)$

其中，Cat(⋅)表示特征拼接操作。

同样地，给定来自卷积阶段第 $\ell$ 层的跳跃连接特征 $X_{\ell}$ 和来自最后一个上采样模块的特征 $X_{\ell+1}^{'}$ ，第ℓ个上采样模块的输出特征 $X_{\ell}^{'}$ 为：

$X_{\ell}^{'}=\mathrm{Cat}\Big(X_{\ell+1}^{'},(X_{\ell})\Big)$

四、试验结果-图像分割

（一）数据集

BUSI、GlaS、CVC-ClinicDB

BUSI数据集由描绘正常、良性和恶性乳腺癌病例的超声图像及其相应的分割图组成。在我们的研究中，本文使用了647张代表乳腺良性和恶性肿瘤的超声图像。所有这些图像都调整为256 × 256的尺寸。该数据集提供了一个全面的图像集合，有助于检测和区分各种类型的乳腺肿瘤，为医学专业人员和研究人员提供了有价值的见解。

GlaS 数据集由31个序列的612帧标准清晰度（SD）图像组成。每帧的分辨率为384×288，来自23名患者。该数据集与位于西班牙巴塞罗那的临床医院相关联。数据集中的序列是使用 Olympus Q160AL 和 Q165L 等设备结合 Extra II 视频处理器录制的。按照常见的做法，本文特别选用了 GlaS 数据集中的165张图像，并将它们调整为512×512的尺寸。

CVC-ClinicDB 数据集，通常简称为 "CVC"，是一个公开可获取的结肠镜视频中的息肉诊断资源。该数据集包含总共612张图像，每张图像的分辨率为384×288，精心从31个不同的结肠镜序列中提取。这些帧提供了各种不同的息肉实例，对于息肉检测算法的开发和评估非常有用。为了确保本研究中使用的不同数据集之间的一致性，CVC-ClinicDB 数据集中的所有图像均被统一调整为256 × 256的尺寸。

（二）试验结果

试验指标为IoU和F1分数。

（三）效融试验

1、KAN Layer数量

在U-KAN中引入KAN层通过有效的embeddings ，促进了更精细的分割细节建模。本次消融实验的目标是评估不同数量的KAN层对模型性能的影响。将KAN层的数量从1层修改至5层，实验表明，包含3层KAN层的配置表现最优。这些结果证明，在U-KAN中合理地集成适量的KAN层可以有效地捕捉与分割相关的复杂细节。

2、KAN Layer和MLP相比较

为了验证KAN层在提升模型性能中的作用，在这些实验中，将引入的KAN层替换为传统的多层感知器（MLP），以观察这种替换是否会导致性能下降。该方法能够更直观地理解KAN层在提升模型整体表现中的重要性。首先，修改了一个已经包含KAN层的模型，将部分或全部KAN层替换为标准的MLP。随后，使用相同的数据集和训练参数，重新训练修改后的模型，并记录其在不同任务中的性能。实验结果显示，当KAN层被MLP替换时，特别是在需要强大特征提取和表达能力的复杂任务中，模型性能明显下降。这些发现进一步证明了KAN层在增强模型表达能力和提升整体性能中的关键作用。

3、模型大小

还对U-KAN的不同模型规模进行了消融研究。具体来说，试验了U-KAN的不同配置，分别称为小型和大型模型。这些模型的主要区别在于它们的通道设置，分别由第一到第三KAN层的不同通道数量（C1-C3）决定，详见表6。小型模型的通道设置为64-96-128，而大型模型的通道数量设定为256-320-512。与之相比，设置默认模型的通道数量配置为128-160-256。可以看到，较大的模型与更高的性能相关，这与集成KAN的模型所表现出的扩展法则特性一致。最终，为了在性能和计算成本之间取得平衡，本文在实验中采用了默认的基础模型。

4、解释性

进一步通过分析激活模式来探索KAN层的可解释性，如下图所示。当使用MLP层（第一列）时，模型难以识别合适的激活区域，且其Plausibility IoU（可信度IoU）较差，这是一种衡量激活图与真实掩码（GT masks）之间IoU的指标（数值越高越好）。相比之下，集成了KAN层后（第二列），模型在准确定位兴趣区域方面表现出显著的改进，激活边界更接近于真实情况（第三列）。这进一步证明了KAN层在增强深度模型解释性决策中的关键作用，尤其是在掩码预测任务中，与KAN的其他观察结果一致。

五、结论

本文介绍了 U-KAN，并展示了 Kolmogorov-Arnold 网络（KANs）在提升 U-Net 等主干网络在各种视觉应用中的巨大潜力。通过将 KAN 层集成到 U-Net 架构中，可以构建一个在准确性、效率和可解释性方面都表现出色的视觉任务网络。我们在多个医学图像分割任务下对该方法进行了实证评估。此外，U-KAN 的适应性和有效性也凸显了其作为扩散模型中噪声预测的优越替代方案的潜力。这些研究结果强调了探索 KAN 等非传统网络结构对于推进更广泛的视觉应用的重要性。