The Effect of Personalization in FedProx: A Fine-grained Analysis on Statistical Accuracy and Communication Efficiency

Authors: Xin Yu, Zelin He, Ying Sun, Lingzhou Xue, Runze Li
Primary area: learning theory
Keywords: Personalized Federated Learning, Statistical Complexity, Communication Complexity, Minimax Optimality

Abstract:
FedProx is a simple yet effective federated learning method that enables model personalization via regularization. Despite remarkable success in practice, a rigorous analysis of how such a regularization provably improves the statistical accuracy of each client's local model hasn't been fully established. Setting the regularization strength heuristically presents a risk, as an inappropriate choice may even degrade accuracy. This work fills in the gap by analyzing the effect of regularization on statistical accuracy, thereby providing a theoretical guideline for setting the regularization strength for achieving personalization. We prove that by adaptively choosing the regularization strength under different statistical heterogeneity, FedProx can consistently outperform pure local training and achieve a minimax-optimal statistical rate. In addition, to shed light on resource allocation, we design an algorithm, provably showing that stronger personalization reduces communication complexity without increasing the computation cost overhead. Finally, our theory is validated on both synthetic and real-world datasets and its generalizability is verified in a non-convex setting.

Decision: reject
Reviewer scores: 1, 5, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1729934444371

thread_type:
official_review

<thread_object>
creator: reviewer_han2

<content>
summary:
This paper has studied how personalization in FedProx impacts the statistical error of FedProx and  the trade-off between local computation cost and communication cost. First,  this paper provided sharper statistical error bounds by adaptively tuning the regularization parameter $\lambda$ in FedProx based on the degree of statistical heterogenelity parameter $R$. Then this paper provided a new algorithm to solve FedProx and analyzed the rate of convergence on both local models and global model, showing personalization can increase the computational cost per iteration but also reduce the number of communications. Finally, numerical experiments have been conducted to validate the theoretical findings.

strengths:
This paper studied a fundamental problem that how statistical heterogenelity impacts the process of personalized federated optimization,
gave new statistical error bounds under different levels of statistical heterogenelity,
and established a new algorithm that shows an interesting trade-off between local computation cost and communication cost.

weaknesses:
(1) The authors claimed that their statistical error bounds improve previous results which I do not fully agree with. The main weaknesses of this paper include the unfair comparisons over previous results and the limited improvement in statistical errors over previous results. There are three reasons.

(i) The statistical error bounds in this paper omitted the optimization error of algorithm solving FedProx.
In other words, this paper just analyzed the statistical error of ERM estimators.
However, previous results [1] incorporate the optimization error of algorithm solving FedProx,
that is, the statistical error of approximate ERM estimators.

(ii) The statistical error defined in this paper is slightly different from previous work [1].
While the two definitions are closely related, there exists a gap represented by a problem-dependent constant.

(iii) Except for problem-dpenednt constants,  this paper only improves previous results [1] by a factor of $O(1/\sqrt{m})$ 
in the case of $R\leq \frac{1}{m\sqrt{n}}$. 

(2) This paper aims to provide contributions in theory. In addition to the theoretical results, 
it is necessary to explain the technical challenges on the theoretical analyses and the corresponding technical controbutions, both of which are currently unclear to me.
Besides, there are many typos in the current version.

(i) In Eq. (29), (30), (32) and so on, the expectation operator $\mathbb{E}[\cdot]$ are omitted.

(ii) In Eq. (33), (35), (36), I think $\tilde{w}^{(i)}$ should be corrected as $w^{(i)}_*$.

References

[1] Chen et al. Minimax estimation for personalized federated learning: An alternative between fedavg and local training? JMLR, 2023.

questions:
Could the authors elucidate the technical challenges and contributions regarding the theoretical analyses?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730359001943

thread_type:
official_review

<thread_object>
creator: reviewer_5muf

<content>
summary:
The paper investigates the effect of regularization in FedProx, providing a theoretical guideline for setting the regularization parameter for achieving personalization. Experiments conducted on synthetic and real datasets validate the theoretical analyses.

strengths:
1. It is intuitive that GlobalTrain is more suitable for i.i.d. data, while LocalTrain (and personalized FL) is more appropriate for non-i.i.d. data, and FedProx represents a hybrid approach. This paper substantiates this intuition and formally provides theoretical guidelines to demonstrate how the regularization term in FedProx affects model convergence.
2. In comparison to similar works, the theoretical results in this paper are more precise with fewer assumptions. Furthermore, the experimental results perfectly validate the theoretical analyses.

weaknesses:
1. Figures in the paper could be larger for better readability.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700572055

thread_type:
official_review

<thread_object>
creator: reviewer_bmut

<content>
summary:
This paper proposes a theoretical guideline for adaptively setting the strength of the regularization under different levels of statistical heterogeneity to achieve personalization in FL. The authors demonstrate that by adaptive tuning the regularization to strengthen the personalization, communication complexity can be reduced without increasing the computation cost overhead.

strengths:
1. This paper addresses a practical issue that commonly exists in the FL framework the statistical heterogeneity is prevalent. The motivation for optimizing regularization strength in Fed Prox is compelling. 
2. The paper is well-written and has a logical structure. The mathematical derivations are rigorous and precise. The authors also provide code to facilitate reproducibility.

weaknesses:
1. The experiments in this paper use only one real dataset, MNIST, which may not comprehensively illustrate the effect of personalization on statistical accuracy under various levels of heterogeneity. In addition, the model used in the non-convex case is relatively simple. Including more datasets and exploring more complex models could better demonstrate the main contribution of this paper.
2. The experimental results primarily focus on the model's error. However, it would also be beneficial to include accuracy to show the model performance with different lambdas. 
3. The descipriotns and analyses of the experiments lack clarity and it would be better to improve the content. For example, L. 492-493 references Fig.4 in the appendix as displaying error-related results, but Fig. 4 (a) actually shows accuracy variation. Improving these experimental explanations would make the results more interpretable for readers. 
4. I am curious whether this adaptive regularization approach can be extended to scenarios where clients are sampled in each communication round.

questions:
Refer to the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703154177

thread_type:
official_review

<thread_object>
creator: reviewer_kgrk

<content>
summary:
The FedProx algorithm enables personalization through penalization on the distances between local models and the global model. 
This paper theoretically analyzes the regularization strength influence the statistical rate and show that FedProx can consistently outperform pure local training and achieve a nearly minimax-optimal rate by properly choosing the penalization quantity based on the relative size of heterogeneity $R$ and local sample size $n$.  Experiments on synthetic and MNIST datasets are given.

strengths:
- The rates are reasonable and established under common assumptions in federated learning.
- The paper is written in a way that is easy to follow.

weaknesses:
- It is claimed that FedProx achieves a nearly minimax-optimal statistical accuracy but it is not compared to other works like Scaffold where heterogeneity $R$ does not appear in the upper bound.

- Empirical results are only given for synthetic datasets and MNIST on logistic regression. Other personalization methods are not compared. 

- The adaptive strategy of choosing of $\lambda$ relies on unknown quantity and therefore limits its contribution.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730879251223

thread_type:
official_review

<thread_object>
creator: reviewer_iwvj

<content>
summary:
The paper studies the FedProx learning rule for multi-distribution learning. It establishes statistical rates as a function of first-order heterogeneity, showing how the regularization (to form consensus) should be set in a heterogeneity-aware manner. The authors also provide an algorithm to solve the empirical FedProx objective, providing its convergence analysis in the noise-less setting. Some experiments are also provided to verify the theory's predictions.

strengths:
The paper identifies an interesting problem in federated learning: the issue of personalizing more v/s less as a function of data heterogeneity.

weaknesses:
**Incorrect units**: Theorem 1 seems wrong, as stated. In particular, the units of the convergence rate are incorrect, i.e., the r.h.s. does not have the units of the squared parameters. For instance, if we blow up both the parameter space and function space by the same factor but very large factor, then r.h.s will grow much more quickly because of the $C_2$ and $C_4$ terms that will blow up due to the additional $\sqrt{1+2\mu}$ factor as opposed to say $\rho/\mu$ which will grow with the same rate as the l.h.s. To put it succinctly, the result is not scale invariant, and the authors are likely either hiding some boundedness assumptions or omitting important problem parameters in the Theorem statement. Even the choice of $\lambda$ is incorrect. $\lambda$ should have the same units as strong convexity, but it does not. It is unclear what the indicator function does in the choice of $\lambda$ because comparing $R$ with a unitless/scaleless quantity such as $1/\sqrt{n}$ makes no sense. For another sanity check, note that when each client's objective is the mean estimation objective (i.e., mean squared error of the estimator), $\mu=1$ and only dependence on the right in (4) should be $\rho^2$ and not $\rho^4$. I request the authors to correct this. 

In light of the above issue, it is hard for me to interpret the convergence result, which is the main result of this paper, but I will try to disambiguate what the correct convergence rate should have been. 

**Vacuous guarantees**: Let's consider (5) for the particular case of mean estimation so that $\mu=1$. I claim that the upper bound in (5) is worse than the best of two trivial algorithms: (i) estimate the mean using only the samples on machine $i$ (i.e., pure local training), (ii) estimate the mean using all the samples across all the machines (i.e., consensus optimization). To see this, consider the first and third terms of (5), which, when combined, simplify to:$$\min(\frac{\rho^2}{mn} + R^2, \frac{\rho^2}{n}),$$
which is precisely the best of simple algorithms (i) and (ii). Since (5) has another term in the upper bound, the guarantee of Theorem 1 can not beat this trivial baseline and is thus vacuous. Furthermore, the best of (i). (ii) is extremely simple to implement: do the usual federated learning and pure local training side by side by splitting the samples in half. Then, compare these two solutions for each machine and pick the best one. 

**Noiseless setting**: It is tough to justify the noiseless setting of the optimization result because, in the absence of noise, there is no benefit of collaboration whatsoever, and even the Theorem statement in the paper suggests setting $\lambda=0$ to do purely local training. Thus, studying optimization with a complex procedure in such a setting is a moot point. The authors brush this away, speaking as if dealing with noise were trivial and techniques like variance reduction can be added without much effort to their algorithm. That does not seem to be the case.  

In light of the above issues, I strongly recommend rejecting the work, as it does not meet the rigor requirements for publication at ICLR.

questions:
1. Why is $w_\star^{(g)}$ the appropriate true global solution in the heterogeneous setting? It seems the authors are conflating between the average of optimas and the optima of the average function. In federated learning, the aim is usually to recover the latter, and both these solutions can be further apart. To see the difference between these two solutions, see the section on fixed point discrepancy in [Patel et al.](https://arxiv.org/abs/2405.11667). 

2. It is unclear what is $\tilde w^{(i)}$; what do the authors mean by local models? Do they use this to refer to each machine's ERM solutions, which are unique given the strong convexity?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732259880704

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to All Reviewers

comment:
We sincerely thank the reviewers for their thoughtful and constructive feedback. Below, we summarize the key improvements made to our manuscript. 

1. **Tight Upper Bound and Minimax Optimality** : We have refined the theoretical analysis and established that with a proper choice of the tuning parameter $\lambda$, the solution of FedProx achieves a matching upper bound with the minimax lower bound of the estimation problem, thereby demonstrating the *minimax optimality* of the FedProx method. See Theorem 1 for the new upper bound and Lemma 1 for a discussion on the technical contribution. Additionally, we have refined the choice of $\lambda$ to ensure that the tuning parameter and the resulting statistical bounds are both interpretable and expressed with correct units.


2. **Improved Numerical Studies** : We have expanded our experiments to include more datasets, applied more complex models, and provided a comparison with more baseline methods, including the dichotomous strategy suggested by reviewer $iwVj$, as well as other regularization-based personalized federated learning algorithms. Our results demonstrate the effect of personalization on the statistical accuracy of personalized federated learning method, further validating the theory.

We invite the reviewers to refer to the revised manuscript for further details on these additions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper establishes that, by appropriately setting the tuning parameter, the FedProx method achieves an error upper bound that matches the minimax lower bound. This demonstrates the minimax optimality of the FedProx method under the considered setting (smooth and strongly convex).

The primary criticism is that the dichotomous method, a simple baseline, can also achieve the same minimax optimal error bound. One reviewer argued that this undermines the novelty of the paper, as it does not explain why FedProx performs better than the dichotomous method in numerical experiments. The reviewer suggested that the authors explore a different setting where FedProx theoretically outperforms the dichotomous method in terms of the upper bound.

The authors counter that proving the minimax optimal upper bound for FedProx is valuable in itself, as it addresses a gap in understanding the theoretical guarantees of this widely used federated learning method.

Both the reviewer and the authors make valid points, making this decision challenging. I believe the paper does contribute to the field, as the authors claim. However, the fact that the same bound can be achieved by a baseline method diminishes the significance of the result. **If the paper fills an existing gap, it should better justify why this gap is important or interesting.** Specifically, the smooth and strongly convex assumption is relatively strong from an optimization perspective and may not align with the practical scenarios where FedProx is applied and has demonstrated success. Why is this setting a meaningful or interesting one for analyzing the error bound of FedProx? I recommend that the authors either provide a stronger justification for this setting or consider alternative settings where FedProx outperforms the baseline.

One additional minor concern I have is that the minimax optimal error bound was proven in the revised version rather than in the original submission. While the authors have highlighted the changes in blue, I am concerned that the review team may not have had sufficient opportunity to thoroughly evaluate these updates.

Overall, I am slightly inclined toward rejection. A good theory should ideally help explain practical phenomena. However, in this paper, the contribution may not be significant enough to guarantee acceptance. That said, I believe the score 1 given by one reviewer is unfairly low, as the paper does have merit.