U-KAN Makes Strong Backbone for Medical
Image Segmentation and Generation

Arxiv 2024

Chenxin Li1*, Xinyu Liu1*, Wuyang Li1*, Cheng Wang1*, Hengyu Liu1, Yifan Liu1, Zhen Chen2, Yixuan Yuan1

1The Chinese University of Hong Kong    2Centre for Artificial Intelligence and Robotics, Hong Kong   

Paper

Code

BibTex

Abstract

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.


Highlight

  • We present the first effort to incorporate the advantage of emerging KANs to improve established U-Net pipeline to be more accurate, efficient and interpretable.
  • We propose a tokenized KAN block to effectively steer the KAN operators to be compatible with the exiting convolution-based designs.
  • We empirically validate U-KAN on a wide range of medical segmentation benchmarks, achieving impressive accuracy and efficiency.
  • The application of U-KAN to existing diffusion models as an improved noise predictor demonstrates its potential in backboning generative tasks and broader vision settings.


Network



Overview of U-KAN pipeline. After feature extraction by several convolution stages, the intermediate maps are tokenized and processed by stacked KAN layers. The time embedding is only injected into the KAN blocks when applied for Diffusion U-KAN.

Segmentation U-KAN

Qualitative Results

Quantitative Results


Diffusion U-KAN

Qualitative Results

Quantitative Results


Citation

@article{li2024ukan,
  author    = {Chenxin Li and Xinyu Liu and Wuyang Li and Cheng Wang and Hengyu Liu and Yixuan Yuan},
  title     = {U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation},
  journal   = {arXiv preprint},
  year      = {2024}
}