Beyond the Federation: Topology-aware Federated Learning for Generalization to Unseen Clients

TL;DR

The paper addresses the challenge of out-of-federation (OOF) generalization in federated learning by introducing Topology-aware Federated Learning (TFL), which jointly learns a privacy-preserving client topology and uses it to regularize robust optimization. TFL alternates between Client Topology Learning (CTL) and Learning on Client Topology (LCT), with CTL deriving a sparse graph from model-weight similarities and LCT focusing on influential clients via a topology-informed uncertainty set and a Lagrangian formulation. Across diverse real-world and benchmark datasets, TFL delivers superior OOF robustness and scalability, outperforming state-of-the-art baselines in OOF and often improving in-federation performance as well. The approach also demonstrates practical benefits in terms of communication efficiency and shows potential for extensions to adversarial settings and privacy-preserving topology inference.

Abstract

Federated Learning is widely employed to tackle distributed sensitive data. Existing methods primarily focus on addressing in-federation data heterogeneity. However, we observed that they suffer from significant performance degradation when applied to unseen clients for out-of-federation (OOF) generalization. The recent attempts to address generalization to unseen clients generally struggle to scale up to large-scale distributed settings due to high communication or computation costs. Moreover, methods that scale well often demonstrate poor generalization capability. To achieve OOF-resiliency in a scalable manner, we propose Topology-aware Federated Learning (TFL) that leverages client topology - a graph representing client relationships - to effectively train robust models against OOF data. We formulate a novel optimization problem for TFL, consisting of two key modules: Client Topology Learning, which infers the client relationships in a privacy-preserving manner, and Learning on Client Topology, which leverages the learned topology to identify influential clients and harness this information into the FL optimization process to efficiently build robust models. Empirical evaluation on a variety of real-world datasets verifies TFL's superior OOF robustness and scalability.

Figures (9)

• Figure 1: Our empirical evaluation of patient mortality prediction using Federated Learning on distributed Healthcare dataset (eICU pollard2018eicu). We observed that a model exhibiting high accuracy with in-federation (IF) data can fail catastrophically when presented with out-of-federation (OOF) data.
• Figure 2: Accuracy vs. wall-clock time on PACS dataset li2017pacs. Wall-clock time is used as a holistic evaluation of scalability (communication and computation costs). We see a tradeoff between OOF robustness and scalability in existing methods.
• Figure 3: Overview of Topology-aware Federated Learning (TFL). TFL contains two alternating steps: client topology learning (CTL) and learning on client topology (LCT). CTL learn the client topology that describes the relationships between local clients. We use model weights as node embedding and construct a $\epsilon$-graph by measuring the similarity of node pairs. LCT leverage the learned client topology to achieve better OOF robustness. We identify the influential client and then use the influential client as prior knowledge to regularize a distributionally robust optimization framework. In this way, the optimization process can balance the "worst-case" client and the "influential" client to avoid overly pessimistic models with compromised OOF-resiliency.
• Figure 4: Model similarity under different metrics. Client models are trained using the same algorithm and hyperparameters. Clients 6 to 10 share similar data distributions. We observe that clients with similar data distribution tend to have more similar models.
• Figure 5: Qualitative results on unseen patients from FeTS dataset. We also show the DSC ($\uparrow$) score. Our method consistently demonstrates superior OOF robustness under diverse local demographics. Additional visualizations can be found in Supplementary A.