Heterogeneous Federated Learning with Prototype Alignment and Upscaling

TL;DR

This paper introduces ProtoNorm, a novel prototype-based federated learning framework that improves class prototype separation through alignment and upscaling, leading to better performance especially in resource-limited settings.

Contribution

ProtoNorm enhances prototype separation in heterogeneous federated learning by combining prototype alignment inspired by physics and prototype upscaling, improving discriminative power.

Findings

Outperforms existing HtFL methods on benchmark datasets.

Improves prototype separation and classification accuracy.

Maintains communication efficiency suitable for resource-constrained environments.

Abstract

Heterogeneity in data distributions and model architectures remains a significant challenge in federated learning (FL). Various heterogeneous FL (HtFL) approaches have recently been proposed to address this challenge. Among them, prototype-based FL (PBFL) has emerged as a practical framework that only shares per-class mean activations from the penultimate layer. However, PBFL approaches often suffer from suboptimal prototype separation, limiting their discriminative power. We propose Prototype Normalization (ProtoNorm), a novel PBFL framework that addresses this limitation through two key components: Prototype Alignment (PA) and Prototype Upscaling (PU). The PA method draws inspiration from the Thomson problem in classical physics, optimizing global prototype configurations on a unit sphere to maximize angular separation; subsequently, the PU method increases prototype magnitudes to enhance separation in Euclidean space. Extensive evaluations on benchmark datasets show that our approach better separates prototypes and thus consistently outperforms existing HtFL approaches. Notably, since ProtoNorm inherits the communication efficiency of PBFL and the PA is performed server-side, it is particularly suitable for resource-constrained environments.