Estimating Connectional Brain Templates with Augmented Federated Learning

1533 

Abstract Submission 

Geng Chen¹, Qingyue Wang¹, Abid Ali¹, Islem Rekik²

¹Northwestern Polytechnical University, Xi'an, China, ²Imperial College, London, United Kingdom

Geng Chen  
Northwestern Polytechnical University
Xi'an, China

Qingyue Wang  
Northwestern Polytechnical University
Xi'an, China

Abid Ali  
Northwestern Polytechnical University
Xi'an, China

Islem Rekik  
Imperial College
London, United Kingdom

Federated learning of Connectional Brain Templates (CBTs) estimates a central brain connectivity across datasets collected by different hospitals [1]. Despite its progress, it suffers from the non-Independent and Identically Distributed (non-IID) issue of multiple domains, where the data distribution varies across hospitals. This issue degrades the centeredness of the learned CBT. To this end, we propose a novel augmented federated learning framework for robust CBT learning across diverse domains. The key innovation lies in generalizing the model by augmenting the data diversity of each hospital, which is achieved through a custom brain connectivity generator. Extensive experiments demonstrate the superiority of our method in learning cross-hospital CBTs from multi-view morphological brain networks.

An overview of the proposed method, called AugFedCBT, is shown in Figure 1. We build upon FedCBT [1], which combines FedAvg [2] with the DGN [3] for federated CBT learning. Our improvement involves the integration of the Multigraph Generator Network (MGN) [4] for local domain augmentation. Specifically, each hospital is equipped with a pretrained MGN-based generator, which generates brain connectivities from subject-specific CBTs obtained with DGN. During the federated learning, local model DGN_k, trained with brain connectivities S_k at hospital k, updates its network weights w_k for the t-th round via:
$$\mathbf{w}_k^{t+1}=\mathrm{DGN}_k(\mathbf{w}_k^t,\mathbf{S}_k).$$
After updating weights across all hospitals, the global DGN on the server is updated by computing the weighted average weights according to the proportion of local samples to the total samples as follows:
$$\mathbf{w}^{t+1}=\sum_{k=1}^K\frac{N_k}N\mathbf{w}_k^{t+1}.$$
Next, a post-federated-learning procedure generalizes the domain by augmenting new brain connectivities from a set of subject-specific CBT C_k perturbed by the same random Gaussian noise matrix n_k via MGN_k:
$$\mathbf{A}_k=\mathrm{MGN}_k(\mathbf{n}_k,\mathbf{C}_k).$$
The network weight learning is then improved by including the augmented domain data A_k as follows:
$$\mathbf{w}_k^{t+1}=\mathrm{DGN}_k(\mathbf{w}_k^t,\mathbf{\hat{S}}_k);\mathbf{\hat{S}}_k=\mathbf{S}_k\bigcup\mathbf{A}_k.$$
The DGN model is trained with an enhanced DGN Loss [3], which considers data S_k and A_k from both source and augmented domains.

Dataset: We evaluate our method using right hemisphere morphological brain networks of 186 normal subjects from the ABIDE-I dataset [5]. Using FreeSurfer [6], we generate cortical surfaces from T1-weighted MR images and further construct six morphological networks based on the Desikan-Killiany Atlas. These networks encode shape relationships between cortical regions in terms of different attributes.
Experimental Results: We compare our AugFedCBT with state-of-the-art (SOTA) models [1,3] and its non-federation-learning ablated version (AugDGN). To obtain stable results, we conduct 10 random data simulations and report the average results. The CBT centeredness results, shown in Figure 2(a), indicate that AugFedCBT significantly (p<0.0001) outperforms the SOTA models and achieves the minimum Frobenius distance in all hospitals. Additionally, it surpasses its ablated version, demonstrating the effectiveness of combining connectivity augmentation with federated CBT learning to overcome the non-IID issue. Figure 2(b) further confirms the superior CBT-centeredness performance of our AugFedCBT, aligning with the observations in Figure 2(a).

We proposed AugFedCBT, an effective CBT learning model to tackle non-IID brain graph data through connectivity augmentation. The effectiveness of AugFedCBT is demonstrated qualitatively and quantitatively by extensive experiments.

Connectivity (eg. functional, effective, structural) ¹

Methods Development ²

Atlasing

Cortical Layers

Design and Analysis

Machine Learning