Prototypical Representation Learning for Multi-Site Domain Generalization in Schizophrenia Diagnosis

Area of Science:

• Neuroscience
• Machine Learning
• Medical Imaging

Background:

• Brain functional networks (BFNs) from fMRI are crucial for diagnosing psychiatric disorders like schizophrenia (SZ).
• Site-induced distribution shifts in multi-site fMRI data hinder the generalization of classification models.
• Existing domain generalization (DG) methods often fail due to the assumption of consistent class structures across sites, overlooking intra-class diversity.

Purpose of the Study:

• To develop a robust domain generalization (DG) framework for classifying schizophrenia using multi-site BFNs.
• To overcome the challenge of site-induced distribution shifts in fMRI data.
• To improve the generalization performance of diagnostic models to unseen domains.

Main Methods:

• A transformer encoder was used to extract discriminative subject-level representations from BFNs.
• A site-independence module with Hilbert-Schmidt Independence Criterion (HSIC) regularization enforced site-invariant feature learning.
• Prototype learning with Sinkhorn matching, exponential moving average (EMA) updates, and maximum likelihood estimation (MLE) loss refined feature-to-prototype matching.

Main Results:

• The proposed DG framework achieved superior classification performance on two independent SZ datasets (88.89%±2.22% and 86.05%±1.64%), outperforming 7 DG, 6 domain adaptation (DA), 6 multi-site, and 6 state-of-the-art methods.
• Ablation studies confirmed the significant contributions of MLE, contrastive, and alignment losses to performance enhancement.
• The method identified discriminative temporal regions, offering insights into the neural mechanisms underlying schizophrenia.

Conclusions:

• The developed transformer-based DG framework effectively addresses site-induced distribution shifts in multi-site fMRI data for schizophrenia diagnosis.
• The integration of prototype learning and site-invariance constraints leads to enhanced generalization and robust classification performance.
• The findings provide valuable insights into the neurobiological underpinnings of schizophrenia and highlight the potential of advanced machine learning techniques in psychiatric research.