Dual-Distilled Heterogeneous Federated Learning with Adaptive Margins for Trainable Global Prototypes

TL;DR

This paper introduces FedProtoKD, a novel federated learning framework that uses adaptive margins and dual knowledge distillation to improve prototype aggregation, addressing margin shrinking and heterogeneity issues in HFL.

Contribution

The paper proposes a dual-knowledge distillation approach with adaptive margins to enhance prototype-based HFL, significantly improving accuracy over existing methods.

Findings

FedProtoKD improves test accuracy by up to 34.13%.

The framework effectively addresses prototype margin shrinking.

It outperforms state-of-the-art HFL methods.

Abstract

Heterogeneous Federated Learning (HFL) has gained significant attention for its capacity to handle both model and data heterogeneity across clients. Prototype-based HFL methods emerge as a promising solution to address statistical and model heterogeneity as well as privacy challenges, paving the way for new advancements in HFL research. This method focuses on sharing class-representative prototypes among heterogeneous clients. However, aggregating these prototypes via standard weighted averaging often yields sub-optimal global knowledge. Specifically, the averaging approach induces a shrinking of the aggregated prototypes' decision margins, thereby degrading model performance in scenarios with model heterogeneity and non-IID data distributions. The propose FedProtoKD in a Heterogeneous Federated Learning setting, utilizing an enhanced dual-knowledge distillation mechanism to enhance system performance by leveraging clients' logits and prototype feature representations. The proposed framework aims to resolve the prototype margin-shrinking problem using a contrastive learning-based trainable server prototype by leveraging a class-wise adaptive prototype margin. Furthermore, the framework assess the importance of public samples using the closeness of the sample's prototype to its class representative prototypes, which enhances learning performance. FedProtoKD improved test accuracy by an average of 1.13% and up to 34.13% across various settings, significantly outperforming existing state-of-the-art HFL methods.