[Paper Review] 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 Techniques4 citations

TL;DR

TransUKAN proposes a computing-efficient hybrid KAN-Transformer architecture for medical image segmentation, replacing traditional MLPs and QKV projections in Transformers with an improved, parameter-efficient KAN (EfficientKAN) to enhance nonlinear modeling and reduce computational load. The model achieves state-of-the-art performance on multiple medical datasets with significantly fewer parameters—20.85M—demonstrating superior accuracy, efficiency, and robustness.

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.

Motivation & Objective

• To address the high parameter count and computational cost of existing U-Net and Transformer-based models in medical image segmentation.
• To overcome the limitations of standard KANs in image processing by introducing a sparsified, efficient variant (EfficientKAN) that reduces memory and computation.
• To integrate KANs into the Transformer mechanism within a U-Net framework to enhance local feature modeling and global dependency learning.
• To achieve state-of-the-art segmentation accuracy while drastically reducing model size and training complexity.

Proposed method

• Proposes EfficientKAN, a sparsified version of KAN that reduces matrix computation complexity by integrating only key activation values via average pooling.
• Replaces fully connected layers and QKV projections in the multi-head self-attention mechanism with EfficientKAN to improve nonlinear fitting with fewer parameters.
• Integrates the enhanced KAN module into a U-Net-like encoder-decoder architecture with skip connections to preserve spatial details.
• Employs a hybrid design combining CNNs, U-Net, Transformers, and KANs to balance local and global feature extraction.
• Uses a learnable activation function in KANs to model complex nonlinear relationships more efficiently than standard MLPs.
• Validates the architecture across multiple medical imaging datasets using cross-validation and ablation studies.

Experimental results

Research questions

• RQ1 Can a KAN-based replacement for MLPs and QKV projections in Transformers improve segmentation performance while reducing parameters and computation?
• RQ2 How does EfficientKAN compare to standard KAN and ReLUKAN in terms of parameter efficiency and performance on medical image segmentation tasks?
• RQ3 Does integrating KAN into a U-Net-Transformer hybrid architecture enhance both local detail extraction and global context modeling?
• RQ4 To what extent does the proposed TransUKAN model generalize across diverse medical imaging datasets with varying modality and anatomy?
• RQ5 Can the model achieve SOTA performance with significantly fewer parameters than existing state-of-the-art methods?

Key findings

• TransUKAN achieves a Dice score of 91.17% on the ISIC dataset, surpassing the baseline TransUNet (86.30%) and other SOTA models.
• On the Kvasir dataset, TransUKAN reaches 87.75% Dice, outperforming the next best model (86.30%) by 1.45 percentage points.
• The model reduces parameters to 20.85M—significantly lower than the 105.32M of TransUNet—while improving performance across all evaluated datasets.
• Ablation studies confirm that EfficientKAN reduces parameter count and training time while boosting accuracy, unlike ReLUKAN, which increased parameters and degraded performance.
• TransUKAN maintains strong generalization, achieving 90.29% Dice on the NKUT dataset and 89.09% on MWT, demonstrating robustness across diverse anatomical structures.
• The model achieves state-of-the-art performance on all seven evaluated datasets, including BUSI, ISIC, Kvasir, and SAM, with consistent improvements over prior methods.