Higher-Order Domain Generalization in Magnetic Resonance-Based Assessment of Alzheimer's Disease

TL;DR

This work tackles domain shifts in sMRI-based Alzheimer's disease classification by targeting single-domain generalization (SDG). It introduces Extended MixStyle (EM), which extends the MixStyle augmentation to include higher-order moments—skewness and kurtosis—via two variants, EM1 (skewness) and EM2 (skewness and kurtosis), integrated into a 3D U-Net backbone. Evaluated on NACC for training and ADNI, AIBL, and OASIS as unseen cohorts, EM achieves consistent cross-domain improvements, with an average macro-F1 increase of $2.4\%$ over state-of-the-art SDG baselines, and realigned feature representations as shown by t-SNE and Grad-CAM analyses. The results demonstrate the practical potential of modeling distributional properties beyond mean and variance to deliver robust AD detection across heterogeneous clinical data, paving the way for more reliable real-world deployment.

Abstract

Despite progress in deep learning for Alzheimer's disease (AD) diagnostics, models trained on structural magnetic resonance imaging (sMRI) often do not perform well when applied to new cohorts due to domain shifts from varying scanners, protocols and patient demographics. AD, the primary driver of dementia, manifests through progressive cognitive and neuroanatomical changes like atrophy and ventricular expansion, making robust, generalizable classification essential for real-world use. While convolutional neural networks and transformers have advanced feature extraction via attention and fusion techniques, single-domain generalization (SDG) remains underexplored yet critical, given the fragmented nature of AD datasets. To bridge this gap, we introduce Extended MixStyle (EM), a framework for blending higher-order feature moments (skewness and kurtosis) to mimic diverse distributional variations. Trained on sMRI data from the National Alzheimer's Coordinating Center (NACC; n=4,647) to differentiate persons with normal cognition (NC) from those with mild cognitive impairment (MCI) or AD and tested on three unseen cohorts (total n=3,126), EM yields enhanced cross-domain performance, improving macro-F1 on average by 2.4 percentage points over state-of-the-art SDG benchmarks, underscoring its promise for invariant, reliable AD detection in heterogeneous real-world settings. The source code will be made available upon acceptance at https://github.com/zobia111/Extended-Mixstyle.

Figures (4)

• Figure 1: Higher-order feature statistics across cohorts. Bar plots of skewness and kurtosis computed from intermediate 3D U-Net feature maps across four sMRI cohorts. For each dataset, channel-wise feature distributions were aggregated over test samples and their higher-order moments were summarized. Error bars denote the 95% confidence intervals (CI). Clear variability in skewness and kurtosis was observed across cohorts, indicating that feature distributions differ beyond mean and standard deviation.
• Figure 2: EM integration into a 3D U-Net encoder and its internal operation. EM receives a batch of feature maps as input, and produces a batch of mixed feature maps as output. In detail, the batch $(\boldsymbol{x}_1, \boldsymbol{x}_2, \ldots, \boldsymbol{x}_B)$ are obtained from Layer 2, and are used to compute per-channel feature statistics $m$. For each sample $b$, the mean ($\mu_b$), standard deviation ($\sigma_b$), skewness ($\gamma_b$), and kurtosis ($\kappa_b$) are computed (depending on the specific variant, $EM_1$ or $EM_2$). These statistics are randomly paired (e.g. between $\boldsymbol{x}_b$ and $\tilde{\boldsymbol{x}}_b$), and mixed through a random sample-specific mixing coefficient $\lambda_b$ producing mixed statistics $\boldsymbol{m}^{mix}_b$. Each feature map is then normalized using its original statistics and renormalized using its corresponding mixed statistics, thus generating mixed feature maps that are forwarded to the next encoder layer.
• Figure 3: Grad-CAM visualizations on 3D sMRI samples across cohorts. The figure presents NC, MCI and AD subjects from ADNI (top row), AIBL (middle row), and OASIS (bottom row). For each cohort, columns show: original sMRI scans, MixStyle baseline, EM1 based on mean, standard deviation, and skewness, and EM2 extending EM1 with kurtosis.
• Figure 4: t-SNE visualizations of sMRI embeddings under different training settings. Data were drawn from four cohorts: NACC, ADNI, AIBL, and OASIS. The vanilla 3D U-Net (a) shows clear cohort-specific clustering, with AIBL forming compact islands and OASIS concentrated in the upper region, while ADNI and NACC remain distinct. EFDM (b) increases inter-cohort mixing, creating a dense shared embedding space though NACC still trends toward the outer edge. EM1 (c) and EM2 (d) further enhance overlap, dispersing cohort-specific clusters and producing a more uniform interleaved structure.