Federated GNNs for EEG-Based Stroke Assessment

TL;DR

The paper tackles privacy-preserving prediction of stroke severity from EEG by combining Federated Learning (FL) with Graph Neural Networks (GNNs) on multilayer brain graphs. It introduces a masked self-attention GNN architecture with EdgeSHAPer explanations and uses MQTT-based FL with FedAvg and SCAFFOLD to handle data heterogeneity across institutions. In experiments with 72 patients across four centers, FL models achieved MAE around 3.22–3.23, comparable to centralized training (MAE 3.22) and approaching human expert performance (~3.0), while providing explanations for edge-level brain connectivity. The work demonstrates that private, explainable, multi-site EEG-based stroke assessment is feasible and highlights directions for scaling and broader validation in clinical settings.

Abstract

Machine learning (ML) has the potential to become an essential tool in supporting clinical decision-making processes, offering enhanced diagnostic capabilities and personalized treatment plans. However, outsourcing medical records to train ML models using patient data raises legal, privacy, and security concerns. Federated learning has emerged as a promising paradigm for collaborative ML, meeting healthcare institutions' requirements for robust models without sharing sensitive data and compromising patient privacy. This study proposes a novel method that combines federated learning (FL) and Graph Neural Networks (GNNs) to predict stroke severity using electroencephalography (EEG) signals across multiple medical institutions. Our approach enables multiple hospitals to jointly train a shared GNN model on their local EEG data without exchanging patient information. Specifically, we address a regression problem by predicting the National Institutes of Health Stroke Scale (NIHSS), a key indicator of stroke severity. The proposed model leverages a masked self-attention mechanism to capture salient brain connectivity patterns and employs EdgeSHAP to provide post-hoc explanations of the neurological states after a stroke. We evaluated our method on EEG recordings from four institutions, achieving a mean absolute error (MAE) of 3.23 in predicting NIHSS, close to the average error made by human experts (MAE $\approx$ 3.0). This demonstrates the method's effectiveness in providing accurate and explainable predictions while maintaining data privacy.

Figures (5)

• Figure 1: Illustration of a FL setup with hospitals acting as nodes. Each hospital processes data locally while sharing model updates (represented by arrows) with a central server.
• Figure 2: Number of patients for each hospital in our federation.
• Figure 3: Illustration of the multi-layer network structure. Each layer corresponds to a specific frequency band, with inter-layer connections helping cross-frequency information flow protani2024explainable.
• Figure 4: Convergence of MAE over training rounds for different setups. Centralized learning shows the fastest and most stable convergence. The federated learning setups converge more slowly but arrive at the same level as the centralized. The isolated learning approaches exhibit higher error and slower convergence, confirming the benefits of collaborative training in federated setups.
• Figure 5: Illustration of Edge Shapley Values for various federated learning setups for the same patient. The color intensity is proportional to the contributions to model predictions. Node sizes are proportional to their weighted degree centrality, adjusted for the number of connections, highlighting each node's significance within the network.