Paper: Understanding Convergence and Generalization in Federated Learning through Feature Learning Theory (ICLR-2024)

Reviewer comment (Novelty):
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 wonder whether this setting is representative enough, could the authors elaborate the technical innovations on this point?

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Writing):
(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.

Author response:
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

Reviewer comment (Writing):
It would be better if the authors explain the rationale behind the two-patch feature generation model again *just below* its definition. Moreover, should the signal patch and the noise patch necessarily have same dimension?

Author response:
Thank you for your suggestion regarding the explanation of the two-patch feature generation model. We have revised our manuscript to include a more detailed explanation of this model just below its definition.

Regarding the dimensionality of the signal and noise patches, they are indeed designed to have the same dimensions in our model. This design choice is particularly relevant for the convolutional neural network framework we employ, where a single weight is used in relation to both signal and noise patches. Having equal dimensions for both types of patches simplifies the convolution operation and aligns with the structure of typical CNNs.

However, it is possible to introduce more patches to adjust the dimensionality ratio between signal and noise patches to explore different model configurations.

Reviewer comment (Theory):
What is the expression of the distribution $P(\mu^{(1)},\mu^{(2)},\cdots,\mu^{(C)})$?

Author response:
$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.

Reviewer comment (Writing):
Page 3: $F_{-1}(W_{+1}, x)$ should be $F_{-1}(W_{-1}, x)$

Author response:
Thanks for pointing them out, we have modified them.

Reviewer comment (Writing):
Page 4: $yf(\bar{W}^{(t)}, x<0)$ should be $yf(\bar{W}^{(t)}, x)<0$

Author response:
Thanks for pointing them out, we have modified them.

Reviewer comment (Theory):
(minor) Theorem 4.3 is a "best-iterate" result. Could it be turned into a "last-iterate" or "average-iterate" guarantee?

Author response:
Our belief is that with detailed analysis, Theorem 4.3 can be extended to a last-iterate guarantee. Specifically, in the second stage of the optimization dynamics, although the loss derivative is non-constant, our preliminary observations suggest that the iterative process of weight adjustment tends to approach the optimal solution progressively. This behavior is indicative of the potential for a last-iterate guarantee.

Reviewer comment (Experiments):
My primary concern with this work is its inability to effectively account for the influence of local training rounds (i.e., $E$ in the paper) on generalization performance. The paper only demonstrates that when $E=\infty$, which signifies no server-worker communication, the generalization performance is inferior compared to the scenario with $E<\infty$. 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 $1\leq E<\infty$.

Author response:
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$.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Writing):
A minor point relates to the clarity of the writing. The paper's notations appear somewhat redundant and disorganized. For instance, the notations $\mu^{(1)}, \dots, \mu^{(C)}$ are not used beyond Section 3.2, and the definition of $\sigma_p$, which plays a pivotal role in signal-to-noise ratios, is missing. In Equation (4), it might be more appropriate to replace $n$ with $n_k$. I recommend that the authors thoroughly review the manuscript to enhance its overall clarity and coherence.

Author response:
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.

Reviewer comment (Writing):
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.?

Author response:
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.

Reviewer comment (Writing):
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.

Author response:
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.

Reviewer comment (Theory):
Why is $y_{k,i}$ Rademacher if $\mathbf y_k$ are labels? Are you assuming a balanced dataset? What is the distr. $P$?

Author response:
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.

Reviewer comment (Theory):
Eq. (6) and (8) are not clear. What is $k'$? And why is there a sum over $k$ present? Where is $k'$ on r.h.s.?

Author response:
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.

Reviewer comment (Writing):
Assum 4.2: what is $\sigma_p$?

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

Reviewer comment (Theory):
What is the motivation for the defn of $\bar{SNR}_k$?

Author response:
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.

Reviewer comment (Theory):
Theorem 4.4: in the following discussion, it is stated that $\xi_k \geq 1$. Since $\{\mu_c\}$ are orthogonal, are you assuming that $K \geq C$?

Author response:
In our analysis, we do assume that $K \ge C$, where $K$ is the number of clients and $C$ the number of classes. Specifically, in the IID case, we assume $K=C$, whereas in non-IID cases, we assume $K>C$.

Reviewer comment (Theory):
In Assum 4.2, $\eta \lesssim 1/d$ is assumed, but in experiments, $\eta=1$ is chosen. Can you explain this divergence?

Author response:
The conditions set forth in Assumption 4.2 were intended to account for the worst-case scenarios, often leading to 'extreme' conditions in a theoretical context. In practical experiments, however, it is reasonable to expect that the results hold true under less strong conditions, such as somewhat larger learning rates and lower dimensions, with a relatively high probability. This divergence between theoretical assumptions and practical experiment settings is common, aiming to ensure robustness in theoretical analysis while remaining applicable in real-world scenarios.

Reviewer comment (Evaluation):
In some rows of both Tables 1 and 2, singleset has better performance than FedAvg, e.g., CIFAR10 with Dirichlet or SVHN in Table 2. Why is this happening?

Author response:
FedAvg aggregates all local models with uniform weights to produce a global model. However, due to the data heterogeneity issues, data distribution among clients may be quite different. Aggregating local models with uniform weights may introduce useless even detrimental knowledge for the local model. SingleSet trains the model locally and will not be interfered by those clients with different data distribution. As a result, in these cases, SingleSet outperforms FedAvg.

Reviewer comment (Presentation):
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.

Author response:
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.

Reviewer comment (Theory):
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 [1]. Also please specify what the are 'statistical properties' that are maintained using Assumption 2.

Author response:
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$.

Reviewer comment (Theory):
The analysis does not discuss the effect of the number of local epochs $E$ on the convergence of FedAvg in Theorem 4.3. For instance, in the i.i.d case when $\mu_1 = \mu_2 = \dots, \mu_K$, we would expect that doing more local steps leads to faster convergence. However, if we substitute the lower bound on $\eta$ from Assumption 3 in Theorem 4.3, then it appears that doing more local steps slows down convergence regardless of the data heterogeneity.

Author response:
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
 
|     E    | B (batch size)  |  R (IID) | R (Non-IID)
|----------|-----|--------------|--------| 
|   1      | 10 |   34 | 350 |
|   5      | 10 |   20 | 229 | 
|   20     | 10 |   18 | 173 |

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.

Reviewer comment (Theory):
I'm not sure if Lemma 5.2 is implying convergence due to the dependence of $\bar{w}^{\*}\_{j,r}$ on $\log(2/\epsilon)$. In my understanding, if we set $\epsilon \rightarrow 0$ then this would imply $\bar{w}^\*_{j,r} \rightarrow \infty $ and hence $\|\|\bar{W}^{(T_1)}-\bar{W}^\*\|\|_F^2 \rightarrow \infty$.

Author response:
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].

Reviewer comment (Experiments):
The authors should add personalization baselines in the experiments on measuring the performance of Weighted FedAvg. I'm not surprised to see Weighted FedAvg outperform FedAvg in Table 1 and Table 2 given that each client has a personalized model for Weighted FedAvg. Also the authors should make it clear in the abstract and introduction that Weighted FedAvg is a personalization method.

Author response:
Thank you for the advice to add personalization baselines in the experiments. Here, we introduced two baselines FedBN [13] and FedPer [14], which separately personalize the Batch Normalization (BN) layers and the Classification Head for each local model. For a fair comparison, we also personalize all BN layers for our Weighed FedAvg. As demonstrated in Table 1 and Table 2 as follows:

| Methods             | CIFAR10 Pathological | CIFAR10 Dirichlet | CIFAR100 Pathological | CIFAR100 Dirichlet |
|---------------------|----------------------|-------------------|-----------------------|--------------------|
| SingleSet           | 20.08±0.13           | 38.60±0.02        | 6.59±0.18             | 6.97±0.11          |
| FedAvg              | 40.16±7.18           | 33.11±7.06        | 13.99±0.79            | 17.34±0.47         |
| FedPer              | 38.09±5.73           | 36.52±7.13        | 29.10±2.42            | 17.63±0.93         |
| FedBN               | 65.71±0.67           | 61.25±3.78        | 42.54±0.53            | 25.64±3.46         |
| Weighted FedAvg     | 60.32±0.31           | 60.52±0.30        | 31.88±0.56            | 18.82±0.36         |
| + Personalized BN   | **67.30±0.31**       | **65.31±1.23**    | **43.51±0.50**        | **26.11±1.77**     |

*Table 1: Experimental Results for FL with Label Distribution Skew on CIFAR10 and CIFAR100.*

| Methods             | MNIST               | SVHN                | USPS                | SynthDigits         | MNIST-M             | Avg                 |
|---------------------|---------------------|---------------------|---------------------|---------------------|---------------------|---------------------|
| SingleSet           | 94.38±0.07          | 65.25±1.07          | 95.16±0.12          | 80.31±0.38          | 77.77±0.47          | 82.00±0.40          |
| FedAvg              | 95.87±0.20          | 62.86±1.49          | 95.56±0.27          | 82.27±0.44          | 76.85±0.54          | 82.70±0.60          |
| FedPer              | 96.21±0.02          | 67.61±0.04          | 96.53±0.02          | 83.88±0.02          | 76.85±0.54          | 81.89±0.03          |
| FedBN               | **96.57±0.13**      | 71.04±0.31          | **96.97±0.32**      | 83.19±0.42          | 78.33±0.66          | 85.20±0.40          |
| Weighted FedAvg     | 95.75±0.18          | 67.82±1.07          | 95.66±0.22          | 84.27±1.06          | 80.22±0.11          | 84.74±0.22          |
| + Personalized BN   | 96.11±0.02          | **75.36±0.04**      | 96.41±0.06          | **84.74±0.03**      | **82.02±0.04**      | **86.93±0.02**      |

*Table2: Experimental Results for FL with Feature Skew on Digits*

From the results in Tabel 1 and 2, we find that our weighted FedAvg exhibits better performance than FedPer in some cases. When we also apply personalized BN layers for our Weighted FedAvg, it outperforms FedBN as well, which means our algorithm can combine with personalized methods and improve their performance. We will illustrate that Weighter FedAvg is a personalized method in the introduction part.

Reviewer comment (Novelty):
What are some of the novelties/challenges of this analysis compared to Cao et al. (2022) and Kou et al. (2023)? It seems that many of the proof techniques such as using coefficient dynamics is similar to these works.

Author response:
Our analysis introduces several novel aspects and addresses unique challenges compared to the work of Cao et al. (2022) [2] and Kou et al. (2023) [3], which primarily focus on single neural network models. Our research extends the feature learning theory to the more intricate setting of federated learning, involving multiple neural networks and the integration of weight averaging operations. Although inspired by the theoretical frameworks in [2,3,15], our approach diverges significantly due to the distinct nature of our problem setting:

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.

Reviewer comment (Theory):
Can the analysis be extended the case where clients have signals of different kinds (say with some bound on the dissimilarity of signals at a particular client)? Restricting the data at one client to belong to only one signal is a bit too restrictive in my opinion.

Author response:
The adaptability of our theoretical framework to scenarios where a client processes multiple signal types is indeed feasible. A key condition for this extension is the orthogonality of the signal vectors corresponding to different classes within a single client. In our current analysis, we treat the case where a client is associated with only one signal as an illustrative but informative example.

Reviewer comment (Writing):
How is $P$ defined in Section 3.2?

Author response:
In Section 3.2 of our paper, $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 manuscrip.

Reviewer comment (Theory):
What is the difference between Lemma 5.3 and Theorem 4.4? It seems they are almost stating the same thing.

Author response:
Lemma 5.3 and Theorem 4.4 in our paper indeed convey similar concepts. We have clarified in the revised version of our paper that Lemma 5.3 is essentially a restatement of Theorem 4.4.

Reviewer comment (Writing):
I would also advise the authors to use different superscripts when referring to local iteration steps versus global averaging steps. One common notation is to use $w_k^{(t,r)}$ for the local models where $t$ refers to the round number and $r$ refers to the iteration number in the $t$-th round. In the current analysis the authors use $t$ for both the local update (Eq. (4)) and global update (Eq. 5) which is a bit confusing.

Author response:
We appreciate the reviewer's advice on improving the clarity of our notation, particularly in differentiating between local iteration steps and global averaging steps. We commit to revising the notation in our manuscript to better distinguish between these two types of steps. A common notation, as suggested, involves using different superscripts to represent the round number and the iteration number within a given round for local models.

Reviewer comment (Writing):
$F_{-1}(W_{+1}, x) $ should be $F_{-1}(W_{-1},x) $ in the line above Eq. (2).

Author response:
We appreciate the thoroughness of the review and the opportunity to improve the accuracy and professionalism of our paper. These modifications have been implemented in the revised version of the manuscript

Reviewer comment (Writing):
$L_{D_k}$ in Lemma 5.3 is not defined. Maybe the authors meant $L_{D_k}^{(0-1)}$.

Author response:
We appreciate the thoroughness of the review and the opportunity to improve the accuracy and professionalism of our paper. These modifications have been implemented in the revised version of the manuscript

Reviewer comment (Writing):
"Provably" and not "Probably" in the heading of Section 4.2.

Author response:
We appreciate the thoroughness of the review and the opportunity to improve the accuracy and professionalism of our paper. These modifications have been implemented in the revised version of the manuscript