← Back to projects

MedViTV2 Notes

KAN-integrated transformers for medical image classification under domain shift.

PyTorch Vision Transformers KAN Domain Shift

Problem

Medical image classifiers often fail under domain shift due to differences in acquisition, staining, and imaging conditions.

What I Built

Reproduced MedViTV2 and trained on BloodMNIST + Czech WBC to study cross-domain behavior and robustness.

Architecture

MedViTV2 architecture diagram
MedViTV2 architecture overview (Local + Global Feature Perception blocks with DiNA and KAN).
  • Local + Global feature perception stages
  • KAN module to enhance non-linear representation
  • Attention design intended to scale while preserving efficiency (DiNA)

Results

Accuracy (BloodMNIST)

98.86%

Accuracy (Czech WBC)

97.68%

F1 (Both)

~97–99%

Confusion matrices show strong diagonal dominance across both domains, with most errors concentrated between a small set of visually similar classes.

Confusion Matrix

BloodMNIST confusion matrix

Key Insights

DiNA in practice

The paper introduces DiNA to scale attention while controlling complexity. In my runs, the KAN-integrated transformer backbone performed strongly on these datasets; DiNA’s benefits may become more pronounced at larger scale and higher diversity.

Why KAN helped

KAN layers use spline-based activations that can model complex non-linearities effectively. On smaller medical datasets, this can contribute meaningful gains even when explicit local attention improvements are limited.