Source code for A foundation-model approach to pediatric headache classification from resting-state fMRI, published at the Machine Learning for Healthcare (MLHC) 2026 conference.
We test whether resting-state fMRI (rs-fMRI) can classify pediatric headache using machine learning, comparing two representations of the same scans:
- Functional connectivity (FC) matrices — the standard neuroscience approach, using pairwise correlations between 48 cortical ROIs (Harvard-Oxford atlas).
- NeuroSTORM embeddings — a pre-trained fMRI foundation model applied directly to the raw 4D fMRI volumes, without any extra feature engineering besides common denoising and normalization preprocessing.
Shallow classifiers (decision trees, logistic regression, random forests, gradient boosting) are trained on each representation using a bootstrapped, one-patient-out cross-validation scheme.
189 rs-fMRI scans from 110 participants (ages 8-22) recruited at Boston Children's Hospital. Any headache prevalence of 74%, with chronic migraine being the most common subtype. Other identified headache types are new daily persistent headache, post-viral headache, and post-traumatic headache. Data were collected on a 3T Siemens Prisma scanner and preprocessed with the CONN toolbox (brain extraction, functional realigment, slice-timing correction, denoising, 8mm smoothing, and normalization to MNI152 space).
Raw data cannot be shared publicly (BCH IRB / patient privacy) but may be requested — see data/README.md.
| Model | AUROC | AUPRC |
|---|---|---|
| Demographic (age + sex only) | 0.61 [0.61, 0.61] | 0.82 [0.80, 0.83] |
| Decision Tree (FC) | 0.53 [0.48, 0.56] | 0.77 [0.75, 0.79] |
| Logistic Regression L1 (FC) | 0.67 [0.67, 0.67] | 0.85 [0.85, 0.85] |
| Logistic Regression L2 (FC) | 0.58 [0.58, 0.58] | 0.79 [0.79, 0.79] |
| Gradient Boosting (FC) | 0.40 [0.38, 0.43] | 0.71 [0.69, 0.73] |
| Random Forest (FC) | 0.54 [0.50, 0.57] | 0.78 [0.75, 0.80] |
| NeuroSTORM (embedding) | 0.82 [0.82, 0.82] | 0.93 [0.93, 0.94] |
NeuroSTORM embeddings (fed into an L1-regularized logistic regression) substantially outperformed every FC-based model at distinguishing headache patients from healthy controls, and performance held up when stratified by sex, suggesting the signal is not driven by demographic confounding.
We then used the same NeuroSTORM pipeline for one-vs-all headache subtype classification (healthy vs. chronic migraine vs. non-chronic headache), obtaining a macro-AUROC of 0.69. The model separated healthy controls (AUROC 0.81) and chronic migraine (AUROC 0.67) reasonably well, but struggled to distinguish other, more heterogeneous headache subtypes (AUROC 0.58) — likely reflecting both diagnostic overlap between subtypes.
conda env create -f environment.yml
conda activate pbrainFC-based experiments (functional connectivity features):
cd experiments
bash run.shFoundation-model experiments (NeuroSTORM embeddings) additionally require the NeuroSTORM codebase checked out at third-party/NeuroSTORM_git and its own environment.
First clone neurostorm into the third-party/NeuroSTORM_git path, install it's environment and extend it with our environment.yml, then:
cd experiments
bash run_foundational.shBoth scripts train classifiers with bootstrapped leave-one-patient-out cross-validation and write results to results/. See the comments in run.sh / run_foundational.sh for options (e.g., --use_age, regression targets, choice of preprocessing pipeline).
Preprocessing is handled by the CONN toolbox, an open-source software for common resting-state fMRI preprocessing steps. After CONN preprocessing, run data/prepare_data (redacted).ipynb to compute the functional connectivity matrices (see data/README.md for details).
| Path | Contents |
|---|---|
experiments/ |
training/evaluation scripts for FC and foundation-model classifiers/regressors |
exploratory/ |
notebooks used for data/model exploration |
post processing/ |
notebooks for reports, embedding exploration, feature importance |
data/ |
data directory (raw data not included, see data/README.md) |