Understanding Convergence and Generalization in Federated Learning through Feature Learning Theory

Thank you for your thoughtful review and the constructive comments on our submission. We are pleased that you found the paper well-structured and the strong motivation and theoretical results. We appreciate your feedback and would like to address your concerns as follows:

---

Weakness 1 The contribution of this paper is still a bit unclear to me, probably since I am not familiar with federated learning: It seems that the optimization dynamics turns to be explicit in the setting of this paper (a certain data-generating model, homogeneous two-layer CNNs and gradient descent optimization), and the results seem direct under the classical analysis of gradient descent. I am wonder whether this setting is representative enough, could the authors elaborate the technical innovations on this point?

While our setting might initially appear explicit with homogeneous two-layer CNNs and gradient descent optimization, the complexity and challenges of our analysis are significantly increased by the inherent non-convexity of two-layer neural networks and the unique dynamics of federated learning. We elaborate on our main contributions as follows:

1. Impact of Weight Averaging:

Unlike single network models, federated learning involves weight averaging across multiple neural networks. This introduces additional complexity, as the training dynamics of each network are influenced by this averaging process. This will lead to the dynamic correlation between neural networks after weight averaging. Besides, despite the fact that weight averaging is a linear combination of weight, its impact on the nonlinear dynamics of individual clients is substantial. We address this by employing a global weight decomposition approach, enabling us to effectively track the dynamics of feature learning within this federated framework.

2. Impact of Data heterogeneity

In federated learning, unlike in single network scenarios, the design of a data model distribution that encompasses both iid and non-iid cases is a significant challenge. To capture the nuances of data heterogeneity and its impact on generalization, we introduce an effective signal-to-noise ratio. This innovative approach allows us to characterize the influence of data diversity on the learning process more accurately.

---

Weakness 2: Another thing that I am concerned is that the current result are strongly related to the neural networks. In my understanding, federated learning is a topic independent of neural networks. Could the authors discuss the probabilities that the theory in this paper can go beyond just using two-layer CNNs to be individual models?

While federated learning is a broad topic that encompasses a variety of models, our research specifically addresses the challenges and dynamics associated with neural networks, which are increasingly prevalent in modern federated learning applications. The urgency to analyze federated learning within the neural network paradigm stems from both their widespread usage and the complexity of their optimization and generalization characteristics, particularly given their non-convex nature. The main motivation for our work is the lack of existing studies that comprehensively reveal the generalization gap between Federated Averaging (FedAvg) and local training within the neural network framework. Our focus on two-layer CNNs is driven by the need to tackle these hard-to-analyze aspects and to provide foundational insights into how federated learning operates in neural network settings.

Regarding the possibility of extending our results to more complex neural network models, we recognize this as a challenging yet intriguing direction for future research. Currently, most studies on feature learning, including our references [1,2,3], concentrate on two-layer CNNs due to their tractable analysis framework. Expanding this research to encompass more intricate neural network architectures presents a valuable opportunity for advancing the field and is an exciting prospect for subsequent work.

---

Weakness 3: (minor) Some symbols and equations should be defined and clearly explained. The authors should check Section 3 and 4 to make sure that all symbols are defined in their occurances.

We appreciate the reviewer's attention to detail in highlighting the need for clearer definitions and explanations of certain symbols and equations in our paper. In these sections, we have now ensured that:

• All symbols are defined clearly at their first point of use.

• Equations are explained with adequate context and clarity to facilitate better understanding

Thank you for your insightful feedback and constructive comments on our submission. We appreciate the opportunity to clarify your concerns and questions as follows:

---

Weakness 1: My primary concern with this work is its inability to effectively account for the influence of local training rounds (i.e., E iin the paper) on generalization performance. The paper only demonstrates that when which signifies no server-worker communication, the generalization performance is inferior compared to the scenario with . This can be partly attributed to the reduced amount of data utilized in client collaboration. It will be more exciting to see any trend within the regime

We are a bit confused about this comment “The paper only demonstrates that when $E= \infty$ …” as we never do that. Our analysis actually focuses on the regime where $0 < E < \infty$. This means that during the communication rounds, the number of local updates $E$ is finite. In particular, we use $T = E R$ to count the total number of local update steps, where $R$ is the number of communication rounds. In summary, all of our theoretical analyses are based on $0 < E < \infty$ instead of $E= \infty$.

---

Weakness 2: The paper establishes that training convergence is independent of data heterogeneity, a finding that appears to contrast with a substantial body of research and empirical observations in FL. It would be greatly beneficial if the authors could provide further elucidation regarding why this particular model exhibits this behavior.

The effect of data heterogeneity on convergence is a critical aspect of our analysis and highlights the significant influence on convergence rates. As demonstrated in Theorem 4.4, we establish condition related to the signal-to-noise ratio (SNR):

$\overline{n} \cdot \frac{\|\boldsymbol{\mu}\_k \|^2\_2}{\sigma^2\_p} \cdot \overline{\mathrm{SNR}}^2\_k = \Omega( 1 )$ with the effective SNR for each climent $k$ given by $\overline{\mathrm{SNR}}\_{k} = \left(\sum\_{k'=1}^K \frac{\langle \boldsymbol{\mu}\_k, \boldsymbol{\mu}\_{k'} \rangle}{ \| \boldsymbol{\mu}_k \|^2\_2} \right) \mathrm{SNR}_k$

In scenarios with increased data heterogeneity, the term $\sum\_{k'=1}^K \frac{\langle \boldsymbol{\mu}\_k, \boldsymbol{\mu}\_{k'} \rangle}{ \| \boldsymbol{\mu}_k \|^2_2}$ tends to decrease. To maintain a consistent SNR level in such cases, it becomes necessary to reduce $\sigma^2_p d$, which, according to the convergence time in Theorem 4.3: $T = \tilde{\Theta}( \eta^{-1} K mn\sigma\_0 ^{-1} \sigma^{-2}\_p d^{-1} + \eta^{-1}\epsilon^{-1} mn \sigma^{-2}\_p K d^{-1})$

results in slower convergence. On the contrary, if we reduce the data heterogeneity, the convergence time will become shorter.

---

Weakness 3: A minor point relates to the clarity of the writing. The paper's notations appear somewhat redundant and disorganized. For instance, the notations are not used beyond Section 3.2, and the definition of which plays a pivotal role in signal-to-noise ratios, is missing. In Equation (4), it might be more appropriate to replace with. I recommend that the authors thoroughly review the manuscript to enhance its overall clarity and coherence.

We greatly appreciate your feedback regarding the clarity and organization of notations in our manuscript.

1. We would like to clarity in section 4.2, we have used $|\mu^{(C)}|$;
   We have introduced a clear definition of $\sigma_p$, which represents the strength of the noise vector. The definition is now presented as follows:

2. Noise vector is distributed as $\boldsymbol{\xi}\_{k,i} \sim \mathcal{N}(\mathbf{0}, {\sigma^2\_p}\mathbf{I}-\sum\_{c=1}^C \boldsymbol{\mu}^{(c)} {\boldsymbol{\mu}^{(c)}}^\top/\| \boldsymbol{\mu}^{(c)} \|\_2^{2})$,

3. Thanks for your suggestion, we have replaced the term $n_k$ in Eq (4). We have conducted a thorough review of the manuscript to ensure that all notations are used consistently and defined clearly at their first point of use.

Thank you for your further question.

Firstly, it is important to note that a smaller $E$ does not necessarily imply better generalization. To explore this relationship further, we conducted additional experiments of FedAvg on CIFAR 10. From Table 1 we can see that the generalization performance of FedAvg is not closely related to the local epoches $E$:

Table 1: Test Accuracy of FedAvg on CIFAR10 with IID setting and 10 clients, learning rate is 0.01

Regarding the local update $E$ there is an implicit relationship with the convergence time. Under our learning rate assumption $\eta \leq \tilde{O}(\frac{K}{E} \min \{\|\boldsymbol{\mu}\|\_{2}^{-2}, \sigma\_{p}^{-2}d^{-1}\})$, increasing $E$ may lead to slower convergence. This observation aligns with the discussion in Weakness 4 with Reviewer vgxi.

The impact of varying $E$ on the required number of communication rounds to achieve adequate convergence is demonstrated in Table 2 of the FedAvg work [10], which compares the communication rounds needed for the same training accuracy on MNIST. The table shows:

This table indicates a clear trend: as $E$ increases, so does the total number of updates ($T=ER$), evident in both IID and Non-IID settings.

Lastly, we would like to emphasize that our theorem does not just simply show that federated learning generalizes better than local training. Instead, our theorem demonstrates that the generalization performance of federated learning relies on the signal-to-noise ratios (SNR) among the clients. SNR effectively captures the similarity in data distribution across different clients. The large SNR values mean less data non-IID, leading to better generalization performance. This assertion also matches the practices of federated learning.

Thank you for your thorough review and constructive feedback. We appreciate your recognition of our paper's strengths and would like to address the concerns you raised:

---

Minor writing 1: Page 2, para 1: "under" (rather than "in") a certain condition. Last sentence in this para is not clear. Didn't you earlier say that local algo has no comm.?

We agree with your suggestion to use the term 'under' rather than 'in' a certain condition, and have amended this in the manuscript.

Regarding the last sentence of this paragraph, we have revised it for better clarity. Our intent was to convey that by 'local training', we refer to the training process that occurs without any communication among the different clients.

---

Minor writing 2: At a few places, mathematical symbols or terms are used before being defined. E.g, sigma_p in Assum. 4.2, weighted FedAvg in the para above Sec. 6.2.

Thank you for highlighting the need for clearer definitions of certain mathematical symbols and terms used in our manuscript. We have introduced a clear definition of $sigma_p$, which represents the strength of the noise vector. The definition is now presented as follows:

Noise vector is distributed as $\boldsymbol{\xi}\_{k,i} \sim \mathcal{N}(\mathbf{0}, {\sigma^2\_p}\mathbf{I}-\sum\_{c=1}^C \boldsymbol{\mu}^{(c)} {\boldsymbol{\mu}^{(c)}}^\top/\| \boldsymbol{\mu}^{(c)} \|\_2^{2})$,

Besides, we have refined the description of the Weighted FedAvg method in Section 6.1.

---

3.2 Data Model: Why is Rademacher if y are labels? Are you assuming a balanced dataset? What is the distr P.

In our analysis, we draw labels from a Rademacher distribution, which are essentially binary variables taking values {+1,-1}. Regarding the balance of the dataset, we assume that with a sufficiently large training sample, the dataset will be approximately balanced. In this context, 'balanced' means that each class is represented proportionately in the dataset, with each class having at least one-third of the samples, as can be provably achieved under our assumptions.

$P$ is defined as a discrete distribution that assigns probabilities to each element in the range from 1 to $C$, where $C$ represents the number of classes or distinct data categories.$P$ is a discrete distribution assigning probabilities to each element from $1$ to $C$. It governs the likelihood of each class or category appearing within the data set of a client. We have added this information in our updated manuscript.

---

4.1: Eq. (6) and (8) are not clear. What is ? And why is there a sum over present? Where is on r.h.s.?

To address the concerns regarding Equation (6) and Equation (8), we offer the following clarification:

The rewritten form of Equation (6) is as follows:

$\mathbf{w}\_{k,j,r}^{(t)} = \sum\_{k'=1}^K ( \alpha^{(t)}\_{k',j,r} \mathbf{w}\_{k',j,r}^{(0)} + \gamma\_{k',j,r}^{(t)} \| \boldsymbol{\mu}\_{k'} \|^{-2}\_2 \boldsymbol{\mu}\_{k'} + \sum\_{i=1}^{n\_{k'}} \rho\_{k',j,r,i}^{(t)} \| {\boldsymbol{\xi}}\_{k',i} \|^{-2}\_2 {\boldsymbol{\xi}}\_{k',i})$

for all $k \in [K]$. The right-hand side of this equation involves a summation over all $k’ \in [K]$, which includes the specific $k$ on the left hand side. This formula represents the iteration considering weight averaging in our model. Crucially, in each iteration, the variable $\gamma^{(t)}\_{k,j,r}$ and $\rho^{(t)}\_{k,j,r,i}$ will be differ for each $k$, leading to distinct growth rates for $\mathbf{w}\_{k,j,r}$ across different clients.

---

4.2 Why is using the same training size for all clients without loss of generality?

We use the same training size for all clients for simplicity and to establish a baseline for comparison. However, our framework allows for generalization to cases with varying client numbers. A more general effective signal-to-noise ratio (SNR) can be derived to account for different training sizes across clients, which adds versatility to our model.

---

Assum 4.2: what is sigma_p

The $\sigma^2_p$ is the variance of the noise vector. We have added this information.

---

What is the motivation for the defn of SNR

The definition of $\overline{SNR}_k$ is introduced to characterize the generalization of Federated Averaging (FedAvg) in the presence of data heterogeneity. While the standard SNR is sufficient for demonstrating the generalization of local training, FedAvg introduces communication among clients, affecting their individual learning processes. Therefore, $\overline{SNR}_k$ measures the effective SNR on client $k$ under FedAvg, providing a more comprehensive understanding of generalization in a federated learning environment.

Thanks again for your thorough review and constructive feedback. As the deadline for the author-reviewer discussion phase is nearing, we would like to respectfully inquire if our rebuttal has effectively addressed the concerns raised in your review.

Your insightful feedback, particularly concerning the clarifications of Eq. (6) and Eq. (8), as well as other notations and the verification of experimental results, has been invaluable in refining our manuscript. We have carefully addressed each of these points in our rebuttal and have made corresponding updates to our manuscript to reflect these changes.

Thank you once again for your valuable contributions to the review process. Your feedback is crucial to the continued refinement of our work.

We would like to thank the reviewer for acknowledging our work on feature learning theory and the significant generalization result achieved. We address your concerns and question as follows:

---

Weakness 1: The presentation of the neural network model can be improved. The authors refer to this as a CNN model but it is not clear immediately what is the channel size, width and stride of this CNN model. It would be nice if authors could illustrate this in a figure either in the main text or appendix. Additionally, the authors also assume the weights of the second layer are fixed, which simplifies the problem.

Our choice to label the model as a CNN is based on the specific characteristics of the data model we used. We adopt a two-patch model, where $\mathbf{x} = [\mathbf{x}^{(1)}, \mathbf{x}^{(2)} ]$. We implement a convolution operation using a single weight across these two patches, akin to applying a filter in traditional CNNs. This is concisely represented as follows:

$f= \frac{1}{m} \sum_{r=1}^m [ \sigma( \mathbf{w}^\top_{k,+1,r} \mathbf{x}^{(1)}) + \sigma( \mathbf{w}^\top_{k,+1,r} \mathbf{x}^{(2)}) ] - \frac{1}{m} \sum_{r=1}^m [ \sigma( \mathbf{w}^\top_{k,-1,r} \mathbf{x}^{(1)}) + \sigma( \mathbf{w}^\top_{k,-1,r} \mathbf{x}^{(2)}) ]$

Here the filter $\mathbf{w}_{k,j,r}$ operates on two patch, with $k \in [K]$ denoting the index of client, $j \in \{+1, -1 \}$ corresponding to the weights value at the second layer, and $r \in [m]$ as the filter index.

Our approach alongside fixing the second layer is consistent with the related works [1,2,3], where similar network structures applying weights across multiple data patches are identified as CNNs. Despite the fixed second layer, the optimization remains non-convex despite the fixed second layer, adding complexity to our analysis.

We have added an explanation in the appendix.

---

Weakness 2: Assumptions need to be justified better. For instance, I don't understand why d has to be large to ensure an over-parameterized setting. In my understanding, just increasing m should be enough to imply over-parameterization as done in [4]. Also please specify what the are 'statistical properties' that are maintained using Assumption 2.

Our approach differs significantly from the "lazy training" or "neural tangent kernel (NTK)" regime [4,5,6]. This is because we consider a smaller initialization scale than that in the NTK regime. In our model, the width of the neural network $m$ is set to be at least polylogarithmic in the dimension, rather than extremely large as typically assumed in the NTK regime [4]. Then, to make sure that the learning is in a sufficiently over-parameterized setting, we set $d$ to be large, this can make sure that the noise vectors are near-orthogonal among each other. Similar conditions have been made in the study of learning over-parameterized models [7,8,9].

Regarding the 'statistical properties' referenced in Assumption 2, these are detailed in Lemma B.1 and Lemma B.2 of our paper. These lemmas require the width of the network and number of samples to be polylogarithmic in $d$. Under this condition, our theoretical results hold with a high probability of at least $1-1/d$.

---

Weakness 3: There either seems to be a major typo/mistake in Equation (6) (and consequently Equation 8) or I am misunderstanding something. The LHS of Eq. (6) has a k’ but the RHS has no dependence on k’. Effectively this is implying that w is going to be the same for any k’, which does not seem to be correct.

We appreciate the reviewer's attention to the details of our mathematical formulation. To address the concerns regarding Equation (6) and Equation (8), we offer the following clarification:

The rewritten form of Equation (6) is as follows:

$\mathbf{w}\_{k,j,r}^{(t)} = \sum\_{k'=1}^K ( \alpha^{(t)}\_{k',j,r} \mathbf{w}\_{k',j,r}^{(0)} + \gamma\_{k',j,r}^{(t)} \| \boldsymbol{\mu}\_{k'} \|^{-2}\_2 \boldsymbol{\mu}\_{k'} + \sum\_{i=1}^{n\_{k'}} \rho\_{k',j,r,i}^{(t)} \| {\boldsymbol{\xi}}\_{k',i} \|^{-2}\_2 {\boldsymbol{\xi}}\_{k',i})$

for all $k \in [K]$. The right-hand side of this equation involves a summation over all $k’ \in [K]$, which includes the specific $k$ on the left hand side. This formula represents the iteration considering weight averaging in our model. Crucially, in each iteration, the variable $\gamma^{(t)}\_{k,j,r}$ and $\rho^{(t)}\_{k,j,r,i}$ will be differ for each $k$, leading to distinct growth rates for $\mathbf{w}\_{k,j,r}$ across different clients.

Weakness 4: The analysis does not discuss the effect of the number of local epochs on the convergence of FedAvg in Theorem 4.3. For instance, in the i.i.d case when, we would expect that doing more local steps leads to faster convergence. However, if we substitute the lower bound from Assumption 3 in Theorem 4.3, then it appears that doing more local steps slows down convergence regardless of the data heterogeneity.

We thank the reviewer for highlighting the need to clarify the effect of the number of local epochs on the convergence of FedAvg as presented in Theorem 4.3. In our analysis, the term $T$ in Theorem 4.3 represents the total number of local updates, defined as $T = ER$, where $R$ is the number of commutation rounds and $E$ is the number of local epochs.

We would like to offer further clarification on the impact of increasing the number of local updates ($E$) on the overall training process in federated learning. This relationship is illustrated in the results shown in Table 2 of the referenced seminar work [10], which we have partially reproduced below for clarity

As demonstrated in this table, there is a clear trend that as the number of local updates $E$, the total number of updates also increases. This trend is evident in both IID and Non-IID settings.

Additionally, the effect of data heterogeneity on convergence is a critical aspect of our analysis. As demonstrated in Theorem 4.4, we establish condition related to the signal-to-noise ratio (SNR):

$\overline{n} \cdot \frac{\|\boldsymbol{\mu}\_k \|^2\_2}{\sigma^2\_p} \cdot \overline{\mathrm{SNR}}^2\_k = \Omega( 1 )$ with the effective SNR for each climent $k$ given by $\overline{\mathrm{SNR}}\_{k} = \left(\sum\_{k'=1}^K \frac{\langle \boldsymbol{\mu}\_k, \boldsymbol{\mu}\_{k'} \rangle}{ \| \boldsymbol{\mu}\_k \|^2\_2} \right) \mathrm{SNR}\_k$

In scenarios with increased data heterogeneity, the term $\sum\_{k'=1}^K \frac{\langle \boldsymbol{\mu}\_k, \boldsymbol{\mu}\_{k'} \rangle}{ \| \boldsymbol{\mu}\_k \|^2\_2}$ tends to decrease. To maintain a consistent SNR level in such cases, it becomes necessary to reduce $\sigma^2\_p d$, which, according to the convergence time in Theorem 4.3: $T = \tilde{\Theta}( \eta^{-1} K mn\sigma\_0 ^{-1} \sigma^{-2}\_p d^{-1} + \eta^{-1}\epsilon^{-1} mn \sigma^{-2}\_p K d^{-1})$ results in slower convergence. On the contrary, if we reduce the data heterogeneity, the convergence time will become shorter.

---

Weakness 5: I'm not sure if Lemma 5.2 is implying convergence due to the dependence of $\bar{w}^\ast_{j,r}$ on $log⁡(2/\epsilon)$. In my understanding, if we set $\epsilon \rightarrow \infty$ then this would imply $|| \bar{w}^\ast_{j,r}||\_2→\infty$

We appreciate the opportunity to clarify the implications of Lemma 5.2, particularly in the context of classification problems where logistic loss is applied. In such scenarios, the logistic loss function is defined as $\ell = \log(1+\exp(-fy))$

In the case of logistic loss, it is a well-established property that as the loss tends to zero, the optimal weight vector tends to grow to infinity as training time goes on. Therefore, in the context of Lemma 5.2, the expectation is that the norm of the optimal weight vector $|| \bar{\mathbf{w}}_{j,r}^\ast ||\_2$ will increase over time, reflecting the model's convergence towards an optimal solution by minimizing the logistic loss. This behavior aligns with the principles of logistic regression models in machine learning and has been similarly observed and established in related works [8,11,12].