ConvNeXt-KAN: a Kolmogorov-Arnold Network head for robust multi-class staging of Alzheimer's disease from brain MRI

Alzheimer's disease (AD) is a progressive neurodegenerative disorder for which accurate, stage-wise classification from structural magnetic resonance imaging (MRI) can support earlier intervention. Convolutional neural networks (CNNs) dominate this task, but their classifier heads rely on fixed activation functions and linear weights, which constrain the modelling of the subtle, highly non-linear atrophy and texture patterns that separate dementia stages. In addition, the widely used Open Access Series of Imaging Studies (OASIS) collection is severely imbalanced, and two common practices-undersampling every class to the smallest one, and splitting at the image (slice) level-respectively discard most of the data and inflate reported accuracy.

We propose ConvNeXt-KAN, a hybrid model that couples a transfer-learned ConvNeXt-Tiny feature extractor with a Kolmogorov-Arnold Network (KAN) classification head whose edges are learnable spline activations implemented efficiently with Gaussian radial basis functions. The model is trained with a robustness-oriented recipe (MixUp and CutMix, label smoothing, decoupled weight decay, a cosine learning-rate schedule with warm-up, and test-time augmentation) and a data strategy that caps abundant classes and oversamples only the training partition, while evaluating on a fixed, balanced, leakage-controlled test set.

On the held-out test set the model attains 94.5% accuracy, a macro-averaged F1 of 0.944, and a mean one-vs.-rest ROC-AUC of 0.991, with an ablation confirming that pre-trained features and expanded data are the main drivers of performance. A head-to-head comparison under a matched protocol shows that the KAN and MLP heads are statistically indistinguishable (McNemar p = 0.52), while the proposed image-space resampling strategy outperforms focal loss, class-balanced cross-entropy and feature-space SMOTE. Evaluated without any fine-tuning on an independent four-class Alzheimer's MRI collection distributed on Kaggle, test-time-augmented accuracy is 61.6% with a macro-averaged F1 of 0.565.

The learnable-activation head improves the discrimination of the difficult non-demented vs. very-mild-demented boundary at negligible runtime cost, and the efficient formulation keeps the approach practical in standard deep-learning frameworks. A class-wise analysis of the external evaluation shows that moderate dementia continues to be recognized perfectly while the two clinically intermediate stages absorb most of the domain-shift error. We further quantify the impact of the evaluation protocol and argue that subject-level validation and multi-site external testing are essential before clinical claims.