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

AI summary

Overview

Research area: Medical image segmentation; deep learning architecture design, specifically combining Vision Transformers (Swin Transformer) with Kolmogorov–Arnold Networks (KANs).

Technical level: Advanced. The paper assumes familiarity with U-Net architectures, self-attention, Swin Transformer shifted-window attention, Kolmogorov–Arnold Networks, and medical image segmentation benchmarks.

Scope in one sentence: The paper introduces UKAST, a U-Net-like architecture that replaces the MLP feed-forward blocks in a Swin Transformer encoder with rational-function-based Group Rational KANs, and evaluates it on four 2D and 3D medical image segmentation datasets under both full-data and reduced-data settings.

The work is published in the Proceedings of IEEE ISBI 2026 (arXiv:2511.04084v2, cs.CV, 26 Jan 2026), by Nishchal Sapkota, Haoyan Shi, Yejia Zhang, Xianshi Ma, Bofang Zheng, Fabian Vazquez, Pengfei Gu, and Danny Z. Chen, with affiliations at the University of Notre Dame and the University of Texas Rio Grande Valley. Code is available at https://github.com/nsapkota417/UKAST.

What This Paper Is About

Medical image segmentation — outlining anatomical structures and pathological regions in scans — needs high accuracy, but medical datasets are small because expert annotation is expensive and privacy restrictions are common. CNN-based models such as U-Net capture local detail well but have limited receptive fields for long-range dependencies, while Transformers capture global context but are "data-hungry" and computationally expensive. The paper's goal is to build a Transformer-based segmentation model that is both more expressive and more data-efficient by swapping the standard MLP feed-forward blocks in a Swin Transformer encoder for learnable rational-function KANs.

Key Contributions

1. A new architecture, UKAST, which integrates rational-function-based KANs directly into Swin Transformer encoders as feed-forward layers, overcoming the computational, memory, and training inefficiencies of vanilla spline-based KANs. The authors state this is the first pure Transformer-based encoder with KANs for medical image segmentation.
2. State-of-the-art segmentation results reported on both 2D and 3D benchmarks spanning multiple imaging modalities, with the largest reported gains on the 3D datasets.
3. A systematic evaluation of Transformers with KANs across different training-data scales (10%, 25%, 50%, and 100% of training data), rather than only at full data.
4. A demonstration of data efficiency under limited annotations, showing that KAN-enhanced Transformers remain competitive with CNNs in data-scarce settings where standard Vision Transformers degrade.

Main Findings

• Full-data comparison (Table 1, Dice scores with 100% training data): UKAST uses 7.2M parameters and scores 81.7 on Kvasir, 79.9 on ISIC, 71.2 on BCV, and 80.8 on MMWHS, compared with its SwinUNETR + RC baseline at 7.2M parameters with 81.9, 78.9, 68.9, and 80.4 respectively (reported deltas: -0.2, +1.0, +2.3, +0.4). The largest reported gains are on the 3D datasets, BCV and MMWHS.
• UKAST slightly underperforms on the Kvasir dataset in Table 1, at 81.7 versus 82.3 for UNeXt, 82.2 for U-Net v2, and 81.9 for SwinUNETR + RC, while improving on ISIC, BCV, and MMWHS.
• UKAT, a constructed baseline that replaces the ViT encoder in UNETR with a KAT encoder, uses 4.8M parameters and reaches 70.0 on Kvasir, 75.3 on ISIC, 55.0 on BCV, and 78.0 on MMWHS — outperforming the larger UNETR (8.3M parameters, 63.7, 74.7, 52.7, 70.3).
• Data efficacy (Table 2): on the 3D BCV dataset, UKAST scores 63.9, 64.6, 67.2, and 71.2 at 10%, 25%, 50%, and 100% of training data, versus SwinUNETR + RC at 60.1, 62.7, 63.1, and 68.9 (reported deltas of +3.8, +4.9, +4.1, and +2.3). On the 2D ISIC dataset, UKAST scores 71.4, 72.1, 75.5, and 79.9 versus 69.8, 71.8, 74.3, and 78.9 (deltas +1.6, +0.3, +1.2, +1.0).
• UNETR collapses in low-data regimes on BCV, scoring 33.1, 33.7, 34.5, and 52.7 across the four data fractions; UKAT reports much larger deltas over it at the low end (+12.0, +15.9, +16.1, and +2.3).
• Complexity (Table 3): replacing MLPs with GR-KANs lowers FLOPs slightly while adding only a marginal parameter increase. Without residual convolutions, SwinUNETR with MLP is 1.2500 GFLOPs / 6,302,228 params versus UKAST at 1.2467 GFLOPs / 6,302,836 params. With residual convolutions, SwinUNETR is 1.4419 GFLOPs / 7,183,508 params versus UKAST at 1.4386 GFLOPs / 7,184,116 params.
• Ablation (Table 4), average Dice across two 2D and two 3D datasets: ViT + MLP scores 69.2 (2D) and 61.5 (3D); ViT + GR-KAN scores 72.7 and 66.5; SwinT + MLP scores 77.1 and 70.5; SwinT + GR-KAN scores 79.2 and 71.3; SwinT + RC + MLP scores 80.4 and 74.7; and SwinT + RC + GR-KAN scores 80.8 and 76.0.
• Ablation summary as stated by the authors: (1) GR-KANs in ViT encoders improve performance by +3.5% (+5.0% in 3D); (2) switching from ViT to SwinT gives +7.9% (+9.0%); (3) within SwinT, GR-KANs add +2.1% (+0.8%); and (4) with residual convolution added, GR-KANs over MLPs add +0.4% (+1.3%).

Methodology in Plain English

The model keeps the overall shape of a U-Net: an encoder that compresses the image into feature representations, a decoder that upsamples back to pixel-level predictions, and skip connections that carry encoder detail across to the decoder.

The encoder is built from Swin Transformer blocks, which split the image into local windows and compute self-attention within them, then shift the windows in the next layer so information can cross window boundaries. Each block contains attention followed by a feed-forward network. In a standard Swin Transformer that feed-forward network is an MLP with a fixed activation such as GELU or ReLU. This paper replaces it with a Group Rational KAN.

A KAN swaps fixed activations for learnable functions on each connection. The problem is that vanilla KANs use spline base functions, which are computationally inefficient, memory-intensive, and hard to train at depth. The authors instead adopt the rational base functions and GR-KAN design from the Kolmogorov–Arnold Transformer (KAT). Each learnable function is a ratio of two polynomials, P(x) over 1 + |Q(x)|, with an added learnable scaling factor w. Writing the denominator as 1 + |Q(x)| avoids blowing up numerically when Q(x) approaches zero. The authors empirically found polynomial orders m = 3 and n = 4 to be the best-performing configuration.

The GR-KAN step addresses parameter cost: a standard KAN needs a separate activation function for every input–output pair, giving d_in × d_out functions. GR-KAN instead splits the input channels into g groups that share the same rational function parameters, reducing the number of unique functions needed from d_in × d_out to just g. The authors tested group sizes 1, 4, and 8, observed only minor performance differences, and chose 8 for better parallelization and faster runtime.

The architecture also adds residual convolution blocks in the encoder stages (the paper refers to the "SwinUNETR + RC" variant). The image is split into non-overlapping patches, each projected by a linear layer into an embedding space, and partitioned into M × M local windows. The decoder uses deconvolution layers followed by 3 × 3 convolutions with batch normalization and ReLU, and a final 1 × 1 convolution produces the pixel-wise predictions. For 3D volumes, the model operates on 2D slices.

Experimental setup: Four datasets were used — Kvasir-SEG (polyp tissue endoscopy), ISIC 2017 (skin lesion dermoscopy), BCV 2016 (abdominal CT), and MMWHS 2017 (cardiac MRI). Training sets contained 700 images, 2000 images, 24 volumes (1,982 2D slices), and 16 volumes (3,834 2D slices), respectively; test sets contained 300 images, 600 images, 6 volumes (603 2D slices), and 4 volumes (973 2D slices). Baselines were U-Net, U-Net v2, UNeXt, UNETR, SwinUNETR, and UKAN, all in their smallest or comparable configurations for fairness. Models were trained from scratch using MONAI with the AdamW optimizer (initial learning rate 0.0002, weight decay 0.001) and a cosine annealing schedule, using a combined Dice and cross-entropy loss on a single NVIDIA A10 GPU for 400 epochs with batch size 24. Training used 320 × 320 random crops augmented with horizontal and vertical flips, random 90° rotations, and Gaussian noise. At inference, predictions used overlapping patches with 50% overlap.

Why This Matters

The paper argues that the fixed activations in Transformer MLP blocks limit their ability to capture complex functional relationships in medical images, and that learnable functional expansions can both increase expressiveness and improve behavior when annotated data is scarce. Because medical datasets are typically small, any architectural change that preserves accuracy at reduced data has direct practical value. The work also reports reducing FLOPs slightly while keeping parameter counts nearly identical to SwinUNETR, which matters for clinical deployment where compute budgets are constrained.

Real-world applications:

• Delineating anatomical structures and pathological regions in abdominal CT for diagnosis, quantitative analysis, and treatment planning.
• Polyp tissue segmentation in endoscopy for colorectal screening support.
• Skin lesion segmentation in dermoscopy for dermatological assessment.
• Cardiac structure segmentation in MRI for cardiac analysis and planning.

Industry relevance: The model is trained from scratch rather than requiring large-scale pretraining, uses a single A10 GPU and 400 epochs, and is built on the MONAI ecosystem, all of which lower the barrier to adoption in clinical research pipelines. The 7.2M parameter count is comparable to existing SwinUNETR deployments, so it does not demand new hardware. The strong performance at 10% and 25% of training data is relevant to hospitals and vendors that have small, private, or newly collected annotated datasets.

Future Directions

• Improving the 2D results on Kvasir. UKAST scored 81.7 there versus 82.3 (UNeXt) and 82.2 (U-Net v2), so the conditions under which KAN-enhanced Transformers beat CNNs on polyp endoscopy are unresolved by this work.
• Extending KAN integration beyond the feed-forward pathway. The paper notes that prior hybrid designs placed KANs only at bottlenecks or used hybrid MLP–KAN designs; whether KANs can usefully replace or augment other components of the encoder is left open.
• Testing on 3D volumes directly. The paper states that for 3D volumes it operates on their 2D slices; evaluating true volumetric modeling could change the reported 3D gains.
• Understanding why data efficiency improves. The paper demonstrates the effect across 10%, 25%, 50%, and 100% training data but does not offer a theoretical account of why rational-function KANs alleviate the data-hungry behavior of Vision Transformers.

Target Audience

This paper is most useful for graduate students and researchers working on medical image segmentation architectures, particularly those already familiar with U-Net, SwinUNETR, and attention-based dense prediction. It is also relevant to practitioners evaluating whether to replace MLP feed-forward layers in their own Transformer models, and to engineers concerned with parameter and FLOP budgets in clinical or resource-constrained deployment. Readers unfamiliar with self-attention, shifted-window attention, or KANs will find the paper's notation-heavy methodology section demanding without prior background.

Authors’ 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