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

TL;DR

U-KAN augments the standard U-Net backbone with tokenized Kolmogorov–Arnold Networks to enhance nonlinear feature modeling and interpretability in medical image segmentation. The architecture combines a Convolution Phrase with a Tokenized KAN Phrase, enabling efficient, expandable processing near the bottleneck and skip-connected decoding. Experiments on BUSI, GlaS, and CVC-ClinicDB show superior segmentation accuracy with competitive efficiency, and diffusion-backbone experiments demonstrate improved noise-prediction and image generation metrics. Ablation studies confirm the benefit of multiple KAN layers, highlight the advantage over MLP baselines, and reveal improved explainability through KAN activations, suggesting broad applicability in vision tasks beyond segmentation.

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. Project page:\url{https://yes-u-kan.github.io/}.

Figures (4)

• Figure 1: Overview of U-KAN pipeline. After feature extraction by several convolution blocks in Convolution Phrase, the intermediate maps are tokenized and processed by stacked Tok-KAN blocks in Tokenized KAN Phrase. The time embedding is only injected into the KAN blocks when applied for Diffusion U-KAN.
• Figure 2: Visualized segmentation results of the proposed U-KAN against other state-of-the-arts over three heterogeneous medical scenarios.
• Figure 3: Generated images by proposed Diffusion U-KAN in three heterogeneous medical scenarios.
• Figure 4: Explainability of U-KAN with channel activation.