UnifiedFL: A Dynamic Unified Learning Framework for Equitable Federation

TL;DR

UnifiedFL is a novel federated learning framework that enables collaborative training across clients with different neural network architectures and data distributions, improving model generalization and performance in heterogeneous settings.

Contribution

The paper introduces UnifiedFL, which uses a shared graph neural network to unify heterogeneous models and addresses both statistical and domain heterogeneity in federated learning.

Findings

UnifiedFL outperforms existing methods on MedMNIST and hippocampus segmentation tasks.

UnifiedFL effectively handles architectural heterogeneity among clients.

UnifiedFL improves model robustness across diverse data domains.

Abstract

Federated learning (FL) has emerged as a key paradigm for collaborative model training across multiple clients without sharing raw data, enabling privacy-preserving applications in areas such as radiology and pathology. However, works on collaborative training across clients with fundamentally different neural architectures and non-identically distributed datasets remain scarce. Existing FL frameworks face several limitations. Despite claiming to support architectural heterogeneity, most recent FL methods only tolerate variants within a single model family (e.g., shallower, deeper, or wider CNNs), still presuming a shared global architecture and failing to accommodate federations where clients deploy fundamentally different network types (e.g., CNNs, GNNs, MLPs). Moreover, existing approaches often address only statistical heterogeneity while overlooking the domain-fracture problem, where each client's data distribution differs markedly from that faced at testing time, undermining model generalizability. When clients use different architectures, have non-identically distributed data, and encounter distinct test domains, current methods perform poorly. To address these challenges, we propose UnifiedFL, a dynamic federated learning framework that represents heterogeneous local networks as nodes and edges in a directed model graph optimized by a shared graph neural network (GNN). UnifiedFL introduces (i) a common GNN to parameterize all architectures, (ii) distance-driven clustering via Euclidean distances between clients' parameters, and (iii) a two-tier aggregation policy balancing convergence and diversity. Experiments on MedMNIST classification and hippocampus segmentation benchmarks demonstrate UnifiedFL's superior performance. Code and data: https://github.com/basiralab/UnifiedFL