Spectral Graph Sample Weighting for Interpretable Sub-cohort Analysis in Predictive Models for Neuroimaging

Abstract

Recent advancements in medicine have confirmed that brain disorders often comprise multiple subtypes of mechanisms, developmental trajectories, or severity levels. Such heterogeneity is often associated with demographic aspects (e.g., sex) or disease-related contributors (e.g., genetics). Thus, the predictive power of machine learning models used for symptom prediction varies across subjects based on such factors. To model this heterogeneity, one can assign each training sample a factor-dependent weight, which modulates the subject's contribution to the overall objective loss function. To this end, we propose to model the subject weights as a linear combination of the eigenbases of a spectral population graph that captures the similarity of factors across subjects. In doing so, the learned weights smoothly vary across the graph, highlighting sub-cohorts with high and low predictability. Our proposed sample weighting scheme is evaluated on two tasks. First, we predict initiation of heavy alcohol drinking in young adulthood from imaging and neuropsychological measures from the National Consortium on Alcohol and NeuroDevelopment in Adolescence (NCANDA). Next, we detect Dementia vs. Mild Cognitive Impairment (MCI) using imaging and demographic measurements in subjects from the Alzheimer's Disease Neuroimaging Initiative (ADNI). Compared to existing sample weighting schemes, our sample weights improve interpretability and highlight sub-cohorts with distinct characteristics and varying model accuracy.

Fig. 1.

We first create a population graph using pre-defined sociodemographic, genetic, or environmental factors that share associations with specific brain diseases. Next, a machine learning model is trained to predict symptom outcome given imaging and non-imaging data for each subject. During training, the classification loss is weighted by $w$, which is a linear combination of the graph eigenbases $\mathbf{E}$ with a learnable vector $\mathbf{a}$. The learned weights highlight sub-cohorts that share common characteristics and achieve higher predictive power.

Fig. 2.

Comparison of learned weights across cohorts by graph factors for NCANDA and ADNI. The statistical difference of the weights across cohorts is measured with the Mann-Whitney U-test. **:p<0.001, *:p<0.05, ns:p>0.05. The BACC for each sub-cohort is shown under or above the box.

Fig. 3.

Impact of choice of neighbors and centering hyperparameter for NCANDA and ADNI. We measure the absolute difference in % of BACC between the cohorts with high vs. low weights to compare the ability of different models to create distinct and highly separable sub-cohorts based on the learned weights.