Paper: Understanding the Theoretical Generalization Performance of Federated Learning (ICLR-2024)

Reviewer comment (Novelty):
the contribution beyond existing results is unclear

Author response:
General Response about contribution: We appreciate the author’s comments on numerous existing papers in online learning, model averaging, and classic convex optimization. However, it is essential to highlight that federated learning constitutes an independent subarea due to its intricate interplay between heterogeneous data and local update steps throughout the learning process. The distinctive settings of federated learning, characterized by the diversity of data sources and the incorporation of local update steps, differ significantly from other contexts. Our objective is to explicitly quantify the influence of data heterogeneity, local update steps, and the total number of communication rounds on the generalization performance of FL. This area remains under-explored in federated learning. More specifically:

1). The unique setting of federated learning in deep learning (over-parameterization) is significantly different compared to works in online learning, model averaging, and classic convex optimization [2, 5, 8, 9, 10, 13].

2). While many works study the convergence of federated learning, including those mentioned by the reviewer [4, 7, 11, 12, 14, 15], our focus is on *generalization*. 

3). The only paper mentioned by the reviewer that includes generalization analysis is [6]. This generalization extends from centralized learning using the Probably Approximately Correct (PAC) framework, which falls into the first class of related works in our paper. Please see Page 1 in our paper (“From a theoretical perspective, there has been relatively limited studies in addressing this question. We can categorize existing explorations into two classes. The first line of work employs the traditional analytical tools from statistical learning, such as the “probably approximately correct” (PAC) framework. These works focus on the domain changes due to the data and system heterogeneity…”). However, it's important to note that these works, including [6], do not provide an explicit relationship showing how critical factors in FL, such as the local update process, the number of communication rounds, and data heterogeneity, collectively affect FL's generalization. In other words, existing works have not fully explored the influence of these factors on the generalization performance of FL, which is the primary goal of our paper.

Reviewer comment (Novelty):
relevant related work is not discussed

Author response:
Related works: We appreciate the reviewer’s comment for various previous works studying the convergence of federated learning, and we will include them in the revision for a self-contained introduction of federated learning. It is essential to note, however, that our primary emphasis lies on the generalization of federated learning, for which we have meticulously presented an in-depth exploration of related works in our submission.

Reviewer comment (Presentation):
the results are presented in a convoluted and unintelligible way with an unnecessary function $\mathcal{F}$ that hides the actual result and is not interpretable to me

Author response:
Response: We explain the meaning of  $\mathcal{F}$ in Section 2.5. In short,  $\mathcal{F}$ corresponds to the general-term formula for a linear recurrence relation. Notice that we consider FL in multiple rounds, which is essentially a recurrence relation. Meanwhile, to help the readers understand our results, we provide expressions of the results in the simple case that does not contain $\mathcal{F}$.

Reviewer comment (Experiments):
there are no experiments (e.g., on simulated data) to evaluate the tightness of the bounds

Author response:
Response: Our Figs. 1, 2, and 3 are the numerical results for three cases to verify our three theorems. As shown in these figures, each marker is the average of 20 simulation runs and the curves are theoretical values from corresponding theorems. All these experiments show the tightness of our theoretical results.

Reviewer comment (Writing):
the proofs provided in the Appendix are not well presented

Author response:
Response: Thanks for your comments. However, it is unclear to us which parts of the proofs in the Appendix the reviewer referred to are not well presented. We would highly appreciate it if the reviewer can pinpoint to a certain proof, and we would be more than happy to make improvements.

Reviewer comment (Novelty):
How do the results for $K>1$ in this paper relate to the results on model averaging for linear regression in [5]?

Author response:
Response: Thanks for your question. Although both works study the linear regression model, there exist two major differences between our paper and [5]: 1) we study the over-parameterized setting (p > n) while [5] studies the under-parameterized case (see Page 5 in [5]). 2) we study the generalization of iterative processes of multiple model averages (i.e., the federated learning process) while [5] studies exactly one-step weighted averaging (see Abstract and Sec 3 in [5]).

Reviewer comment (Theory):
Why do you need the assumption of standard normal features? While Gaussian noise is a common assumption in the analysis of linear regression, this assumption appears too strong to me. Even the analysis in [1] only assumes that data is generated by a process using a covariance operator and a random variable whose components are independent and subgaussian. Shouldn't such an already strong assumption suffice?

Author response:
Response: Thanks for your question. We use the assumption of standard normal features to make our analysis and results more concise and tractable. Specifically, with standard normal features, our theories give exact values of the generalization performance, instead of upper or lower bounds (which is less precise than the exact value). Compared with [1] where the authors used sub-Gaussian instead of standard Gaussian distributions, the results in [1] are only the upper bounds of the generalization error.

Reviewer comment (Evaluation):
Why do you not use (a variant of) the notion of (shifting-)regret to analyse the non-stationary case? For this case (i.e., online learning), the (shifting-)regret is typically used as a success measure, since it captures the nature of the task better than the loss [3].

Author response:
Response: Our work proposes a unified framework to analyze the generalization of general federated learning, which includes the case of non-stationary data for each client. However, the goal for federated learning is to find a single solution given all the data from all clients. That is, the goal of federated learning is to find a solution that can approximate the global solution x*. Thus, the difference between the solution x_t and x* is used as the metric for generalization. This is different from online learning, in which regret is used to quantify how well an algorithm adapts to changes in the environment over time. Meanwhile, the regret can be easily calculated based on our current results by summing up the expressions over $t$.

Reviewer comment (Novelty):
The K=1 case of FL is equivalent to simple distributed SGD, i.e., it can be considered to be a centralized SGD with larger batch size and smaller learning rate [cf. Prop. 2 in 8]. For convex problems, this has been extensively studied [2]. How does the case in this paper differ?

Author response:
Response: Since we considered heterogeneity and non-stationarity, the $K=1$ case of FL is not necessarily equivalent to a simple distributed SGD. Since we want to know how system parameters such as the number of rounds and heterogeneity affect the generalization performance, our analysis and results are much richer compared to and are not included by [2].

Reviewer comment (Novelty):
Since the paper investigates convex problems, standard convergence results even for non-convex FL (where, e.g., it is shown that $||\nabla L|| = 0$) imply that $w^* = w$. This includes, e.g., [4, 12, 14]. Please elaborate on the contribution beyond these works.

Author response:
Response: Thanks for your comment. The major contribution in our work is that our focus is on the generalization of federated learning while these prior works only studied the convergence of federated learning.

Reviewer comment (Novelty):
The problem of heterogeneous local data in FL and its impact on convergence has also been extensively studied [7, 11, 15]. Please elaborate on the contribution beyond these works, in particular to Lemma 3 in [6] which seems to be a generalization of the results in this paper for deep learning.

Author response:
Response: Thanks for your comments. We note that our focus is the generalization of federated learning, while these papers study the convergence of federated learning. In [6], the authors studied the generalization of federated NTK (it seems that we couldn’t find Lemma 3 there, if we didn’t miss anything). The generalization of [6] is the extension of centralized learning, where it can be seen from the proof (Sec C in their appendix) that their results heavily rely on [R2]. From a high level, this work belongs to the first class of generalization as summarized in the introduction of our paper (Page 1 - 2). These works do not provide an explicit relationship to show how the critical factors in FL, (e.g., the local update process, the number of communication rounds, and data heterogeneity) affect the generalization of FL in general, which is the goal of our paper.

[R2] Arora, Sanjeev, et al. "Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks." International Conference on Machine Learning. PMLR, 2019.

Reviewer comment (Theory):
The results for the overparameterized case are not discussed with respect to benign overfitting for linear models [1]. The results of [1] should at least be discussed for the K=1 case, where FL boils down to distributed SGD.

Author response:
Response: Thanks for your comments. However, we do have comparisons with the benign overfitting of classical (single-task) linear models in Eq. (25) for $K=\infty$ case, which leads to the insight 4 in the paper (FL can alleviate the null risk compared to the classical single-task/agent benign overfitting). Notice that $K=1$ is the case of one-step local gradient update in each round, which does not overfit the training data set (and thus does not belong to benign overfitting).

Reviewer comment (Theory):
The model, i.e. the consideration of a linear model with i.i.d. Gaussian data, is very simple and unrealistic. However, as mentioned by the authors, this is now a commonly used model as a first step in understanding the theory.

Author response:
Thanks for your insightful comments regarding our model setting. For better readability, we also structure our response to your question accordingly as follows:

**Simultaneously minimized risk and generalization erros:** We respectfully disagree that both empirical risk and the generalization error are simultaneously minimized. In the cases of $K=1$ and $K<\infty$, the empirical risk is not minimized because the SGD/GD update is stopped after K steps (i.e., early stop). In the case of $K=\infty$, the empirical risk is minimized but the generalization error is not, since Eq. (6) suggests a zero empirical risk but the generalization error shown by Theorem 3 is still positive.

**“Model error”:** As we explained in the footnote of Page 4, the model error is equal to the expected test error for noise-free test data. This quantity is widely used to indicate the generalization performance in a linear model setup.

Reviewer comment (Experiments):
In particular, there is an inherent alignment between empirical risk and generalization error in the considered setup. That is, they are somehow simultaneously minimized, which is in sharp contrast to most real-world scenarios. In fact, throughout the paper, the authors study "model error," a quantity that one would also study for empirical risk. For this reason, I am not sure if the findings on generalization behavior using this setup can be extended or useful for realistic setups.

Author response:
Thanks for your insightful comments regarding our model setting. For better readability, we also structure our response to your question accordingly as follows:

**Simultaneously minimized risk and generalization erros:** We respectfully disagree that both empirical risk and the generalization error are simultaneously minimized. In the cases of $K=1$ and $K<\infty$, the empirical risk is not minimized because the SGD/GD update is stopped after K steps (i.e., early stop). In the case of $K=\infty$, the empirical risk is minimized but the generalization error is not, since Eq. (6) suggests a zero empirical risk but the generalization error shown by Theorem 3 is still positive.

**“Model error”:** As we explained in the footnote of Page 4, the model error is equal to the expected test error for noise-free test data. This quantity is widely used to indicate the generalization performance in a linear model setup.

Reviewer comment (Theory):
In addition, the authors stated that the model error can be shown to be equal to the expected test error for noise-free data. What happens if the data is not noise free? (I guess this is the interesting case, right?). How does the "model error" relate to the generalization performance in this case?

Author response:
Thanks for your question. The difference in the expected test error between the case of noise-free test data and the case of noisy test data is only the noise level (a constant that is irrelevant to the learning process). In the revision, we added Lemma 6 to make this claim rigorous.

Reviewer comment (Theory):
Continuing the point 2, a good “generalization bound” typically decreases with the total number of used symbol (here, it would be $m\times n \times t$ at the end of iteration $t$, in the simple case). How does the provided bounds behave with these quantities?

Author response:
Thanks for your question. It appear sthat there are some misunderstandings in here. First, we want to clarify that our characterization of the generalization error is not a bound. Rather, it is the **exact value** of the expected generalization error. The correctness of our characterization is also validated by Figs. 1-3 that analytical values perfectly match simulation values. Second, the expressions in Theorems 1~3 contain $m$, $n$, and $t$ explicitly. For example, in Eq. (14), both H and G have $mn$ in their denominator and thus decrease with respect to $mn$. When the learning rate is small, we have $H<1$, which suggests the first term of Eq. (14) also decreases with $t$. The second term of Eq. (14) (noise term) increases with $t$ due to the accumulation of the noise effect over time.

Reviewer comment (Evaluation):
All plots provide the numerical evaluation of the bounds on the model errors. It would be useful to plot the estimated model error as well as the estimated generalization error in the plots to observe how well the provided bounds capture the correct behavior of both the model error and the generalization error.

Author response:
Thanks for the suggestions. Note that we have already done so in the current figures. In Figs 1~3, the points (markers) are the numerical evaluations, and the curves (without the markers) are the analytical values. As these markers are very near to the curves (the curves almost pass through points exactly), the correctness of our analytical results is indeed verified.

Reviewer comment (Novelty):
Some relevant references are missing, e.g. arXiv:2306.03824 and arXiv:2306.05862.

Author response:
We appreciate the reviewer’s comment about the references, and we will include them in the revision. These two works use traditional analytical tools from statistical learning, such as stability, PAC-Bayes and rate-distortion, to establish the bounds. However, we note that these traditional tools *cannot* explicitly quantify the influence of data heterogeneity, local update steps, and the total number of communication rounds on the generalization performance of FL, which is part of the motivation of our work.

Reviewer comment (Theory):
Theorem 1 (and other results) suggest that increasing the learning rate $\alpha$ increases the model error. While this behavior makes sense from an empirical risk point of view, it's often the opposite for generalization behavior: larger learning rates, at least in some cases and for the classical centralized setup, lead to better generalization performance.

Author response:
We respectfully disagree that Theorem 1 suggests increasing the learning rate $\alpha$ always increases the model error (i.e., generalization error). In Eq. (14), a larger $\alpha$-value leads to a larger $G$ (the term for noise and heterogeneity). However, $H$ is a quadratic function with respect to $\alpha$. When there is no noise and no heterogeneity (i.e., $G=0$), the optimal $\alpha$ that can minimize the generalization error (i.e., minimize $H$) should equal to $\frac{mn}{mn+p+1}$. In this situation, when the learning rate increases from 0 to $\frac{mn}{mn+p+1}$, the generalization performance becomes smaller. Thus, there is no contradiction between our theoretical results and empirical observations.

Reviewer comment (Theory):
Furthermore, Theorem 1 suggests that more heterogeneity leads to higher generalization error. This is also counterintuitive to me. After all, this is exactly what I would expect for empirical risk, but I would expect the opposite for generalization.

Author response:
Thanks for your comments. In general, heterogeneity has been observed to be a negative factor for generalization. For example, the paper mentioned by the reviewer in Comment 7 (arXiv:2306.03824) shows the same conclusion that heterogeneity leads to higher generalization error (Corollary 1 in arXiv:2306.03824).

Reviewer comment (Theory):
If I'm not mistaken, equation (16) (and Proposition 1) is developed for the "constant learning rate" scenario. While for this scenario the existence of an optimal $K$ makes sense, I was wondering what would be the behavior with respect to $K$ for the carefully tuned decreasing learning rates?

Author response:
Thanks for your question. We agree that it is an interesting topic to further investigate whether an optimal $K$ exists when the learning rates are carefully tuned for different $K$. Meanwhile, we also note that the use of constant learning rates is the most common practice in training. Since our expression of the generalization performance contains $K$ and learning rate explicitly, we can replace $\alpha$ by a function of $K$ and then analyze the monotonicity of the whole expression. Depending on how $\alpha$ changes with $K$, the conclusion might be different. Due to the goal of the current paper is to take a first attempt to understand the generalization error of federated learning, we leave this to future work.

Reviewer comment (Experiments):
The bound of Theorem 3 interestingly suggests a double descent phenomenon. However, the paper would greatly benefit from a numerical verification of this.

Author response:
Thanks for your suggestions and questions. In Figs 1~3, the markers/points are simulated values while the curves are analytical values. In other words, the curves are not drawn by connecting the markers/points but are directly calculated by the analytical expressions.

Reviewer comment (Presentation):
It is written that "In Figure 3, we plot the model error against p...". However, the caption of Figure 3 states that "The curves are theoretical values from Theorem 3". It is a bit confusing whether the exact "model error" is plotted or the established bound on it? It would be useful to plot both (as well as the generalization error).

Author response:
Thanks for your suggestions and questions. In Figs 1~3, the markers/points are simulated values while the curves are analytical values. In other words, the curves are not drawn by connecting the markers/points but are directly calculated by the analytical expressions.

Reviewer comment (Theory):
The proofs are very long and technical, and it takes a lot of time to go through them. I think it would be helpful (and perhaps even necessary) to give a proof sketch in the main text, explaining the main steps, ideas, and tools used from the literature, and what the authors have added. In the current format, it is not clear what technical novelty the authors bring to this work.

Author response:
Thanks for your suggestion. Due to the page limit, it is difficult to provide proof sketches for all three main theorems in the main body. The main technical tools we used are listed in Appendix A. Specifically, Lemmas 2~4 are from the literature, while Lemma 5 is the new tool we created. To help the readers understand the proof of Lemma 5, we also provide a geometric interpretation in Fig. 4.

Reviewer comment (Presentation):
The supplementary material (in particular Appendix A) may be more appropriately organized, with a table of contents to guide the reader to the various sections.

Author response:
Thanks for your suggestion. In the revision, we have provided a table of contents to guide the reader to the various sections of the supplemental material. We have also added one paragraph at the beginning of Appendix A to make Appendix A more organized as follows:
"Moreover, in Appendix A, we have provided some useful lemmas. Specifically, Lemma 1 is used to support the claim of the convergence speed in Insight 4. Lemmas 2 to 4 are some results about the Gaussian random matrices that can be found in the literature. We want to highlight that Lemma 5 is part of our technical novelty, which gives the exact values of terms related to the projection formed by each agent’s training inputs."

Reviewer comment (Theory):
A major concern is the parameter defined by Equation (3). It does not seem like a suitable measurement for heterogeneity. Consider the case where $w_{(i), t} = w_{(j),t} = w^* + v$ with $\Vert v \Vert$ very large. In this case, all agents’ ground truth are the same, *i.e.* the data are homogeneous. However, the *level of heterogeneity* is very large by the definition in the paper.

Author response:
Response: Thanks for the comments. We note, however, that if every agent shares a very large and common $v$, then this quantity $v$ should be considered part of $w^*$ to accommodate such bias. Indeed, the choice of $w^*$ depends on all $w_{(i),t}$ (e.g., their average). In general, $w^*$ should be regarded as the “largest” overlapped part of $w_{(i),t}$. Therefore, our measurement for heterogeneity is still suitable.

Reviewer comment (Theory):
Is there a reference for Equation (25)?

Author response:
Response: Thanks for your question. Yes, Eq. (25) can be found in [R1] (with changes of notations). Eq. (25) is also a special case of Eq. (23) in our Theorem 3 by letting $m=1$ and $t=1$. 

[R1] Mikhail Belkin, Daniel Hsu, and Ji Xu. Two models of double descent for weak features. SIAM Journal on Mathematics of Data Science, 2(4):1167–1180, 2020.

Reviewer comment (Theory):
The paper claims to be general enough to provide results for the non-stationary setting. However, it is unclear if measuring the distance from \( w^\star \) is useful in the non-stationary setting. In particular, it is never discussed what \( w^\star \) is and why the machines would want to recover it in a general setting, as opposed to minimizing regret with respect to a fixed best model in hindsight. Due to this reason, it seems that the simple setting for presenting results, where \( w^\star \) is the average optimum of different machines, is the only reasonable setting. As a result, the generality of the results in this paper is not very useful, without further motivation.

Author response:
Response: Thanks for your comments and suggestions. While letting $w^*$ be the average of different machines seems to be the most reasonable choice, there are some other possibilities in reality. For example, we may want to let $w^*$ to be the “smallest” (in certain criteria such as l2 norm) among different machines when the whole system’s performance depends on the machine with the smallest $w$.

Reviewer comment (Novelty):
In the most challenging setting, that is, for a general \( K < \infty \), \( \bar{||\gamma||}^2 \) is conveniently assumed to be zero. This implies that all the machines have the same optimum as \( w^\star \). We already know the min-max optimal algorithms for quadratic functions in the homogeneous setting due to [Woodworth et al.](https://arxiv.org/abs/2002.07839). Thus, Theorem 2 is not interesting.

Author response:
Response: Thanks for your comments. It appears that there are some misunderstandings. First, our paper focuses on offering a theoretical understanding of the generalization performance of FL, not min-max optimal algorithms for quadratic functions. Second, Theorem 2 not only provides the results for the simple setting but also provides the general results with heterogeneity and non-stationarity. Third, we believe that results in Theorem 3 are non-trivial and the key part of the proof is Eqs. (78)~(80), which depends on our Lemma 5. The proof of Lemma 5 is complex and novel, which by no means can be obtained by standard arithmetic. We also use Figure 4 to give a geometric interpretation of the proof.

Reviewer comment (Novelty):
Theorem 3 is interesting but it is very easy to obtain using standard arithmetic.

Author response:
Response: Thanks for your comments. It appears that there are some misunderstandings. First, our paper focuses on offering a theoretical understanding of the generalization performance of FL, not min-max optimal algorithms for quadratic functions. Second, Theorem 2 not only provides the results for the simple setting but also provides the general results with heterogeneity and non-stationarity. Third, we believe that results in Theorem 3 are non-trivial and the key part of the proof is Eqs. (78)~(80), which depends on our Lemma 5. The proof of Lemma 5 is complex and novel, which by no means can be obtained by standard arithmetic. We also use Figure 4 to give a geometric interpretation of the proof.

Reviewer comment (Novelty):
Overall, the paper does not appear novel enough and does not convincingly add to the existing theory of local update methods.

Author response:
Response: We would like to clarify and highlight that the main contribution and novelty of this paper is to derive exact expressions (instead of upper or lower bounds) of the generalization performance of FL for various settings of $K$, which provide meaningful insights. To the best of our knowledge, our work is the first of its kind in literature. That being said, we would highly appreciate it if the reviewer provides pointers to existing works that we are unaware of, and we would be happy to include and compare them in the revision of this paper.

Reviewer comment (Experiments):
How are the step sizes tuned in all the experiments? In particular, in Figure 2, are step sizes tuned separately for each experiment?

Author response:
Response: Thanks for the question. The (fixed) step sizes are set to 0.02 in all experiments.

Reviewer comment (Theory):
Equation (23) also has null risk, depending on the order of taking limits. For instance, if \( p \to \infty \) before \( t \to \infty \), then there is a residual null risk. As a result, Insight 4 seems wrong. Figure 3 also seems to show this. What is the insight here then?

Author response:
Response: Thanks for your question. Again, it appears that there are some misunderstandings in how to interpret Insight 4. We note that Insight 4 only says the null risk is *alleviated* (not completely removed) by using more communication rounds and we explicitly discuss the issue of the speed w.r.t. $t$ and $p$ in the paragraph below Eq. (25). Hope this clarifies the confusion.