Understanding Convergence and Generalization in Federated Learning through Feature Learning Theory
Authors: Wei Huang
Ye Shi
Zhongyi Cai
Taiji Suzuki
Venue: ICLR 2024 poster (2024)
Keywords: Federated Learning
Feature Learning Theory
FedAvg
Primary area: societal considerations including fairness, safety, privacy

Abstract: Federated Learning (FL) has attracted significant attention as an efficient privacy-preserving approach to distributed learning across multiple clients. Despite extensive empirical research and practical applications, a systematic way to theoretically understand the convergence and generalization properties in FL remains limited. This work aims to establish a unified theoretical foundation for understanding FL through feature learning theory. We focus on a scenario where each client employs a two-layer convolutional neural network (CNN) for local training on their own data. Many existing works analyze the convergence of Federated Averaging (FedAvg) under lazy training with linearizing assumptions in weight space. In contrast, our approach tracks the trajectory of signal learning and noise memorization in FL, eliminating the need for these assumptions. We further show that FedAvg can achieve near-zero test error by effectively increasing signal-to-noise ratio (SNR) in feature learning, while local training without communication achieves a large constant test error. This finding highlights the benefits of communication for generalization in FL. Moreover, our theoretical results suggest that a weighted FedAvg method, based on the similarity of input features across clients, can effectively tackle data heterogeneity issues in FL. Experimental results on both synthetic and real-world datasets verify our theoretical conclusions and emphasize the effectiveness of the weighted FedAvg approach.

Published as a conference paper at ICLR 2024
UNDERSTANDING CONVERGENCE AND GENERALIZA- TION IN FEDERATED LEARNING THROUGH FEATURE LEARNING THEORY
Wei Huang
Ye Shi*
RIKEN AIP
Shanghaitech University
wei wei.husr huang wr@riken. ip
abiyy@@shanghaättchheeee
Zhongyi Cai Shanghaitech University
Taiji Suzuki
The University of Tokyo RIKEN AIP tai taijitmist.. ac- ip
cai salkhyeshannhaatechh.ee cn
ABSTRACT
Federated Learning (FL) has attracted significant attention as an efficient privacy- preserving approach to distributed learning across multiple clients. Despite ex- tensive empirical research and practical applications, systematic way to theo- retically understand the convergence and generalization properties in FL remains limited. This work aims to establish a unified theoretical foundation for under- standing FI. through feature learning theory. We focus on a scenario where each client employs two-layer convolutional neural network (CNN) for local training on their own data. Many existing works analyze the convergence of Federated Averaging (FedAvg) under lazy training with linearizing assumptions in weight space. contrast, our approach tracks the trajectory of signal learning and noise memorization in FL. eliminating the need for these assumptions. We further show that FedAvg can achieve near-zero test error by effectively increasing signal-to- noise ratio (SNR) in feature learning, while local training without communication achieves a large constant test error. This finding highlights the benefits of commu- nication for generalization in FL. Moreover, our theoretical results suggest that a weighted FedAvg method, based on the similarity of input features across clients, can effectively tackle data heterogeneity issues in FL. Experimental results on both synthetic and real-world datasets verify our theoretical conclusions and emphasize the effectiveness the weighted FedAvg approach.
INTRODUCTION
Federated Learning (FL) (McMahan al., 2017: Yang al., 2019) distributed learning frame- work that involves collaboratively training a global model with multiple clients while ensuring pri- vacy protection. The pioneering work, Federated Averaging (FedAvg) (McMahan al., 2017), learns the global model by aggregating the local client models. It achieves satisfactory performance with both IID and non-lID client data. Since then, with the growing importance of privacy protec- tion, federated learning emerged hot area of research and application (Bonawitz et al., 2019; Li et al., 2020a;b; 2023b). Numerous technologies and methods began to surface, encompassing a wider range of fields and application scenarios, including smartphone applications (Li et al., 2019: Yang et al., 2021), Internet of Things (IoT) devices (Nguyen et al., 2021; Khan et al., 2021), and more.
Despite federated learning's success, there has been limited research on understanding its conver- gence and generalization properties when neural networks used client models. From theoret- ical perspective, many recent studies focus on the convergence of FedAvg using wverparaaettrizeed neural networks. For example, Song et al. (2023) show FedAvg converges linearly to a solution that achieves (almost) zero training loss through a special initialization strategy and network overparam- eterization. Besides, Huang al. (2021) present class of convergence analysis for FL with the
"Corresoonding author

---

Published as a conference paper at ICLR 2024
help of a neural tangent kernel. Similarly, Deng al. (2022) prove that FedAvg can optimize deep neural networks with ReLU activation function in polynomial time. However, these analyses often rely on linearizing assumptions in the weight space. Moreover, current generalization analyses in FL. draw from classical theories like VC dimension (Hu et al., 2022), PAC Bayes (Sefidgaran et al., 2023), and convex loss (Chen al., 2021). These theories often require strong assumptions, such as Lipschitz continuity and (strong) convexity. As result, they may not fully capture the success of FL. The above works analyze either the convergence or generalization properties of FL, lacking unified theoretical framework for FL with the neural network.
To fill the theoretical gap in FL, we delve into the feature learning theory (Allen-Zhu & Li, 2023; Cao et al., 2022) to establish a unified theoretical framework to analyze both the convergence and generalization properties of FedAvg. We introduce a two-level data distribution that simulates data collection in FL. For the input feature, we use a two-patch model that mimics the latent structure of an image, containing both signal patches and noise patches. The signal patch is correlated to the label while the noise patch represents the background of the image and has nothing to do with the label. We then study the dynamics of a two-layer convolutional neural network in each client with FedAvg by comparing the learning speed of signal and noise. Once we show that the neural network in FedAvg achieves near zero training error, we then compare the generalization ability between FedAvg and the local algorithm where no communication is allowed. Under certain con- dition, FedAvg can converge to the global minimum with small test errors, while the local training achieves a high test error as it memorizes the noise from the training data. This demonstrates the advantage of federated learning over standard local learning (without communication) with the help of communication between clients. Our contributions are summarized follows:
We demonstrate the convergence properties of the FedAvg algorithm from a feature- leatning per- spective. This result goes beyond the conventional lazy training approach, which relies lineariz- ing assumptions in weight space to ensure network parameters remain close to their initialization values throughout training. In contrast, our analysis does not depend on such assumptions and reveals that FedAvg can achieve a near-zero training error within polynomial time.
To the best of our knowledge, we are the first to rigorously establish the exact generalization gap between FedAvg and local algorithms by leveraging effective signal-to-noise ratio in feature learning theory. Specifically, under the same conditions, we demonstrate that the test error of FedAvg exhibits asymptotic behavior of o(1), whereas local training converges to @(1) test error. Additionally, our theoretical result reveals the importance of cooperation in FL based on the similarity between input data among clients from feature learning theory.
We have conducted comprehensive experiments on both synthetic and real-world datasets. The simulation results validate our theoretical assertion regarding the superiority of weighted FedAvg. a personalization method, and provide empirical support for our insights into the cooperative aspects of this approach.
2 RELATED WORK
Federated Learning In Fl., the global model generated by weighted aggregation of the up- loaded local models, where the weight is proportional to the local data sizes (McMahan et al., 2017; Yang et al., 2019). However, when the data among clients is not independently and identically dis- tributed, the global model produced by this aggregation mechanism may fail to fit either the global data distribution or the local data distribution (Zhao et al., 2018; Li et al., 2022; Cao et al., 2023; Cai et al., 2023; Li et al., 2024). To address issues arising from data heterogeneity, various aggre- gation methods have been proposed. For instance, Li et al. (2022) leveraged validation sets to learn aggregation weights tailored to each client's 's local data distribution. In a different approach, Xu al. (2023) addressed an optimization problem for each client, identifying an optimal linear combination of local classifier heads serve as the new local classifier. Mcanwhile, Li al. (2023c) utilized proxy dataset to learn both the aggregation weights and the global weight shrinking factor. Ad- ditionally, Ye et al. (2023) adjusted aggregation weights based on the disparities between local and global category distributions to enhance performance.
Feature Learning Theory Different from neural tangent kernel (Jacot et al., 2018) or lazy training Chizat et al. (2019), where weight updates during gradient descent training are relatively small,
2

---

Published as a conference paper at ICLR 2024
feature learning theory (Allen-Zhu & Li, 2023; Cao et al., 2022; Yang & Hu, 2021) states that weight updates can be more substantial during gradient descent, allowing the network to learn patterns from the data. The feature learning has been widely explored to provide interpolation to various neural networks, such as graph neural network (Huang et al., 2023), convolutional neural network (Cao al., 2022; Kou et al., 2023), vision transformer (Jelassi et al., 2022; Li et al., 2023a); and different training algorithms like Adam (Zou et al., 2023a), momentum (Jelassi & Li, 2022), OOD Chen al. (2023), and Mixup (Zou et al., 2023b; Chidambaram et al., 2023). Among the works related to feature learning, Kou et al. (2023) investigate the optimization and generalization of two-layer ReLU CNNs in centralized training scenarios. In contrast, our research delves into the application of feature learning theory in Federated Learning (FL), introducing additional techniques tailored for the FedAvg algorithm.
Theoretical Analysis of FL The convergence analysis of federated learning has been gaining increasing attention. Initially, researchers explored convergence in FL without employing neural network models for clients. For instance, Li et al. (2020c) analyzed the convergence of FedAvg on non-iid data, establishing convergence rates for strongly convex and smooth problems. Wang al. (2020) provided a comprehensive framework for analyzing the convergence of federated opti- mization algorithms, considering heterogeneous local training progress at clients. Subsequently, the advent of over-parameterize neural networks sparked a new wave of convergence analysis for FL (Huang et al., 2021; Deng et al., 2022; Song et al., 2022). These convergence analyses may highly rely on linearizing assumptions in weight space to ensure network parameters remain close to their initialization. In terms of generalization in FL, much of the existing research has been grounded in classical theories that often require strong assumptions, such as Lipschitz continuity and (strong) convexity. For example, works like (Mohri al., 2019; Chen al., 2021; Masiha al., 2021; Hu et al., 2022; Sun et al., 2023) have presented generalization results in FI. based on these assump- tions. Another research strand delved into variance reduction strategies for heterogeneous FL, as exemplified by works like (Woodworth et al., 2020; Murata & Suzuki, 2021; Oko et al., 2022). In contrast to these approaches, our work focuses on both the convergence and generalization of neural networks in FL. based on the feature learning theory, offering unified theoretical framework for FL.
3 PRELIMINARY
3.1 NOTATIONS
We use lower bold-faced letters for vectors, upper bold-faced letters for matrices, and soan-bold-fccee letters for scalars. For vector v. its f2-norm is denoted as ||vlee For a matrix A. we use ||A||2 to denote its spectral norm and || | || for its Frobenius norm. When comparing two sequences, we employ standard asymptotic notations such as 0(.). o().. 2(()), and 0(-) to describe their limiting behavior. We use Ö(), ((.), and ((-) to hide logarithmic factors in these notations respectively.
Lastly, sequences of integers are denoted as [ni]
3.2 DATA MODEL
We assume there are K clients, and each client is associated with distribution Dr. In this work, we adopt a two-level distribution framework for FL (Yuan et al., 2021; Hu et al., 2022; Wang et al., 2021). In the first level, we select a signal vector f4 for client h through pa ((u(1) (2) (C), where P is discrete distribution assigning probabilities each feature vector (e)) for [C]. For simplicity,
assume every Mrc for E [C] are orthogonal to each other. In the second level, the input feature associated to each client are sampled based on μ4g In particular, utilize two-patch model for feature generation (Allen-Zhu& & Li, 2023; Shen al., 2022; Huang et al., 2023; Kou et al., 2023; Cao et al., 2022). Conditional on the Rademacher random variable yk,t E (-1, 1). for each client ls e [K]. the input with index is generated from:
X = [ x
(1)
pyr Ex,el,
(1)
where noise vector Ex. ~ N(0, of oI - E pl() p (1)) with op being the strength of
noise. Note that Σ (0) p(c) / is introduced to ensure noise vector orthogonal to signal
3

---

Published as a conference paper at ICLR 2024
vector. We recover IID by setting μ4 Ml MK. Note that two-patch model that mimics the latent structure of an image. The signal patch is correlated to the label while the noise patch represents the background of the image and has nothing to do with the label.
3.3 NEURAL NETWORK MODEL
We introduce two-layer CNN model, denoted Sk. for every client & [K], which utilizes
ReLU activation function, defined a(=) = max (0. =). Given the input data x, the CNN's output
represented as f&( (WW,)) F+1 BetWW+1. F+1(W+1.x* x1 F-1 (--1W-11 -1. The sign is associated with the neuron of second layer which are fixed as either +1 or -1. Furthermore, we define F41 W1..))
and x) follows:
(2 j. x)
Σ o
(2)
where is the width, and Wk. Rd refers the weight vector of the first layer client k. The symbol W collectively represents the model's 's weights. Moreover, each weight in the first layer is initialized from random draw of Gaussian random variable, Wj. (O, o Laxa) for all r E [m] and {-1,1}, with regulating the initialization magnitude for the first layer's weight.
3.4 OBJECTIVE FUNCTION AND FEDAVG
We denote the training set as S₁ - {xi,i 3 The for A E [K], and overall training set as S A {S We aim to minimize the empirical cross-entropy loss function:
1 Σ Lf xk,));
(s(W)=
Ls. (w)) with Ls. (W&)
T
where (yf(WW x)) log(1 + exp(-f((W x)y)) is the logistic loss. Next, we introduce the FedAvg training algorithm, which comprises two types of updates: gradient descent updates and weight averaging updates. In each iteration, we first apply gradient descent by E epochs to the CNN fü for client & [K]:
(++1) kij.r =
wk.
V We Ls (s() = Wkj.r 7&m
kit o' ((w (1) k.j.r (x.)))
1J
(1t ki
((w
k.r Va.s px))juu
(4)
where we define the loss derivative as lh A ll (JK, = ssp(-ve,ree.l update through weight averaging can be expressed follows:
Then the server
((
(0) Rjrrt
(5)
K where Wj.r is the weight server, Ek=1 n& is the total number of training samples. After a certain number H of training loops until convergence, we examine the training through test error defined as LD-1 (w(()) (y x)<0).
MAIN RESULTS
4.1 COFFFICIENT DYNAMICS FOR FEDERATED LEARNING
In this section, we introduce our key theoretical findings that demonstrate the convergence and generalization of federated learning. Through the application of the gradient descent rule outlined
in Equation (4) and Equation (5), we observe that local weights w (!) kkj.r for all client k [K] and
4

---

Published as a conference paper at ICLR 2024
server weights w j.r are linear combination of the random initialization W Mij.er the signal vectors Mg, and the noise vectors Êx, for E [mx] and & E [K]. Consequently, for [m], and [K], the decomposition of weight vector iteration can be expressed:
(0)
Be
(t) k Wir + Th. Th.. wlw 122 pk
(0)
Pkck Jna E (6)
R.j.r
K
w j.r
Riji kii +
pe + Pha alluell 2
> 1=1
j.r. Th and Pk. j.r. Pk.jor for k., [K] j E (-1,1), E [m], [mx]. serve coefficients. We refer to Equation (6) and Equation (7) as the signal-noise decomposition
where (0)
(0)
(0)
j.r
of (1)
"ki.r" The normalization factors || and ||Ex,el1² are introduced to ensure that (2)
"e,k' Jir
(wk)... Rigar px)). and Pk,k' 2 (w Evr,). To facilitate a fine-grained analysis for the evolution
of coefficients, we introduce the notations (( J.r,' - (0) j.r.' I (Pk,e 2 0), kk,k'j.. j.n Pk,k' J.n [( iint S 0). Consequently, we further express the vector weight decomposition (6,7) as:
-
nga
(a
(0) Lak jor Wk Kjr +
(t)
Tk. for | 2e +
(0)
Me
>
(wbl Jaa + RR sna 1l (w
k=11
(8)
nk mu -' + Pk.s (ll Ex. 2x. 2
(a k=1
W (0) k.j.r + 1e
(1)
Our first key result is to turn the dynamics of weights into the dynamics of coefficients. To analyze the feature learning process of federated learning during gradient descent training, introduce an iterative methodology. based on the signal-noist decomposition in (8) and (9) and gradient descent update (4) and (5). The following lemma offers us means monitor the iteration of coefficients:
Lemma 4.1. The coefficients (1) (1
"I j.r WkK jir 800 R.M. in decomposition (8), for all k, E [K],
follow the update rules governed by:
(0)
(0)
j.r' Vk,k' J.r.s k. jor.
0,
(10)
J.r = (t)
(t+1)
- n&m
"
k2 o' ((w
k.j. yk : 11 12 = k);
(11)
Pk don, = Wk..
(1+1)
(1)
n
(()) o' ((w k.j.r
k.
(vk.s j)1(k' = k),
(12)
n&m
Pk.k J... LK. jir.
(1+1)
(t)
(() k.i
0' d' k.j.r & ER (vn.e [K) k).
(t)
(13)
nym
Furthermore, the coefficients Tke Viim Ikiiin in decomposition (9), for all k E [K], are up- dated according to the following rules:
711)
(t)
K
K
K
K
(1) Tk..r
Tkk
"K 1( px), Vkijon
Σ W=1
ny (0) W.k...
-0
Σ
Pkiia
1=1
W=l
(14)
The proof of Lemma 4.1 can be found in Appendix A. Lemma 4.1 provides us new approach to track the dynamics of neural network in FL. In particular, we study the learning speed ratio between signal and noise, represented by Equations (10-14).
4.2 FEDAVG PROVABLY BENEFITS FROM COMMUNICATION
Before stating the main results, we demonstrate our setting and assumption for federated learning. First, without loss of generality, we the same training size for all clients and fo norm of signal
5

---

Published as a conference paper at ICLR 2024
vector is the same, namely π n1 n2 nk. and || 12 = || 11:00111100 10
||ll|2. Therefore n = π.K Our analysis will be made under the following set of assumptions:
Assumption 4.2. Suppose that:
(1) The dimension d is sufficiently large: d= ((max(nn , k
(2) The size of training sample and width of CNNs m adhere to n. = ((polylog(d)), m ((polylog(d); All network model share the same initialization.
(3) The learning rate yn satisfies S O
K min(llulll .07 o,72d-1)). The standard deviation of Iullz'):
Gaussian initialization og is chosen such that ((n)/(o,d) S ou S Õ(1) (0100(mi(() viil((apV))= (
Brief explanations for each assumption are provided: (1) The requirement for high dimension in our assumptions specifically aims to ensure that learning occurs in an adequately wver-paameeerizee setting. (2) It's necessary for the sample size and neural network width be least polylogarithmic in the dimension d. This condition guarantees that certain statistical properties of the training data and weight initialization are maintained. (3) The condition on n and initialization dp is to ensure that gradient descent can effectively minimize the training loss.
Finally, we introduce a critical quantity that plays an important role in our analysis. First, we define the local signal-to-noise ratio (SNR) for each client as SNR; = (M,ll2//o, and the effective
signal-to-noise ratio in FedAvg, represented as SNR+
k'=1 (1 SNR,
VX
where we define X = - 1|
(es) The signal-to-noise ratio serves as a measure of the relative learning speed between the signal and noise. Given the aforementioned assumptions and definitions, present our main result on feature learning of FedAvg as follows:
Theorem 4.3 (Convergence of FedAvg). Suppose e > O, and let T = - Knno, + --'m'mmo,"" "). Under Assumption 4.2, then with probability at least d d=1, there exists 0 Osts T such that the training loss of FedAvg converges to i.e., T.s(W(!)) <.
Theorem 4.3 demonstrates the global convergence of FedAvg on the training dataset by tracking the signal learning and noise memorization trajectory. We demonstrate the detailed proof sketch for convergence through feature learning theory in the next section and leave the complete proof in Appendix B. Based on results in Theorem 4.3, we illustrate the generalization of FedAvg through the following theorem:
Theorem 4.4 (Generalization of FedAvg). Under the same assumption as Theorem 4.3, there exists (1(1), the trained model by FedAug achieves a small test loss:
S t T such that if 11x2111111
Koyd
(w()) S exp (-c -x211111 Koja
(pe Here, Xk = =E#= 11 1 Otherwise, if ni millell = 0(1, the ofa
trained model by FedAvg achieves a large constant test error Lb-. (D-((W()) = 211.
Theorem 4.4 further demonstrates the separation condition based on the effective SNR for the gen- eralization of FedAvg. In this result, the signal-to-noise plays a significant role. When the number of sample vi and effective SNR SNR; is strong enough, FedAvg can achieve a small test error. Con- versely, if these conditions not met, FedAvg may not generalize effectively to the test samples.
Based on theorem 4.4, we draw a comparison between FedAvg and local training. Note that the SNR for FedAvg has an additional factor X = =#=1 1| similarity of input features across clients and can also be interpreted ity. In the context of local training, X = 1, which means the effective SNR simplifies to the vanilla 7x2 [a l1 ?(1) and Tlll ogd
[M,) This factor measures the overall
metric for data heterogene-
SNR. Given the conditions
O(1), FedAvg achieves a small test
Koğd loss, while local training has large test error.
Lastly, we interpret the effective SNR. In the case of the IID setting, we find that X = K for A E [K]. which can boost the performance of FedAvg mostly. On the other hand, when we consider the non-IID, the smallest value that Xk can achieve is 1, when all the input features are different from each other. In this case, FedAvg performs worse than local training. These results point out the importance of communication in FedAvg, suggesting that weighted FedAvg can further improve the performance of FL., we examine this understanding in the experiment section.

---

Published as a conference paper at ICLR 2024
5 PROOF SKETCH
In this section, we provide proof sketch for Theorem 4.3 and Theorem 4.4. The key idea is analyze the coefficients update function for l E [K]. To achieve our optimization result, employ a two-stage dynamic analysis. Next, we demonstrate the generalization analysis based on the expected test error.
5.1 ITERATIVE ANALYSIS OF COEFFICIENTS FOR CONVERGENCE
We adopt two-stage dynamics analysis, inspired by Cao et al. (2022); Kou et al. (2023), to track the behavior of these coefficients in FL based on the behavior of the loss derivative.
Stage 1 In the first stage, we adopt small initialization for weights such that the neural network at initialization has constant level cross-entropy loss derivatives for all training data and clients, namely 2.00 f (WW()) xx,:)l ((1) for all i [7]] and & [K]. This is guaranteed based on assumption (4.2) regarding dn- Consequently,
replace the loss derivative ( by their constant upper and lower bounds in the analysis of the coefficients dynamies in Equations (10 14). This simplifies the analysis significantly, and then we present our main conclusion in Stage 1:
Lemma 5.1. Under assumption 4.2, there exists Ti = ( "Timo, 2d such that
max, Tkiir O SNR)) E {+1), and ls E [K].
(Ti)
Wkggnil (() some j [=1}, r [m]) E [n] and all k E [K].
* maxy,r Ōk.j.. = 0(1) for all i e [n] and k E [K].
The proof can be found in Appendix B.1. Lemmas 5.1 leverages the period of training when the derivatives of the loss function are of constant order.
It's important note that weight averaging plays significant role in differentiating the learning speeds between signal learning and noise memorization during this initial stage. In each client's 's local update via gradient descent, the learning speeds for signal learning and noise memorization are determined by || l2 and ||E||, respectively (Cao et al., 2022). Our theoretical analysis further reveals that, after the application of weight averaging, these learning speeds are approximately deter- mined by X || and || la. respectively. This distinction arises because, during communication, clients with identical signal types can share their learning signal and preserve the signal magnitude. In contrast, the independence of noise vectors across clients means that their magnitude has sig- nificant decrease. Consequently, communication in FedAvg slows down the learning speed noise memorization, thus enabling FedAvg to focus on signal learning.
Stage 2
Building on the results from the first stage, we study the second stage of the training process. During this phase, the derivatives of the loss no longer remain constant. We demonstrate that the training eтrоr can be minimized to an arbitrarily small value. Importantly, the scale differ- ences established during the first stage of learning continue to be maintained throughout the training process, and the convergence of the FedAvg in all clients can be guaranteed:
Lemma 5.2. Suppose 0, let T, Tı be defined in Theorem 4.3 and Lemma 5.1 respectively. Let W be the collection of CNN parameters Wi, = w j.r + log(2/e) Σkn n. SNR T +
E. 1 yk.i)
St.i 1 [1
Then under the same conditions as Theorem 4.3, for any t E [Ti, T).
Eiem Ls (w("))
(WT)) W 11
holds that: +1
+
W +1)
Lemma 5.2 provides a convergence result for FedAvg. It suggests that as the training time increases, the training loss approaches arbitrarily small value.
7

---

Published as a conference paper at ICLR 2024
5.2 TEST ERROR ANALYSIS FOR GENERALIZATION
Finally, we consider new data point (Xx, yt) ~ Dx drawn from the distribution for client k. Theorem 4.4 indicates that FedAvg can achieve low test emror when both the sample size and effective SNR are sufficiently large. However, in the absence of these conditions, the test error achieved by FedAvg may exceed certain constant. The proof can be found in Appendix B.3.
EXPERIMENTS
In this section, we first validate our theoretical findings through numerical simulations using syn- thetic data. Then we demonstrate the effectiveness of FedAvg real-worke datasets.
6.1 FEDERATED LEARNING ON SYNTHETIC DATA
In this section, we conduct experiments to substantiate our theoretical findings using synthetic data drawn from the two-level distribution in both IID and non-lID settings. For these experiments, the client number is set to be K 20 and all of them selected in each communication round. Each client is equipped with a two-layer ReLU CNN descried in Section 3.3, where width m = 50.
Then, establish our synthetic dataset. The training data size is set to Atrain 100 and the testing data size is set to Rlest = 2000 with instance dimension to d = 1000 for each client. We further set the signal strength ||u|| 2 and noise variance Sp 1 for all clients. We trained models for E 5 epochs the local training phase and conducted R = 100 communication rounds in total. All models trained with the SGD optimizer whose learning rate 1) = 1.
Feature Learning in IID setting. Under the IID) setting, all clients share identical signal vectors, denoted as, μ4 = #2 ... 4420- Figure displays the training loss, training accuracy, test loss, and test accuracy for both FedAvg and SingleSet (local training) the dataset of the first client. Notably, both FedAvg and local training successfully achieve small training loss and high train accuracy, verifying results presented in Theorem 4.3. Furthermore, FedAvg displays superior test accuracy compared to local training, validating theoretical results in Theorem 4.4.
Train Loss
Train Accuracy
Test Loss
Test Accuracy
1.5
FadAug Singleset
1.0
1.0
0.8
0.8
0.5
0.6
0.5
0.6
Figure 1: Convergence behavior comparison of train loss, train accuracy, test loss, and test accuracy on synthetic data. Both local training and FedAvg demonstrate convergence on the training set. FedAvg outperforms significantly on the testing set.
6.2 WEIGHTED FEDAVG ON REAL- WORLD DATASET
To verify our theoretical conclusion in FL with data heterogeneity issues, we conducted experiments in real-world datasets. In this case, the conventional aggregation strategy, which assigns model ag- gregation weights based on local data sizes, may not be optimal due to the diverse signals across clients. To address this, we enhanced collaboration between clients that exhibit similar signals. Par- ticularly, we employ prototype similarity as a measure to approximate the signal similarity between clients. This prototype similarity is then utilized to determine the aggregation weights for each client, ensuring a more effective aggregation process.
Weighed FedAvg For each client k. we learn a well-trained model fa that is composed of a feature extractor ge and linear classifier eg. Supposing there are overall M classes, a prototype of the j-th class in client & is defined as P/= | Ex,y)ssu, gi (X). We concatenate all class prototypes in client get the client prototype P& concatue[M] [P]]] This client prototype is updated after
8

---

Published as a conference paper at ICLR 2024
Table 1: Experimental Results for FL with Label Distribution Skew on CIFARIO and CIFAR100.
CIFARIO 10
CIFARION
Methods
Pathological
Dirichlet
Pathological
Dirichlet
SingleSet
20.08+0.13
38.60=0.02
6.59+0.18
6.97#0.11
FedAvg
40.1647.18
33.11=7.06
13,99:0.79
17.3440.47
FedPer
38.0945.73
36.52+7.13
29.10#2.42
17.6340.93
FedBN
65.71+0.67
61.25+3.78
42.5420.53
25.6443.46
Weighted FedAvg
60.32+0.31 67.30+0.31
60.52+0.30
31.88+0.56
18, 02.422
+ Personalized BN
65.31+1.23
43.51+0.50
26.11+1.77
Table 2: Experimental Results for FL with Feature Skew on Digits
Digits
Methods
MNIST
SVHN
USPS
SynthDigits
MNIST-M
Avg
SingleSet
94.38+0.07
6..25=1.07
95.16.2.12
80.31
77.77+0.47
82.00+0.40
FedAvg
95.87+0.20
12.86=1.49
95.56:1.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.6620.22
84.27+1.06
80.2240.11
84.7410.22
+ Personalized BN
96.11#0.02
75.3620.04
96.4120.06
84.7420.03
82.02+0.04
86.93+0.02
the local training phase of each involved client. Different from the normal model aggregation in FedAvg, we personalize the aggregation weights for each client model. Specifically, we determine the aggregation weight $g.k' of each involved client k for client k' based on the client prototypes, exp(sim(ps. Pp))T) where
i.c., 5k,k' Σe exp(sim(Pe.PPi/7
0 is a temperature parameter and sim(-) is the cosine similarity function. The default value of 7 is set as 1. Then the server conducts personalized weighted aggregation to produce weight Wt for [K] with the formula W
Σ
x==1= $k.k W.r kr-
Experimental settings We conducted experiments on three image classification datasets: CI- FARIO, CIFAR 100 (Krizhevsky al., 2009), and Digits. These were chosen to address FL scenarios with label distribution imbalance and feature shift data heterogeneity issues. For FL with label dis- tribution skew using CIFARIO and 100, clients possess distinct object category distributions. We adopted similar "Pathological" and "Dirichlet" data partition strategies similar to (Marfoq al., 2022; Li et al., 2023b) to make the class distribution among clients distinct. In the "Pathological" partition, we sampled eight/twenty classes for each client. And, in the "Dirichlet" partition, we em- ployed a symmetric Dirichlet distribution with the default parameter a = 0.8 to divide the dataset into sub-datasets. Conversely, for FL with feature skew using the Office dataset, each client has a roughly similar category distribution but originates from diverse domains. We introduced two baselines FedBN (Li al., 2021) and FedPer (Arivazhagan et al., 2019), which separately person- alize the Batch Normalization (BN) layers and the Classification Head for each local model. For fair comparison, we also personalize all BN layers for our Weighed FedAvg. For more in-depth exploration of this experiment setting, please refer to Appendix C.
The experimental outcomes for FL with label distribution skew are provided Table 1, while the results for FL with feature skew are showcased in Table 2. The results confirm that FedAvg can outperform local training in almost all datasets and Weighted FedAvg can further enhance the performance by increasing the communication efficiency, aligning with our theoretical predictions.
Results
CONCLUSION
This work establishes a unified theoretical framework for the analysis of convergence and gen- eralization Federated Learning by leveraging feature learning theory. Based on a signal-noise decomposition for weights, we can track the dynamics of signal learning and noise memorization, providing convergence results. Furthermore, we demonstrate the generalization gap between Fe- dAvg and local training through effective signal-to-noise ratio. This result reveals how FedAvg benefits from communication. As a pioneering study in feature learning within the context FL, theoretical framework is limited to examining FedAvg with gradient descent optimization and a certain data generalization model. Future work can extend our framework to consider stochastic gradient descent and more advanced settings in FI..

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported by JSPS KAKENHI Grant Number 24K20848 and RIKEN Incentive Re- search Project 00847720230106201 We thank the anonymous reviewers for useful suggestions improve the paper.
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. The Ninth Inter- national Conference on Learning Representations, 2021.
Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. The Eleventh International Conference on Learning Represen- tations, 2023.
Manoj
Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Fed- erated learning with personalization layers. arXiv preprint arXiv: 1w/.//99.0000.. 2019.
Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konečný, Stefano Mazzocchi, Brendan McMahan, et al. Towards federated learning at scale: System design. Proceedings of machine learning and systems, 1: 374-388, 2019.
Zhongyi Cai, Ye Shi, Wei Huang, and Jingya Wang. Fed-COg: Cooperation of online and offline models for severe data heterogeneity in federated learning. Advances in Neural Information Pro- cessing Systems, 36, 2023.
Yu-Tong Cao, Ye Shi, Baosheng Yu, Jingya Wang, and Dacheng Tao. Knowledge-aware federated active learning with Non-IID data. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 22279-22289, 2023.
Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convo- lutional neural networks. Advances in neural information processing systems, 35:25237-25250 2022.
Shuxiao Chen, Qinqing Zheng, Qi Long, and Weijie J Su. A theorem of the alternative for person- alized federated learning. arXiv preprint arXiv:2 2021.
Yongqiang Chen, Wei Huang, Kaiwen Zhou, Yatao Bian, Bo Han, and James Cheng. Towards un- derstanding feature learning in processing systems, 2023.
generalization. Advances in neural information
Muthu Chidambaram, Xiang Wang, Chenwei Wu, and Rong Ge. Provably learning diverse features in multi-view data with midpoint mixup. In International Conference on Machine Learning, pp- PMLR, 2023.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Local SGD optimizes over- parameterized neural networks in polynomial time. In International Conference on Artificial Intelligence and Statistics, pp. 6840-6861. PMLR, 2022.
Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In International conference on machine learning, pp. 1180-1189, PMLR, 2015.
Xiaolin Hu, Shaojie Li, and Yong Liu. Generalization bounds for federated learning: fast rates, participating clients and unbounded losses. The Eleventh International Conference on Learning Representations, 2022.
Baihe Huang, Xiaoxiao Li, Zhao Song, and Xin Yang. Fl-ntk: A neural tangent kernel-based frame- work for federated learning analysis. In International Conference on Machine Learning, pp- 4423 4434. PMLR, 2021.
10

---

Published as a conference paper at ICLR 2024
Wei Huang, Yuan Cao, Haonan Wang, Xin Cao, and Taiji Suzuki. Graph neural networks provably benefit from structural information: A feature learning perspective. arXiv preprint arXiv:2306. 13926, 2023.
Jonathan J. Hull. database for handwritten text recognition research. IEEE Trunsactions on pattern analysis and muchine intelligence, 16(5):550-554, 1994.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernl:: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning, pp. 9965-10040. PMLR, 2022.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37833 2022.
Latif U Khan, Walid Saad, Zhu Han, Ekram Hossain, and Choong Seon Hong. Federated learning for internet things: Recent advances, taxonomy, and open challenges. IEEE Communications Surveys & Tutorials, 23(3):1759-17999 2021.
Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting for two-layer relu networks. International Conference on Machine Learning, 2023.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11) (6(11):2278-2322 1998.
Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. The Eleventh International Conference Learning Representations, 2023a.
Hongxia Li, Zhongyi Cai, Jingya Wang, Jiangnan Tang, Weiping Ding, Chin-Teng Lin, and Ye Shi. FedTP: Federated learning by transformer personalization. IEEE Transactions on Neural Net- works and Learning Systems, 2023b.
Hongxia Li, Wei Huang, Jingya Wang, and Ye Shi. Global and local prompts cooperation via optimal transport for federated learning. arXiv preprint arXiv:2403.000 00041, 2024.
Li Li, Haoyi Xiong, Zhishan Guo, Jun Wang, and Cheng-Zhome Xu. SmartPC: Hierarchical pace control in real-time federated learning system. In 2019 IEEE Real-Time Systems Symposium (RTSS), pp. 406-418. IEEE, 2019.
Li Li, Yuxi Fan, Mike Tse, and Kuo-Yi Yi Lin. A review of applications in federated learning. Com- puters & Industrial Engineering. 149:106854, 2020a.
Shuangtong Li, Tianyi Zhou, Xinmei Tian, and Dacheng Tao. Learning to collaborate in decentral- ized learning of personalized models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9766-9775, 2022.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE signal processing magazine, 37(3):50-60 -0, 2020b.
Xiang Li. Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang. On the convergence of FedAvg on Non-IID data. International Conference on Learning Representations, 2020c.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. FedBN: Federated learn- ing on NID-IID features via local batch normalization. The Ninth International Conference on Learning Representations, 2021.
Zexi Li, Tao Lin, Xinyi Shang, and Chao Wu. Revisiting weighted aggregation in federated learning with neural networks. International Conference Machine Learning, 2023c.
11

---

Published as a conference paper at ICLR 2024
Othmane Marfoq. Giovanni Neglia, Richard Vidal, and Lactitia Kameni. Personalized federated learning through local memorization. In International Conference on Machine Learning. pp- 15070-15092. PMLR. 2022.
Mohammad Saced Masiha, Amin Gohari, Mohammad Hossein Yassaee, and Mohammad Reza Aref. Learning under distribution mismatch and model misspecification. In 2021 IEEE International Symposium on Information Theory (ISIT), pp. 2912-2917. IEEF, 2021.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effcieent learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1288 PMLR, 2017.
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In Interna- fional Conference on Machine Learning, pp. 4615 PMLR, 2019.
Tomoya Murata and Taiji Suzuki. Bias-variance reduced local sgd for less heterogeneous federated learning. International Conference on Machine Learning, 2021.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Dinh C Nguyen, Ming Ding, Pubudu Pathirana, Aruna Seneviratne, Jun Li, and H Vincent Poor. Federated learning for internet of things: A comprehensive survey. IEEE Communications Sur- veys & Tutorials, 23(3):1622-1658, 2021.
Kazusato Oko, Shunta Akiyama, Tomoya Murata, and Taiji Suzuki. Versatile single-loop method for gradient estimator: First and second order optimality, and its application to federated learning. 14th International OPT Workshop on Optimization for Machine Learning in NeurIPS2022, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Milad Sefidaran, Romain Chor, Abdellatif Zaidi, and Yijun Wan. Federated learning you may communicate less often! arXiv preprint arXiv:2306. 05862, 2023.
Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypemnetorkss. In International Conference on Machine Learning pp. 1489-95022 PMLR, 2021.
Ruoqi Shen, Sébastien Bubeck, and Suriya Gunasekar. Data augmentation as feature manipulation. In International conference on machine learning, pp. 19773-19808 PMLR, 2022.
Bingqing Song, Ioannis Tsaknakis, Chung- Yiung-YYu Yau, Hoi-To Wai, and Mingyi Hong. Distributed optimization for overparameterized problems: Achieving optimal dimension independent com- munication complexity. Advances in Neural Information Processing Systems, 35:6147-6160, 2022.
Bingqing Song, Prashant Khanduri, Xinwei Jinfen Yi, and Mingyi Hong. Fedavg converges to zero training loss linearly for overparameterrized multi-layer neural networks. International Conference on Machine Learning. 2023.
Zhenyu Sun, Xiaochun Niu, and Ermin Wei. Understanding generalization of federated learning via stability: Heterogeneity matters. arXiv preprint arXiv:2306.03882, 2023.
Chong Wang, Weiqiang Ren, Kaiqi Huang, and Tieniu Tan. Weakly supervised object localization with latent category learning. In Computer Vision-ECCV 2014: 13th European Conference, Zurich. Switzerland, September 6-12, 2014, Proceedings, Part VI 13, pp. 431-445. Springer, 2014.
Janny Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7621 1-7623, 2020.
12

---

Published as a conference paper at ICLR 2024
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al. A field guide to feder- ated optimization. arXiv preprint arXiv:2107.06917, 2021.
Blake E Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch local SGD for heteroge- neous distributed learning. Advances in Neural Information Processing Systems, 15:6281-6292 2020.
Jian Xu, Xinyi Tong, and Shao-Lun Huang. Personalized federated learning with feature alignment and classifier collaboration. The Eleventh International Conference on Learning Representations, 2023.
Chengxu Yang, Qipeng Wang, Mengwei Xu, Zhenpeng Chen, Kaigui Bian, Yunxin Liu, and Xu- anzhe Liu. Characterizing impacts of heterogeneity in federated learning upon large-scale smart- phone data. In Proceedings of the Web Conference 2021, pp. 935-999999 946, 2021.
Greg Yang and Edward Hu. Feature learning in infinite-widie neural networks. International Conference on Machine Learning, 2021.
Qiang Yang. Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions on Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
Rui Ye, Mingkai Xu, Jianyu Wang. Chenxin Xu, Siheng Chen, and Yanfeng Wang. Feddisco: Federated learning with discrepancy-awarn collaboration. International Conference on Machine Learming. 2023.
Honglin Yuan, Warren Morningstar, Lin Ning, and Karan Singhal. What do we mean by general- ization in federated learning? The Tenth International Conference on Learning Representations, 2021.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with Non-IID data, arXiv preprint (rXiv:1806000582 2018.
Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. Understanding the generalization of Adam in learning neural networks with proper regularization. The Eleventh International Conference on Learning Representations, 2023a.
Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. The benefits of mixup for feature learning. Proceedings of the 40th International Conference on Machine Learning, 2023b.
A COMPLETE PROOF FOR COFFFICIENT ITERATION
In this section, we provide the proof for Lemma 4.1. This lemma plays crucial role by transforming the perspective from tracking weights to learning feature coefficients. The main idea combines weight decomposition with gradient update rules. Before presenting the proof for Lemma 4.1, introduce several preliminary results concerning the inner product between noise vectors and signal vectors, both inter- and intra-client:
Signal vector versus signal vector Between different clients, namely and k' according to our assumption, μ4 is either orthogonal to pa or identical to pr.
Noise vector tion, specifically, Éx, N/ (N - ( p(()) (1)))
Noise vector It known that all vectors drawn from Gaussian distribu-
We employ the following lemma to demonstrate that noise vectors are almost orthogonal to each other:
13

---

Published as a conference paper at ICLR 2024
Lemma A.1 (Cao et al. (2022)). Suppose that ö > 0. and d log(4n/6), then with probability least - 8, we have
J02 - || S 2002d.
3
(E aw - Vog((nn /6)d.
Remark A.2. The total number of samples in this work is K= Thus the union bound is taking over n1
Signal vector versus Noise vector We consider the inner product between the noise vector and the signal vector. According to our setting of the noise vector, namely, Eki No o
I p(()) 13), the noise vector strictly orthogonal to the signal vectors.
The above statements collectively demonstrate that noise vectors and signal vectors are (nearly) orthogonal to each other. They not only imply that the coefficient decomposition is guaranteed to be unique with high probability but are also rooted in theoretical analysis. Now, are ready to demonstrate the Lemma of coefficient iteration.
Lemma A.3 (Restatement of Lemma 4.1). The coefficients (0) " ki Jina dnn in decom-
(0
(1)
position (8), for all k k E follow the update rules governed by:
J.r Kk, j.ki
(0)
(0)
(0)
(15)
mh
((+1) (1)
mmm Σ
(1) R.
(8) ( W. = -k);
(16)
(++1) (1)
((0) d' d'((( k,
(yK. j)i k),
(17)
n&m
Kk, M.r. KK
(1+1)
k,t
d'
R...t Ex. 12 -j]i( = k).
(18)
Furthermore, the coefficients Tkgj.r Vajia Pkiia in decomposition (9), for all is [K] are up- dated according the following rules:
K
K
K ny W Wkrii Okjina
K
(0) Tkgjr
ny
WW."" S. 1(p px),, Vkiinna
-0
n
(19)
Proof of Lemma A.3. First, we have proven that the vectors are both linearly independent and or- thogonal, within each client and across different clients, with high probability. Consequently, this ensures that the decomposition, as expressed in Equation (9), remains linear independent as the dimensionality increases
Then, demonstrate the dynamics concerning the gradient update for each client
(1) Kkj.r lgg.r
(t-1)
1) V Vwn.. Ls. (W(!))
(1) ag
o d'((w!!) (t)
(( Exii)) En,
leg.r
mm
o' ((w
Wiiir Wa.s pk)) juk
(t)
Tigm
At the same time, the weight decomposition can be expressed as:
Cok j.r War + Tk. allllwww 2 pn P amall
(8)
(2)
Wkj.r
==1
14

---

Published as a conference paper at ICLR 2024
By compared with the coefficients of w Weg.ir 14x, and Ex. on above equations, for all k k'
(0)
[K], we arrive the following results:
Th
(0) Th Tk
((+1)
(D) ki o'
=k),
m
1=1
(0+1)
"I
k,
EA. E (yn. 5) k))
ny
kk,k' Ok,k'
(++1)
(0)
n&m k,
p(0)
o'
d'
k... E (UK,I = - k),
(++1)
(+)
d.r ak. .i.""
Note that the gradient descent update (4) does not change the coefficients of w Wi.j.r However, will see later, the weight averaging operation will mix the initializations of weights among different clients. Furthermore, observe that G by the definition of the cross-entropy loss. Therefore,
(t+1) the signs of kk,.j... and Kk.kjj.r. can be determined by and yk,i-
(t++))
Lastly, illustrate the dynamics of weight averaging among clients. Examining the weight average update (5) and weight decomposition (9), we find that:
(1)
nkk (c)
W k.j.r
(0)
+ TkJ.r B.s.rllulll
(1))
1xll22 pa +
>
Eohhon + Pk..n,) Ex. 2
1=1
By comparing the coefficients of w kig.r pk: and Exis in the above two equations, we derive the following results:
(0)
K
Tkkk
(8)
kig.r
kig.
n
K
[()
n2y (t)
i .j. (
mk),
Jkjjr
n
#"=11 k==1
(()
Tiky (0)
Vkdin
Vi k.j...
n
K
-0
Ti
@kji.
B COMPLETE PROOF FOR COEFFICIENT DYNAMICS
B. STAGE ONE: INCREASE OF SIGNAL LEARNING AND NOISE MEMORIZATION
According to the behavior the loss derivative, we divide the entire training dynamics into two stages. In the first stage, the loss derivative is close to a constant. Thus, we can simplify the analysis of coefficient iteration dynamics by replacing the loss derivative with the corresponding constant value. Before we conduct the proof for the main lemma in stage one, list several useful preliminary lemmas for random initialization.
Lemma B.1 (Cao al. (2022)). Suppose that 8 then with probability least 8, we have
(0) 1( k.j.m (0)) (wh.r Evii) /log(4mKn//)ddoopz
Voog(4mKC/0jaul l|2.
15

---

Published as a conference paper at ICLR 2024
With Lemma B.1 hand, we ready demonstrate how the effect of initialization changes control the variance of the weights' initialization.
Lemma B.2. Denote that: 3 = 2 maxk,i..r (lww k.j.r Hx>,,
(0) k.j.r (k,))) Suppose that do S
ht//lo((KKm///// min {Ilulla', (oyvd)-1} then with probability least 1 8, we have 1.
Proof. By Lemma B.I, with probability at least 6, we can set an upper bound for 3 as follows:
B S log(4KmnC/0) Oo max maxx|1mu12,0,,,v
According to the exact condition for do. which is also part of Assumption 4.2:
oo [4vlog( (0log(4mmCC))) (4KmnC/)]]
mmin min {Ial15" (apvd)
(20)
conclude to verify the following inequality:
B S Mllog(4Kmmll) oo max( op Vy 1
(21)
L
Next bound the loss derivative in the first stage.
Lemma B.3. Suppose that (t) J.. O(1) and (t)
[K], and i for i [i]. then we that j. 0(1), d00 j.r. O(1) with k.
C
-(()) <
(22)
where C1 is a positive constant.
Proof. According to the assumption that for 0 St Ti. we have max Kk.k'.j.r O(1 and (0)
(0)
(1 K.ki O()) and max j.r 1kkkk' = O(1) for all i [n]], the output function
j.
satisfies:
F)(W()) xx) To x x(1) + ow x()))))(()))
(w kwggrr p1), (1)
(0)
Tkgir (w 10)
wkg.r E) (v(()) (e 1obg
max
s, J. 10 1 g
(t)
max
0(1)
where (a) by the weight decomposition (6), (b) is by the definition of (3, and (c) by Lemma B.2.
Next, we bound the loss derivative for t E [0, Tij as follows:
(()) exp( fwy)
1+ exp(exffp O(1)
(23)
Finally, we show the lower bound for - (())
(()) exp(-Sitk)
1 expe (11y))
1 + exp (frv)
1 1+ +O1))
Similarly, we can obtain
5 - (0)) (W) < 1,
by the relation w -K
((()) (0) kgj.
and K O())
16

---

Published as a conference paper at ICLR 2024
With the above results at hand, ready present the outcomes of the first stage.
Lemma B.4 (Restatement of Lemma 5.1). Under assumption 4.2, there exists total number of local update T R, E (nn "Kno,, ?dd such that
(Ti) (n SNR for j E {+1), and ls E [K].
max, T Tkgj.
Whegonal = (())for some j E {+1], r [m]], i E [n]] and all k E [K]
maxjjr Ok.j,r.3 = O(1) for all i E [x] and k E [K].
Proof of Lemma B.4. We first consider noise memorization. Iet
Define that for each i E [m]] and (e)
l E [K].
( K. max,,r PPk,kj.r jora max,,r (v J.r. "i J.r.). Similarly, we define
(() the coefficient after weight average (1) K, max,,r PR, jr,, = maxj,r Whik'snn P joad.
Clearly we have 4 (0) k = 0 for all E [n]] and ls k E [K] by definition. Then by Equation (12)
and Equation (13) we have:
V (1+1) = (1) +
(() kii
1..1--o' d' (wob) En.i) +
K
(e) R
(EW Ek.
R.K.s
mm
EE 112
k==1d Mml
RE ExAii
||Ex,el12
M=11==1
K
KEN Ek.i)l 1E 113
o
(w K. Ek,s) +
|1Ex.
k,k,i
m&m
k=11==1
1
(1)
11Ex,,112
k.k.
n&m
1
(8)
(24)
ngm
where the first inequality (a) follows from ( < 1, as stated in Equation (22) within Lemma B.3; the second inequality (b) is due to the non-linear activation (xx max(0, x}; and the last inequality (c) derived from Lemma A.I.
Before the first step of averaging weights, by taking a telescoping sum over = 0, 1...... E, we obtain:
V E
E 3 ngm od
2
At the same time, by Equations (17) and (18), we obtain V R 0 with h: * k for i E [n] and t E [E]. After ł = E steps, we perform weight average operation in server. Then by Equation (19)
we have:
E
E K m 2 3d
Note that noise vectors εx. are almost independent of each other. Thus, during the first weight averaging, the effect is directly reduced by K times. Besides, every client will store this information after accepting the averaged weight from the server.
In the next gradient descent steps at each client, by applying (24), we further obtain:
V
(2E) k.ks
<
Knom o F 3
F
n&m 2 o
K +
7/ 3 0 dd
K
ngm
(2E) k Kngm
<
2d pd
k' k
17

---

Published as a conference paper at ICLR 2024
In the second round of local updates, there two sources of noise. The first originates from the last return of weight averaging, and the second comes from the gradient descent update in the local client. From this result, find that the growth of noise memorization is akin to only E steps of gradient descent update when K is a large number.
Again, after gradient descent, apply second weight average operation in server and obtain the following result:
(2
K-1 K2
n 3 7mm
K - K2
D 3 Tg. 20 K2
2K
D 3 Tigm 2 ogd.
o
During the second weight averaging operation, calculate the coefficient of noise memorization by averaging all weights. The first term K+1 76
n,m 2 2od results from averaging the coefficient on client k. Morcover, the second term originates from the other K- clients who stored the coefficient during the first weight averaging operation. Together, applying Equation (19) yields
K²
(2E) N
2
nF 3 nmm 2 o7dd
K
which is equivalent to being divided by K compared to the local update without weight aggregation. Similarly, we repeat the computation procedure and obtain the following results for the third round of local updates plus the weight averaging operation:
Vp
(3F)
2
nE 3 ogd
d
3 m
E
o
K+2
DE 3 n&m 2 o.d.
K
Tum 2
K
(3E) V
2
E 3 d
W + k.
K
R&m
2
(3F)
K +
2
nE
3
2(K 2 (K-1)
K2 Tim 2 20 3 nE 3
DE
71gm 2 odd.
K2 nm
not hard observe that after the third round of local updates and weight averaging, the noise memorization is equivalent to the noise from local training divided by K. Using the same technique, we can summarize that, given R 2 times of communication, at the end of stage one with R Ti / round of communication, the noise memorization has an upper bound:
BBq E k,k
S
K + (R1
1)
mm 2 odd
F 3
K
B)
R1 1
nE 3
o k' k,,
K
7gm
V
k,
B)
Rı
K
ngm
F
3
2 p
2d.
Note that K+R,-1
R1 given that K > 1.
-(0) On the other hand, we establish the lower bound for f...... and (t) gradient descent update before weight average, we have:
Pkijon We show for uk = j with
K
(w kgrrt Ex,a) = (w/)). Cn,s) +
K. Jon IE 1 (Ex.s.
K
R.k Jona E 22 (Ex. Ewr,r)
log (2/42///0 15)
(0) k.j.r En.i) + Vkkjjaa -4v
R,k jon + k.
d
R²'
P
log (4n2
(0) Ex,)) + Pkkaas
4C2m
where C2 is positive constant, (a) is by Lemma A.I and (b) is by ((() j.r. S C2.
K 0 when Me. j and
18

---

Published as a conference paper at ICLR 2024
Let v(( ", = maxxj=wa Wa.""" (w (w EvrR) + whe Phana - /
1og(4n2 /6)
At initialization,
d
is easy to check that:
! k,k', 2ooo d/4 10204 4C log(4n2/04
(e)
2 0.
where (a) by the following condition:
ou Σ Cany log(4n') (6n2 (opd).
(25)
Then, compute the growth of (t)
follows:
KER, Ek,i)
(t-1)
(1()
"
(()) ki
o' (w k.jorr Ex.a) +
11Ex.i112
k,k,s
k,
n&m
11E
#=10=1
(1)
1/ 06
(w (0)
kijino Sare)
K
(1)
KEM,e, Ek.r
1E 11Ex.i112
k,
nmm
k=1P=1
C
1/202d
mmm
where (a) by Lemma B.3 stating the lower bound of minus loss derivative, (b) is by Lemma A.1 and Equation (25). Before the first step of weight average, take telescoping sum t = 0, 1,..., E, then gives
EE C1 mmm 2
V
(E) R.R
ozd.
At the same time, we claim (1( and (17). After t = E steps, perform weight average operation in server. Then by Equation (19)
k with k + k for E [n] and t [F] through Equations (18)
we have:
EE 1 n C1
K ngm 2 ofd.
The next E gradient descent steps yields:
EE C1
D C
K K+1
D C1 7gm 2 K
y (2E)
o.d+
0.
odd.
k.k.
Knamn 2
7mm 2
DE
C1
22
Knom- o,d,
k' k
k.
In the second round of local update, there are sources. The first source is from the previous weight average update, and the second is from the gradient descent update in the local client. From the result, we find the growth of noise memorization is close to only E steps of gradient descent update when K is a large number. Similarly, after gradient descent, we apply second weight average operation server and obtain the following result:
(2 RR
K+1
E
C1 2
K-1
K
DE
2 C1 o 2
-d
K TigHT 2 odd.
EE C1
K2
TigT
2
K2
It is not hard to observe that after the third round of local updates and weight averaging, the noise memorization is equivalent local training divided by K. Using the same technique, we can sum- marize that, given R₁ times of communication, at the end of stage one with Rí = Tí T E E z round of communication, the noise memorization has lower bounds:
K + (R1 - 1) EE C
(R
7gm 2 odd,
K
E)
R1
nF Ci
ond, W k
K
n&m
k.
F)
R1
K
E
C1
2
0
2d
19

---

Published as a conference paper at ICLR 2024
Finally, confirm that
(TL)
r/2i
Ci
oz k. Ex,) + 4C2m log 1og(4n²//0 /0)
J 3 Weii
max Waira
ngmK 2
IV3 IVE
rili
C o
Wijirr Exii) + C3
m 2
C4.
where the inequality (a) is by 4C2n/ Iog (4n /8)
I C3 which holds when 2 Cs log(4n2 /6)n² and the inequality (b) by taking the value of Tı and Lemma B.2.
Ing(4n2
A the same time, we calculate the growth of " for and (() Th.j.r Using the iteration equation for the coefficient of signal learning under local gradient descent in the first round, we have:
(1)
(++1)
(1) kk.....r
o' (( k.j.er (a. ( M 1
Vigm
The next step to provide an upper bound for signal learning at the first round of local updates.
Kkk.j.r I "kkjjr m
((+1) (a) (0)
"
1
where (a) is by Lemma B.3. Before the first step of weight average, taking telescoping sum over t == 0, I,... E then yields:
1
After E steps, we perform the weight averaging operation on the server through Equation (19). Then, we have:
K
(pk m
n m
px 12
K
12
k=11
Note that Ek=11 He counts the total similarity between signal vectors among clients. In the case of i.i.d., this sum approaches the maximum value of 1. the extreme non-i.i.d case, where no client shares the same class, this sum equals π. For ease of derivation, denote this sum (Mxig)) In the subsequent E gradient descent steps each client, by applying as X A Ex=1 || Equation (24), we further have:
K
(2E)
E
(pk mP) | 12
||
7F ||
kk.k.j.r
m
1
+ +1)
| l1
K
||
11
m
X
E
|
K
m
In the second round of update, there are two sources. The first one arises from the last round of weight averaging, and the second comes from the gradient descent update in the local client. Again, after performing gradient descent, we apply second weight averaging operation the server and obtain the following result:
(2 K.j.r
<
Xk
+
X* K
D
2
Xk: ME K
|
K
K
K
Here, the first term
X
1)
nE mh
|
results from averaging the coefficient across X clients
|| 1l originates from other K X clients who retained the coefficient during the first weight averaging operation.
who share the same signal vector. Moreover, the second term (1 X X aE
m
20

---

Published as a conference paper at ICLR 2024
Similarly, repeat the computation procedure and obtain the following results for the third round of local update plus the weight averaging operation:
(3E)
S 2 Xk
nE
|
11
T
11 (2
Xk
+
1)
7E
p 12:
K
m
m
K
m
Ph..r 1
(3E)
Xk X +
(2
nE
x (11 | 1 (2
Xh
K
F
|
Xk F || 112.
K
m
m
K
m
We repeat the computation process and find that, given Rı times of communication, at the end of stage one, with R T rounds of communication, the noise memorization has an upper bound:
(Rg E) Kk.j.r
((B) - x)x² +
n m
|
K
(R, E) Tkji
R1
Xk
B m
1
K
Lastly, provide the lower bound for signal learning in the first stage. It is known that
(w())
(1)
Wer.r 14x)
(0) (w
k.j.r> 14))
(0) Ear (Eur,
snalllall
k. jon 1l
2 (Ex...
Rujnt pk),
where (a) by the orthogonal relation between feature vector and noise vector.
is known that there exists some r for which (w (0) gradient update, we have:
kui M. 0. Thus, for the first round of
Rk
(1+1) (0)
n
(1)
d d'w 3i(le)) ||
m&m
" C1 |
+
m
where (a) follows from the fact that C1 < - C S 1. Before the first step of weight averaging,
taking a telescoping sum over f = 0, l..... E' yields:
n/ E
(B) Nk,k.j
m
After E steps, we perform the weight averaging operation on the server. Then, we have:
C (pe
m 1
K
1
counts the total similarity between the signal vectors among clients. In the next E gradient descent steps on each client, by applying Equation (24) we further have:
Recall that XA A E
(2F)
K Σ
(pL.
mP) 12
7 F 11
Vk.k.j.
mK
==1
Xk
11 E
1)
K
Again, after gradient descent, we apply the second weight averaging operation on the server and obtain the following result:
((2))
Xk Xk
1)
+
(1-
X XA
C
E
|
12
K K
K
K
m
2 1
|| 112
K m
21

---

Published as a conference paper at ICLR 2024
During the second weight averaging operation, calculate the coefficient of noise memorization by averaging all weights. Similarly, we repeat the computation procedure and obtain the following results for the R -th round of local update plus weight average operation:
(R Ki
((R) - 1) X +1 1 K
|
m
B)
R1
XA
7 F
K
m
Finally, taking the value of into the above inequality, we have:
Crnxe
Cmn
| 112
((TTSSRR))
k.j.
1126d
m
B.2
STAGE TWO: CONVERGENCE ANALYSIS AND FEATURE LEARNING SCALE
In this section, demonstrate how the learning scale of signal learning and noise memorization at convergence completes the characterization feature learning initiated in the first stage.
B.2.1 FEATURE LEARNING SCALE AT CONVERGENCE
In the second stage, we aim to prove the following scaling behavior feature learning
(())) Wkina S Vkkdiin $ log(T");
(26)
(0)
0 Pkjjr 2 Pk.kd.r 2 - llog(?"),
(27)
where RE for an integer R and T > Tı.
The proof is by induction. Before proceeding with the main lemma, it is necessary to introduce several technical lemmas that are fundamental to our argument.
Lemma B.5. Under Condition 4.2, suppose Equation (26) and Equation (27) hold at iteration t.
Then
(0) Pkj.r
8n
1og(4n2/00 (0)
log(T")) (w(1) j.r
(t)
w Ek,))
J( Pleg.r. + 8n
log (n2 18)
(T*)
200 Ok.j.
/ny
log(4n2/14
log log(T")
Sur
Ex.i) S Pkiji + +sn (0)
log( 4n (4n2/8) log(T"),
jor
d
for all r [m], [+1), i E in] and [K]. Here we define wher
>
Re (0)
n kjj."
Proof of Lemma B.5. It is known that, when yk. = j:
K ay
(w! (1)
J.r J.r Ex.i)
w (U)
Vk danee
(t)
(Ew.y En.i)
W=1
P=1
k J.ra 1Er
log 4n log(4n²
/8)
(t)
log(An2
(1)
(()
41
Wkinnll + 41
sinee vl + Vkiii
d
d
k=1 3'
(6)
log (og(4n2/)) log(T"),
1
+ /8n//
d
22

---

Published as a conference paper at ICLR 2024
where the first inequality (a) is by Lemma A.1 and the last inequality by 1V Pelesrar Pkk.jrre - log( in Equation (26). Similarly, we can show that when ys + j:
(2)
K
(t)
Jr Ekii)
(0) Viire 1l122 WE Ex,a)
A
Pk Jina 1E 1122 (En, Eka))
(4)
log log(4nd (n2² 15)
1og(4n2/8
K
(0)
(1) FKuJirs
4V
R seral + 41
OK
shree +
d
d
k==1 zf
log 442/// log(T");
Oki.r + V
d
where the first inequality is by Lemma A.1 and the second inequality by Pegrel Pkaian log(T*) in equation 26, which completes the proof.
Lemma B.6. Under Condition 4.2, suppose Equation (26) and Equation (27) hold at iteration
Then we have:
(1)
log(4n² Kki.a /
(0) log(T") S (Wij.r , Vkk.ijn + Sny
(0)
log log(4n2/0m 15)
(og(
d
d
1og(4n2/0) /0)
log( /5) (T"( S (wk.r -W Ex,i) Okikiiit + 8m
(1) Okkiin 8n
(t)
(1)
log(T"),
d
d
for all r E [m], j [+1), ie in] and E [K]. Here we define w w Et=r w
(0)
k.j.""
Proof of Lemma B.6. It known that, when yk.i j:
K
(w kig ww j.r En.
(1)
=
Ve da
#=1r=
Rig
kf K.jia Me (Sur,,En.))
(c)
W=1==1
(0)
log 1og(4n2/80 18)
log (on442n0/ 18)
VV
W Joaar + +AV
l5((0 Wereir VkRjja
M=1Fzi
k=
(b) k.kj.. + +8m
Tog(4n2 lTg(
d
where the first inequality (a) is by Lemma A.1 and the last inequality is by W s.arh 10 j.xll log(?") in Equation (26). Similarly, we can show that when ya + j:
K ng
Kighr - (0)) En,s)
k. ai ae IE (Ew,,, Ek.s)
K
Kk Jow 1u (Evr,,EK..
W=1f=1
(a)
log(4n2 18)
1og(4n2/0 /0)
(0) Weddmel + +4
( 1k anvl + Oikidin
(0)
I 4
d P
(b)
logí 4n222/2//
< Pkkijoo 8m
log log(T*))
d
where the first inequality is by [v(())
Lemma A.I and the second inequality is by
s.rar s log(T") in equation 26, which completes the proof.
(0(0)
23

---

Published as a conference paper at ICLR 2024
Lemma B.7. Under Condition 4.2. For any E, it holds that
(w) J.r
W j.r p4x) = - (1)
(w(1)
j.r ma) = ".
ligh
for all r [m], j {=1}, and & [K].
Proof of Lemma B.7. For any time t > E, we have that
K
(w k.j.r
w (0) j.r: 1b) j Je
Rak anellell (Ew,e,))
k=lr=1
R. Jina 1l, -Ex, (E)
= K.j.
where the equation is by our orthogonal assumption between feature vector and noise vector. Similarly,
K
(w kor - 500
(ex) =
sirt
-j.
+
Vkinn Ils. (Ew
@h Jun MEM (E
Lemma B.8.
Under
Condition 4.2, for <
S
T where T
77 " |||l
p 0 n, m, d) = R E, we have that
0 s kj.r S Kkk.jjrr. S log(T"))
(1)
(1)
0 2 0 Kjj.. 2 Pk.jr 2 -log(T"),
0 S Jajj S 1k,kj.r S & log(T"))
(1)
for all r E [m]) E [11], k E [K] and i E [n].
Proof of Lemma B.8. The proof relies on induction. At t = 0, the results are straightforward, given that all coefficients are zero. We assume that there is time R,E < T for which the Lemma B.8 is valid for all moments 0 < 1 < (R. - 1)E. Our goal is to demonstrate that the result also stands true for E.
16n log(4n"/0) lT'( We first prove that Equation (27) holds fort RgE, i... Ok.... 2 -3- -3-166n
-00
d
for t = E,E, [m]. j {+1), k E [K]. and i E [n]. Notice that @kj.r. Vj yk.i Therefore, we only need to consider the case that j + yk.i-
(()--1)) When 0k.jnn
< -0.5. Sny
log(T"), by Lemma B.5 we have that log(4n2/1 18)
1) (w j,r
Ex.i) <
Ok.j.r. + (w j.r Ex.s) / d
(RR, 1)E)
(0)
log(4n2/0
log(T") S 0,
Then we have E steps of local update, and we find that, for E [E]:
(w)) k.
(1)++1)))+
Ex.x) < (w)) jrr
(R,-1)))
Ek.))
24

---

Published as a conference paper at ICLR 2024
Therefore, have,
E 7
(R,E))
((R,)1)E)
((R d'((w ( ((R,11)++7) Ex,s)) 1yk. = --]]l
Okjjn, = -ki + K ==1
kijir
m
1))
(a)
log /62 log(T");
=Pkini
-3 16n
where the last inequality (a) is by induction hypothesis.
When Pk.jo,
(()))
z -0.53 Sn
V
log(An2 (8)
log(T"), we have that
E
(R.BB)
=
((RA-11 1)E)
+
7
(((Ra-))B)
o'((w
((Ro (1)++))) ,k,ii) ,x,ttt) 1 EA,1
1V
log (og(4n2// /6)
n -0
log(4n2/0)
-0.53 /
0' (0.58 + sny
log (T()
d
nymK)
d
log(4n²/1 /0)
Fo
-0.53 /
log(T")
d
nymK
- 16 log( (og(4n2/0) /0)
log(
where we use k.i S 1, and [|Ex; (ozed) by Lemma A.I, and Lemma B.6 in the first inequality
(a), the second inequality (b) is by 0.5,3 + 8n 1og(4n2/8) log(T*) $ 1, and the last inequality (c) is by n ((ozK/((ZdE)) in Assumption 4.2.
((t) Next we prove Equation (26) holds for R.E. Direct computation leads to Pk.j....
E
Vk.j.r. = Vk.j.r.
(RE) (E(R-1)) n
BKKKK ki
(ER-E+T))
d'(w k.j.r
(RE-E Etr) (k,,)) 1 - 511
(t) T that Wkfer S (.5log(T** Then we have that
Let 4, = RGE to be the last time
E
Wharr =
Vasina
(1))
(latT)
o'( kij
((str) ,Ex,i)) 1 1ykk, = 11)
ny.KK
T=1
Σ
Σ
e ((RETT) 0 ((W!Ret7
(RE+T) Ek.i)) 1 (yA,i =R),
T=1 RRSRRRR ngMK
2 Fo
D n
eep((=o((w))) (WW.r En,i)) + 1) - d'(ww En,i)) ||E
nmK
R.SRST.
0.25 (T") + 0.257 exp( - (T) log (**)
S 0.5 lTg(
where the first inequality (a) is by Lemmas B.6 and A.I, the second inequality (b) is by " S (Eo))(Eofd) log(T"))
(()) Rkkjann > 0.5 log((") and (w(0) k.jor Ex.i) 2 -0.53 due to the definition of to and B. (c) is by B s log and / 100(( /6) log(T") $ log
(2)
d
Finally, we can prove that S - I Vkkkj.rr S SNR
(t)
Direct computation yields:
(1+1) ((1)
Σ
TIK
Σ
((-r) o -''(( ((+T)
(WWkr Uk. Mx)) 11
Tkj.r kkj.r
ngmK
25

---

Published as a conference paper at ICLR 2024
Let T. R&E be the last time
T that x(1) Tkg. I 0.5llg(TT)) FSNNR
Then we have that
F
(T.)
(T.)
TIX
(T++))
d' Rij. :Be,
(Te+T)
Tk.j.
E
E
((E+T)) o' d'( (RE+T) R3
((wk3et7)
ngmK
R.RReR.
(0) (T.)
IXK
E
(ToTT) d'((W
(Te-T) -.KK)44))
m&m
T= i=1
E
NX
exp( exp(-o((w)
(RE+T) kij. -UKSSEE) + 1) o' d( (RE+T) -.k.amm))
n&m
Re<R<RR
(6)
Rk.. 0.25 S log(?") + 0.257 (Tx log(T" STX&SSRR
where the first inequality (a) is by the sign of f,i the second inequality (b) is by Lemma B.3, prop- erty of ReLU activation, mK/ (E) M (12). Furthermore, we have used the following inequality,
W
Jan
(4) + Tkj.r
IV8 IV3
0.53 + 1 log(T")
0 log(T"),
where the first inequality (a) by Lemma B.7, the second inequality (b) is by 7k..r 0.5 log(T SX and (w j.r px) 2 -0.33 due to the definition T. and 3, the last inequality (c) is by B I SNR log (T")." Similarly, for To < l < To and y/s j. we can also upper bound (w(() Jun (a) as follows,
(f)
(0)
MR) + (2)
mk)
W5 1A
0.5. + 0.5m% SNR log(?")
0.5 log(
where the first inequality (a) is by Lemma B.7, the second inequality (b) is by induction hypothesis (1)
Tk..r I S log(T"); the last inequality (c) is by 8 S S.1 log(T*).
B.2.2 CONVERGENCE ANALYSIS
In this section, our proof is highly inspired by Cao et al. (2022) and Kou et al. (2023).
Lemma B.9. We choose the solution of FedAvg W follows:
K
j .. SNR 2
Ek
W f.r +
log log(2/E)
10 = mk.ss)
Under the same conditions as Theorem 4.3, we have that
|
W s δ m 1/2 Ko - 12).
Proof of Lemma B.9, Recall that in the first stage and time step Ti we know that:
(T))
(0)
(T))
Ex. 1
Ex
J.r
R.j.r
PAR
1
i=1
26

---

Published as a conference paper at ICLR 2024
Then the distance between the two weights can be calculated as follows:
W
l| (a)
S
W
(T)
w
|L +
W (0)
WF
(TT)
(Tt)
(Tt)
max 3.r
Phij.r
mm max 3.r
vm max 3.r
Ok.jara
| ll2
E 112
E l12
+ O0mm 1/2 (og(1/))
n SNNR 1 + Vn max [|sx.ilE")
(6)
S 50 1/2
| 112 17 111111111
SNR + +(KTm 1/2 +OKKTm
K
+ O(m lo////)g((/(((CC))
SNR 2
||1|12² + Kn 1 -t121122
1=1
50 1/2 (T
SNR
||1| + ( 2 1/2 VVTKo 1/2),
where the first inequality (a) is by triangle inequality, the second inequality (b) is by the decompo- sition of (T)) and the definition of W", the third inequality (c) is by Lemma B.4, and the last
inequality is by direct derivation, the last inequality (d) is by that of 2 K²n p [ Assumption 4.2.
Lemma B. Under the same conditions as Theorem 4.3, we have that
Ve, (V, xx,)), W log(2/e) for all i [m]] k E [K] and Tı < <ISTo T
Proof of Lemma B. Recall that f(w xk.i) (I/m)E (r(WW,rrBessauu + a(W j.r: {,,, so we have
XF.), W o' (ww j.r Uk. (ME
yk.
1
1
d'((W( (0)
jort (...
for
1
o' (t) j.r j.. (lo. F. SNR 1
o' d'(WW!! Jim (H.
m
m
(((w j..j ( m
ΣΣ
a d'w( joro (n.
m
dr
1
o' d'(ww!!) (t) W (w SNR 2 m
1
(t) (.. (ww.:
Σ
m
((W)) (000 m 1= m
Σ
d'((w())
(a. C..) E1..)-1E...1 12.211
m
J.r
a'(W!! f. H42 SNR
1 Σ
(1) (' (1000000l0l
m
m
j.
jor
d'((w To ER. (0))), m
Σ
d((((() f. Coweooeea. logg)))
m
ji
==
(28)
m
Vi J.r
27

---

Published as a conference paper at ICLR 2024
where the first inequality (a) is by Lemma B.I, and the definition of Wj.r and the second inequality (b) is by Lemma A.I, and we have used the following inequality:
K
ng
K
Jirt Ek,i)
=
Σ
(1) Wk
S 2 (EK. Ew,e)
+
@k Pina Eu (Ex.i, Eur,)
(1)
K=10==1
(w
j.r Ex,s) + Wkisnn 8V
(0)
(()
15)
log (e(4nn2/0)
K
(t)
8V
d
(6)
j.r
(0)
Ex,))
8C2mV Kigo
(4)
log(4n2 /6)
d
2 2 2 1og (log(mmm/0))) 00o, Vd V 1og(4n2/0
(9)
(4)
20
1,
(29)
where C2 is positive constant, (a) is by Lemma A.1 and (b) is by Pkd.. 0 when yk,, = j and Veg... = O() for all E [z).. (c) is by Lemmas B. and (d) is by conditions on and do stated in Assumption 4.2.
Plugging Equation (29) into Equation (28) gives:
y V x),,),
> 2lo((2/e) Ö(oollmmlle) - Ö(00op
log(2/c),
where the last inequality (a) is by do S Õ(1) min || in Assumption 4.2. Lemma B.11 (Cao al. (2022)). Under Assumption for 0 S 1 1 T", the following result holds for all k E [K].
VLs, (W s O (max [11 o2d)) L.s., (
Lemma B.12 (Cao et (2022)). Under the same conditions as Theorem 4.3, we have that
W 1/
W nLs((()))) me w (++1)
for all Ti TT
Lemma B.13. Under the same conditions Theorem 4.3, let T Ti | (T) W
2ne
Ti + Ömmm le -nK22 2d 2). Then we have Ls W () for some Ti StST.
Proof of Lemma 1.1.. By Lemma 12, for any & [Ti, T]. we have that
W w w | (+t1) - 1 > nLs ( TE
holds for Taking summation, obtain that
W (Ti)
P + ne(1 - T, 1) W
Σ Ls
(30)
for T, S T. Dividing (I Tj + 1) on both side of Equation (30) gives that
W T1) W
1
Σ
Ls W
+ c.
t - Tj +
qit T + 11)
+= T
28

---

Published as a conference paper at ICLR 2024
Then can take T and have that
W (1)) W F
Ls (W
+
T - 71 +
11T Ti + 1)
m m/2 k -
+ < 2c,
- -i+
W
where the second inequity is by Lemma B.9. Therefore, our choice that T Ti + Because the mean is smaller than 6, we can conclude that there exist Ti S f
27e
T such that
Ls( <.
B.3 POPULATION LOSS
Consider a new data point (xx, 3) ~ Dx drawn from the distribution of client Moreover, by the signal-noise decomposition, the learned neural network has parameter:
(0) 1og(2/e)
j.n. SNN,,
MR
Σ
n
vk.e)
Ék.i E 1
1
for j {+1} and E [mm] The calculation of the test follows the results established in Theorems E.I and E.3 from Kou et al. (2023). Lemma B.14 (Restatement of Theorem 4.4). Let T be defined in Theorem 4.3. Under the same conditions as Theorem 4.3, there exists 0 t S T such that if axcllall Koyd
22(1) then the test error
satisfies LD, (w( S exp Kac 1
Koya
Proof of Lemma 1.14. Recall that the test error is define as P( (v * f(w(() xx)) It is equiv- alent to calculate that
P(x,))Da (v + f(W!)) x)) = P( (Ds (w ,) < 0).
is essential to calculate the output function of neural network which follows:
ws(w(()
a((w vpa)) + Σ
(ow En.i))
3.
3.
We first calculate the signal learning part:
(w)) v1n) (w (0) j.r y1un)
juTk... (t)
Σ> k=1i=1
(0) (hasina) (E wpa)) ( Ex,
Next, we calculate the noise vector part. It is known that (w (1) mean zero. Let w w!O and
j.m Ex,e) is a Gaussian distribution with
Σk j Tki."
then have that (w(() (t)
En.x) (w()) Ex.i)
T
| 1o s01/
(31)
As result, by the condition
Thg.r E,dh Whiter
A 11 ), we have that: nogd
Pxxy)-P,uff xx) < <0) < Pix,y)>Da
Σ a @a( Jum Bmm))) 2 [o((w() ( Exii))
ToX 1l
< exp
K Kofd
29

---

Published as a conference paper at ICLR 2024
C SUPPLEMENTS FOR EXPERIMENTS
C.I EXPERIMENTAL DETAILS FOR WEIGHTED FEDAVG ON TEAL-WORLD DATASET
Following the footsteps of the prior works (Acar et al., 2021; Shamsian et al., 2021), A ConvNet LeCun al. (1998) with two convolutional layers and three fully-connected layers is adopted for experiments on CIFARIO and CIFARIO0 datasets. The real-world dataset comprises five distinct sub-datasets exhibiting feature shift: SVHN (Netzer al., 2011), 1), USPS (Hull, 1994), SynthDigits (Ganin & Lempitsky, 2015), MNIST-M (Ganin & Lempitsky, 2015), and MNIST (Wang et al., 2014), where each domain serves client.
As for the Digits dataset, we utilized another ConvNet LeCun et al. (1998) similar to Li et al. (2021) with four convolutional layers and three layers. For all experiments in this part, models were trained with 300 communication rounds and all tasks were optimized with the SGD optimizer. For experiments on CIFAR CETARIO/CIFAROO we set the learning rate 0.03 and randomly sampled half of the clients participating in each communication round. For experiments on Digits, the learn- ing rate set to 0.01 and all clients are involved each communication round. Algorithms were implemented on PyTorch Paszke et al. (2019) with an RTX 3090 Ti GPU. The learned models are tested on the private test data of each client. To be specific, we run five times trials and re- port their mean and std. Codes are available at https: /anonymous 4open. science/se / fed- Teeture-learning
C.2 FEATURE LEARNING PROCESS IN FL
We provide an experimental evidence to validate our theoretical analysis of feature learning for FL, shown in Figure 2. In this experiment, engage two clients under the IID setting. Both the data and neural networks are synthesized in accordance with our theoretical framework. Both data and neural networks are synthetic by our theoretical setting. In the plot, we illustrate the learning trajectory for client for client 1 in terms of signal learning represented by maxj,r "/r and noise memorization represented by maxj,r Pj.r- We observe that, with the help of communication, the signal learning keeps increasing without any degradation. In contrast, noise memorization exhibits a marked decline with each weight averaging. This observed behavior further reinforces our theoretical analysis.
Feature Learning
Loss
2.5
train_userl
0.6
2.0
train_userf _
1.5
0.4
1.0
maxp).r.i
0.5
0.2
maxY),r Jrr
0.0
250
500
750
1000
250
500
750
1000
Feature learning trajectory of FedAvg on synthetic data. The signal learning (blue curve) keeps increasing without any degradation. In contrast, noise memorizations (red curves) exhibit marked decline with at each weight averaging.
Figure 2:
C.3 ADDITIONAL EXPERIMENT WITH HIGH DIMENSION SETTING
We further examine our theoretical result by experimental simulation with high dimension setting.
The training data size is set to Tirrain = 100 and the testing data size is set to Neest = 2000 with instance dimension to d 8000 for each client. The result is shown in Figure 3. These additional experiments confirm that Federated Averaging (FedAvg) can indeed achieve better results than local training methods, aligning with our theoretical predictions.
30

---

Published as a conference paper at ICLR 2024
Train Loss
Train Accuracy
Test
Test Accuracy
Fedwig Singleset
1.0
0.7
0.6
0.8
2
0.4
0
200
400
200
400
200
400
200
400
Figure 3: Convergence behavior comparison of train loss, train accuracy, test loss, and test accuracy on synthetic data. Both local training and FedAvg demonstrate convergence on the training set. FedAvg outperforms significantly on the testing set.
Performance on each participant
Accuracy
0.6
0.4
Client IDs
SingleSet
FedAvg
Weighted FedAvg
Figure 4: Comparison results on each client under non-IID setting.
C.4 FEATURE LEARNING IN NON-IID SETTING
To simulate the non-lID scenario, we split 20 clients into four groups. Within each group, the signal vectors identical. However, signal vectors from different groups orthogonal to each other. Specifically, the client groupings are as follows: {*) Besides, we implement the Weighted FedAvg algorithm, where each client only communicates with those clients in the same group. The test accuracy on each client demonstrated in Figure 4. The findings indicate that as the number of members in group increases, the performance of both FedAvg and Weighted FedAvg improves. This aligns with our theoretical insights regarding the generalization of Federated Learning (FL).
{k) ke[s,10)- and (k) [e])ep11
D ADDITIONAL DISCUSSION
D.1 NAME OF EMPLOYED MODEL AS CNN
Our choice to label the model as a CNN based on the specific characteristics of the data model
we used. We adopt two-patch model, where X = [x (1) (2)] = [uK. Ex.i). We implement k,
convolution operation using single weight across these two patches, akin to applying filter in traditional CNNs. This is concisely represented as follows:
m Σ (W (1) + (W
(2)]
To(w k, x(!) + OW x(!1
31

---

Published as a conference paper at ICLR 2024
Here the filter Wk,ir operates two patch, with & [K] denoting the index of client, j {+1, -1} corresponding to the weights value at the second layer, and r E [m] as the filter index. Our approach alongside fixing the second layer is consistent with the related works [Cao et al. (2022); Kou et al. (2023)), where similar network structures applying weights across multiple data patches are identi- fied CNNs.
D.2 RELATION BETWEEN CONVERGENCE AND DATA HETEROGENEITY
The effect of data heterogeneity convergence is critical aspect of our analysis and highlights the significant influence on convergence rates. demonstrated Theorem 4.4, we establish condition related to the signal-to-noise ratio (SNR):
m 2(1)
Koyd
with the effective SNR for each climent given by SNR, = EE
(pe- SNR,. In scenarios 11
with increased data heterogencity, the term X = Σx=1 consistent SNR level in such cases, becomes necessary reduce ozd. which, according the convergence time in Theorem 4.3: T = O(1 Kmño, 2d + n "='mmoo'''' ( results in slower convergence. On the contrary, reducing data heterogeneity will shorten the convergence time.
===l llaall12 tends to decrease. To maintain
(pea)
D.3 EFFECT OF THE NUMBER OF LOCAL EPOCHS ON THE CONVERGENCE AND GENERALIZATION OF FEDAVG
In analysis, the term T in Theorem 4.3 represents the total number of local updates, defined as = RF, where R is the number of commutation rounds and E is the number of local epochs.
Under our assumption on the learning rate, where 7) S Ō min - op---1))))) increasing E may lead to slower convergence. This observation aligns with the empirical results reported in McMahan et al. (2017). The impact of varying E on the required number of communication rounds to achieve adequate convergence is demonstrated in Table 2 of McMahan et al. (2017), which
E
compares the communication rounds needed for the same training accuracy on MNIST.
Besides, it important to note that smaller E' does not necessarily imply better generalization. To explore this relationship further, we conducted additional experiments of FedAvg on CIFAR 10. From Table 3 we can see that the generalization performance of FedAvg is not closely related the number of local epochs.
Table 3: Relation between local epochs and test accuracy.
EIR
20
40
48.18
66.28
66.24
64.98
64.55
10
54.17
64.6
64.48
63.44
64.5
15
54.34
64.1
63.91
63.44
64.33
32


Reviews:
[{"review_id": "efUaw94NYw", "reviewer": "Reviewer_vgxi", "summary": "The paper studies the convergence and generalization properties of FedAvg for two-layer convolutional neural networks (CNNs). To do so, the authors consider a setting where the input data is divided into a signal part and noise part, with the signal being different across different clients and the noise being independently sampled from a Gaussian distribution. The analysis then tracks the coefficients of the signal and noise in the model updates done by clients and also in the global averaging step done in FedAvg. Based on this analysis, the authors show that FedAvg can provably benefit from the communication across clients by increasing the signal to noise ratio (SNR) in the global model updates. In some regimes, this can lead to FedAvg achieving near-zero test error while individual local training at clients achieves a large constant error. These theoretical results are then verified experimentally on a synthetic dataset. The authors also propose a weighted FedAvg approach which is tested empirically on some real-world datasets and shown to outperform FedAvg.", "questions": "* 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.\n\n* 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.\n\n* How is $P$ defined in Section 3.2?\n\n* What is the difference between Lemma 5.3 and Theorem 4.4? It seems they are almost stating the same thing. \n\n* 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.\n\nTypos\n\n* $F_{-1}(W_{+1}, x) $ should be $F_{-1}(W_{-1},x) $ in the line above Eq. (2)\n* Summation should be over $n_k$ in the Eq. (4) and not $n$\n* $<$ should be outside the bracket in the definition of $L_{D_k}^{(0-1)}(\\bar{W}^{(t)})$\n* $L_{D_k}$ in Lemma 5.3 is not defined. Maybe the authors meant $L_{D_k}^{(0-1)}$\n* \"Provably\" and not \"Probably\" in the heading of Section 4.2\n\n\n\n**References**\n[1] Du, Simon S., et al. \"Gradient descent provably optimizes over-parameterized neural networks.\" arXiv preprint arXiv:1810.02054 (2018).", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The analysis clearly establishes feature learning for this particular class of NNs, unlike most existing work on NNs which are focused on the lazy training regime and hence are not really doing feature learning.\n\n* The analysis shows that FedAvg provably benefits from collaboration among clients, by increasing the SNR compared to local training. This in turn leads to better generalization guarantees for FedAvg.", "weaknesses": "* 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.\n\n* 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. \n\n* 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_{k',j,r}^{(t)}$ is going to be the same for any $k'$, which does not seem to be correct.\n\n* 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. \n\n* 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$. \n\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2ufiGFj0gO", "reviewer": "Reviewer_u4St", "summary": "The paper does joint generalization and convergence analysis of FL, using 2-layer NNs.", "questions": "3.2 Data Model:\n- Why is $y_{k,i}$ Rademacher if $\\mathbf y_k$ are labels? Are you assuming a balanced dataset? What is the distr. $P$?\n\n4.1:\n- 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.?\n\n4.2:\n- Why is using the same training size for all clients without loss of generality?\n- Assum 4.2: what is $\\sigma_p$?\n- What is the motivation for the defn of $\\bar{SNR}_k$? \n- 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$?\n\nExperiments:\n- In Assum 4.2 it was assumed that $d$ is larger than $\\bar{n} m$. But, in experiments in Sec 6.1, that's not the case ($d=1000 < \\bar{n}_{train}.m=5000$). Why?\n- In Assum 4.2, $\\eta \\lesssim 1/d$ is assumed, but in experiments, $\\eta=1$ is chosen. Can you explain this divergence?\n- 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is largely well-written and seems to be a good work. I have some basic questions which I have asked below.", "weaknesses": "Minor writing issues:\n- 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.?\n- 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.\n- Questions are listed below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "f0YzTdNP2e", "reviewer": "Reviewer_bJ2V", "summary": "This paper studies a scenario where each\nclient employs a two-layer CNN for local training\non their own data by tracking the trajectory of signal learning and noise\nmemorization in FLs. The paper shows\nthat FedAvg can achieve near-zero test error by effectively increasing the signal-to-noise ratio in feature learning, while pure local training without communication\nachieves a large constant test error. Inspired by the theoretical results, the paper proposes a heuristic weighted FedAvg which leverages the similarity of local representations into mixing weights.", "questions": "NA\n\n\n=================================\n\nI raised my score after the rebuttal.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a fresh perspective and innovative techniques. It successfully demonstrates fruitful results, offering valuable insights into the generalization of FedAvg.", "weaknesses": "1. 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$. \n\n2. 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.\n\nI am open to raising my score if the two points can be adequately addressed.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Syx4Un7xbj", "reviewer": "Reviewer_aCq5", "summary": "This paper studies the optimization dynamics and generalization behavior of FedAvg several neural networks in federated learning. Inspired by feature learning theory, they prove that FedAvg converges with gradient descent optimization under a certain data-generating model.", "questions": "- 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? \n- What is the expression of the distribution $P(\\mu^{(1)},\\mu^{(2)},\\cdots,\\mu^{(C)})$?\n- Page 3: $F_{-1}(W_{+1}, x)$ should be $F_{-1}(W_{-1}, x)$\n- Page 4: $yf(\\bar{W}^{(t)}, x<0)$ should be $yf(\\bar{W}^{(t)}, x)<0$\n- What is $w_{k',j,r}$ in eq.(6) and (8)? It seems that the right-hand-side is independent of $k'$.\n- (minor) Theorem 4.3 is a \"best-iterate\" result. Could it be turned into a \"last-iterate\" or \"average-iterate\" guarantee?\n- The discussion below Theorem 4.4 could be made more accessible. Specifically, when justifying the superiority of FedAvg compared to local training, could the authors elaborate more, rather than just discussing on two equations on SNR?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is generally well-written, making me easy to follow the logic flow. The motivation is strong, and the theoretical results are intuitive. Moreover, the authors honestly discuss the limitations of the paper in conclusion.", "weaknesses": "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?\n\nAnother 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? \n\n(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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)