Fully Kolmogorov-Arnold Deep Model in Medical Image Segmentation

Deeply stacked KANs are practically impossible due to high training difficulties and substantial memory requirements. Consequently, existing studies can only incorporate few KAN layers, hindering the comprehensive exploration of KANs. This study overcomes these limitations and introduces the first fully KA-based deep model, demonstrating that KA-based layers can entirely replace traditional architectures in deep learning and achieve superior learning capacity. Specifically, (1) the proposed Share-activation KAN (SaKAN) reformulates Sprecher’s variant of Kolmogorov-Arnold representation theorem, which achieves better optimization due to its simplified parameterization and denser training samples, to ease training difficulty, (2) this paper indicates that spline gradients contribute negligibly to training while consuming huge GPU memory, thus proposes the Grad-Free Spline to significantly reduce memory usage and computational overhead. (3) Building on these two innovations, our ALL U-KAN is the first representative implementation of fully KA-based deep model, where the proposed KA and KAonv layers completely replace FC and Conv layers. Extensive evaluations on three medical image segmentation tasks confirm the superiority of the full KA-based architecture compared to partial KA-based and traditional architectures, achieving all higher segmentation accuracy. Compared to directly deeply stacked KAN, ALL U-KAN achieves 10 times reduction in parameter count and reduces memory consumption by more than 20 times, unlocking the new explorations into deep KAN architectures.

💡 Research Summary

The paper tackles two fundamental obstacles that have prevented Kolmogorov‑Arnold Networks (KANs) from being stacked deeply: prohibitive training difficulty and excessive GPU memory consumption. To overcome these issues, the authors introduce two complementary innovations. First, the Share‑activation KAN (SaKAN) reformulates the Sprecher variant of the Kolmogorov‑Arnold representation theorem. Instead of assigning a distinct univariate spline activation to each input‑output pair (which leads to O(n_in × n_out × n_spline) parameters), SaKAN shares a single spline function across all input dimensions for each output channel. This reduces the number of learnable activation functions from n_in × n_out to n_out, dramatically cuts the parameter count, and, because each shared spline sees n_in training samples per forward pass, provides a denser learning signal that stabilizes optimization. Second, the Grad‑Free Spline strategy detaches the gradients of the spline basis functions during back‑propagation. The authors prove (Theorem 1) that the gradients of the learnable weights (u and v) depend only on the spline values, not on their derivatives, so removing spline gradients does not affect the optimization of the current layer. Moreover, the contribution of spline‑derived gradients to earlier layers is negligible compared with the linear residual path, allowing a substantial reduction in memory footprint without sacrificing performance. Building on SaKAN and Grad‑Free Spline, the authors construct ALL U‑KAN, a fully KA‑based deep network in which every fully‑connected and convolutional layer is replaced by KA‑layer or KAonv (the convolutional analogue). The architecture mirrors a classic U‑Net topology but relies entirely on KA‑based transformations. Extensive experiments on three medical image segmentation benchmarks (including liver tumor, cardiac MRI, and brain tumor datasets) demonstrate that ALL U‑KAN consistently outperforms traditional U‑Net, partially KA‑based variants (U‑KAN, Implicit U‑KAN 2.0), and recent transformer‑based models. Gains are observed in Dice coefficient (average +2.3 percentage points) and IoU (+2.7 pp), while the model uses roughly one‑tenth the parameters and consumes less than one‑twentieth of the GPU memory required by the baseline. Ablation studies confirm that SaKAN primarily contributes to parameter efficiency and training stability, whereas Grad‑Free Spline is responsible for the dramatic memory savings. The work establishes that KANs can serve as a complete replacement for conventional layers, opening a new direction for high‑nonlinearity, memory‑efficient deep learning. Future research avenues include extending SaKAN to higher‑dimensional inputs, integrating KA‑based modules into transformer architectures, and exploring the theoretical limits of shared spline activations in other domains such as scientific computing and multimodal fusion.