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

Abstract

It is essential to be able to combine datasets across imaging centres to represent the breadth of biological variability present in clinical populations. This, however, leads to two challenges: an increase in non-biological variance due to scanner differences, known as the harmonisation problem, and, data privacy concerns due to the inherently personal nature of medical images. Federated learning has been proposed to train deep learning models on distributed data; however, the majority of approaches assume fully labelled data at each participating site, which is unlikely to exist due to the time and skill required to produce manual segmentation labels. Further, they assume all of the sites are available for training. Thus, we introduce UniFed, a unified federated harmonisation framework that enables three key processes to be completed: 1) the training of a federated harmonisation network, 2) the selection of the most appropriate pretrained model for a new unseen site, and 3) the incorporation of a new site into the harmonised federation. We show that when working with partially labelled distributed datasets, UniFedproduces high-quality segmentations and enable all sites to benefit from the knowledge of the federation. The code is available at https://github.com/nkdinsdale/UniFed.