Learn What You Need in Personalized Federated Learning

TL;DR

Learn2pFed introduces an unrolling-based framework for personalized federated learning that learns per-parameter participation degrees via a diagonal regularization matrix $\Lambda_i$ and participation weights $p_i$. By reformulating the problem as a bi-level optimization and unrolling an ADMM-based solver into a deep network, it automatically tunes how much each local parameter collaborates during federation, improving robustness to data heterogeneity. The method includes a comprehensive theoretical convergence analysis and a federated implementation, and demonstrates superior performance across polynomial regression, time-series forecasting, and image classification, while reducing communication and memory costs. This approach advances personalized FL by enabling fine-grained, data-driven collaboration patterns and opens avenues for model compression and selective sharing in distributed learning.

Abstract

Personalized federated learning aims to address data heterogeneity across local clients in federated learning. However, current methods blindly incorporate either full model parameters or predefined partial parameters in personalized federated learning. They fail to customize the collaboration manner according to each local client's data characteristics, causing unpleasant aggregation results. To address this essential issue, we propose $\textit{Learn2pFed}$, a novel algorithm-unrolling-based personalized federated learning framework, enabling each client to adaptively select which part of its local model parameters should participate in collaborative training. The key novelty of the proposed $\textit{Learn2pFed}$ is to optimize each local model parameter's degree of participant in collaboration as learnable parameters via algorithm unrolling methods. This approach brings two benefits: 1) mathmatically determining the participation degree of local model parameters in the federated collaboration, and 2) obtaining more stable and improved solutions. Extensive experiments on various tasks, including regression, forecasting, and image classification, demonstrate that $\textit{Learn2pFed}$ significantly outperforms previous personalized federated learning methods.

Figures (7)

• Figure 1: Three federated ways of local model parameters in (personalized) federated learning: sending (a) full parameters; (b) partial parameters with binary decision; (c) adaptive partial parameters. We aim to determine the part and the degree of local model parameters that participate in federated collaboration.
• Figure 1: Regression performance w.r.t. three personalized FL settings on synthetic data. The proposed $\textit{Learn2pFed}$ achieves the best performance in all the three personalized FL settings.
• Figure 2: Illustration of the $\ell$-th cell in $\textit{Learn2pFed}$. It unrolls (\ref{['eq:layer1']})-(\ref{['eq:layer4']}) into one four-layer cell of the deep network. Black lines indicate the flow of intermediate variables, e.g., $\{w^\ell,\alpha^\ell, z_i^\ell, v_i^\ell\}$ in the $\ell$-th cell. Blue line indicates the local data flow, which however, will not be shared across clients.
• Figure 2: Learnable parameters ablation study w.r.t. three personalized FL settings on synthetic data: we leave the check marks under the learnable parameters and the blanks under the non-learnable parameters. The experiments with more learnable parameters perform better.
• Figure 3: Diagonal values of $\{\Lambda_i\}$. All clients share the same initialization of $\{\Lambda_i\}$ (on the left). The right shows the learned $\{\Lambda_i\}$ in five clients by $\textit{Learn2pFed}$ under Setting 1.

Theorems & Definitions (2)

• Lemma 1
• Theorem 1