PaperlayerView exact revision on arXiv

Source-linked AI summary

Medical Image Classification with KAN-Integrated Transformers and Dilated Neighborhood Attention

Omid Nejati Manzari, Hojat Asgariandehkordi, Taha Koleilat, Yiming Xiao, Hassan Rivaz

arXiv:2502.13693v3cs.CV

TL;DR

Clinical image corruptions linked to multi-center imaging variation and MedViTV1’s scaling weakness motivate a more robust medical classifier. MedViTV2 combines KAN layers, adapted DiNA, and a hierarchical local/global design; it reports state-of-the-art performance across evaluated benchmarks with approximately 44% lower computational complexity than its predecessor.

  • Problem

    Medical-image corruptions related to imaging equipment and acquisition conditions challenge robust classification, and MedViTV1 accuracy drops as parameter count rises.

    Source: §1 · p. 2

  • Method

    MedViTV2 integrates KAN layers into the Transformer feed-forward pathway and combines adapted DiNA with a hierarchical strategy balancing local and global feature perception.

    Source: §1 · pp. 3–4

  • Results

    MedViTV2 achieves state-of-the-art performance on all evaluated medical benchmarks and reduces computational complexity by approximately 44% compared with its prior version.

    Source: §6 · p. 22

  • Takeaways & Limitations

    The authors conclude that the model is highly suitable for real-time deployment and federated-learning applications.

    Source: §6 · p. 22

Abstract

from arXiv · show

Convolutional networks, transformers, hybrid models, and Mamba-based architectures have demonstrated strong performance across various medical image classification tasks. However, these methods were primarily designed to classify clean images using labeled data. In contrast, real-world clinical data often involve image corruptions that are unique to multi-center studies and stem from variations in imaging equipment across manufacturers. In this paper, we introduce the Medical Vision Transformer (MedViTV2), a novel architecture incorporating Kolmogorov-Arnold Network (KAN) layers into the transformer architecture for the first time, aiming for generalized medical image classification. We have developed an efficient KAN block to reduce computational load while enhancing the accuracy of the original MedViT. Additionally, to counteract the fragility of our MedViT when scaled up, we propose an enhanced Dilated Neighborhood Attention (DiNA), an adaptation of the efficient fused dot-product attention kernel capable of capturing global context and expanding receptive fields to scale the model effectively and addressing feature collapse issues. Moreover, a hierarchical hybrid strategy is introduced to stack our Local Feature Perception and Global Feature Perception blocks in an efficient manner, which balances local and global feature perceptions to boost performance. Extensive experiments on 17 medical image classification datasets and 12 corrupted medical image datasets demonstrate that MedViTV2 achieved state-of-the-art results in 27 out of 29 experiments with reduced computational complexity. MedViTV2 is 44\% more computationally efficient than the previous version and significantly enhances accuracy, achieving improvements of 4.6\% on MedMNIST, 5.8\% on NonMNIST, and 13.4\% on the MedMNIST-C benchmark.

1. Introduction

MedViTV2 targets medical-image corruptions and scaling limits by combining KAN layers, adapted DiNA, and a hierarchical design balancing local and global feature perception. Across the reported benchmarks, its variants achieve state-of-the-art performance on most datasets.

1. Introduction

  • MedViTV2-T/S/L improve average accuracy by 2.6%, 2.5%, and 4.6%, respectively, over MedViTV1-T/S/L across MedMNIST datasets.The comparison is framed as average accuracy versus FLOPs.Source: Figure 1 · p. 2
  • Medical-image corruptions can arise from acquisition protocols, data handling, post-processing, and imaging-equipment differences, motivating robust classification across modalities.MedViTV1 also shows declining accuracy as parameter count increases, limiting scalability.Source: §1 · p. 2
  • The adapted DiNA block addresses feature collapse and strengthens sparse global feature competition, while a hierarchical hybrid strategy balances local and global perception.The hybrid strategy is presented as an efficient way to improve performance.Source: §1 · pp. 3–4
  • The paper introduces KAN layers in the Transformer feed-forward pathway, reporting 44% lower computational complexity alongside enhanced performance.The authors describe this as the first such integration for medical image classification.Source: §1 · p. 3
  • Across 17 medical image datasets and 12 corrupted benchmarks, MedViTV2 variants achieve state-of-the-art performance on most datasets.The reported variants are tiny, small, base, and large.Source: §1 · p. 4

2. Related Works

Prior medical image classification work spans CNNs and Transformers, while KAN research explores varied function bases and applications, including medical imaging.

2.1. Medical Image Classification.

  • CNNs have long been widely used in medical image classification, while Transformers offer long-range dependency modeling and attention maps that can highlight image regions.The literature also discusses ViTs' adaptability to diverse inputs and multimodal applications.Source: §2.1 · p. 4

2.2. Kolmogorov–Arnold Networks

  • KAN research spans multiple domains and function-basis designs, and medical applications include segmentation, Transformer-U-Net hybrids, and uncertainty-aware prediction.Examples include RBF-based approaches and Bayesian-KAN.Source: §2.2 · pp. 4–5

3. Methods

MedViTV2 combines local feature perception with global feature perception in a hierarchical design, incorporating dilated attention and KAN layers. Its stages repeat LFP blocks with a GFP block to integrate local and global information.

3.1. Overview of MedViTV1

  • The methods section describes MedViT overview, LFP and GFP designs, the hierarchical hybrid strategy, and multiple model sizes.The overview frames MedViT as the proposed model before detailing its components.Source: §3 · p. 5

3.2. Dilated Attention Block

  • DiNA confines attention to nearby pixels while retaining shifted-window attention’s network complexity and parameter count, and overlapping windows preserve translation equivariance.LFP combines DiNA with LFFN to capture local and global features.Source: §3.2 · pp. 5–6

3.3. Kolmogorov–Arnold Networks (KANs)

  • KAN layers replace traditional MLP weight matrices with parameterized functions and learnable edge activations; MedViTV2 uses efficient RSWAFs based on Gaussian RBFs approximating third-order B-splines.GFP combines efficient MHSA, MHCA, and KAN outputs, then concatenates the residual and KAN branches.Source: §3.3 · pp. 6–7

3.4. Hierarchical Hybrid Strategy

  • The hierarchical hybrid strategy repeats LFP blocks and adds a GFP block at each stage after the first, targeting global-feature learning in shallow layers and integration across stages.The pattern is (LFP × N + GFP × 1) × L.Source: §3.4 · p. 9
  • Stacked (LFP × N + GFP × 1) groups outperform a single group at similar latency, and N = 2 provides the adopted performance-complexity trade-off.For the small model, L = 3 and N = 3 gave marginal improvement over L = 2 and N = 2 with significantly greater complexity.Source: §3.4 · p. 10

3.5. MedViTV2 Architectures

  • MedViTV2 has four variants—T, S, B, and L—with stage strides and channels specified in its architecture configurations.The described configuration also sets GFP spatial reduction ratios to and fixes channel shrink ratio r at 0.75.Source: Table 3 · p. 8, §3.5 · p. 10

4. Experiments

The experiments evaluate MedViTV2 on publicly available medical image datasets and a corruption benchmark. The implementation uses separate training configurations for MedMNIST, NonMNIST, and MedMNIST-C.

4.1. Datasets

  • 17 publicly available, multi-center medical image datasets were used to evaluate MedViTV2, alongside MedMNIST-C, which covers 12 datasets and 9 imaging modalities.The MedMNIST collection spans multiple imaging types and classification tasks, with dataset sizes ranging from about 100 to over 100,000 images.Source: §4.1 · pp. 10–11
  • The dataset collection also includes Fetal-Planes-DB, CPN X-ray, Kvasir, and PAD-UFES-20, covering ultrasound, X-ray, endoscopy, and skin-lesion images.Fetal-Planes-DB contains over 12,400 images from 1,792 patients; the CPN X-ray collection contains 5,228 images.Source: §4.1 · p. 11
  • MedMNIST-C extends the MedMNIST+ test set with modality-aware and task-specific corruptions, including noise, blur, color changes, and task-specific distortions.Figure 3 shows four artifacts simulated on four datasets, each at five severity levels.Source: §4.1 · p. 12, Figure 3 · p. 12

4.2. Implementation Details

  • MedViT variants were trained separately for each dataset; MedMNIST used 100 epochs and batch size 128, whereas NonMNIST used 150 epochs and batch size 64.MedViTV2-S was employed for the MedMNIST-C datasets.Source: §4.2 · pp. 12–13

BreastMNIST BloodMNIST TissueMNIST OrganAMNIST OrganCMNIST OrganSMNIST

MedViTV2 performs strongly across MedMNIST, NonMNIST, and corrupted-image evaluations, with additional analyses of evaluation metrics and model attention. The reported results include dataset-specific gains and qualitative comparisons of feature focus.

4.3. Evaluation Metrics

  • Evaluation uses ACC and AUC for MedMNIST, six classification metrics for NonMNIST, and balanced accuracy and error measures for MedMNIST-C.MedMNIST-C's rBE measures performance drop relative to the clean test set; the overall rBE averages across corruptions.Source: §4.3 · pp. 14–15

4.4. Performance on MedMNIST

  • MedViTV2 variants improve AUC and accuracy across the MedMNIST2D datasets; MedViTV2-small gains 3.6% AUC and 16.6% accuracy over ResNet-50 on OCTMNIST.MedViTV2-tiny gains 3.0% AUC and 5.2% accuracy over MedMamba-T on PneumoniaMNIST; average accuracy ranges from 86.6% for tiny to 88.8% for large.Source: §4.4 · p. 15

4.5. Performance on NonMNIST

  • MedViTV2-tiny leads nearly all PAD-UFES-20 metrics, with 63.6% accuracy and 87.7% AUC at the lowest reported computational complexity and parameter size.It also reports 63.6% precision, 62.5% sensitivity, 91.7% specificity, and a 61.2% F1-score.Source: §4.5 · p. 16, Table 6 · p. 16
  • On CPN X-ray, MedViTV2-small has the best reported metrics; on Kvasir, it leads in overall accuracy and AUC, including a 0.7% AUC gain over DeiT-small.Its overall-accuracy gains on CPN X-ray are 0.9% over MedMamba-S, 1.4% over VMamba-S, 2.8% over Swin-S, and 2.6% over ConvNext-S.Source: §4.5 · p. 16, Table 7 · p. 17
  • MedViTV2-base achieves the best overall accuracy and AUC on Fetal-Planes-DB while maintaining the lowest FLOPs.Its overall-accuracy gains are 0.9% over MedMamba-B, 1.5% over VMamba-B, 6.1% over Swin-B, and 6.2% over ConvNext-B.Source: §4.5 · pp. 16, 18, Table 8 · p. 17

4.6. Robustness Evaluation

  • MedViTV2-small achieves the best reported robustness on MedMNIST-C with 32.3 million parameters; digital and noise corruptions affect it least, while color corruptions widen the clean-to-robust performance gap.The table reports bACC, rBE, and BE averaged across 12 datasets, with BE also broken down by corruption category.Source: §4.6 · p. 18, Table 9 · p. 18

4.7. Heatmap Visualization

  • Grad-CAM visualizations compare the final GFP, LGP, and normalization layers across tiny and large MedViTV1 and MedViTV2 models.MedViTV2-L shows improved attention quality and focuses on more relevant image regions compared with its smaller version.Source: §4.7 · p. 19

5. Ablation Study

The ablations examine component combinations, attention receptive fields, dilation settings, KAN alternatives, and feature diversity. They report that DiNA and KAN support strong clean and corrupted-dataset performance, while MedViTV1 shows feature collapse.

5.1. Impact of Components

  • On corrupted TissueMNIST, MedViTV2's Dilated and KAN combinations achieve the best clean and robust accuracy; KAN consistently improves clean accuracy across tested combinations.LFFN attains the second-best large-model robust balanced accuracy, but improves clean accuracy less effectively than KAN.Source: §5.1 · p. 19

5.2. DiNa vs. Attention Methods

  • On TissueMNIST, DiNA exceeds Swin and NAT by 4.8% and 1.6%, respectively; on TissueMNIST-C, it exceeds them by 14.8% and 5.4%.DiNA's receptive field can grow from NAT-like linear expansion to k^ℓ with progressively increasing dilation; Table 11 contrasts this with linear growth for NA and Swin.Source: §5.2 · p. 20, Table 11 · p. 21

5.3. Dilation Values

  • At 224 × 224 resolution, (8, 4, 2, 1) is the upper limit of feasible dilation for the classification task; higher dilation values could be used at higher resolutions.The study compares uniform (1, 1, 1, 1) dilation with MedViTV2's varied sequence and adopts a gradual scheme.Source: §5.3 · pp. 20–21, Table 12 · p. 21

5.4. Impact of KAN Block

  • The proposed KAN balances resource use and accuracy across both benchmarks, outperforming FC-KAN by 0.3% on TissueMNIST and 7.8% on TissueMNIST-C while using about 73.33% less memory.The comparison holds LFP and GFP construction settings constant across the tested blocks.Source: §5.4 · p. 21, Table 13 · p. 21

5.5. Feature Collapse Analysis

  • Feature cosine-distance analysis of 5,000 PathMNIST validation images shows that pre-trained MedViTV1 tends toward feature collapse, motivating DiNA's use to promote feature diversity.The reported collapse in MedViTV1 is concentrated in its deepest layers; a higher pairwise cosine distance indicates greater feature diversity.Source: §5.5 · p. 22, Figure 6 · p. 22

6. Conclusion

The conclusion presents MedViTV2 as a hybrid model family combining enhanced Transformer blocks with KAN. It reports improved performance across medical benchmarks and reduced computational complexity.

6. Conclusion

  • MedViTV2 combines enhanced Transformer blocks with KAN and improves performance when scaled across medical benchmarks.The paper reports a balance between clean and robust accuracy on corruption benchmarks and state-of-the-art performance on evaluated medical benchmarks.Source: §6 · p. 22
  • MedViTV2 reduces computational complexity by approximately 44% compared with its prior version.The conclusion identifies real-time deployment and federated learning as potential applications.Source: §6 · p. 22
  • Future work will apply the architecture as a backbone for medical segmentation and 3D medical imaging.Source: §6 · p. 22

Related papers

Loading 2502.13693v3…