Implicit U-KAN2.0: Dynamic, Efficient and Interpretable Medical Image Segmentation
Chun-Wun Cheng, Yining Zhao, Yanqi Cheng, Javier Montoya, Carola-Bibiane Schönlieb, Angelica I Aviles-Rivero
TL;DR
Implicit U-KAN2.0 tackles medical image segmentation by delivering a dynamic, memory-efficient, and interpretable U‑Net variant. It introduces the SONO block with second-order NODEs, expressed as $\\ddot{x}(t) = f(\\boldsymbol{x}, \\dot{\\boldsymbol{x}}, t, \\theta_f)$, and a velocity-enhanced formulation with adjoint-based memory savings, combined with a tokenized MultiKAN layer for richer interactions. A theoretical analysis shows that MultiKAN's approximation error is independent of input dimensionality and depends on the residual rate. Extensive 2D and 3D experiments demonstrate state-of-the-art performance across datasets with improved Dice and boundary accuracy, while maintaining constant memory cost and robustness to noise, highlighting potential for clinical deployment.
Abstract
Image segmentation is a fundamental task in both image analysis and medical applications. State-of-the-art methods predominantly rely on encoder-decoder architectures with a U-shaped design, commonly referred to as U-Net. Recent advancements integrating transformers and MLPs improve performance but still face key limitations, such as poor interpretability, difficulty handling intrinsic noise, and constrained expressiveness due to discrete layer structures, often lacking a solid theoretical foundation.In this work, we introduce Implicit U-KAN 2.0, a novel U-Net variant that adopts a two-phase encoder-decoder structure. In the SONO phase, we use a second-order neural ordinary differential equation (NODEs), called the SONO block, for a more efficient, expressive, and theoretically grounded modeling approach. In the SONO-MultiKAN phase, we integrate the second-order NODEs and MultiKAN layer as the core computational block to enhance interpretability and representation power. Our contributions are threefold. First, U-KAN 2.0 is an implicit deep neural network incorporating MultiKAN and second order NODEs, improving interpretability and performance while reducing computational costs. Second, we provide a theoretical analysis demonstrating that the approximation ability of the MultiKAN block is independent of the input dimension. Third, we conduct extensive experiments on a variety of 2D and a single 3D dataset, demonstrating that our model consistently outperforms existing segmentation networks.