FedMAP: Unlocking Potential in Personalized Federated Learning through Bi-Level MAP Optimization

We thank the reviewer for their careful consideration of our paper and detailed comments. We address each of their concerns below.

1. On reflection, we agree that Section 2, in particular, could be more accessible. In the revised version, we will update the section to provide a high-level overview and build intuition before introducing the technical details.

2. We agree that theoretical analysis should be included in the main paper. In the revised manuscript, we will integrate the Theorem 1 and its implications directly into the main text.

3.1. We agree that evaluating on more standard benchmarks is valuable. We chose Office-31 as it is commonly used in approaches that address feature distributions and naturally introduces real-world data heterogeneity through its three distinct domains, which includes realistic distribution shifts in feature distributions, quantities, and label distributions. Moreover, we used synthetic datasets that allow controlled analysis of specific non-IID scenarios to confirm FedMAP's effectiveness. We are currently performing the experiments on larger datasets such as FEMNIST and GLUE.

3.2. Our initial experiments focused on methods addressing similar challenges and were integrated within the Flower framework to ensure fair comparison and reproducibility. We agree that comparing against additional personalized FL approaches would strengthen our evaluation, and we are extending our comparisons to include meta-learning based approach FedL2P[1] in the revised version.

Reference:

[1] Lee, R., Kim, M., Li, D., Qiu, X., Hospedales, T., Huszár, F., & Lane, N. D. (2023). FedL2P: Federated Learning to Personalize. Proceedings of the Thirty-seventh Conference on Neural Information Processing Systems.

We thank the reviewer for their time and detailed review. They raised a critical oversight in our omission of pFedMe, we were unaware of this paper and acknowledge that at an initial look, it seems to be very similar to our work and should have been referenced. Below, we highlight in detail how pFedMe is a special case of FedMAP, rather than being equivalent, and that FedMAP is a significantly more general framework which we would exploit in a revised version.

Q1: The oversight in citing pFedMe was a good-faith error, we completely agree that this should have been cited as it is an important contribution, also being a special case of FedMAP, it is entirely relevant for citation and will be corrected in the revised manuscript. While we acknowledge the mathematical similarity between pFedMe and FedMAP, we stress that this similarity only arises for a specific case of FedMAP. Indeed, the regularizer $\mathcal{R} (\cdot, \gamma)$ in FedMAP is a parametrized function with parameter $\gamma$. When $\mathcal{R} (\cdot, \gamma)$ is taken as the negative log-likelihood of a Gaussian distribution, where the parameter $\gamma$ is the mean, the regularizer is of the form $\mathcal{R} (\theta, \gamma) = \alpha \|\theta - \gamma\|^2$, for some $\alpha>0$ fixed. In this case, FedMAP coincides with pFedMe. We admit that in the present version of the paper, we mainly considered this choice, and therefore it is not clear what are the differences between pFedMe and FedMAP. However, one can consider other parametrizations for the regularizer, which offer significant advantages over the quadratic case with a fixed hyperparameter $\alpha$.

For instance, one of the key considerations when using a quadratic regularizer is the choice of $\alpha$, which is related to the variance of the parameters of the model across the clients. When the model, $\Phi( \cdot; \theta)$ has $p$ parameters, i.e. $\theta\in \mathbb{R}^p$, one can consider the FedMAP approach with a regularizer of the form

$

\mathcal{R} (\theta, (\gamma_1, \gamma_2)) = \sum_{i=1}^p \alpha (\gamma_2^{(i)} ) \left( \theta_i - \gamma_1^{(i)} \right)^2. \tag{1}

$

where $\alpha: \mathbb{R} \to \mathbb{R}^+$ is a function that has to be suitably chosen in a way that the function $(\gamma_1, \gamma_2) \mapsto \mathcal{R} (\theta, (\gamma_1, \gamma_2))$ is convex for all $\theta$. With this choice, the parameter $\alpha$, associated with the variance of each model parameter, is updated during the aggregation step according to the local models obtained from the different clients. This is important since some parameters might be similar across the local models whereas others might be needed to be different. This represents a significant improvement over the pFedMe approach.

In a revised version of the paper, we will now consider the more general regularizer, as shown in Eq.~(1), instead of the quadratic case for the parameterized regularizer, and describe precisely how to choose the function $\alpha(\cdot)$.

Q2: The current experiments show the effectiveness of FedMAP in handling non-IID challenges in two ways. We used synthetic datasets that were specifically constructed to show the controlled combinations of non-IID issues (feature, label, and quantity skew). We also used Office-31 datasets which naturally introduced the feature distribution skewness across three different domains despite their size being small. However, as recommended by the reviewer, we will evaluate FedMAP with large datasets FEMNIST and GLUE to strengthen our findings.

Q3:In this paper, we focused on the foundational method such as FedAvg to establish baseline performance. FedProx was chosen because it particularly addresses client drift with regularization. We included FedBN because it specifically addresses feature shift challenges, which are key issues FedMAP aims to address. However, on reflection, we agree with the reviewer that comparing with pFedMe, Ditto, and FedDyn would provide a more comprehensive evaluation. We have already started these experiments and in the revised manuscript, we will add these additional baselines to our experiments and report the results.

Q4: In the current version, Algorithm 3 represents the aggregation step for the specific case when $\mathcal{R} (\theta, \gamma) = \| \theta - \gamma\|^2$. In a more general setting, this should be replaced by a gradient step as in equation (11) in the paper. Convergence of the FedMAP iterations is guaranteed since, as shown in Theorem 1, these iterations correspond to gradient descent applied to a convex function. We will make sure this is clear to the reader in the revised version of the manuscript.

Q5: This has been addressed in the response to Q4 above.

Q6: While Office-31 is common for demonstrating feature shifts, we agree to include larger datasets (FEMNIST and GLUE) in our revised experiments.

We would like to thank the reviewer for their thoughtful comments. Below are our detailed responses to each point raised.

Q1: We would like to clarify the choice of the prior distribution. In Section 3, we present an implementation of FedMAP using a normal distribution as the prior, in which the parameter is the mean. This choice is made for practical demonstration because of its simplicity. However, FedMAP is designed as a more general framework. In Section 2, we present our method for any parameterized family of probability measures over the parameter space. The theoretical foundations in Section 2 and Theorem 1 are established for this general setting, with the requirements being that the convexity condition in Assumption 1 holds. The Gaussian implementation was chosen only for its simplicity and interpretability, but the framework can accommodate non-isotropic priors, non-Gaussian parametric families, and task-specific prior distributions that better capture domain knowledge. In the revised version, we will consider, as an example, a more flexible parameterization of a Gaussian prior. Namely, we shall consider a Normal distribution, in which the parameters will be both the mean and the variance. This choice lifts the requirement of having to manually choose the variance of the prior as a hyperparameter.

Q2: FedMAP has the same communication complexity as FedAvg. In each round, the server broadcasts the parameters $\gamma$ defining the current prior distribution, and each client sends back their local parameters $\theta_k$ and weight $\omega_k$. FedMAP indeed introduces additional computational overhead in the local optimization step through MAP estimation. However, this computational cost is offset by improved communication efficiency. As shown in Theorem 1, the FedMAP updates correspond to gradient descent iterations for a strongly convex function $M(\gamma)$, which ensures faster convergence with fewer communication rounds. In non-IID settings, client drift often requires more rounds using traditional methods. The trade-off between increased local computation and reduced communication rounds is especially beneficial in FL where communication is often the primary bottleneck. We agree that analyzing the trade-offs is important for practical deployment and plan to include detailed empirical analysis in the revised version of the manuscript.

We thank the reviewer for their thoughtful comments and efforts toward improving our manuscript. We are glad that you found our proposed algorithm novel and appreciated the suggested improvements. Below, we would like to address your main concerns:

W1: We agree that our paper should provide a more technical discussion. We will add a new subsection that provides formal analysis comparing our bi-level optimization to the existing PFL approaches and demonstrate how FedMAP's envelope function and adaptive scheme facilitate better handling of non-IID data. Also, we added a brief discussion in our response to Q1, which shows FedMAP is different from other PFL approaches.

W2: We will move the statements of the key results of the theoretical analysis from the appendix to the main paper. The new section will present the convergence guarantees and connect theoretical results to practical benefits in non-IID settings.

W3: We will update the related work section to include the relevant 2024 publications such as PerAda[1], FedASA[2] and pFedEM[3].

W4: Although our current experimental evaluation includes both conventional and personalized FL baselines, and their comparison results present FedMAP outperforms, we agree that additional baselines will strengthen our argument. We will include more recent and relevant PFL approaches in our evaluations.

Q1: Other PFL approaches mentioned in our related work section include fine-tuning based approaches, which use sequential optimization (first global, then local) without regularization, whereas FedMAP provides continuous adaptation with its learned prior. Layer-wise based methods have constraints on specific layers/components, but FedMAP allows continuous, full-model adaption. Compared to meta-learning based approaches which require complex second-order optimization, FedMAP is more efficient with first-order updates. In contrast to the clustering based which uses discrete clustering, FedMAP provides more smoother knowledge transfer using the adaptive weighing. The key novelty of FedMAP is its bi-level optimization framework. It minimizes

$

M(\gamma) = \sum_{k=1}^q M_k(\gamma; Z_k) = \sum_{k=1}^q \min_{\theta} {L(\theta; Z_k) + R_k(\theta, \gamma)}

$

where the inner optimization personalizes local models and the outer optimization learns a global prior. This structure provides several advantages:

1. The framework offers convergence guarantees for any parameterized regularizer $R_k(\theta, \gamma)$ satisfying suitable convexity conditions (Theorem 1), making it general and flexible for different personalization strategies.

2. The local optimization

$

\theta_k^* = \arg\min_{\theta} L(\theta; Z_k) + R_k(\theta, \gamma)

$

enables client-specific adaptation while maintaining theoretical properties through MAP estimation.

3. The global parameter optimization through $M(\gamma)$ balances local personalization and global knowledge sharing, which is important for heterogeneous data distributions as demonstrated in our example for a linear regression task.

References:

[1] Xie, C., Huang, D.-A., Chu, W., Xu, D., Xiao, C., Li, B., & Anandkumar, A. (2024). PerAda: Parameter-efficient federated learning personalization with generalization guarantees. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 23838–23848.

[2] Deng, D., Wu, X., Zhang, T., Tang, X., Du, H., Kang, J., Liu, J., & Niyato, D. (2024). FedASA: A personalized federated learning with adaptive model aggregation for heterogeneous mobile edge computing. IEEE Transactions on Mobile Computing, 23(12), 14787–14802. https://doi.org/10.1109/TMC.2024.3446271

[3] Chen, S., Liu, W., Zhang, X., Xu, H., Lin, W., & Chen, X. (2024). Adaptive personalized federated learning for non-IID data with continual distribution shift. 2024 IEEE/ACM 32nd International Symposium on Quality of Service (IWQoS), 1–6. https://doi.org/10.1109/IWQoS61813.2024.10682851