Beyond Transfer Learning: A Generative Self-Supervised Framework for fMRI-Based Diagnosis on Small and Imbalanced Datasets
Discuss this preprint
Start a discussion What are Sciety discussions?Listed in
This article is not in any list yet, why not save it to one of your lists.Abstract
Diagnosis of neurological and psychiatric diseases (NeuroPsyD) from functional Magnetic Resonance Imaging (fMRI) poses significant challenges when using Deep Neural Networks (DNNs), especially in cases where datasets are small and imbalanced, leading to severe overfitting and instability. Transfer Learning (TL), whether supervised or self-supervised, partially alleviates these challenges but remains hampered by domain shift, annotation bias, majority-class bias, and limited gains on very small target datasets. We propose an alternative to TL that addresses these limitations without requiring pre-training. We introduce a purely in-domain solution, Boundary-aware Variational Autoencoder + Self-Supervised Mixup (BVAE + SSup-Mixup), which integrates: (1) SSup-Mixup to learn robust, label-free representations via Mixup NT-Xent loss; (2) minority-class oversampling using a multivariate VAE; and (3) a boundary-aware filter that retains only synthetic examples near the decision boundary, enriching the region of greatest decision uncertainty while discarding outliers. Our framework was evaluated on five fMRI-based NeuroPsyD diagnosis tasks, covering both extremely and moderately small, imbalanced datasets. An ablation study assessed the impact of synthetic fMRI sample generation and selection. We then compared ours with supervised- and self-supervised-TL baselines. Across most tasks, our framework significantly improved accuracy, F1, and AUC by at least 3% compared to the stronger TL variant, while maintaining both sensitivity and specificity above 83%. These balanced gains confirm that boundary-aware augmentation provides precisely the minority-class evidence that TL still lacks. The results demonstrate that targeted generative augmentation, coupled with self-supervised feature learning, allows DNNs to generalize reliably even in data-limited fMRI-based classification tasks. BVAE + SSup-Mixup thus offers a proof-of-concept feasibility study suggesting an alternative to TL for data-limited NeuroPsyD studies and can be extended to other imaging modalities.