Prototype Normalization: Optimizing Prototype Separation for Heterogeneous Federated Learning

Heterogeneity in data distributions and model architectures poses significant challenges in federated learning (FL). While various heterogeneous FL (HtFL) approaches have been proposed to address these challenges, prototype-based FL (PBFL) has emerged as a promising framework that exchanges prototypes—per-class mean activations from the penultimate layer—rather than model parameters. However, existing PBFL methods suffer from suboptimal prototype separation in the feature space, which limits their discriminative capacity. We propose Prototype Normalization (ProtoNorm), a novel PBFL framework that systematically addresses this limitation. Inspired by the Thomson problem in classical physics, ProtoNorm optimizes global prototype configurations on a unit hypersphere, maximizing separation between class prototypes. Extensive experimental evaluation demonstrates that ProtoNorm consistently outperforms existing HtFL methods across diverse heterogeneous settings. Importantly, ProtoNorm maintains communication efficiency while operating entirely on the server side with less computation than existing PBFL methods, making it well-suited for resource-constrained environments.