[论文解读] TransUKAN:Computing-Efficient Hybrid KAN-Transformer for Enhanced Medical Image Segmentation

Yanlin Wu, Tao Li+3 more|arXiv (Cornell University)|Sep 23, 2024

Medical Image Segmentation Techniques被引用 4

一句话总结

TransUKAN 提出了一种计算高效的混合 KAN-Transformer 架构，用于医学图像分割，通过改进的、参数高效的 KAN（EfficientKAN）替代传统 Transformer 中的 MLP 和 QKV 投影，以增强非线性建模能力并降低计算负载。该模型在多个医学数据集上实现了最先进性能，参数量仅为 20.85M，展现出卓越的准确性、效率和鲁棒性。

ABSTRACT

U-Net is currently the most widely used architecture for medical image segmentation. Benefiting from its unique encoder-decoder architecture and skip connections, it can effectively extract features from input images to segment target regions. The commonly used U-Net is typically based on convolutional operations or Transformers, modeling the dependencies between local or global information to accomplish medical image analysis tasks. However, convolutional layers, fully connected layers, and attention mechanisms used in this process introduce a significant number of parameters, often requiring the stacking of network layers to model complex nonlinear relationships, which can impact the training process. To address these issues, we propose TransUKAN. Specifically, we have improved the KAN to reduce memory usage and computational load. On this basis, we explored an effective combination of KAN, Transformer, and U-Net structures. This approach enhances the model's capability to capture nonlinear relationships by introducing only a small number of additional parameters and compensates for the Transformer structure's deficiency in local information extraction. We validated TransUKAN on multiple medical image segmentation tasks. Experimental results demonstrate that TransUKAN achieves excellent performance with significantly reduced parameters. The code will be available athttps://github.com/wuyanlin-wyl/TransUKAN.

研究动机与目标

• 解决现有 U-Net 和基于 Transformer 的模型在医学图像分割中参数量过高和计算成本过高的问题。
• 通过引入一种稀疏化、高效的变体（EfficientKAN），克服标准 KAN 在图像处理中的局限性，降低内存和计算需求。
• 将 KAN 集成到 U-Net 框架内的 Transformer 机制中，以增强局部特征建模和全局依赖性学习能力。
• 在大幅减少模型规模和训练复杂度的同时，实现最先进水平的分割精度。

提出的方法

• 提出 EfficientKAN，即 KAN 的稀疏化版本，通过仅整合关键激活值（利用平均池化）来降低矩阵计算的复杂度。
• 用 EfficientKAN 替代多头自注意力机制中的全连接层和 QKV 投影，以更少的参数实现更强的非线性拟合能力。
• 将增强后的 KAN 模块集成到类似 U-Net 的编码器-解码器架构中，并通过跳跃连接保留空间细节。
• 采用混合设计，结合 CNN、U-Net、Transformer 和 KAN，以平衡局部与全局特征提取。
• 在 KAN 中使用可学习激活函数，以比标准 MLP 更高效的方式建模复杂的非线性关系。
• 通过交叉验证和消融研究，在多个医学影像数据集上验证了该架构的有效性。

实验结果

研究问题

• RQ1 能否通过基于 KAN 的 MLP 和 QKV 投影替代方案，在减少参数量和计算量的同时提升分割性能？
• RQ2 在医学图像分割任务中，EfficientKAN 相较于标准 KAN 和 ReLUKAN 在参数效率和性能方面表现如何？
• RQ3 将 KAN 集成到 U-Net-Transformer 混合架构中，是否能同时增强局部细节提取和全局上下文建模能力？
• RQ4 所提出的 TransUKAN 模型在具有不同成像模态和解剖结构的多样化医学影像数据集上，其泛化能力如何？
• RQ5 该模型能否以远少于现有最先进方法的参数量，实现 SOTA 性能？

主要发现

• 在 ISIC 数据集上，TransUKAN 的 Dice 分数达到 91.17%，超过基线模型 TransUNet（86.30%）及其他 SOTA 模型。
• 在 Kvasir 数据集上，TransUKAN 达到 87.75% 的 Dice 分数，比表现第二好的模型（86.30%）高出 1.45 个百分点。
• 该模型将参数量减少至 20.85M，显著低于 TransUNet 的 105.32M，同时在所有评估数据集中性能均得到提升。
• 消融研究证实，EfficientKAN 在降低参数量和训练时间的同时提升了准确率；而 ReLUKAN 反而增加了参数量并导致性能下降。
• TransUKAN 在 NKUT 数据集上保持了 90.29% 的 Dice 分数，在 MWT 数据集上达到 89.09%，展现出在多样化解剖结构上的强大泛化能力。
• 该模型在全部七个评估数据集（包括 BUSI、ISIC、Kvasir 和 SAM）上均实现了最先进性能，且对先前方法保持了一致的性能提升。