Source-linked AI summary
U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation
Chenxin Li, Xinyu Liu, Wuyang Li, Cheng Wang, Hengyu Liu, Yifan Liu, Zhen Chen, Yixuan Yuan
TL;DR
Existing U-Net-based vision models have limited nonlinear pattern modeling and deficient interpretability. This paper redesigns U-Net with tokenized KAN layers as U-KAN, reporting higher accuracy with lower computation cost on medical segmentation benchmarks and exploring it as a diffusion noise predictor.
Problem
Existing U-shaped vision networks are limited in nonlinear pattern modeling and interpretability, motivating a stronger and more explainable backbone.
Method
U-KAN redesigns the U-Net pipeline by inserting a tokenized KAN block near the bottleneck within a convolutional encoder-decoder architecture with skip connections.
Results
U-KAN achieves higher accuracy with lower computation cost than established U-Net backbones on medical segmentation benchmarks and is evaluated as a diffusion noise predictor.
Takeaways & Limitations
U-KAN shows potential as a backbone for medical image segmentation, diffusion-based generation, and broader vision settings.
Abstract
from arXiv · showhide
U-Net has become a cornerstone in various visual applications such as image segmentation and diffusion probability models. While numerous innovative designs and improvements have been introduced by incorporating transformers or MLPs, the networks are still limited to linearly modeling patterns as well as the deficient interpretability. To address these challenges, our intuition is inspired by the impressive results of the Kolmogorov-Arnold Networks (KANs) in terms of accuracy and interpretability, which reshape the neural network learning via the stack of non-linear learnable activation functions derived from the Kolmogorov-Anold representation theorem. Specifically, in this paper, we explore the untapped potential of KANs in improving backbones for vision tasks. We investigate, modify and re-design the established U-Net pipeline by integrating the dedicated KAN layers on the tokenized intermediate representation, termed U-KAN. Rigorous medical image segmentation benchmarks verify the superiority of U-KAN by higher accuracy even with less computation cost. We further delved into the potential of U-KAN as an alternative U-Net noise predictor in diffusion models, demonstrating its applicability in generating task-oriented model architectures. These endeavours unveil valuable insights and sheds light on the prospect that with U-KAN, you can make strong backbone for medical image segmentation and generation. Project page:\url{https://yes-u-kan.github.io/}.
Introduction
U-Net-based vision backbones remain constrained by limited nonlinear modeling and weak interpretability. U-KAN integrates tokenized KAN processing into U-Net, achieving higher segmentation accuracy with lower computation cost and extending to diffusion noise prediction.
- Introduction: U-Net remains a landmark encoder-decoder backbone with skip connections for medical image segmentation and has also supported image translation and diffusion tasks.Diffusion models train U-Net to iteratively predict noise for removal during denoising.
- Introduction: Existing U-shape networks face limited nonlinear pattern modeling because conventional kernels capture spatial dependencies and channel relationships linearly.Medical images contain intricate diagnostic patterns and channels with differing clinical relevance.
- Introduction: Their empirical, heuristic architecture design also leaves existing U-shape models difficult to interpret and explain in clinical decision-making.The passage connects this unexplainability with risks in diagnostic system design.
- Introduction: U-KAN combines a multilayer encoder-decoder and skip connections with a tokenized KAN block near the bottleneck.The block projects intermediate features into tokens and applies a KAN operator to extract informative patterns.
- Introduction: U-KAN improves established U-Net pipelines through KAN-based nonlinear modeling and interpretability, while retaining compatibility with convolutional designs.The contribution list identifies the tokenized KAN block as the mechanism for compatibility.
- Introduction: Medical segmentation benchmarks report higher U-KAN accuracy with lower computation cost than established U-Net backbones.The contribution summary also characterizes U-KAN as more accurate, efficient, and interpretable.
- Introduction: The paper also evaluates U-KAN as an alternative U-Net noise predictor for diffusion models and broader generative vision tasks.This extends the proposed backbone beyond segmentation.
Related Work
Related work establishes U-Net as a central medical-segmentation architecture amid convolutional, nested, and Transformer-based variants. KANs provide the paper’s alternative basis through learnable one-dimensional activation functions and claimed interpretability.
- U-Net Backbone for Medical Image Segmentation: Medical image segmentation is a challenging task where deep learning has achieved substantial recent advances.The literature includes convolutional and Transformer-based approaches.
- U-Net Backbone for Medical Image Segmentation: U-Net uses an encoder-decoder architecture to capture image features, while CE-Net and Unet++ add contextual encoding and nested multiscale fusion.These modifications target receptive-field, semantic-representation, and segmentation-accuracy improvements.
- U-Net Backbone for Medical Image Segmentation: Transformer-based models, including Vision Transformer, Medical Transformer, and TransUNet, have also gained attention for visual and medical-imaging tasks.
- Diffusion Probability Models: Diffusion probability models use a distinct generative paradigm based on a fixed Markov chain that maps latent space.They are contrasted with VAEs, GANs, and vector-quantization methods.
- Kolmogorov–Arnold Networks (KANs): KANs comprise concatenated layers with learnable one-dimensional activation functions and are described as interpretable while approximating complex high-dimensional functions.
Method
U-KAN combines a convolutional encoder-decoder with tokenized KAN blocks near the bottleneck, preserving U-Net skip connections while applying learnable nonlinear functions to intermediate representations. Its diffusion variant adds time conditioning and trains the network to predict noise for image generation.
- U-KAN architecture: U-KAN uses a two-phase encoder-decoder with convolution blocks, tokenized KAN blocks, resolution changes, and skip connections.The encoder begins with convolution operations and the decoder ends with convolution blocks; encoder blocks halve resolution and decoder blocks double it.
- KAN as Efficient Embedder: KAN layers replace linear weight matrices with learnable activation functions on network edges, supporting smaller models and improved interpretability.The paper characterizes KAN layers as compositions of learnable one-dimensional activation functions rather than traditional linear weight matrices.
- Tokenized KAN Phrase: Intermediate convolutional feature maps are reshaped into flattened patches and projected into a latent embedding space before KAN processing.The projection E maps vectorized patches into D-dimensional embeddings; a convolution layer provides the projection and positional information.
- Tokenized KAN Phrase: Tokenized KAN blocks apply KAN layers followed by depth-wise convolution, batch normalization, ReLU activation, residual addition, and layer normalization.The implementation uses three KAN layers per block, with two blocks in the Tokenized KAN Phrase.
- U-KAN Decoder: The decoder concatenates skip-connected features with upsampled features to recover low-level details and produce the final segmentation map.The final output feature map has one channel dimension for the semantic categories and is optimized with pixel-wise cross-entropy loss.
- Extending U-KAN to Diffusion Models: Diffusion U-KAN injects learnable time embeddings into KAN blocks, removes depth-wise convolution and residual connections, and predicts noise from corrupted images.Its noise-prediction objective is optimized with mean squared error, after which DDPM sampling uses the trained model for denoising.
Experiments
Experiments evaluate U-KAN across heterogeneous medical datasets for segmentation and diffusion-based generation. Results report accuracy, efficiency, qualitative outputs, ablations, and activation-based explainability.
- Experimental Setup: Three heterogeneous medical datasets provide evaluation settings with differing characteristics, data sizes, and image resolutions.The datasets support segmentation and generation experiments.
- Segmentation Performance: Across all datasets, U-KAN surpasses compared segmentation methods on IoU and F1 scores.Comparisons include U-Net, U-Net++, Att-UNet, U-Mamba, U-NeXt, and Rolling-UNet.
- Segmentation Performance: U-KAN achieves the best trade-off between segmentation accuracy and efficiency, except that UNext has an efficiency advantage.Efficiency is assessed using parameter volume and Gflops alongside segmentation accuracy.
- Segmentation Performance: Qualitative comparisons show fewer false positives for U-KAN than other methods across the evaluated medical scenarios.CNN-based methods are described as more prone to over- or undersegmentation.
- Generation Performance: Diffusion U-KAN shows superior generative performance to compared state-of-the-art diffusion models across three datasets using FID and IS.The evaluation compares generated and real images with FID and assesses classification-based image quality with IS.
- Ablation Studies: Three KAN layers yield the strongest ablation performance, while replacing KAN layers with MLPs causes declines across multiple tasks.The reported decline is particularly pronounced for intricate tasks requiring robust feature extraction and representational capacity.
- Ablation Studies: Larger U-KAN models correlate with enhanced performance, but the default base model balances performance against computational expense.The Small, default, and Large configurations differ in channel settings across the first three KAN layers.
- Explainability: Activation analyses show KAN layers locating regions of interest and boundaries more closely aligned with ground-truth masks than MLP layers.The comparison uses Plausibility IoU, where higher values indicate closer overlap between thresholded activation maps and ground-truth masks.
Conclusion
U-KAN integrates KAN layers into U-Net backbones to improve accuracy, efficiency, and interpretability across visual applications. The paper evaluates it on medical image segmentation and highlights future extensions to broader settings and higher-dimensional data.
- U-KAN integrates KAN layers into U-Net to strengthen vision-task backbones in accuracy, efficiency, and interpretability.
- The method is empirically evaluated across several medical image segmentation tasks.
- Future work will extend the network operators to broader settings and higher-dimensional formats, including temporal, genomic, and 3D data.