FedPKD+: A Prototype-Based Framework for Efficient and Flexible Heterogeneous Federated Learning

Abstract

Federated Learning (FL) enables collaborative model training across mobile and edge devices without sharing raw data, but its deployment is hindered by <italic>system heterogeneity</italic> and <italic>non-IID</italic> data. Existing FL methods either require homogeneous architectures or suffer from accuracy loss and high communication overhead in heterogeneous settings. To address these issues, we propose FedPKD⁺, a prototype-based knowledge distillation framework that leverages both output-space knowledge (logits) and feature-space knowledge (prototypes) to enhance flexibility and efficiency. Preliminary experiments reveal three key directions for improvement: logit quality, public data utilization, and feature regularization. FedPKD<inline-formula><tex-math notation="LaTeX">$^+$</tex-math><alternatives><mml:math><mml:msup><mml:mrow/><mml:mo>+</mml:mo></mml:msup></mml:math><inline-graphic xlink:href="lyu-ieq1-3710770.gif"/></alternatives></inline-formula> realizes them with four interacting modules: (1) <italic>dual knowledge transfer</italic> shares logits for supervision and prototypes for feature regularization, (2) <italic>heterogeneous prototype alignment</italic> makes cross-model prototypes aggregatable, (3) <italic>prototype-based ensemble distillation and data filtering</italic> use quality-weighted logits and global prototypes to refine server knowledge and filter public samples, and (4) <italic>server knowledge transfer</italic> feeds the refined knowledge back to clients Beyond framework design, we provide a convergence analysis of prototype-based KD in FL, proving both single-round and cross-round convergence under time-varying objectives. Extensive experiments under diverse non-IID settings demonstrate that FedPKD⁺ consistently outperforms state-of-the-art baselines in both learning performance and communication efficiency.