When Swin Transformer Meets KANs: An Improved Transformer Architecture for Medical Image Segmentation

TL;DR

This paper introduces UKAST, a novel transformer-based architecture integrating Kolmogorov-Arnold Networks into Swin Transformer encoders, achieving state-of-the-art medical image segmentation performance with improved data efficiency and reduced computational costs.

Contribution

The paper presents UKAST, a new transformer architecture that incorporates rational-function based KANs into Swin Transformers, enhancing data efficiency and segmentation accuracy in medical imaging.

Findings

Achieves state-of-the-art results on multiple benchmarks.

Performs well in data-scarce scenarios.

Reduces computational complexity compared to baseline models.

Abstract

Medical image segmentation is critical for accurate diagnostics and treatment planning, but remains challenging due to complex anatomical structures and limited annotated training data. CNN-based segmentation methods excel at local feature extraction, but struggle with modeling long-range dependencies. Transformers, on the other hand, capture global context more effectively, but are inherently data-hungry and computationally expensive. In this work, we introduce UKAST, a U-Net like architecture that integrates rational-function based Kolmogorov-Arnold Networks (KANs) into Swin Transformer encoders. By leveraging rational base functions and Group Rational KANs (GR-KANs) from the Kolmogorov-Arnold Transformer (KAT), our architecture addresses the inefficiencies of vanilla spline-based KANs, yielding a more expressive and data-efficient framework with reduced FLOPs and only a very small increase in parameter count compared to SwinUNETR. UKAST achieves state-of-the-art performance on four diverse 2D and 3D medical image segmentation benchmarks, consistently surpassing both CNN- and Transformer-based baselines. Notably, it attains superior accuracy in data-scarce settings, alleviating the data-hungry limitations of standard Vision Transformers. These results show the potential of KAN-enhanced Transformers to advance data-efficient medical image segmentation. Code is available at: https://github.com/nsapkota417/UKAST