UniFed: A unified segmentation framework for partially labelled, distributed neuroimaging data

1995 

Abstract Submission 

Nicola Dinsdale¹, Mark Jenkinson^1,2,3, Ana Namburete¹

¹University of Oxford, Oxford, United Kingdom, ²Australian Institute for Machine Learning (AIML), Adelaide, Australia, ³South Australian Health and Medical Research Institute (SAHMRI), Adelaide, Australia

Nicola Dinsdale  
University of Oxford
Oxford, United Kingdom

Mark Jenkinson  
University of Oxford|Australian Institute for Machine Learning (AIML)|South Australian Health and Medical Research Institute (SAHMRI)
Oxford, United Kingdom|Adelaide, Australia|Adelaide, Australia

Ana Namburete, Professor  
University of Oxford
Oxford, United Kingdom

Deep learning models are powerful tools for neuroimage analysis but their current clinical impact is limited. The required pooling of data to capture clinical populations leads to two major challenges: first, the harmonisation problem created by scanner differences; second, data privacy as medical images are inherently personal. Federated learning (McMahan2016) (FL) has been used with distributed data to protect privacy, but most approaches make infeasible assumptions that data from all sites are both fully labelled and available at training time. Thus, we propose UniFed, a unified FL framework (Fig. 1) formed of three processes: a federated network for partially labelled datasets; model selection for an unseen, unlabelled site, and model adaptation to the unseen, unlabelled site.

We used the ABIDE dataset (Martino2013), split into 16 distributed sites, preprocessed using the FSL ANAT pipeline. We split the sites into four sets: Reference site (NYU), Labelled and Unlabelled sites in the federation, and Unseen sites. We considered segmentation of four subcortical structures (brain stem, thalamus, putamen and hippocampus), using FSL FIRST (Patenaude2011) generated labels. A UNet architecture and Dice loss were used and features from the second to last layers were considered. Dice Score was used to assess performance.

The framework is unified by two key concepts: encoding each learned feature as a 1D Gaussian distribution, enabling feature distribution information to be communicated without violating data privacy, and the use of the Bhattacharyya distance (DB) to calculate the distance between sites in feature space.

Partially labelled FL: For every site in the federation, local training minimises a combination of two loss functions: the main task loss, if labels are available, and the DB between the local and global feature distributions for all data. The local models are aggregated equally (Dinsdale2022) to create the global model; the global feature distribution is the average across supervised sites.

Model selection: Given a model zoo of pretrained models, we hypothesise that the best choice for a new unlabelled dataset has the shortest DB between the source global features (shared already for FL) and the features for the target data for that model. We used the models trained in Partially Labelled FL for the model zoo.

Model adaptation: We follow our previously published SFHarmony (Dinsdale2023), a source-free domain adaptation (SFDA) approach, which adapts the network to the unseen site by minimising the DB between the source and target feature. We considered two source models – the NYU-only model and the UniFed model trained on 6 labelled sites.

Partially labelled FL: We considered increasing numbers of supervised sites (1-6) and percentages of labels available (1,5,10% across 5 sites with NYU fully labelled), and compared to standard FL methods (McMahan2016,Li2018), semisupervised FL methods (Bdair2021) and domain adaptation FL methods (Peng2020,Dinsdale2022). Our approach matched or outperformed existing methods across all degrees of supervision. Fig 2a shows the results for 2 fully supervised sites: improved performance on the Unseen sites is clear, despite only sharing the µ and σ of the features.

Model selection: Each site showed a negative correlation between Dice score and the DB. The slope was significant for each unseen site (p<0.05). Our federated models were selected for 4/5 unseen sites, showing that UniFed models generalise better.

Model Adaptation: For both source models our approach outperformed existing methods (Bateson2022) for SFDA. Model selection is clearly important: although the performance of the NYU-only model increased from 52.02 to 68.65, the performance was much worse than with the UniFed model (83.14 to 85.43).

UniFed enables the training of high-performing models for distributed and partially labelled datasets. The approach is generalisable across segmentation tasks and imaging studies.

Methods Development ²

Segmentation and Parcellation ¹

Subcortical Structures

Databasing and Data Sharing

Machine Learning

STRUCTURAL MRI

Other - Data Privacy