UnifiedFL: A Dynamic Unified Learning Framework for Equitable Federation

TL;DR

UnifiedFL tackles fully heterogeneous architectures and non-IID data in medical federated learning by representing each local network as a model-graph and updating a single, architecture-agnostic GNN parameter vector $oldsymbol{oldsymbol{ heta}}$. A dynamic, Theta-guided clustering mechanism groups clients by optimization trajectories and a two-tier aggregation schedule balances rapid convergence with reduced cross-cluster interference. Empirical results on MedMNIST variants and hippocampus segmentation show UnifiedFL approaching centralized performance while preserving data privacy, outperforming static clustering and other heterogeneous-FL baselines. This framework advances equitable, scalable collaboration across diverse hospital settings in medical imaging.

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

Figures (5)

• Figure 1: Conceptual comparison between conventional federated learning and the proposed unified learning (UnifiedFL) workflow. (a) Conventional FL. Clients deploy heterogeneous backbones (MLP, CNN, Transformer). Each site uploads its native weight tensor to the server (blue dashed arrows). Because the tensors differ in shape, the server cannot perform element-wise aggregation (red cross). (b) UnifiedFL. Each client converts its backbone to a model-graph and trains a shared set of GNN parameters $\boldsymbol{\Theta}$ that rescale the underlying weights. Only the compact $\boldsymbol{\Theta}$ is exchanged. All parameter vectors have identical length, so the server can average them directly (green tick) and broadcast the result back to the hospitals. This mechanism enables architecture-agnostic collaboration without exposing raw images or full model weights.
• Figure 2: Overview of the proposed UnifiedFL workflow. At every federation round $t$ each client (top & bottom rows show examples for two hospitals with distinct image distributions $P_{1}(x)$ and $P_{2}(x)$) converts its private backbone into a model-graph, optimizes the shared GNN parameters $\boldsymbol\theta^{[t]}$ on local data and sends the updated parameters $\boldsymbol\theta^{[t+1]}_{[c_i]}$ (blue dashed arrows) rather than raw network weights $\mathbf W$. The server (center) clusters clients according to graph topology; frequent intra-cluster aggregations every $t_{\mathrm{ic}}$ (dotted red arrows) are complemented by sparser inter-cluster merges every $t_{\mathrm{bc}}$ (dotted yellow arrows). This topology-aware schedule prevents interference between dissimilar architectures while still enabling global knowledge transfer. Aggregated parameters are broadcast back to all clients, where they rescale/shift local weights for the next round, yielding an architecture-agnostic and communication-efficient federated learning process.
• Figure 3: t-SNE visualisations of the raw feature space used to create non-IID client splits. We run $k$-means with $k{=}10$, and project the features to two dimensions for display. Points sharing color belong to the same $k$-means cluster and will be assigned to the same federated client. The four panels correspond to BreastMNIST, PathMNIST, PneumoniaMNIST, and Hippocampus (clock-wise from top left). Well-separated color clouds indicate strong inter-cluster heterogeneity, whereas overlap signals milder shifts; these visual patterns anticipate the non-IID difficulty faced during federated training.
• Figure 4: Quantitative comparison of seven training protocols on three MedMNIST medmnistv2 benchmarks. Columns show datasets; rows show evaluation metrics. Each bar represents the mean of three folds, with error bars denoting one standard deviation. color code: (light-blue) individual training, non-IID split; (red) individual training, random split; (yellow) vanilla uFedGNN, non-IID; (orange) topology-aware uFedGNN, non-IID; (dark-blue) centralized uGNN, non-IID (upper bound); (olive-green) centralized uGNN, random; (dark-green) proposed UnifiedFL. Ten heterogeneous backbones are plotted per metric: three CNNs, one U-Net, and six MLPs. Across datasets UnifiedFL (dark-green bars) consistently attains scores closest to the upper bound (dark-blue bars) and exceeds all federated baselines for the majority of backbones.
• Figure 5: Quantitative comparison of seven training protocols on the Medical Decathlon -- Hippocampus dataset. Rows show evaluation metrics (Micro-F1, DICE score, and IoU score), columns show heterogeneous backbones (three CNNs, one U-Net, and six MLPs). Each bar represents the mean of three folds, with error bars denoting one standard deviation. Segmentations for a randomly selected test sample below the plots illustrate qualitative differences across protocols, with corresponding IoU scores reported underneath.