Generalization of FedAvg Under Constrained Polyak-Lojasiewicz Type Conditions: A Single Hidden Layer Neural Network Analysis

• To establish the linear convergence, an important step is to show that the iterates stay within a neighborhood around the initialization point. However, it seems that the analysis towards this aim is problematic. For example, the proof of Theorem E.1 uses the inequality $\|w_k^{T}-w_k^0\|\leq \eta\sum_\tau\|\nabla \Phi_k\|$. Then, the paper uses the bound of $\|\nabla \Phi_k\|$ given in Lemma E.5. However, Lemma E.5 only holds in expectation. The underlying reason is that Lemma E.5 is based on Lemma E.4 and Lemma E.2. Note Lemma E.2 only gives bounds in expectation. Therefore, the first inequality in the proof of Lemma E.4 should involve an expectation over both sides. Therefore, the arguments in the paper only give bounds on $\mathbb{E}[\|w_k^{T}-w_k^0\|]$. This bound in expectation is not enough to apply the local PL condition to get linear convergence.
• The proof of Theorem 5.4 indicates that the covering number is of order $(3\rho\sqrt{d}/\theta)^d$. However, this covering number estimate seems to consider only a single $w_l$ instead of all nodes. The underlying reason is that the covering number estimate does not involve $m$.
• Below Eq (43), the paper uses the identity $\sigma(w_l^\top x)=\sigma(w_{l_j^\top x})+(w_{l_j}-w_l)^\top\nabla \sigma(w_{l_j^*}^\top x)$. This identity seems to not hold.
• As far as I see, the proof of Theorem 5.4 only gives Rademacher complexity bounds under the condition that $v\in \mathcal{G}_v$. It is not quite clear how to transfer this bound to Rademacher complexity bounds without this condition.

Unfortunately, there are several shortcomings in the paper. Please refer to questions for the detail of the below points.

• There are several inaccurate statements in the paper.

• The paper overlooks the literature on the generalization gap of federated learning algorithms.

• Some of the main claims of the paper are not sufficiently justified and contradict previous work. For example, the claim that the number of clients has no significant impact on the (especially) generalization gap of FedAvg.

• The derived generalization bounds do not reflect the behavior with respect to many previously reported factors, such as the number of clients, rounds, and the degree of heterogeneity. Moreover, the proof technique is not new and consists of classical $\epsilon$-net covering methods to bound the Rademacher complexities of each client; independent of the other clients and the FedAvg algorithm.

• The derived orders for $\rho$ and $m$ cause the NN output to change order-wise with $n$.

• The simulation part is disappointing.

• The paper only specifically analyzes single hidden layer neural networks, and the findings may not directly apply to other model architectures.
• The analysis relies on certain assumptions, such as the scaling of eigenvalues of the Jacobian matrix, which have been experimentally verified but may not hold universally.
• The paper focuses on a theoretical analysis of these conditions and their experimental validation on the MNIST dataset. Further research is needed to explore their applicability to other model architectures and datasets.
• The conditions for achieving linear convergence and good generalization require the radius $\rho$ and the model size to scale with the number of samples. While the paper acknowledge this aspect, choosing a large radius and model size might have practical limitations in real-world applications
• The results need more elucidation.

I detail all the found weaknesses in the paper below. As a result, I believe the paper is not ready for publication yet.

Regarding the theoretical results:

• Regarding Assumption 2.1: It seems that the “$L_k$-smoothness imply the $L-smoothness$ with $L=\frac{1}{K}\sum^K_{k=1}L_k$. If this is the case, why is there a need to introduce two types of smoothness in the assumption when one implies the other?
• Regarding Assumption 2.2: Why is the smoothness definition in Assumption 2.2 so different than Assumption 2.1, i.e., without being expressed as the norm of a difference of gradients?
• Moreover, there is no mentioning of Assumption 2.1 in Theorem 3.1, nor in any of the other results’ statements. Where is it used?
• Also, it is my understanding that we are always dealing with empirical losses when we are trying to provide optimization convergence and not with population losses (population losses are only important for generalization results). Then, why do we really need Assumption 2.1 which is only in terms of population losses (e.g., see the definition of $\Phi_k$ which is in terms of expectations over the true distribution of the data instead of an average of over the data samples)? Why is Theorem 3.1, which should deal with data and empirical measures, considering Assumption 2.1 instead of Assumption 2.2?
• My other concern is with the existence of Assumption 4.4. This assumption is only “used” to explain a little remark about the conditions in equation (11). However, it is not an assumption that is explicitly stated in any of the theoretical results of the paper. I think perhaps the authors had in mind to consider Assumption 4.4 for Corollary 4.6, in which case, this should be made explicit. In fact, I would just remove Assumption 4.4 and place its content inside the Corollary 4.6’s statement. Then, the simulations of Figure 1 can be placed afterwards.
• Corollary 4.6. seems incorrect. According to my calculations, the choices of $\rho$ and $m$ would provide: $\frac{\lambda}{m}>\mathcal{O}\left(\frac{c^2}{n^2}+\frac{1}{c^2}\right)$. If we assume Assumption 4.4, then the left-hand side is a constant. The corollary claims that the above inequality is satisfied for a sufficiently large $c$. However, a large $c$ will make the right-hand side grow unbounded, which could easily become (or be) larger than the constant at the left-hand side and thus violate the inequality. Unless I am missing something or the authors omitted something I am not considering, Corollary 4.6 is incorrect.
• In the third point of Section 5.1, it is claimed that certain scaling occurs “by choosing large enough communication rounds $R$”. However, I don’t see the explicit dependence on $R$ in equation (13), nor in Theorem 5.4. Why are the authors making such claim?

Important concerns regarding the contribution of the paper:

• In the Introduction, works such as Haddadpour et al. (2019), Stich, 2018, Qu et al., 2020, are mentioned to impose strong conditions on the losses because of the type of PL conditions or strong convexity being assumed. However, most of the results of the paper are based on the square loss. Indeed, the paper does not provide concrete examples of other losses used in practice for which Assumption 2.4 holds. Are there any other losses? Also, has the square loss been used in other previous works? This last question is very important because it delimits better the scope of the paper’s contribution, and there is currently no answer to these questions in the paper.
• The contributions and introduction should make explicit that the results using the one-layer neural network and generalization ones are only under the square loss. When I started reading the paper I thought that the classes of losses under which the main results hold were more general, but, as a reader, I was misled. The authors should make the contributions’ details more specific.
• In the paragraph that starts at line 255, the authors show how their work requires two conditions on local and global losses, as opposed to one condition as in the work by Chatterjee (2022). Can the authors provide intuition as to why two conditions are needed in their approach? Why isn’t just a condition on the local loss enough?
• The simulation is lacking important plots to strengthen the theoretical findings. For example, I suggest setting $m=c_1\cdot\frac{nK}{d}$ and $d=c_2\cdot n$ for some constants $c_1,c_2\geq1$, and then choose three or four values of $n$ (e.g., 20, 200, 2000) and plot how training and testing losses change across sample sizes and see if all of them converge to the global minimum (zero training loss), which is what I would expect from the theoretical results.

There are many notation problems that may have strong repercussions in the understanding of the paper and which show that paper needs to have further proofread.

• Around line 138 the author introduced a notation for the “empirical loss function (…) using sample $i$”, however it makes no sense to talk about an empirical loss function if it doesn’t include all the samples from the sample set. Besides, the authors do not provide a formula for $\Phi_{k,i}$ which they should in order to clarify what they mean by this expression.
• What is $\alpha_g$ in Assumption 2.4 and Theorem 3.1? This symbol has not been defined anywhere in the paper (the symbol $\alpha$ has been introduced, but not $\alpha_g$).
• In Remark 1. Shouldn’t it be $\Phi_{k,i}(\mathbf{w}^*)$ instead of $\Phi_{i}(\mathbf{w}^*)$?
• In Theorem 3.1’s statement and in the third point of the Section 5.1: Why is “$R>0$” being stated? Shouldn’t it be “$T>0$” according to Algorithm 1, i.e., the number of communication rounds? The symbol $R$ is not defined.
• In parts of the paper, such as after equation (6) and Lemma 4.3, the authors use the notation “$v_j\overset{i.i.d.}{\sim}\\{-1,+1\\}$”; however, this is not a correct notation since “$\\{-1,+1\\}$” is not a distribution, but a set. The authors must define that they want a distribution whose support is the set $\\{-1,+1\\}$. Another problem with this bad notation is that it is used in Definition 5.1 when defining the Rademacher complexity: the problem with this is that for such complexity we want a Rademacher distribution, not just any distribution over $\\{-1,+1\\}$.
• The symbol $\mathbf{H}$ after equation (10) has not been introduced before. It must be introduced.

Things to clarify:

• Assumption 2.4 deals with population losses, thus, its statement should be more explicit about giving the definition of $\Phi_k$ again, to have in mind that it is dependent on the underlying data distribution. In fact, the paper does not stress how these types of assumptions are dependent on the data distribution. As an alternative, I suggest including in the same equation (1) the definition of $\Phi_k$ because it is highly important: currently, the definition of $\Phi_k$ is buried in the text after equation (1).
• Following the previous point: do previous works also make conditions that are dependent on the data distribution?
• The fact that many results are under square loss is almost buried in the paragraph after Assumption 4.1. The authors should place the assumption of the square loss in either a new assumption environment or put it with Assumption 4.1. This is an important assumption, and the paper currently seems to just belittle it but not giving it the proper attention. This new assumption of square loss should be explicitly considered in the statements of all theoretical results that use it.
• What initialization is used for $v_i$ in the linear scaling experiments (the ones described in Figure 1)?

Other issues:

• Line 215: it should say “we will show” instead of “we have shown”.
• First paragraph of Section 3: it should say “Assumption 2.4” instead of “equation 2.4”.
• The term $\max_{k}l’_k$ can be removed since it doesn’t really add anything to the theorem’s statement.
• Please, double check where citations should be within parentheses and when it shouldn’t be.
• Sourav is a male name. Line 256 should say “his work” instead of “their work”.
• In the paragraph after Assumption 4.1, the tanh function is said to be $\sigma(0)\neq 0$: this is not true.
• Corollary 5.3 seems more important than Theorem 5.2 because it states a result comparing the losses, i.e., the generalization result in itself. Thus, I suggest “framing” this corollary instead of Theorem 5.2.
• $\Phi_S$ used in Theorem 5.2 needs to be defined.