Predicting Infant Brain Connectivity with Federated Multi-Trajectory GNNs using Scarce Data

TL;DR

The paper tackles predicting infant brain connectome trajectories across multiple imaging modalities when data are scarce. It introduces FedGmTE-Net, a federated GNN framework that forecasts several modality trajectories from a single baseline graph, augmented by a topology-preserving loss and an auxiliary regularizer. FedGmTE-Net+ adds KNN-based precompletion imputations and the regressor-based regularizer, while FedGmTE-Net++ adds an imputation refinement step using modality- and hospital-specific similarity regressors and subsequent fine-tuning. Across real IID and non-IID and simulated datasets, FedGmTE-Net++ yields the strongest performance on key metrics such as the graph MAE, PCC, and JD, demonstrating robust, privacy-preserving, data-efficient multi-trajectory brain graph prediction with potential generalization to other isomorphic graph domains.

Abstract

The understanding of the convoluted evolution of infant brain networks during the first postnatal year is pivotal for identifying the dynamics of early brain connectivity development. Existing deep learning solutions suffer from three major limitations. First, they cannot generalize to multi-trajectory prediction tasks, where each graph trajectory corresponds to a particular imaging modality or connectivity type (e.g., T1-w MRI). Second, existing models require extensive training datasets to achieve satisfactory performance which are often challenging to obtain. Third, they do not efficiently utilize incomplete time series data. To address these limitations, we introduce FedGmTE-Net++, a federated graph-based multi-trajectory evolution network. Using the power of federation, we aggregate local learnings among diverse hospitals with limited datasets. As a result, we enhance the performance of each hospital's local generative model, while preserving data privacy. The three key innovations of FedGmTE-Net++ are: (i) presenting the first federated learning framework specifically designed for brain multi-trajectory evolution prediction in a data-scarce environment, (ii) incorporating an auxiliary regularizer in the local objective function to exploit all the longitudinal brain connectivity within the evolution trajectory and maximize data utilization, (iii) introducing a two-step imputation process, comprising a preliminary KNN-based precompletion followed by an imputation refinement step that employs regressors to improve similarity scores and refine imputations. Our comprehensive experimental results showed the outperformance of FedGmTE-Net++ in brain multi-trajectory prediction from a single baseline graph in comparison with benchmark methods.

Figures (9)

• Figure 1: Brain graph multi-trajectory evolution including different imaging modalities with a single baseline brain graph as input. We span the full multimodal space of 4D trajectories from a single brain connectome.
• Figure 2: Hospital-specific population graph where nodes symbolize distinct subject brain graphs within the population, and edges denote their connections. The use of various colours for the connections showcases the non-binary nature of the relationship between two subjects.
• Figure 3: Pipeline of proposed FedGmTE-Net+ for infant brain graph multi-trajectory evolution prediction from baseline.(A) KNN-based imputation Imputation technique to complete the missing graphs from each hospital's local training set by utilizing the similarities of subjects at the baseline timepoint. (B) Multi-trajectory evolution network Each hospital's network uses a single input modality to generate multiple trajectories spanning different imaging modalities. (C) Federated learning paradigm A decentralized learning paradigm that allows hospitals to collaborate with each other without sacrificing data privacy. (D) Auxiliary regularizer The auxiliary regularizer improves network performance by utilizing the entire local training dataset across all timepoints.
• Figure 4: (E) Imputation refinement step for the proposed FedGmTE-Net++ network. Following the initial training of the network, we proceed to train the similarity regressors. Their purpose is to update the similarity scores for subject pairs across various timepoints. The updated scores are used to improve the nearest neighbours' selection for each subject, consequently refining the initial KNN imputation. Subsequent to this step, a series of fine-tuning rounds are executed, capitalizing on the improved dataset.
• Figure 5: PCA comparison between real and simulated datasets. Top: Comparison using the morphological connectome. Bottom: Comparison using the functional connectome.