TransUKAN:Computing-Efficient Hybrid KAN-Transformer for Enhanced Medical Image Segmentation

The advent of deep learning technologies, particularly CNNs, has significantly advanced the field of medical image segmentation. U-Net [29] is one of the most classic CNN-based network architectures for medical image segmentation. The success of U-Net has inspired the development of numerous improved models. For instance, UNet++ incorporates nested encoder-decoder sub-networks and redesigned skip connections to improve feature fusion and gradient flow. UNet3+ leverages full-scale skip connections to combine multi-scale features and employs deep supervisions to learn hierarchical representations. 3D UNet and V-Net extend the traditional UNet architecture to three dimensions, making them suitable for volumetric image segmentation. Attention UNet and R2UNet introduce attention mechanisms to focus on relevant features, enhancing segmentation accuracy. Residual UNet and ResUNet++ incorporate residual connections to address the vanishing gradient problem, facilitating the training of deeper networks. Dense UNet and U-Net with Dense Blocks apply dense connectivity to improve feature reuse and network efficiency. Furthermore, UNetGAN combines UNet with generative adversarial networks (GANs) to improve segmentation quality through adversarial training.

Furthermore, to address the lack of long-range dependency modeling capabilities in CNNs, researchers introduced Transformers into the field of computer vision, proposing the Vision Transformer (ViT). Subsequently, networks such as TransUNet, TransFuse, and UCTransNet emerged, combining Transformers with CNNs and fully exploring the extended capabilities of Transformers. Although Transformers can capture global information, they have a large number of parameters, high computational complexity, and are difficult to train. Additionally, Transformers are less effective at extracting local information, which can impact segmentation accuracy.

Kolmogorov–Arnold Networks (KAN) is a neural network architecture based on the Kolmogorov–Arnold super approximation theorem. This theorem states that any multivariate continuous function can be represented as a finite superposition of univariate continuous functions. This theory provides a solid mathematical foundation for the representation and computation of complex functions. In KAN, by constructing a multi-layer network structure and utilizing the combination and superposition of univariate functions, complex nonlinear relationships can be efficiently approximated and represented. The core advantage of KAN lies in its ability to model nonlinear features efficiently with fewer parameters, making it highly effective in handling high-dimensional data and complex tasks. Specifically, KAN captures nonlinear patterns in input data through a series of linear transformations and nonlinear activation functions. Compared to traditional CNNs, KAN has advantages in terms of the number of parameters and computational complexity, making it highly applicable in scenarios requiring efficient models.

In computer vision, ConvKANs [4] adapt KANs into a convolutional architecture by integrating nonlinear activation functions from KANs into the convolutional layer. This integration effectively reduces parameter count while maintaining high accuracy levels. Graph-based applications also benefit from KAN [20, 36, 5], which replaces traditional MLPS in graph neural networks with KAN layers. By using learnable spline-based functions instead of fixed activation functions, this substitution enhances the model’s ability to capture complex relationships in graph-structured data.

KANs offer significant advantages in accuracy and interpretability, positioning them as a promising alternative to traditional neural network models. However, the application of KAN in image processing has not been fully explored, especially in challenging areas such as medical image processing, where issues of high computation and memory usage have not been adequately addressed.

In our research, KAN has been improved and introduced into the task of medical image segmentation. By integrating it with the Transformer, we fully leverage its strengths in modeling nonlinear relationships. The improvements and introduction of KAN not only reduce the number of model parameters and lower computational complexity but also enhance segmentation accuracy. This combined approach effectively addresses the shortcomings of Transformers in local information extraction and significantly improves overall model performance and training efficiency.

As shown in Fig. 1, the overall structure of TransUKAN combines the advantages of CNN, Transformer, and KAN. First, the input image is processed by a CNN to extract features, generating feature maps. These feature maps are reshaped and linearly projected into a high-dimensional feature space. Then, these embedded features serve as the input to the Kansformer encoder, processed through multiple Kansformer layers to capture contextual information. Each Kansformer layer includes Layer Normalization (LN), Multi-Head Self-Attention mechanism (MSA), and EfficientKAN, ensuring the integration of global and local information in the features.

To restore the spatial resolution of the image, the encoded features are progressively upsampled through a cascaded upsampler. It consists of multiple upsampling blocks, each containing a 2x upsampling operation, a 3×3 convolutional layer, and a ReLU activation function. During the upsampling process, the encoded features are fused with high-resolution feature maps from the CNN encoding path through skip connections, enhancing the recovery of low-level spatial information for precise image segmentation. Finally, the fused feature maps are further upsampled to the full resolution of the original image to generate the final segmentation mask.

Since its inception, KAN has been optimized compared to the MLP. Its powerful nonlinear representation capability and training stability enable it to better capture the complex features of data. Therefore, an intuitive improvement approach is to directly replace MLP with KAN. Additionally, to maintain the continuity of nonlinear representation, we also replaced the QKV mapping matrices with a single-layer KAN. This allows the powerful nonlinear representation capability of KAN to be utilized in the self-attention mechanism, ensuring that the model maintains efficient nonlinear representation at each stage of feature extraction and processing, thereby improving overall performance.

The overall structure of Kansformer is shown in Figure 1(a), primarily composed of LN, EfficientKAN layers, and MSA. The output of the l-th layer can be expressed as follows:

Here, $\text{LN}(\cdot)$ denotes the Layer Normalization operator, and $z_{l}$ represents the feature map encoded by the Kansformer at layer $l$.

By gradually replacing the MLP and QKV mapping matrices, the model benefits from KAN’s parameter compression and nonlinear representation at each step. After replacing the MLP, the reduction in parameter count and increase in computational efficiency provides a solid foundation for further replacing the QKV mapping matrices. When the QKV mapping matrices are also replaced with KAN, the overall computational complexity is further reduced, and parameter efficiency is further improved. This synergistic optimization ensures that the model achieves optimal performance at each stage.

The core idea of KAN can be expressed by Eq. (3), which states that a high-dimensional function can be represented as a composition of a finite number of one-dimensional functions:

Specifically, assuming the input vector $x$ has a length of $n$, the computation of the output $y$ can be expressed as follows:

Here, $\phi_{i,j}$ is referred to as the inner function, and $\Phi_{i}$ is referred to as the outer function. Specifically, the inner and outer functions can be expressed in the form of linear combinations and B-spline functions as follows:

Here, $B_{i}(x)$ is a B-spline function, $w_{b}$ and $w_{s}$ are weight parameters, and $c$ is a control coefficient for shaping the B-spline. B-spline functions are a set of bell-shaped functions used to represent any univariate function on a finite domain. These functions have the same shape but different positions.

Due to the computational complexity of B-spline functions, KAN is limited in utilizing the parallel processing capabilities of GPUs. This complexity results in significant limitations in processing speed and scalability, especially in fine-grained classification tasks such as medical image segmentation. The ReLU-KAN architecture simplifies the basis function by adopting ReLU and dot product operations, optimizing the computational process as follows:

Here, $e_{i}$ and $s_{i}$ represent the upper and lower bounds of the basis function, respectively, controlling the range and position of the basis function through these two parameters. The factor $\frac{16}{(e_{i}-s_{i})^{4}}$ in the basis function plays a role in normalization and scaling. Its main purpose is to ensure that the activation value remains within a reasonable range, avoiding overflow or underflow. Subsequently, the computational process of KAN is optimized through convolution operations. This simplified basis function has higher computational efficiency, making it more suitable for GPU processing.

Through our experiments and observations, we found that although using convolution can make the KAN more favorable for GPU computation during the activation value integration, this method integrates not only the activation values of different neurons but also the activation values of individual neurons, treating all activation values as a whole. This approach can result in significant information redundancy and introduce a large number of parameters when handling high-dimensional data, leading to substantial computational resource usage during backpropagation. Specifically, assuming the input $X$ has dimensions $(B,C_{in})$, where $B$ is the batch size and $C_{in}$ is the number of neurons, the network activated based on the ReLUKAN principle will generate $X_{1}$ with dimensions $(B,(G+K),C_{in})$, where $G$ and $K$ are hyperparameters used by KAN to generate the number of grids and also represent the height and width of the feature map and convolution kernel. The relationship between the activation values of each neuron calculated through convolution operations can be expressed as follows:

Here, the size of the convolution kernel is $(G+K,C_{in})$, and the dimension of the integrated feature map $X^{\prime\prime}$ is $(B,C_{out})$. This method introduces significant computational and parameter redundancy when handling high-dimensional data because the convolution kernel needs to integrate all neurons at each position:

Here, $A_{ij}$ represents the activation values of different neurons in the input $X$. The same row represents different activation values of the same neuron, and the same column represents different neurons. $B$, $C$, and $D$ are convolution kernels that integrate the activation values of $X$. ReluKAN replaces the original KAN activation value integration operation with convolution operations using kernels of the same size as the input feature map. Therefore, when the input feature map dimension increases, it leads to a significant increase in the number of parameters and computational load. To reduce computational and parameter redundancy, we limit the integration of activation values to within a single neuron and implicitly learn the relationships between different neurons only during backpropagation, as follows:

By using this method, the activation value processing matrix is maximally sparsified. To further reduce parameters and computational load, the activation value processing matrix is further simplified such that $a=k=1/j$. In practical implementation, this can be replaced by an average pooling operation. Additionally, to further enhance the nonlinear fitting capability of KAN, a square calculation is introduced after the integration of activation values.

ISIC [11]: This dataset contains 2594 skin lesion images captured from real patients using a dermatoscope equipped with a digital camera. Each image has been annotated by a professional physician to mark the area of the skin lesion, and all data have been reviewed and managed by professional dermatologists with knowledge of dermatoscopy.

Kvasir [19]: Kvasir-SEG is an open-source dataset manually annotated and verified by an experienced gastroenterologist. The dataset contains 1000 images of polyps and their corresponding masks, with image resolutions ranging from 332 $\times$ 487 to 1920 $\times$ 1072. The aim of creating this dataset is to promote the development and progress of polyp detection tasks.

BUSI [1]: The BUSI dataset collects ultrasound images of breasts from women aged between 25 and 75 years old. The dataset includes 780 images from 600 female patients, which are divided into three categories: normal, benign, and malignant. Among them, there are 133 normal cases, 437 benign tumors, and 210 malignant tumors. The average image size is 500 $\times$ 500 pixels.

NKUT [39]: NKUT is a specialized dataset designed for the segmentation of pediatric mandibular wisdom teeth (MWT) from Cone Beam Computed Tomography (CBCT) images. This dataset comprises 133 CBCT volumes, representing over 53,000 slices, with patient ages ranging from 7 to 22 years, and an average age of 13.2 years. The dataset includes detailed pixel-level annotations created by pediatric dentistry experts, covering bilateral MWT germs, second molars (SM), and partial alveolar bones (AB).

The images in the dataset were uniformly preprocessed and adjusted to 256x256 pixels to meet the model’s input requirements. The dataset was divided into training, validation, and test sets in an 8:1:1 ratio. Data augmentation techniques, including random cropping, rotation, and flipping, were widely applied during training to increase data diversity. Model training was conducted on a single NVIDIA A6000 GPU.

During training, we used the Adam optimizer with an initial learning rate set to 1e-4 and employed weight decay to prevent overfitting. The total training epochs were set to 200, with the first 10 epochs as a warm-up phase, using linear learning rate growth, followed by cosine annealing learning rate decay. Each training batch size was set to 8. For binary classification, the loss function is a weighted sum of cross-entropy loss and dice loss, while for multi-class classification, only cross-entropy loss is used. During the validation, we use DICE, IOU, and accuracy as evaluation metrics, and we also record model parameters as well as inference times to comprehensively evaluate the model’s performance.

We conducted a comprehensive evaluation of our method against several state-of-the-art models across multiple datasets. The results are presented in Table LABEL:table2, it can be seen that our proposed method has significant competitiveness in medical image segmentation. Our model outperformed its counterparts across all six medical image segmentation tasks. These experimental results provide strong evidence that our method has effectively improved the capabilities of SAM in medical image segmentation. Furthermore, our method achieves the objective of employing a single model for segmenting multiple medical images, while consistently delivering excellent performance.

In this section, we assess the impact of the proposed components on segmentation performance across external and internal datasets. All these models are based on TransUNet. The experimental results are shown in Table 2. The baseline model, comprising 105.3 M, demonstrates moderate performance on the external dataset with a DICE score of 91.8%, 88.51% and 75.58%, respectively.

After replacing all the fully connected operations in the Transformer with KAN, the number of parameters was reduced to 21.23 M, approximately five times lower. The results on external datasets remain highly competitive, demonstrating the potential of the KAN in medical image segmentation. However, after improving KAN to ReLUKAN, the introduction of numerous convolution operations led to an increase in model parameters to 233.2 M. Despite the increase in parameters, the model’s performance did not significantly improve and even showed a downward trend, demonstrating that the ReLUKAN approach is not suitable for medical image processing tasks. Finally, using EfficientKAN to build the TransUKAN reduced the number of parameters to 20.85 M. This significantly lowered memory usage while accelerating training and improved the model’s testing accuracy on external datasets, demonstrating the effectiveness of EfficientKAN.

Additionally, on the internal CBCT validation set, TransUKAN still achieved excellent performance, demonstrating its generalization and robustness.