FedSAM: Hierarchical Structure Adaptation For Model-Heterogeneous Federated Learning

In real-world Federated Learning (FL) settings, clients often differ in both computational resources and local data characteristics. Model-heterogeneous FL addresses this issue by assigning clients sub-models of different-sized network structures partitioned from a single large neural network. However, existing methods typically rely on heuristic partitioning criteria and do not allow clients to customize their sub-models to local data's unique statistical properties. In this paper, we propose FedSAM, a novel hierarchical Bayesian framework that enables probabilistic inference on client-specific network structures while sharing collaborative knowledge across the local models in the hierarchy. Specifically, in the framework we adapt a local model's neural network structure by modeling network depth and width as beta and conjugate Bernoulli processes. We further specify a global-level beta process as a hyper-prior to aggregate and share local models' structural information at central server. We develop a structured variational inference scheme for efficient computation. Extensive empirical studies demonstrate that by adapting local models' structures to non-IID data and sharing the structural information, our framework yields state-of-the-art performance.