PFedAtt: Attention-based Personalized Federated Learning on Heterogeneous Clients

Zichen Ma (The Chinese University of Hong Kong, Shenzhen)*; Yu Lu (The Chinese University of Hong Kong, Shenzhen); Wenye Li (The Chinese University of Hong Kong, Shenzhen); Jinfeng Yi (JD AI Research); Shuguang Cui (The Chinese University of Hong Kong, Shenzhen )


In federated learning, heterogeneity among the clients' local datesets results in large variations in the number of local updates performed by each client in a communication round. Simply aggregating such local models into a global model will confine the capacity of the system, that is, the single global model will be restricted from delivering good performance on each client's task. This paper provides a general framework to analyze the convergence of personalized federated learning algorithms. It subsumes previously proposed methods and provides the principled understanding of the computational guarantees. Using insights from this analysis, we propose PFedAtt, a personalized federated learning method that incorporates attention-based grouping to facilitate similar clients' collaborations. Theoretically, we provide the convergence guarantee for the algorithm, and empirical experiments corroborate the competitive performance of PFedAtt on heterogeneous clients.