Federated Multi-trajectory GNNs Under Data Limitations for Baby Brain Connectivity Forecasting

Abstract

Building accurate predictive models to forecast the trajectory evolution of baby brain networks during the first postnatal year can provide valuable insights into the dynamics of early brain connectivity development. While emerging studies aimed to predict the evolution of brain graphs from a single observation, they suffer from two major limitations: (i) they typically rely on large training datasets to achieve satisfactory performance. However, longitudinal infant brain scans are costly and hard to acquire, and (ii) they adopt a uni-trajectory approach, lacking the ability to generalize to multi-trajectory prediction tasks, where each graph trajectory corresponds to a particular imaging modality (e.g., functional) and at a fixed resolution (graph size). To address these limitations, we propose FedGmTE-Net*, a federated graph-based multi-trajectory evolution network. Given a small dataset, we leverage the power of federation through collaborative model sharing among diverse hospitals. This approach not only enhances the performance of the local generative graph neural network (GNN) model of each hospital but also ensures the preservation of data privacy. To the best of our knowledge, our framework is the first federated learning framework designed for brain multi-trajectory evolution prediction. Further, to make the most of the limited data available at each hospital, we incorporate an auxiliary regularizer that modifies the local objective function, for more effective utilization of all the longitudinal brain connectivity in the evolution trajectory. This significantly improves the network performance. Our comprehensive experimental results demonstrate that our proposed FedGmTE-Net* outperforms benchmark methods by a substantial margin. Our source code is available at https://github.com/basiralab/FedGmTE-Net.