Widening the Network Mitigates the Impact of Data Heterogeneity on FedAvg
Authors: Like Jian
Dong Liu
Venue: ICML 2025 poster (2025)
Keywords: Federated learning
overparameterized neural networks
data heterogeneity
Primary area: deep_learning->theory

Abstract: Federated learning (FL) enables decentralized clients to train a model collaboratively without sharing local data. A key distinction between FL and centralized learning is that clients' data are non-independent and identically distributed, which poses significant challenges in training a global model that generalizes well across heterogeneous local data distributions. In this paper, we analyze the convergence of overparameterized FedAvg with gradient descent (GD). We prove that the impact of data heterogeneity diminishes as the width of neural networks increases, ultimately vanishing when the width approaches infinity. In the infinite-width regime, we further prove that both the global and local models in FedAvg behave as linear models, and that FedAvg achieves the same generalization performance as centralized learning with the same number of GD iterations. Extensive experiments validate our theoretical findings across various network architectures, loss functions, and optimization methods.

Widening the Network Mitigates the Impact of Data Heterogeneity on FedAvg
Like Jian'
Dong Liu
Abstract
convergence and generalization (Li et al., 2019; Zhao et al., 2018).
Federated learning (FL) enables decentralized clients to train a model collaboratively without sharing local data. A key distinction between FL and centralized learning is that clients' data are non-independent and identically distributed, which poses significant challenges in training a global model that generalizes well across hetero- geneous local data distributions. In this paper, we analyze the convergence of overparameterized Fe- dAvg with gradient descent (GD). We prove that the impact of data heterogeneity diminishes as the width of neural networks increases, ultimately vanishing when the width approaches infinity. In the width regime, we further prove that both the global and local models in FedAvg be- have as linear models, and that FedAvg achieves the same generalization performance as central- ized learning with the same number of GD itera- tions. Extensive experiments validate our theoret- ical findings across various network architectures, loss functions, and optimization methods.
To address the challenges posed by non-IlI data in FL, numerous research efforts have emerged, including client regularization (Li et al., 2020), adaptive optimization frame- works (Karimireddy et al., 2020; Reddi et al., 2021), person- alized model architectures (Jeong & Hwang. 2022: T Dinh et al., 2020), and etc. While these approaches have demon- strated empirical success, they often require intricate hyper- parameter tuning restrictive assumptions about convexity, data similarity, or gradient boundedness, limiting their ap- plicability in practical highly heterogeneous environments.
In parallel, overparameterizedd neural networks have gar- nered prominence in centralized learning for their remark- able ability to achieve strong generalization despite non- convex optimization landscapes, underpinned by theoretical frameworks such as the neural tangent kernel (NTK) (Jacot et al., 2018). These networks exhibit implicit regularization properties, enabling interpolation of complex data distribu- tions while maintaining robust generalization (Neyshabur et al., 2015; 2019a;b; Lee al., 2018), motivating a piv- otal question: Can increasing the network width inherently mitigate the effects of data heterogeneity in FL?
1. Introduction
In this work, we analyze the convergence of FedAvg with gradient descent (GD) for multi-layer overparameterized neural networks and establish that the impact of data het- erogeneity can indeed be reduced by widening the network. Further, we prove that the network width approaches in- finity, both global and local models behave as linear models. Strikingly, in this regime, FedAvg and centralized GD yield identical model parameters and outputs under matched iter- ations, achieving equivalent generalization performance. To the best of our knowledge, this is the first work to provide a quantitative analysis explicitly linking the width of neural networks to the impact of data heterogeneity on both FL training and generalization. Our key contributions
Federated Learning (FL) is a distributed machine learning paradigm that enables collaborative model training across distributed clients while preserving data locality (MeMa- han et al., 2017), a critical feature for privacy-sensitive domains such as healthcare, finance, and mobile computing, where regulatory or infrastructural constraints prohibit data centralization. However, fundamental challenge in FL arises from the intrinsic non-independent and identically distributed (non-IID) nature of client data (Li et al., 2021b), where local datasets exhibit significant distributional shifts due to user-specific behaviors, geographic variations, or device-specific usage patterns. Such statistical heterogene- ity leads to divergent local optimizations, degrading model
Theoretical guarantees for heterogeneity reduction: We prove that the model divergence is bounded and decreases inversely proportional to the square root of the network width asymptotically, without relying on restrictive assumptions on convexity or gradient simi- larity/boundednes... This bound vital in the conver- gence analysis of FedAvg, revealing that the impact
'School of Cyber Science and Technology, Beihang Uni- versity, Beijing, China. Correspondence to: Dong Liu < bdaa.cdu.cn>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Widening the Network Mitigates the Impact of Data Heterogencity on FedAvg
of data heterogeneity slows the convergence rate, but vanishes when the width approaches infinity, allowing the convergence rate to recover linearity.
and analyzed its convergence using NTK. However, their analysis is limited two-layer networks, limiting the model capacity. Under the same two-layer assumption, Jiang et al. (2024) proposed a local-global update mixing method and analyzed its convergence via NTK, while Huang et al. (2021) showed that overparameterized FedAvg converges to the global optimal solution with linear convergence rate.
Bridging federated and centralized learning: We extend the NTK theory from centralized learning to FL with multi-layer networks, showing that infinite net- work width induces constant global and local NTKs, further linearizes both global and local models. No- tably, we prove the equivalence between infinite-width FedAvg and centralized GD, thereby achieving the same generalization performance, bridging decentral- ized and centralized learning paradigms.
Moving beyond the depth constraint on networks, Deng et al. (2022) proved that overparameterized FedAvg with ReLU activation converges in polynomial time with stochas- tic gradient descent (SGD). Fed- reedeenseblee (Shi et al., 2024) employs model ensembling to enhance the generalization of FL. supported by an NTK --baeed convergence analysis. Yet, data heterogeneity is not considered. To address data het- erogeneity, Yue al. (2022) proposed to let clients transmit Jacobian matrices rather than weights or gradients, provid- ing more expressive data representation. Yu et al. (2022) observed that the learning of final layers in FL. is strongly in- fluenced by non-convexity and propose the train-convexify- train (TCT) method to alleviate these issues.
Empirical validation:
We conducted numerous exper- iments on MINST and CIFAR-10 10 datasets, spanning diverse network structures, loss functions, and optimiz- ers to validate our theoretical analysis.
2. Related Work
Data Heterogeneity. While prior works have provided convergence analyses of federated learning with non-IlD data, many rely on restrictive assumptions. For example, some studies assume convex loss functions (Cho et al., 2020; Khaled et al., 2019; Li et al., 2019), others require bounded gradient dissimilarity (Li et al., 2020; Zhang et al., 2023; Wang et al., 2020b), and some assume bounded gradients (Li et al.. 2019; Cho al., 2020). These conditions are often diffieult to satisfy in practice. In this paper, we analyze the convergence of FedAvg via NTK without imposing convex- ity, bounded gradients, or bounded gradient dissimilarity.
Among the most relevant works, Song et al. (2023) estab- lished that overparameterized FedAvg achieves linear con- vergence to zero training loss, and empirically observed that wide neural networks achieve better and more stable performance in FL. By contrast, we theoretically prove that the model divergence in FL caused by data heterogeneity is bounded by Onn where n is the network width, thereby establishing the first quantitative relationship between net- work width and the mitigation of heterogeneity. Beyond this, unveil the generalization performance of overpa- rameterized FL by proving that infinite-width FedAvg and centralized learning yield identical model outputs under matched training iterations, while the theoretical analysis in (Song et al., 2023) solely focused on training loss.
To address data heterogeneity, various techniques have been proposed. Regularization-basee methods (Li et al., 2020; Durmus et al., 2021) introduce a regularization term dur- ing local updates to mitigate client drift. Client selection approaches (Cho et al., 2020; Goetz al., 2019; Zhang et al., 2023) choose subset of clients whose aggregated updates approximate those of all clients. Personalized fed- erated learning (Jeong & Hwang, 2022; Jiang et al., 2024) allows clients to leverage aggregated knowledge while fine- tuning on their local data. Other methods, such as SCAF- FOLD (Karimireddy al., 2020) and FedNova (Wang et al., 2020b), correct optimization bias between clients and the global model to improve convergence. FedMA (Wang et al., 2020a) constructs the global model layer by layer to dimin- ish the impact of heterogeneous data, and MOON (Li et al., 2021a) compares local and global models to correct client drift.
3. Notations and Problem Formulation
In this section, we establish the basic notations, formulate the FL problem, and introduce the metric quantifying the impact of data heterogeneity.
We consider a standard FL setup consisting of a server and M clients. The local training dataset of client i is denoted by D, with X; = [=|(z, y) E P.} and ); {y/ (z, y) E D.} representing the of inputs and labels of client é, respec- tively, where E In and y E B². The global dataset is defined as the union of all clients dataset, i.e., D = UM D with X' = U and y = D
Suppose that each client trains a L-layer fully-connected neural network (FNN), where the width of the I-th layer is
Overparameterized FL. Recent works have made sub- stantial progress in overparameterized federated learning. For instance, Li et al. (2021b) proposed FedBN, which ap- plies local batch normalization to mitigate heterogeneity,
denoted by nt. Then, the output of the 1-th layer of client i's
2

---

Widening the Network Mitigates the Impact of Data Heterogencity on FedAvg
model can be expressed as
are the local Jacobian matrix and error vector, respectively. Similarly, we define the global Jacobian matrix and error vector respectively, as
1= 0 0<l< (1)
fia(2) - a (Wa,kt-1(()+ + bi,),
J (0) _ VoJ E RUXMD)
(8)
(iLfiL-1(T)) + bi 1=L
9 (0) f0) vec())).
(9)
where Way E 1R and b, E R are the weight and bias of the (-th layer of client i, respectively, and (-) is the activation function.
In the (++ 1)-th global round, the model is aggregated by FedAvg, i.e.,
M Σ 1=1
We define 0g - vec ((Wa,,battE)) E I" as the vector of all trainable parameters of client i's model. Then, the model output of a single input sample x can be expressed as fz 0.) and the concatenated output of all samples in X; denoted by [i, 0,) = vee (a(rr, 0a)a ) E & P.I. Similarly, we further define the global model f having the same structure as Sa but with parameter 0 E R Analo- gously. its concatenated output of all samples in X' is de- noted by f(X', 8) = vee (f 0)) ) E RAIPI To sim- plify the notations, we use the short hand (0.) A O.)
O(++))
P
(10)
where we let IT and + 1)T to denote the time in- stants before and after the (t 1)-th global aggregation, respectively. Then, the aggregated parameters ((f+1)r are broadcast to all clients, which yields o((+1)=- @(t-1)r Vi. Consequently, the relation between the global and local Ja- cobian and error at the (-th global round can be described by
9: (0,") Pig (0'))
(11)
and f0) e f(x: 0) in the following.
M
We consider the mean square error (MSE) loss function, and hence the loss function of client i is expressed as
J (0 = Σ Jj (0T)) P. t=1
(12)
1
Σ (D.
(z, 0,) 2
(2)
where P; E x is projection matrix defined as
2|D
0
1
0
0
0
0
1
...
The goal is to minimize the global loss function defined by [P:1_
P
(13)
= E P,, where p. = [D]
At initialization, each client's model parameters are sampled from the Gaussian distribution as follows
kIp, (ELL))
APD x Apdi |P3
whose operator norm
llop 1.
Wh, 1.1 69 = = 62. ~ N (0,o?)
Who M1 (0, Fin
(3)
To facilitate convergence analysis, the following notations are also introduced:
(4)
f (B/Ttr) ff (x,
(14)
Upon initialization, each client updates its local model min- imizing the loss function by GD for iterations. For every local iterations, each client uploads its local model to the server for model aggregation, and then the server broadcasts the aggregated model to each client for the next round. Let t denote the number of global rounds. Then, the model pa- rameters of client i after the r-th (1 < r S r) local iteration in the t-th global round can be denoted by B!T+..
J (0/TTT)= Vogt+-J I (077*),
(15)
g (O/+t))= f (0fr+1) vec()))
(16)
where / (0!T+") denotes the output of the global model when its parameters are replaced with client i's parameters in round tr + r.
In the (1 + 1)-th global round, the degree to which client i's model deviates from the global model is characterized by
Specifically, during the 1-th and (1++ 1)-th global round, say in the (tr + r + 1) total iteration, the model parameters are updated by GD as
| 1r
| - g((+1)r
(17)
Therefore, we use ΣM' pel Ao 11 to quantify the de- gree of data heterogeneity and term it as model divergence. Apparently, when the data is IID, EM P: proaches zero as the number of local data increases. By contrast, when the data is non-III, >M" p:|| (allAe((t))) ||
07 grr+++
ppprt
11 (0,7 9i (B/++)) |D,l"
(5)
| 1|, ap-
where 1) is the learning rate and
Jr (0.) = fo (0.) E Ru
(6) remains non con- and increases with the degree of data het- erogeneity.
gi (0.) = fa (0.) - vec()r)
3

---

Widening the Network Mitigates the Impact of Data Heterogencity on FedAvg
4. Convergence Analysis
Theorem 1. Under Assumptions 1 4, for any small
6g O, there exist Ro > 0, N and Ci O, such that for any Z N, the following holds with probability at least (1 - 60) over random initialization:
no > 0, C >
In this section, we analyze the convergence of overparame- terized FedAvg. We derive the bound on the model diver- gence explicitly and analyze how it influences the conver- gence rate and error.
M Σ
Ps
AA
(g(t+tir
- _
Coo
(22)
We first introduce several notations regarding overparam- eterized neural networks. Let π min(n1,72. n22 :.. n and define the global NTK matrix (Lee et al., 2019) in the 1-th global round as
vn - -)
| )112 S & + TC Ros (1 - g¹)
(23)
(1 (1-9)²
Co (1 - -9))
θeT -²==- - ( J (0*T)-
100 - 10 S Vni (1 - -q)
(24)
(18)
n
2110 3 Ro (1 g³)
| 001 F
(25)
Meanwhile, the analytic NTK matrix is defined as
vn(1 -q)
|, - 6%1lF
+
lim O". 2
(19)
2 C RovK
2n Ro (1 - 9¹) V ViDi
+
(26)
Analogously, the local NTK matrix of the standard parame- terization in the (tr + r)-th iteration is defined as
Vn(D:
(1 q) /n
where q = 1 Dor 3
1h 2 e
θ rr 1 (B!TT)) T Ja (0/T+1) P. n
2
(20)
The detailed proof is provided in appendix B and we present the proof sketch in the following.
The following assumptions are made to facilitate the conver- gence analysis.
Proof Sketch. We use mathematical induction to prove Theorem and the induction hypotheses are (23) and (24). It is trivial that (23) and (24) hold when L = 0, and our aim is to prove (23) and (24) for L+
Assumption 1. The minimum width among all hidden layers n sufficiently large such that the terms of order O(n and higher are omitted.
(Step 1] We first present several essential lemmas in Ap- pendix A, including proving the Lipschitz continuity of the global and local Jacobians and some properties regarding the Taylor series expansion.
The analytic NTK θ is full rank, i.c., the minimum eigenvalue Am of θ satisfies Am > O.
Assumption 2.
Assumption 3. The norm of every input data is bounded, i... ||x|||2 < 1.
[Step 2] Prove induction hypothesis (24) holds for t 1. Due the small learning rate r) dh for large n, we treat time as continuous and use gradient flow to approximate GD, which yields
Assumption 4. The activation function a satisfies
[o(x) - (xx')
(o()),, Ilo' sup xfxx x x'|
00
dg (0,T*T) dr
0 tr
[D, 9 (0/Tt)).
(27)
Assumptions are common in analyzing the overpa- rameterized neural network (Lee et al., 2019; Shi et al., 2024).
Applying the mean value theorem of integral, we can obtain
The learning rate is set to 17 and To is a constant inde- pendent of n, which results in infinitesimally updates during each gradient descent step when n is sufficiently large. Con- sequently, we adopt gradient flow as an approximation of gradient descent, which can be expressed as
y (0/ = 19 e
(0 F. E (0.
(28)
Based on (28) and the induction hypotheses, the local model parameters variation can be bounded by:
|10|T -0 S ToTC [lg (0'T) )lla
+ +110 1100T-601/, (29)
Vn(Da
d0 dr
7 D.
J. (B/Tt)) (o) (21)
Notably, (29) captures the relationship between the dynam- ies of local model and global model, based on which can obtain the recursive relationship between ||0'T - 00 and |[(++1)r - -0 from
Next, we present the main theorem regarding the bound of model divergence and the convergence of overparameterized FedAvg-
4

---

Widening the Network Mitigates the Impact of Data Heterogencity on FedAvg
M Pt ||0 - -8011
linear convergence rate and zero training error as shown in (23).
||0(tt1/T - 10 -
1=1
Remark 3 (Lazy training), Inequality (24) shows that as the network width increases, each global update in overpa- rameterized FedAvg remains confined within an increusingly smaller neighborhood of size O(n b) around its initial- ization, thereby extending the tazy-training phenomenon observed in centralized settings (Chizat et al., 2019) to FL settings.
C lg (0'T)]12 + +100T= +100--00|| (30)
-11112
vn
Using (30), we can prove (24) holds for +++ 1.
[Step 3] Based on the inductions hypotheses and a variation of (29), we are able to obtain (22).
Remark (Constant global and local NTKs). Inequation (25) shows that as the network width increases, the global and local NTK experiences less variation during training. When the width approaches infinity, the both the global and local NTKs are constant, which extends the findings in centralized learning (Jacot et al., 2018) to FL settings.
(Step 4] Prove induction hypothesis (23) holds for 1+ 1. By taking the Taylor series expansion of Di= Pig (0)++)) at g(t+1)r, we are able to derive
M
M Σ pig (of
M Σ
+
p.l2
(31)
Next, we will investigate the training dynamics of FedAvg in the infinite-width regime and compare it with central- ized learning to further investigate the generalization perfor- mance of overparameterized FedAvg.
where S2 represents the remainder terms. Based on the results of [Step 1] ~ [Step 3), we can further bound || + ")22 and Σ pil2,l, respectively, as
M Σ Pig (0))
S s la (0:5)l2
(32)
5. Generalization Performance
In this section, we analyze the training dynamics and gener- alization performance of overparameterized FedAvg as the network width n - co. First, we prove that both the global and local models behave as linear models during the train- ing process. Then, we derive the closed-fowe expression of those linear models and establish the equivalence between infinite-winn width FedAvg and centralized GD.
M
21 2 RoS
Σ
P
(33)
(1 ②)
Plugging (32) and (33) into (31), (23) holds for + 1.
[Step 5] Based on the Lipschitzness of the global and local jacobians as well as the results in [Step 2]~([Step 4], we prove (25) and (26).
We define the linear models flin (6*)) and f!in (B/T-T) as the first-order Taylor expansion of the global model f (OCr) and local model fi (B!++)). respectively:
Remark 1 (Bound on the model divergence). Inequa- tion (22) establishes an upper bound Ç on the model diver- gence caused by data heterogeneity. Since (= O(n ), increasing the network width can reduce the effect of data heterogeneity. Note that we do not impose any strict as- sumptions on the convexity of the loss function (Cho et al., 2020: Khaled et al., 2019; Li et al., 2019), the bound of local gradients (Li al., 2019; Cho et al., 2020) or the di- vergence between local and global gradient (Li et al., 2020; Zhang et al., 2023; Wang et al., 2020b). Instead, we prove that model the divergence is indeed bounded as long as the network is sufficiently wide.
flin (0*T) / (09) + J (00) (0'T -00)
(34)
f!in (0) - fo (8") + Jr (80)7 (0 - -89) (35)
Our main results are as follows.
Theorem 2. Under Assumptions 2 to 4, when n - 0o, we
have
sup ) ( =((0=)]===0 On- 120
(36)
Remark (Impact of data heterogeneity on the conver- genee rate). Inequation (23) characterizes the evolution of training error across the global aggregation rounds. Differ- ent from Song et al. (2023); Huang et al. (2021), the pres- ence of Ç >> 0 here slows down the convergence, making the convergence rale no longer linear and the convergence error no longer zero. Recalling that Ç = O(n 1), widening the network enhances the convergence rate by mitigating the model divergence. When n > 30, we have Ç -+ and the impact of data heterogeneity vanishes, resulting in a
sup 1|5 (0!T**) - fo (0/T+T) = 00(=="nn -, Vi (37) 120,
15rST
The detailed proof is provided in appendix C.
Remark 5 (Infinite-width FedAvg induces linearized global/ocal models ). Theorem 2 suggests that as the net- work width approaches infinity, the global and local models become linear models. This extends the findings of Shi et al. (2024), which demonstrated that the global model can be
5

---

Widening the Network Mitigates the Impact of Data Heterogencity on FedAvg
approximated by a linear model, to show that both local and global models can be well approximated by linear models.
loss functions, and optimization methods. Specifically, we evaluate the impact of data heterogeneity under different network widths, verify that both the local and global models of overparameterized FedAvg can be well approximated by linear models, and demonstrate that overparameterized FedAvg generalizes the same as centralized learning. The number of clients our experiments are set to M = 10, and the dataset as well as the model settings are provided as follows.
Therefore, we can analyze the training dynamic of those linear models instead. The main theorem describing their training dynamics is presented as follows.
Theorem 3. Under Assumptions 2 to 4, when n - 0o and NoT is sufficiently small such that the terms of O (nâT²) and higher are neglected, the linear model has closed-fomt ex- pressions for the global parameters and oulputs throughout the training process:
Non-IID Data Generation. 1) Standard dataset: We con- duct experiments on two widely used image classifica- tion datasets: MNIST (LeCun et al., 1998) and CIFAR-10 (Krizhevsky et al.). To partition the datasets among different clients and generate non-IID data, we follow the approach proposed by Hsu et al. (2019), which employs the Dirichlet distribution with concentration parameter a controlling data heterogeneity. Specifically, smaller value of in- dicates a higher degree of data heterogeneity and we set a = 0.1 throughout our experiments.
o - -10))
notr 00) TeT
0(00) + 00, (38)
(1, (x, 60)
wotr go) - (0) (1-6
9(00), (39)
where (xx) A J (0)TJ(00)
The detailed proof is provided in appendix D.
Suppose there is model having the same structure that trains the global dataset D via centralized GD, whose model parameters at the -thh GD iteration is denoted by 0 and the model output is fxx (. When the initialization of Oeen and 0 are the same, i.e., 0 = 00, the following can be obtained according to Lee et al. (2019, Equations (8), (10), (11)):
2) Small dataset: To facilitate MSE loss minimization using gradient descent (GD) as required by our theoretical deriva- tions, we also use the mini-MNIST and mini-CIFAR-10 10 datasets for binary image classification tasks. The mini- MNIST dataset is created by randomly selecting two classes from the MNIST dataset, followed by randomly sampling 50 images from each class for the training set and 10 im- ages from each class for the test set. A similar approach is used to generate the mini-CIFAR- 10 dataset, which contains 500 training images and 100 test images. To generate non- IID data, inspired by McMahan et al. (2017); Zhang et al. (2021), we assign each class exclusively to specific clients: half of the clients receive all images from one class, while the remaining clients receive all images from the others.
((00) I-- (10
8% = n
g(0
+ 80
(40)
fcea (I, f (=x((110) 80)
- (00 (6(0)(60)-(11---
g(0%). (41)
When 8' = tr, by comparing (38), (39) with (40), (41) and employing Theorem 2, we can obtain
Experimental Models. We employ three types of models: FNN, convolutional neural networks (CNNs), and residual networks (ResNets).
d''n= = 0"T f(x, 0'en) = f(),
(42)
The structure of FNNs detailed in Table I of Appendix E, where width factor k is introduced to adjust the network width. By setting k = 1, 2, 4, 16, we construct networks of varying widths, named FNN1, FNN2, FNN4, and FNN16.
1) FNN:
Remark 6 (Infinite-width Aed generalizes the same as centralized GD). Equations (42) suggest that when the total number iterations of centralized GD and FedAvg are the same, both models share the same model parameters in the infinite-widie wiith regime, thereby producing the same output for an arbitrary test input and achieving the same generalization performance. This means that the impacts of data heterogeneity on the generalization performance vanishes.
We adopt the approach described by Park et al. (2019) to obtain CNNs with varying widths having the base architecture in (LeCun et al., 1998, Figure 2). The details of the architecture are presented in Table 2 of Appendix E, where the width factor i is used to scale the channel size. By setting k = 1,2, 8, 32, we generate CNN1, CNN2, CNN8, and CNN32.
2) CNN:
6. Numerical Experiments
In this section, we verify our theoretical findings by numer- ical experiments spanning various network architectures,
Codes reproduce the main results are available htips: Mggthb..comkkuugg//CMM...
6

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
1.0
0.9
0.9
0.9
0.8
0.7
0.6
0.7
0.5
0.6
Tea Accuracy
0.6
0.5
0.5
5 0.4
0.4
0.4
0.3
0.3
0.3
0.2
0.2
0.2
500
1000
1500
2000
0.1
2000
4001
8000
8000
10000
0.1
2000
4000
8000
10000
Global Round
Global Round
FNNi_aid
Fud4-aid
CMN1 i
CNh8 id
--. CNN1 rid
CANE nid
Reelettlasz
Ruelaets
Radies
FNN2 sid
FNNIG rid
FAN2 1d
FNN16
CNK2.8
CNN32J0
CMN2, CNM32.0d
RecNe12aae 2id -n Redtetlizzdd did
Renker? RestenTIAl
Figure 1. Test accuracy of different network families. Each global round consists of r = 5 local SGD iterations. The left figure shows the accuracy FNNs both IID and non-IID MNIST datasets. The middle and the right figures show the accuracy of CNNs and ResNets, respectively, on both IID and non-IlD CIFAR-10 10 datasets.
1.0
Text Accurancy
0.8
Accurancy
0.6
0.4
0.4
0.1
2000
4000
6000
8000
10000
20000
40000
GO000
80000
10000
Global Round
Global Round
FNNI
FMK1.j
FNNSPJM
CNN1 ad
CANI.ii
CMN32.6 CNN32.wi
Figure 2.
Test accuracy of large networks. Each global round consists of = 5 local SGD iterations, ow = 1, 0.1. The left figure shows the test accuracy of the FNN32 and FNN1 both IID and non-IID MNIST datasets. The right figure shows the test accuracy of the CNN32 and CNN1 on both IID and non-lID CIFAR- 10 datasets.
3) ResNet: The network architecture is based on the work of Zagoruyko (2016), shown in Table of Appendix E. The parameter d represents the number of blocks in each group, which is set to vi = 1 in our experiments. The channel size is fixed at 16, and we vary the width factor & 1, 4, 16 to obtain ResNetl, ResNet2, ResNet4, and ResNet16.
pact data heterogeneity diminishes as the network width increases.
To further verify that the impact of data heterogeneity van- ishes the network width approaches infinity, we set the learning rate with ysis. As shown in Figure 2, the convergence rate and final accuracy of FNN32 are nearly identical for both IID and non-IID data, and similar trend is observed for CNN32. In contrast, noticeable gap exists in FNN1 between IID and non-IID data, which is also evident in CNN1.
in line with our theoretical anal-
6.1. Impact of Non-lID Versus Network Width
Although our theoretical analysis is based on GD with learning rate 7 = On- to show our conclusions can be extended to more practical settings, we use SGD with batch size 64 and set common learning rate 7) = 0.1 with a weight decay of 0.0005. Moreover, for MNIST and CIFAR-10, stead of the MSE loss required in the theoretical analy- sis. As shown in Figure 1, in the non-IlI cases, the test accuracy of FNN1, FNN2, FNN4, and FNN16 decreases by 17.4%, 9.5%, 6.3%, and 2.0%, respectively, compared to the IID cases. Similarly, the test accuracy CNN1, CNN2, CNN8, and CNN32 drops by 44.9%. 26.7%. 5.1%, and 2.4%, while the test accuracy of ResNet1, ResNet2, ResNet4, and ResNet16 decreases by 44.6%, 29.1%, 18.7%, and 14 14.8% respectively. These results verify that the im-
Additionally, to monitor the evolution of model parameters and the global/local NTKs during training, we train FNNs on the non-III mini-MNIST dataset using GD for binary classification with MSE loss. As shown Figure 3, in- creasing the network width diminishes the variations in both the global/local NTKs and the model parameters. For suffi- ciently wide networks, the global and local NTKs well as the model parameters remain nearly constant, exhibiting lazy training behavior.
the practical cross-entropy loss in-

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
2-2
2.
10-22
1us (15/ Timmm
2-34
10"33
2-16
2:6
10"5
2-h
1
100
2010
300
100
300
500
1000
1500
200
Global Round
Total Iteration
Global Round
FNNLD
FN864
FNN256
FMN15
FMN6s
FNN255
FNNI6
FNNtos
FMN226
Figure 3. Training dynamic of NTK and model parameters. Each global round consists of т = local GD iterations, = 1, ow = 1.5,
ab 0.1. The left figure shows the variation in the global NTK, the middle figure show the variation in randomly chosen local NTK. while the right figure shows the model parameters' update during the training process.
3.0
train_con_ FNN5122
train FNN512
10P
Train_fed,_r FNNN512
train fed FNN512
10
10-1
FNN512
test
FNN512
1.5
1.5
19-2-
10-4
FNN512
fed.
FMN512
test fed. FNN512
F6
1.0
1.0
10-9.
FNN512
10-4
0.5
0.5
100
230
3000
100
500
300 1000 1100 XX00 2000
0
40
80
120
160
200
100
200
300
400
500
Gloabl Round
Total Iteration
Total Iteration
Total Iteration
Figure 4. Output difference between FedAvg and linear model. Each global round consists of local GD iterations, yoo ow 1.5, = 0.1.
Figure 5. Training and testing loss of FodAvg and centralized Iearn- ing. to = 0.1, dw 1.5, σ 0.1. The global round of FedAvg consists of r = 2 and r = 5 local GD iterations in the left and right figures, respectively.
6.2. Linear Approximation of FedAvg
f(x, (TT) on both the training and testing dataset of mini- CIFAR-10, by ensuring e' tr for fair comparison. can be observed that the outputs of FedAvg and centralized learning are almost identical under the same number of GD iterations, empirically confirming Theorem that overpa- rameterized FedAvg generalized the same as centralized learning.
To show that overparameterized FedAvg can be well ap- proximated by linear models, we train FNN16 and FNN512
on the non-lID Mini-MNIST dataset with MSE loss for bi- nary image classification using GD. In Figure 4, we analyze the difference in the outputs between the global model f in FedAvg and the global linear model flim throughout the training process. Additionally, a randomly selected local model examined by comparing its outputs with those of the corresponding linear model f!in.
7. Conclusion and Future Directions
In this work, we established a quantitative relationship be- tween neural network width and the impact of data het- erogeneity in FedAvg. We proved that the impact of data heterogeneity on the convergence of FedAvg diminishes at a rate of Onn with increasing network width n and vanishes entirely in the infinite-widd limit. In that regime, we extended NTK theory from centralized learning to FL, showing that both the global and local models in FedAvg linear and have constant NTKs. Furthermore, we derived closed-form expressions for the model outputs of FedAvg, revealing the equivalence between infinite-width FedAvg and centralized GD in both training dynamics and gener- alization performance, under matched training iterations.
As shown in Figure 4, we can observe that, for FNN512, the outputs of the global model f and the linear model flin remain nearly identical throughout training. By contrast, the narrower FNN16 exhibits noticeable difference between f and flim, A similar trend is observed for the local models Ji and flin These findings confirm Theorem 2 that wider networks enable linear approximations to align more closely with the dynamies of FedAvg.
6.3. Comparison of FedAvg with Centralized Learning
To compare overparameterized FedAvg with centralized learning, Figure we evaluate the loss of feou (2, 8C) and
8

---

Widening the Network Mitigates the Impact of Data Heterogencity FedAvg
Extensive experiments on MINST and CIFAR- 10 datasets validated our conclusions across different network architec- tures, loss functions, and optimization methods.
analysis. In International Conference on Machine Learn- ing, pp. 4423 4434. PMLR, 2021.
Jacot, A., Gabriel, F,, and Hongler, C. Neural tangent ker- nel: Convergence and generalization in neural networks. Advances neural information processing systems, 31, 2018.
These theoretical findings provide valuable insights for prac- tical federated learning. Notably, the linear dependence between model outputs and parameters suggests a potential commumication-efiee
FL strategy: clients may transmit only the model oulputs instead of the complete model pa- rameters for aggregation. This approach may significantly reduce the communication overhead of FL and deserves further investigation. Another essential direction for future research involves extending these analyses to more realistic FL settings by relaxing the idealized assumptions, such as infinite network width and continuous-ting gradient flow.
Jacot, A.. Gabriel, F., and Hongler, C. The asymptotic spectrum of the hessian of DNN throughout training. In International Conference on Learning Representations, 2020.
Jeong, W. and Hwang. S. J. Factorized-FL. Personalized federated learning with parameter factorization & similar- ity matching. Advances in Neural Information Processing Systems, 35:35684 35695, 2022.
Acknowledgments
Jiang M., Le, A., Li, X., and Dou, Q. Heterogeneous personalized federated learning by local-global updates mixing via convergence rate. In International Conference on Learning Representations, 2024.
We sincerely thank the anonymous reviewers for their valu- able comments and suggestions. This work is supported by National Natural Science Foundation of China under Grant 62301015.
Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S., Stich, S., and Suresh, A. T. Scaffold: Stochastic controlled averag- ing for federated learning. In International conference on machine learning, pp. 5132-5144 5123. PMLR, 2020.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Khaled, A., Mishchenko, K., and Richtárik, P. First anal- ysis of local GD on heterogeneous data, arXiv preprint 2019.
References
Krizhevsky, A., Nair, V., and Hinton, G. CIFAR-10 (cana- dian institute for advanced research). http://www.cs. toronto. wtt.ee////i/iffffa..nlllll
Cho, Y. J., Wang, J., and Joshi, G. Client selection in feder- ated learning: Convergence analysis and power-of-choice selection strategies. arXiv preprint arXiv:2010.01243, 2020.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient- based learning applied to document recognition. Proceed- ings of the IEEE, 86(1 1) 1):2778-2322, 1998.
Deng, Y., Kamani, M. M., and Mahdavi, M. Local SGD optimizes overparameterized neural networks in poly- nomial time. In International Conference Artificial Intelligence and Statistics, pp. 6840-6861. PMLR, 2022.
Lee, J., Bahri, Y., Novak, R., Schoenholz, S. S., Pennington, J., and Sohl-Dickstein, J. Deep neural networks as gaus- sian processes. In International Conference on Learning Represenaatons, 2018
Durmus, A. E., Yue, Z., Ramon, M., Matthew, M., Paul, W., and Venkatesh, S. Federated learning based on dynamic regularization. In International conference on learning representations, 2021.
Lee, J., Xiao, L., Schoenholz, S., Bahri, Y., Novak, R., Sohl- Dickstein, J., and Pennington, Wide neural networks of any depth evolve linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Goetz, J., Malik, K., Bui, D., Moon, S., Liu, H., and Kumar, A. Active federated learning. arXiv preprint arXiv: 1909.1 12641, 2019.
L., Q. He, B., and Song, Model-contrastive federated learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp- 10713 10722, 2021a.
Hsu, T.-M. H., Qi, H., and Brown, M. Measuring the effects of non-identical data distribution for federated visual clas- sification. arXiv preprint arXiv: /
2019.
Huang, B., Li, X., Song, Z., and Yang, X. FL-NTK: A neu- ral tangent kernel-base based framework for federated learning
Li, T. Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, Federated optimization heterogeneous

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
networks. Proceedings of Machine learning and systems, 2:429-450, 2020.
T Dinh, C., Tran, N., and Nguyen, J. Personalized federated learning with moreau envelopes. Advances in neural information processing systems, 33:21394-21405, 2020.
Li, X., Huang, K., Yang, W., Wang, S., and Zhang, Z. On the convergence of FedAvg non-IID data. In International Conference on Learning Representations, 2019.
Wang. H., Yurochkin, M., Sun, Y., Papailiopoulos, D., and Khazaeni, Y. Federated learning with matched averaging. In International Conference Learning Representations, 2020a.
Li, Jiang, M., Zhang, X., Kamp, M., and Dou, Q. FedBN: Federated learning on non-llD features via local batch normalization. In International Conference on Learning Representations, 2021b.
Wang. J., Liu, Q., Liang, H., Joshi, G., and Poor, H. V. Tack- ling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 13:7611-7762 2020b.
McMahan, B., Moore, E., Ramage, D., Hampson, S., and Arcas, B. A. coammunication--ffcicee learning of deep networks from decentralized data. Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
Yu, Y., Wei, A., Karimireddy, S. P., Ma, Y., and Jordan, M. TCT: Convexifying federated learning using bootstrapped neural tangent kernels. Advances in Neural Information Processing Systems, 35:30882-30833 2022.
Neyshabur, B., Tomioka, R., and Srebro, N. In search of the real inductive bias: On the role of implicit regularization in deep learning. In International Conference Learning Representations, 2015.
Yue, K., Jin, R., Pilgrim, R., Wong, C.-W., Baron, D., and Dai, H. Neural tangent kernel empowered federated leamn ing. International Conference on Machine Learning. pp. 17783-25803. PMLR, 2022.
Neyshabur, B.,
Li, Z., Bhojanapalli, S., LeCun, Y., and Srebro, N. The role of over-parametrizatton in general- ization of neural networks. In International Conference on Learning Representations, 2019a.
Zagoruyko, Wide residual networks. arXiv preprint 2016.
Zhang, J., Li, A., Tang, M., Sun, J., Chen, X., Zhang, F., Chen, C., Chen, Y., and Li, H. Fed-CBS: A heterogeneity- aware client sampling mechanism for federated learning via class dass-imbalance reduction. International Confer- ence on Machine Learning, pp. 41354-4136 PMLR, 2023.
Neyshabur, B., Li, Z., Bhojanapalli, S., LeCun, Y., and Srebro, N. Towards understanding the role of over- parametrization in generalization of neural networks. In International Conference on Learning Representations, 2019b.
Park, D., Sohl-Dickstein, J., Le, Q., and Smith, S. The ef- fect of network width on stochastic gradient descent and generalization: empirical study. In International Con- ference on Machine Learning, pp. 5042-5051. PMLR, 2019.
Zhang, X., Hong. M., Dhople, S., Yin, W., and Liu, Y. FedPD: A federated learning framework with adaptivity non-llD data. IEEE Transactions on Signal Processing, 69:6055-6070, 2021.
Zhao, Y., Li, M., Lai, L., Suda, N., Civin, D., and Chandra, Federated learning with non-IlID data, arXiv preprint 1rx6v01806.00582, 2018.
Reddi, S., Charles, Z., Zaheer, M., Garrett, Z., Rush, K., Konečný, J., Kumar, S., and McMahan, H. B. Adaptive federated optimization. In International Conference on Learning Representations, 2021.
Shi, N., Lai, F., Kontar, R. A., and Chowdhury, M. Fed- ensemble: Ensemble models in federated learning for improved generalization and uncertainty quantification. IEEE Transactions on Automation Science and Engineer- ing, 21(3):2792--880. 2024. doi: 10. 10.1109/AASE.2023. 3269639.
Song. B., Khanduri, P., Zhang. X., Yi, J., and Hong, M. FedAvg converges to zero training loss linearly for over- parameterized multi-layer neural networks. In Inter- national Conference on Machine Learning, pp. 32304 32330. PMLR, 2023.
10

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
A. Essential Lemmas
In this section, we introduce the necessary lemmas used deriving Theorems
Lemma 1 (Local Lipschitzness of the Jacobian). There exists constan! C > 0, such that for any C² > 0, with high probability over random initialization the following holds:
1 |1(00 - ((0)) s010- v0,0' E B(0o, C'n
(43)
|1(0))le SC
where B(0o R) & 10: 1 0ol12 R}.
Lemma has been proved by Lee et al. (2019, Lemma 1) and we will apply it directly.
Lemma 2 (Local Lipschitzness of the Local Jacobian). There exists constant C' > 0, such that for any C' > 0, with high probability over random inirialization the following holds:
[J:(0) JF(0') -010001l,
Vn
v0,0 0' E B(fo, -t
(44)
||1:(0)|l= SC.
where B(0. R2) 2 (0: 110 0oll2 < R}.
Proof. Since (0) J(f") is the concatenation of J:(0) - J(0 for 1,..,, ,M, we have
1 |],(0) J.('')ll
1 Vn [I](0) J(0))l s C 0'1
(45)
where the last step applies Lemma 1. Similarly, since J(Ø) is the concatenation of all (4),), we have
|,00)) ||(()]]> SC
(46)
Lemma 3. For square matrix A whose norm satisfies |All S PA. the remainder term Qie (e) sum of second-ord -oddr and higher terms in the Taylor series expansion of -A satisfies
(-
4 RT. iee the
A) < 2 eA 2
(47)
Proof. By taking the norm on both sides of (()=A) = Σズ (-1) 2tt R..
we can readily obtain
A
||A||
|A|
k-2
2 |lll
2 eAs
|| 1|
<
<
(48)
h
2
(k
2)!
2
2
2=2
Lemma 4. For any non-negative aggregation weights combination { M. ToT sufficiently small such that the terms of ) and higher are neglected, the following equation holds;
M
Pre
7o70. -00 DMM P
(49)

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
Proof. Employing the Taylor series expansion both sides of (49) yields
M
Pie 7o76,
Σ Pr [I - 107 0 (1ā+2)] =1 7/77
pi O
(50)
and
-o Σ p:::: = - n0T Σ p,o, + 0 01 7
p p
(51)
respectively. Comparing (50) with (51), we can obtain
Σ Pie 9o78, = Eis P,O.
(52)
B. Proof of Theorem 1
Since the width of network large, with the initialization described Section 3, the output f(6)) converges Gaussian process (0, X(X,, X)) according the central limit theorem, where K((X', X) lim E [(())) Therefore, for n-+00
arbitrarily small 6 O, there exist constants Rg > 0 and n' such that for any n nt', with probability at least 6g over random initialization, we have
||g (0") (0') S Ro.
(53)
In the following, we prove Theorem by mathematical induction, where the induction hypotheses are
|=(00)))12 S SgoRo 200 RoC (1 g¹) 110 -0112 < vn 9)
(54)
(1 (1-9)²
ToTRRo (1 -
(55)
When / = 0, (54) and (55) trivially hold. Then, we aim to prove the hypotheses hold in the (L + 1)-th global round, i.e.,
200oTCC (1 11--++1 g²+1))
(0(0
1 Ro +
(56)
(1- g)²
CoTCRo ((1
(57)
vn
Referring to (5), the local update step of client in the (tr + r + 1)-th iteration is
g
05T J. (0!TT)) gi (0/TT)..
(58)
D2
Since 77 is small, we can approximate the local update with gradient flow by making time continuous, yielding
d
11 dr [D,
J; (B!T+)) 9: (B!T+)).
(59)
Further applying the chain rule, we can obtain
dg (0,
(ofrtr)² do
dr
dr
" J (0) Ja (B.T+T) g: (0(TT))
|D,|
=n|pi" (0!T+T)² Ji (0,T4)) Pig (0,
170
12

---

Widening the Network Mitigates the Impact of Data Heterogeneity on FedAvg
m [D:] 9 (O!T*))
(60)
where the last step follows from the definition of the local NTK in (20). By integrating both sides of (60) from to r and applying the mean value theorem for integrals, we have
9 (0, =e 137l
9 (0 fg (0.r)
(61)
Replacing r with 7 yields
9 (0)
18
no 6 t
g ( 8g (0,7)
(62)
B.1. Proof of (57)
We first prove (57). To bound | 00)
we proceed the following derivations.
d 0 09112
d0
dr
dr
" || (0!T+") 9: (0! 112
|Dj
"I
|| (0/ Pig (0,+**) l12
|D;
"
|| (0/T
IIP 119 (l(
[Dal
1/Cv
||g (0!++)) ll2
[Di
CoC
moz tt
g (0*T)
Co
||g (02))
Dn
op
CoC
|/g (f*r)
I|
(63)
Dn
where the first step is obtained by applying the chain rule and Cauchy-Schwaed inequality, step (a) holds because of 1. and the last step holds because 077+ is not full rank from its definition (20), yielding
Lemma 2 and P op
2n le op
2o 12. (9
I
1. Integrating from O to on both sides of (63) yields
MorC Dn|
10 00112 s
||g (ATT) ll2 + -|00T 8 101
(64)
Further replacing r with 7 yields
Cn
|1 - 000112 -0012 S
||g (0*T) 112 + +110r - -00 00
(65)
/n DD
According to (10), we can obtain
lolt23r
00
Σ ps 10 - 800
M
pi |
80112
i=1
(a) M
noTC
|19 (((T)
+ -110---001 80) 11112
D.
13

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
MroTC
||g (02T)12 + +|00T 10
Vn|Di
MC
||g (0'T))12 + +1100T - -00112
Vn
NoTC
200TCC, Roc
20 RoC
To (1-9*)
g¹ Ro
+
Vn
(1 (1-9)²
(1-9)²
9)²
(1(1-g)
No 11-9- (1-gtt1)
(66)
V 9)
where step (a) comes from (65), step (b) holds because (D] > M, step (c) applies the induction hypothesis (54) and (55).
and the last step omits S = On according to Assumption 1. Therefore, (57) is proved.
B.2. Bounding the Model Divergence
To prove (56), we first bound || | 1r12. According to (64), we have:
110 - -0 <
|1g (0tT) -1100- -0%12 Dn|
C
8' Ro 200 RoS
2no CC,
Co (1 g9)
+
Vn|Dal
(1-4) 9) CoTC (1 g²)
(1-9) -)
v (1-1)
mr Ro
+
(67)
Vn|Di]
Vn (1
where the second step employs the induction hypotheses (54) and (55), and the last step omits (17), we have
O(n Referring
ge
g((+1)r
(0 89) - ((0t1)--9) 10011
M
(0 09) Σ
Pi (0+ - 000
i=1
10
10) + Σ Pi 11 00)
(a)
oo Ro
CoTCRo (1-9') 9²)
M
ToTCq' Ro
VnlD. + NoTC'oo (1 g¹)
Σ
+
Vr|D.
Vi(1 (1-19)
Vi 9)
i=1
Cr Ro
CoTCRo (1- (1-1))
Cq Ro
CoTCRo (1 9²)
Vr(Dil
vπ -9)
Vn|D|
VTi 9)
(b) vorCy' Ro
CoT (1 - (-9
OoTCy Ro vn
NoTCRo (1 - 9")
Vn
vn (1 ②)
vū(1 ②)
2noTCoo
1 1- q'
9
Vn
q
2n (1 1
Vn
②)
20ooTRo
< /n (1 -1)) q)'
(68)
14

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
where the step (a) is from (67), step (b) holds because [ 2 and [D] M, and the last step hold because require
0 S q < 1.2 Therefore, the data heterogeneity term can be bounded by
Σ pb
|
2017 CRo Pr
2np C-
(69)
vn (1 [19)
vn q)
Notably, we have proven (22).
B.3. Proof of (56)
Finally, we prove (56) to finish the induction. Taking the Taylor series expansion of (0)7+)) f(1+1)r yields
9 (0,
(D((4))r) + J (Bl(+1)r) A + +2,
(70)
where S2 represents the remainder terms of order two and above. Taking the sum of both sides yields
Σ
M
Pig
(0, =
Σ Pig
M
(e(ct1)r)
Σ P.J (Blit-ir)² T
Agle+I3r
Σ
M
p.
i=1
Σ pig
(o((+1)r)
Σ p:S2
M
= (o(+t))r)
+ Σ pile,
(71)
where the second step holds because E DiJ (D()+1)r)* T Aott+ = 0 according to the model aggregation (10). Rewrit- ing (71), we have
M
M
g (
Pig (0/
prlla
(72)
Taking the norm of both sides yields
M Σ i=1
|1o (oi++1>7)1
11,
Pig (0,
+
Σ p.S
(73)
In the following, we bound
| Σ Pig (0/
||, and />ما/// M
p,l2
respectively.
1) Bounding
EM Pig (o/+tr)l
According to (62), we have
Σ Pig (0!T+T)
=
D
P.e
1
9
(077)
2
M
)
Pi
1
NOT
th:
e + Ω
10
upy
(0
D
Pi
I
nOT
+ Ω
e
18
119 (85T) 112
De
lop
Σ
Pb
T
noT
(
|/g (0T7))2
M Σ
Ps
S2
e
2.
0 18
)
||9 (077)22
[D
op
op
We prove that there exists Tjo > such that 0 < q < at the end Section B.3.
15

---

Widening the Network Mitigates the Impact of Data Heterogencity on FedAvg
I
1OT
M
+
|1 (0:7
+
Σ
M
Pi
52
c
mus 12,
o
t
)
|/0 (877)12
(74)
lop
lp
where the second step is obtained by employing the Taylor series expansion and Qi () = 2k=z (-1)* A Then, we
bound |
T 2
i=
M
op -i=l Ps
M and
e
respectively. op
I
M
can be rewritten as
op
1
7a+
M
0
707
7
+
7o7
107 [D]
M
+
D
[D]
|D]
[D]
[D]
M
NoT
+
7OT D
0
gT++
[D]
[D]
[D]
[D]
op
op
op
?o Am [D]
+
707 [D]
||0
e
10T D
op
i=1
op
M
<
7oT Am
+
1OT
|
1oT
O
θ
+
(75)
[D]
D
D
t=1
op
According to Lee et al. (2019, Section G.1), there exists n², such that the following event
Am 3
|100 6 S
(76)
with probability at least 1
5e hold for any n 2 n². As for
EM 0r
60| op
according to the definition of local
NTK (20), we obtain
M
M
0
00
Σ
J (0/++)))" T Jj (0!T+)) P.
Σ J (80) T. J; (09) P. i=1
op
op
M
1
J ( T Js (B( P. J (0") J, (6.) P.
n
op
M
J (o.T T Ji (0! - J (00) J, (0%)
||P!llop
lpp
1 >
| (o Ja (0 - J (0 J (09)
op
M
( " [] (.) Je (89) + [] (0 - (09) J. (09)
|| (0 1lap || (0)T - J. (89) + = (0/T+m)==-1))
- J (B!! 7|
|| (8 ll.p
lp
|| (B!T+):)) | J| (0TT+)) J. (8!) (r>)>> + | (0/T+t))" J ((9))" J (8.) 11
M
< 20
Σ)
|| 10
(77)
1=1
16

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
where the last step holds because of Lemma and Lemma 2. Plugging (67) into (77) yields
M
2n0F.
T Ro (1 x¹)
Σ
θ +
O
VTlD.
vū (11-9 -)
2 Ro
2 Ro
vn
vn 1 9)
4 - Ro (2 (2-9) vni (1 9)
Am
(78)
3
where the second step holds because R, 2 7, [D:] 1 and 0 S q 1. and the last step holds when
n² with Then, plugging (76) and (78) into (75), we can obtain
n" A4 (2(2- Ao. (1-
1
7
707 7
5 1
(79)
3/D]
To bound E pi 12(e
app we first bound
| (T ll.p According to (20), we have
| FP
J (o, Jj (0! P. up
op
< n
|| (0 | 1| (0 ) || P.P:llop
|| (o)+*))) || || (0)++)))) |l
< C²
(80)
where the last step holds beacuse of Lemma and Lemma Then, according to Lemma 3, we can obtain
M
Pt
e
[2.
o
18
Pi
JoT
C2
1..TT
C2
Dii
M
7h
72
aurc? 1PT
2 24
2/D
1=1
70
2
n C
CT
(81)
2
Further plugging (79) and (81) to (74), we have
Σ
M
pig (0
S 511
NoT A Am
+ enorea)
[|9 (0/7) e llg (0(T)
(82)
3/D
2
2) Bounding >M Pi Roll2=
By defining
[ = g(0(++1) + 310
),8 E |0, 1]. we can obtain
r (0) (0L+22))).
I' r'(0) Vy (p(c+1)r) A0 [ y (
(83)
17

---

Widening the Network Mitigates the Impact of Data Heterogencity on FedAvg
We further define
" (1-)) (8).
(84)
Then, we can obtain
11 (1) - u (0)1l2 - (I) (0)
|g (0 9 (g(++1rr) - Vv (Bltt))r) A 2
= ||0ill2-
(85)
According to mean value inequality, we have
I|u (1) u (0)122 s sup || (8)122 (1 - 0).
(86)
05331
Combining (85) and (86) yields
||2ll22 I sup us951
|| (8)1l2-
(87)
Taking the derivative on both side of (84) yields
u' (3) = r (8) (1)) (8)
= (1 (1- B) [" (9)
= (1 - 3) (20 V' vig (oR+UI++ + soo(tturr) AQlit1T
(88)
By plugging (88) into (87), and then (87) into || E/ p:02:12- we can obtain
M
Σ
Pe
110,l12
M
$
Σ
Pr sup
|| (9)1l2
05951
1=1
M
(AO T V2g (ol--Dr + (
( s
Pt (1 - 3)
sup
05651
M
Σ pi (1 - 8) /|/
Vyg (pl+1)) (3
sup 2 2
M
Σ Pi (1 - B) AQ (
Civ (((411)1
(6)
M
Coo
C CivR- 2n1oTCoo
Σ
Pi (1 B))
Vnl (1 - 9)
VT(11 ④)
1=1
(6)
4n6 C2CC
Vn 9)²
20, RoC
(89)
(1-4)
where step (a) applies V V (0) S Ci, which is proved by Jacot al. (2020, Lemma 1) under Assumption 4, step (b) comes from (68), step (c) holds because 3 E [0, 1], and the last step comes from (69). Plugging (82) and (89) into (73):
|| (o(t-1+)) < sgllg (0tT) 11, + 20 RoS
(1 9)
18

---

Widening the Network Mitigates the Impact of Data Heterogencity FedAvg
20,
200
21 CC2 RoC
q qp
Ro
+ (11
(11-0) - g)
(1 (1-9)²
(1 (1-4)²
(1-1) 1)²
200TCC,
20
g
Ro
+
(90)
(1-9)²
(1 (1-9)²
where step (a) applies induction hypothesis (54). At this point, we have proven (56). Recall that (57) have also been proven and we require 2 nn' n", n" during the derivations, and hence the induction hypotheses (54) and (55) hold for N A max(n', .n""..
Noting that we also require 0 S < 1 to complete the proof, we prove there exists no > 0 such that 0 S q < 1 in the following. Referring to (82), we rewrite q as a function of 17o as
q (71) = DoTAA Am
1h -4
-00
+
(91)
3/DI
2
Taking the derivative of (no) with respect to 7Jo, we obtain
d' (no)
-Am
+ "
+
C6
Barc
(92)
3/D]
Further taking the second derivative, we have
C8 +
d" (no) = 72C + Ch
70
>
(93)
Therefore, o' (n1) monotonically increasing. According to Assumption 2, we have Am > 0 and hence obvious that (0) 0 and lim o' (77) 0, and hence there exists 1/6 such that o' (Th)
holds. Consequently, q (m))
70000
is monotonically decreasing on (0, vol. and monotonically increasing on (%6 oo). Additionally, from (91), we have lim 9 (17) = 1. Consequently, q (7%) 2 0, 0 S q (77) < 1 holds for 0 < 7o S % Otherwise, if q (%) < 0, there 70
exists no (0, m6) such that q (n6) 0 holds. Then, 0 S (1o) < 1 holds for 0 no S n6- To sum up, long as 0 < Too min 16. 11',. 0 < (7o) holds.
B.4. Bounding the Variation on Global and Local NTKs
We continue prove (25) and (26) in Theorem 1. Referring to (18), we have
| e0|
1
J (BiT) T J (0'T) - J (80) "(00"""((( (69)
n
1
|| [] (BrT)= T - (80)²1 J (B7) J (80) T [J (0+)) - J (8"))
n
1
|J (D'T)- J (09) || (l+) + 1|
1 (0+T) - J (l))
20 110 -001/2 21/0 C Ro (1 - q²) vn (1 -
(94)
where the fourth step holds because of Lemma 1 and the last step holds because of (55).
Referring to (20), we have
| 0%
1
J (O!T+M) T Ja (0tTt) P. J (09)T J. (0") p.
1
J (O!Ttr) J. (D/Tt))= J (0.)²
T J. (0%) |PPll
| (o! J, (0/ - J (0!)² T J. (09)
19

---

Widening the Network Mitigates the Impact of Data Heterogencity on FedAvg
D
(0/rtr)" ( J. (0/ + (69)" [J. (0) - J1 (0?) F
n
(0 J (09) || J| (0, + k|Da |J (0 || || (0) - Js (0.) F
(2)
20 /E|D. 110TT110TTT> -00
-
2nord' C RoVE 2 Ro (1 - q-)
(95)
(1-9) vn
where step (a) of Theorem I.
employs Lemma and Lemma 2, and the last step holds because of (67). We have now completed the proof
C. Proof of Theorem 2
To prove Theorem 2, we first bound | (0((+1)r) 9 (0(ct1)r) 1|2- According to (34), we have
Jin (0!T+)) = (89) + J (09)² (0tT+r - 89)
(96)
Taking the derivative with respect to -- on both sides yields
flin (0/T**) = J (89)
(97)
Then, consider
d
(e1
6
[9 (0. - -g (0.Tt))]
dr
0
elPa
our
lin (QTT)) - 9 (
Di
Tlo
- J (0?) a.th (0, + ر (o]Tt))T Je (0, (0
+
[
e
Vn O
e 1PT
(g lm (0e g (0,
[Di]
1/0
up
6
1
1 J (0T J. (0,0tt)) Pig (0/-t11) ( (0%) P.g ii (0) +
+
e
1P7
[D=]
n
Jur 0
ig (0T+)T) 9 (0,
+
nu
e p.ol
Jur o
[-0 (0/ + 07r y (0[T+T)
D
flo
nor R.
e
(0)
0 g (0,
(98)
[D:
By integrating both sides from 0 to r, we obtain
jin (B(+T)) (B! e 7 e
[g Jin (Q!T) 9 (0TT) +
7o
2 3 0 / e
(0/Tt)- -09) g (B(TtT) dr.
e
D
(99)
The aggregation of linear local models yields
M Σ 1=1
y Jin (o(ct127)
Pig ii (o/TtT).
(100)
Further considering equation (72), we have
M
M [ lin (0)(t)) (B]r+)] + Σ p.
bin (o(l-1))) (B((t2)))=
= Σ
Pi
(101)
2m1
20

---

Widening the Network Mitigates the Impact of Data Heterogeneity on FedAvg
By taking the norm on both sides, we obtain
9 Jin (p( (D(++1177) 11
M
M
Ps [9 (D/TtT) -9 9 (0/T+T)]
+ Σ
p,
2
M
M
Σ
Pi
[9
(dur+))= 9 (0/
+
i=1
i=1
M
Σ
p,e
1871
0
س (B(T)- (8(T)]
+
Σ
Pt
110
[2:] 6
00 (0 0.) g (0/T+)) dr e e|b;7
D
2
M
+
Σ
2
(102)
where the last step is obtained by substituting (99), and E P:still2 2 is bounded in (89). In the following, bound |
Pie 68 (glin (0*T) - 9 (8))) ll2 and || MM
Ps no e 18 0 So Jee ",," 8 (0 - -Q) g (0)**)) drllop re- p..
spectively, and finally bound || سنرر (o(rt))- f (0!T-))l as well as ||!!!m (0:-T) fa (0,T*))lg 112
1/1
C.1. Bounding
E P.e
no o (g (0tT) - g (0*²)) l1 io
By taking the Taylor series expansion of 10,1
2oz o obtain
Σ
(g (0*T) - 9 (0*T))
Pie
12
M
<
Σ
P&e
no 6.
||9 (BTT) - ((*²)]
M
7
0 + (--10
1190th (0PT) - (02))12
Ps
[D;
op
7/0 7
|| (0t²) - 9 x00t)
11 + Σ pi 12 (o 188 0P
M
11 110 lin (OeT) 9 (12)
(103)
D
Next, we bound
n D
and
Σ 1=
12 c
no 12,
6.)
respectively.
up
no 7
OT
OT
07
D
(D)
[D]
op
OT
OT
0%
[
(D)
[D]
NoT Am
DO
07T
[D]
D
[D]
1p
7oTAm
noOT
[D]
D
[D]
(3)
T
Au
[D]
3/D]
Au
(104)
3 3/D]
21

---

Widening the Network Mitigates the Impact of Data Heterogencity FedAvg
where step (a) applies (76). To bound
Ex=
llpp' we first bound |6 as
i=
|
)( (.) P
op
n
op
1
|( |1 |1.(0))) 11=
PPPillop
71
1
|
|1.(0))
n
C2,
(105)
where the last step applies Lemmas and 2. Then, we have
M
00
M
Pi 110
2
ps
Ω E
18,7
e Dj
[D;]
M
Th T
more2
e 12.1
DDD
M
2
i=1
=
2
r²
N0TC2
(106)
2
where the first step employs Lemma 3. Further, plugging (104) and (106) into (103) yields
M
Σ
pie
(g (0²7) - g (077)
S s 1|18 lin (o²²) -g g ll
(107)
C.2. Bounding 1 Pt
noT oo Tea
So
nor of (0 0.9) 9 (0!T**) drll.p
[D.
Considering p. =
we can obtain
D
M
ToT 6 [P21
nor
oe (0/ (0rtt--99] -0%) g (0 dr
P1
elP,
1p
M
770
120g
8
mut o.
24
(0/ - -00) (r.
dr 10p
'D
M
ToT
(0)
- -0?)
e
iPa1
9 (T:)
D
M
No7
[1
6.9 llpp
no --)
0
||9 (0/T+)))l)
op
(b)
A
NoT
[(
0.09 1| |19 (o.T+F))))
[D]
M
r+++
[
- 6,9 |1
e
|/g (8/
e
op
22

---

Widening the Network Mitigates the Impact of Data Heterogeneity on FedAvg
M
7l7
- lp 19 l1:
(108)
where step (a) holds according to the mean value theorem of integrals and P. (0, 7), step (b) holds because @ is not full No (+ |Psl
rank according to the definition of local NTK (20) and hence ||e
op S
on(Fe-)) [Pal
Amta (m)) 1, step (c) comes
from (61) and r E (0, ft). the last step holds because 6 ( is not full rank. Then, we proceed to bound || - 0 'lop' which can be derived as
-
| (of T Jo (0, J (09)² J. (0.) 10p
|| (o+ - -=-+""""=== (0.) "JJ, (DoTt + J (0?)² [] (O,T+Fi)) - J. (0971110)
= J (o T - J -090 7 |J, (0,Tttt || + | (09)) 7| || (OSTtFi) - J: (0))
s 202 1 - -101
( 2n. Rod
+ 2n 8 Ro (1- 9²)
Vn(D,]
vn (1 9))
C Ro
2 C Ro (1 - 9¹)
+
(109)
Vn
vn (1 (1-9)
where step (a) comes from (67) and the last step holds because f; S T and [D;] 2 1. Plugging (109) into (108) yields
Σ
P
7/0
no
el
ugr 06
(0) g (0 dr
D2
MrOT
2 o Ro
+ 20 Ro (1 - q¹)
|19 (0!-)
[D]
Vπ (1 -- 9)
M
2a
2 Ro (1 - -0)
2 200 Ros
2no CC2 MDTTCCRRo
9¹
Ru
+
[D]
Vn
vū 9)
(1 - g)²
(1 - g)²
M
2n
2 Ro (1 - g²)
9¹ Ro
[D]
Vn
Vū 9)
2M , C2 Rgg' (1 (1- /n|D| (1 ④)
21 20 R/g' (1 - +1)
(110)
Vπl (1
where step (a) comes from (23) in Theorem and step (b) omits S0 = O(n
C.3. Bounding س (0+++)- - f (0(++) [|| and ||l!" (0/++))= - Se (0,7777 |12
Plugging (110) and (107) into (102) yields
S q سلواا (gLT) - 9 (0tT) 0 q-1
M
| (0l4+2)77) - (o(+z)=)|| ||
+
Σ i=1
p.lp
Vn (1
2n 3 R3g" (1 (11)
7Cp Ros
S q |g (8th) g (0tT) 112
+
(111)
Vnl (1 ④)
(1 q)
23

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
where the last step comes from (89). By recursively employing (111) and considering the fact ||glh (00) 9 (00) obtain
| (olet1)9)- -9 9 (Dlcttrr)
|| S 1 r
2 C R (1-
2n CC2 7 RoC
V (1 -9)
(1- 9)
21 2 Rig' (1- -1
+ 2010 Ros (1 (1-9
(112)
Vπ () 9)²
(1 - q)
Considering Ç
Rg and replacing t with t, we have Vn(1-g)
sup 9 Jin (0*T) 9 10 sup ||اا (0tr) - (0 ((n-)
(113)
120
Finally, we bound || (0!T+)) - l: (0(T+r))l as
I (0)) fi (0,7+)
2
|10!"h (0 +r) 9i (0/r+1) 11
Pa ti (0, Pig (0 (0:-7) (0++)11 l12
& Ps Pillopp ||g Ji (0:-*) (0)+)) 111
||g bn (0eT+r) 9 (0,++))) l12
2n 2 3 Rhg (1-9tt1)
q ||g lin (0*)) (0*)) +
vn - -9)
2n 2 C* 18g 1- (1 - -22 PRTTCCRooSI (1 - ('))
+ 2 RBg² (1 (1-9t+1)
<
Vn (1 - 9)
(1-9)²
vn(1 9)
r 28 RBg (1 - q¹)²
2 C0o RoCy (1 - -²) +
2n C C RBg" (1 (1- 1(1+)
Vn (1 9)²
(1-9)²
vn (1 - ②)
(114)
where step (a) holds because of (99), (107) and (110). The third inequation holds because of (112), Further considering 2no Rg we have
va(1- (2ll-n)
sup ||5t (0 - Jo (0) = On) Vi 120, 1SrST
(115)
Thus, Theorem is proved.
D. Proof of Theorem 3
The Jacobians of the linear global and local models are
flin (0t)) = J (88).
J.hn (0) = Je (0!)
(116)
The local update process of the linear model can be expressed as
g 0CT - n1 (0P) gfth (0!Tt)).
(117)
Via continuous time gradient flow, we obtain
7
J, (0%) g. (0.++))
J: (0.)) tig (0.
(118)
dr
D.
[D.""
Employing the chain rule yields
dg lin (0,7+))
dr
J (8!)T
dr
dr
24

---

Widening the Network Mitigates the Impact of Data Heterogencity on FedAvg
(0.) T
Ja (0.) Pi (0/
Dil
170 6 lin (0%Tt)).
(119)
[D:
Integrating from 0 to r on both sides yields
lin (0 = 6 - o. g (0
(120)
Plugging (120) into (118) yields
d0tr+r +r
27 [Di]
J (09) P.e z.
ee 9 Jin (0'T)
(121)
dr
Integrating from 0 to T on both sides yields
11 [D.] Je (0 P. (1"
0T+
8tT
E p
60 dr
gtو (0TT)
(122)
where S0 e Di 6. dr can be further derived as
dr
You Dii
0 dr
K!
k=
1
?
0
T
k!
[D,
++1
k=
|Pal
1
(-
h+1
(0!)
0
170
(k + 1)!
[D:
k=
-PPal (0!) -1
TOT [D.]
0p
110
h!
k=1
[Dil (6!) -1 (1 - e 18,7 09)
(123)
1/0
Plugging (123) into (122) yields
0 ger
11 Ja (09) P.
|Dil
(0!) -1 (I 11- i88i g (BtT)
[Di]"
T/0
1 (09)2) -1
J ( J. (0.) P. (0%) -1 (1
1o (1-)-10i0)) g (0tT)
(J (109))) -1 (1 (I-e v 0
g Jin (ger)
(124)
Considering the model aggregation process, we have
M
((++1)r O² Σ
Pi (0!T+T-000
-0- 0²T)
i=1
-1 As
(009))
Σ Pi (1 (1- 19 00) gli (0TT)
1=1
(3)
(100))
1 -
ΣM 7 6P) nm (0+T)
-1 (I(0))))
[ 1- e
788 00) g (8TT).
(125)
25

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
where step (a) employs Lemma 4. Considering the aggregation process of the linear model, have
M lin
g Jin (1
(0!
Pig
M
up 6 lin (0,T)
Pie
0
1=1
NoT
E m,
6 &gggin (0!")
B 0 ساو (07).
(126)
where step (a) employs Lemma 4. By cursively employing (126), obtain
9 Jin (o(++1)7) = 1r e (6))
(127)
mr Replace / + with / yields ساو (0tr) e- 2 gاو (09), plugging which into (125) further yields
g(++1)r ger (009))) (1 1
nox Di
00
e
lin (0")
()( -ytt
Na (0(
6a
g lin (09)
(128)
Employing (128) iteratively yields
-0===1((007) -1 (1 1-- 1o(
0(++1)= -0 00=
lin (09)
(129)
Replacing & with f yields
8tT - 8 - (89) (09) (69)-1 -1 (I-O- mo0 glin (
(130)
Therefore, for an arbitrary input x, we can obtain the closed form of glin (x,, 0CT) and flin (s, 0+0) as
g jim (z. BtT) = g (r, 80) (x. 84) T (geT - 80)
= 9 (x, (0) o (x) (69) -1 (1 (1-e- 101 9 (00)
(131)
and
plin = - -0) (09)- (x)) 1 1-
) (x
(132)
respectively.
26

---

Widening the Network Mitigates the Impact of Data Heterogeneity FedAvg
E. Network Details
Table I. Fully-Connected Network
KAYER
INPUT
IUTPUT
FCI FC2
INPUT SIZE
8 X x
8 k
FC3
8 k
OUTPUT
Table Convolution Network
LAYER
CHANNEL
KERNEL
OUTPUT
STRIDE
CONV1
6 k
5 x
28 x 28
POOLING
6 k
2 2
14 x 14
2
CONV2
16 k
5
x
POOLING
16 &
2
5 x
FCI
120 x k:
FC2
84 1:
FC3
10
Table Residual Network
LAYER
OUTPUT
BLOCK TYPE
CONV1
32 x 32
3 3, CHANNEL SIZE X kj
3 CHANNEL SIZE k
CONV2
32 x 32
3 3, CHANNEL SIZE X k
3 CHANNEL
k X
CONV3
16 X 16
3 CHANNEL SIZE
3 CHANNEL
A X
CONV4
8 X 8
3 CHANNEL SIZE & [8 x 8)
AVG-POOL
1 1
27


Reviews:
[{"review_id": "OFq0kQPe1n", "reviewer": "Reviewer_6HsR", "summary": "This paper analyzes the performance of FedAvg in infinite-width fully connected neural networks trained with gradient descent. It provides theoretical results on both convergence and generalization, revealing new insights—such as the diminishing impact of data heterogeneity as network width increases.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for acknowledging our contributions. Our response to your question is provided below.\n\n**What are the main technical challenges in extending the NTK analysis from centralized learning (e.g., [Jacot et al., 2018]) to FedAvg?**\n\nThe main technical challenges in extending NTK analysis from centralized to FedAvg come from *multiple local updates between global aggregations*  and *heterogeneous (non-iid) data distribution*.\n\nIn centralized learning setting, the dataset is centralized and the classic NTK analysis (e.g., [Jacot et al., 2018]) relies on global (centralized) gradient descent with infinitesimally small learning rate such that the parameter update can be approximated by gradient flow. Consequently, the dynamics of parameter updates in the infinite-width network is smooth and can be characterized by a linear ordinary differential equation. However, in FedAvg, the situation is more complex:\n\n1. **Multiple local updates between global aggregations:** In FedAvg, rather than a single global update step per iteration (as assumed in classical NTK analysis), each client performs multiple local gradient steps on its own data before sending updated parameters back to the server. The cumulation effects of multiple local updates and the global aggregation (averaging) process will lead to jump update of parameter which breaks the smooth dynamics. In our work, we resolve this challenge by establishing the relationship between the parameters before and after aggregation via Taylor expansion and bounding the reminder terms using the mean value theorem.\n2. **Non-iid data distribution:** Federated learning involves heterogeneous local datasets whose distributions differ across clients. Client-specific updates exhibit model-drift (divergence) due to data heterogeneity, complicating the establishment of unified bounds that accommodate all local models within a global training round. In our work, we resolve this challenge by deriving the upper bound of model divergence and show that it decreases with the network width $n$ as $\\mathcal{O}(n^{-\\frac{1}{2}})$. This allows us to analyze the impact of network width on data heterogeneity and extend the NTK conclusions to FedAvg.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No.", "experimental_designs_or_analyses": "No.", "relation_to_broader_scientific_literature": "This work studies how the network width affect the FedAvg performance in theory.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "The theoretical results are strong and novel, providing a quantitative analysis of how neural network (NN) width impacts the training and generalization performance of FedAvg. The main theorem (Theorem 1) is clearly stated, accompanied by a proof sketch. Notably, the results suggest that model divergence in FL due to data heterogeneity is bounded by $O(n^{-1/2})$ where $n$ denotes the width of NN, offering valuable insights into the effects of overparameterization in FL.\n\nA minor limitation is that the analysis is restricted to the asymptotic case (width approaching infinity). However, given that non-asymptotic analysis is significantly more challenging—particularly in the NTK framework—this remains a well-written paper with meaningful theoretical contributions to the understanding of overparameterization in FL.", "other_comments_or_suggestions": "NA", "questions_for_authors": "What are the main technical challenges in extending the NTK analysis from centralized learning (e.g., [Jacot et al., 2018]) to FedAvg?", "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": "4izAVr2y6B", "reviewer": "Reviewer_HFoE", "summary": "This paper mainly studies that in the federated learning scenario, as the width of the neural network continuously increases, the negative impact of non-IID data on the FedAvg algorithm will gradually weaken and finally disappear completely when the width is infinite. Based on the theory of neural tangent kernel (NTK), the authors analyze the training process of multi-layer fully connected networks and prove that when the network is sufficiently wide, the local model and the global model approximate to linear models, and under the same number of centralized training iterations, both can achieve equivalent generalization performance. The authors have conducted experimental verification on datasets such as MNIST and CIFAR-10, using various network structures (such as FNN, CNN, ResNet) and various loss functions and optimization algorithms, and the results are consistent with the theoretical analysis.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive comments. We will clarify your questions and concerns below.\n\n**1. References regarding the convergence behavior of multi-layer NTK in real-world applications**: \n\nAfter carefully reviewing the existing literature again, we were unable to find further studies regarding the convergence behavior of multi-layer NTK in real-world FL applications. If relevant studies become available in the future, we will certainly include and discuss them in our subsequent work.\n\n**2. Weakness of lacking experiments on industrial-scale data**: \n\nDue to our limited GPU resources, we were unable to conduct experiments on substantially larger datasets in the current version, especially for federated learning where multiple clients are involved. However, we have already begun preliminary tests on larger-scale datasets, e.g., CIFAR-100, to further validate our theoretical results. If time permits, we will include the results of these experiments in the next version. \n\n**3. Suggestion 1 — Enhance the paper's heuristic value for applicators**: \n\nThank you for this suggestion. We consider adding the following discussion to our paper regarding the possible further directions to apply NTK theory to the industrial field:\n\n>  Our theoretical analysis confirms that under FedAvg, a sufficiently wide model can be approximated by a linear model whose outputs are linearly related to the model parameters. Consequently, the model parameters can be recovered from the model outputs. This suggests that each client may transmit the model outputs instead of the model parameters for model aggregation, possibly reducing the communication overhead of federated learning in practice, which deserves further investigation.\n\n**4. Suggestions 2 — Provide a more intuitive overview of gradient continuity assumption in the main text**: \n\nThank you for this suggestion. We believe what you mentioned \"gradient continuity assumption\" refers to the *gradient flow approximation* and we would like to highlight that we have already provided a **Proof Sketch** in our main text (**line 178**) where the gradient flow approximation was mentioned in  **[Step 2]** (**line 190**). Nevertheless, we appreciate your suggestion and consider adding a sentence in **line 205** to make it more clear:\n\n> The learning rate is set to $\\eta = \\frac{\\eta_0}{n}$ and $\\eta_0$ is a constant independent of $n$, which results in infinitesimally updates during each gradient descent step when $n \\to \\infty$. Consequently, we adopt gradient flow as an approximation of gradient descent, which can be expressed as $\\frac{d\\theta_i^{t\\tau+r }}{dr} = -\\frac{\\eta}{|\\mathcal{D}_i|}J_i\\left(\\theta_i^{t\\tau+r}\\right)g_i\\left(\\theta_i^{t\\tau+r}\\right).$\n\n**5. Question regarding experiment and theoretical results on extremely high non-iid data**: \n\nThank you for pointing this out. We would like to highlight that *Figures 5* in our paper has already considered an extremely high non-iid case. Specifically, we use a mini-CIFAR-10 dataset containing two image categories, with one category assigned to half of the clients and the other category to the remaining clients. As such, the datasets between these two client groups completely do not overlap. As shown in Figure 5, FedAvg with large-but-limited-width networks performs almost identically to centralized learning. \nMoreover, in the rest experiments (e.g., Figures 1 and 2), we generate highly non-iid data from CIFAR-10 by employing the Dirichlet distribution with a concentration parameter $\\alpha=0.1$, which is sufficiently small to ensure a high degree of non-iid as done by Zhang et al. (2023) and Wang et al. (2020).\n>Zhang et al. Fed-CBS: A heterogeneity-aware client sampling mechanism for federated learning via class-imbalance reduction. ICML, 2023.\n\n>Wang et al. Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization. NeurIPS 2020.\n\nRegarding our theoretical analysis, we do not make any assumption on the non-iid degree of data. Instead, we prove that the model divergence can be bounded by $\\mathcal{O}(n^{-\\frac{1}{2}})$ regardless of the non-IID degree, and this term, i.e., $\\zeta$,  has been already been included in our derived error upper bound in (22).", "claims_and_evidence": "The main argument of the paper is that as the width of neural networks increases, the model differences caused by non-iid data can be mitigated; under the limit of infinite width, FedAvg and centralized training can achieve the same model output and generalization performance under the same number of iterations. The authors provide a complete theoretical derivation, including the analysis of NTK invariance and the conclusion that the magnitude of model parameter updates decreases as the network width increases. The experimental part verifies this point through various networks and diverse data distributions, and the numerical results are consistent with the theoretical inference.", "methods_and_evaluation_criteria": "The author's method is based on NTK theory, characterizing the training dynamics of networks under gradient descent and local multi-round updates, and proving the convergence rate and upper bound of errors using strict mathematical methods. The evaluation criteria mainly include:\n\n1) Convergence metrics during the training process, including the gap between local and global model parameters, as well as the change of training error with iterations.\n2) Generalization performance on the test set, measured by accuracy or MSE, etc.\n3) Various network structures and non-IID partitioning strategies, ensuring the diversity of evaluation.", "theoretical_claims": "The theoretical part of the paper includes:\n\n1) It is proven that under large width, the difference between the local model and the global model can be confined within a range that shrinks as the width increases.\n2) As the network approaches infinite width, the global and local NTK matrices tend towards constant matrices, making the model dynamics equivalent to a linear model, and FedAvg and centralized training yield the same solution.\n\nThe authors provide a detailed derivation process in the appendix, mainly based on the continuous approximation of the gradient flow and the inference of constraints on the NTK variation. The derivation is relatively systematic and lacks obvious logical gaps or mathematical inconsistencies.", "experimental_designs_or_analyses": "1) Compared the convergence speed and model accuracy under different network widths, demonstrating that the convergence and accuracy impact of non-IID data is significantly weakened as the network width increases.\n2) On datasets such as MNIST and CIFAR-10, the theoretical claims were verified to hold for different types of networks (such as FNN, CNN, and ResNet).\n3) Adopted commonly used optimization methods (such as SGD and cross-entropy loss) and selected commonly used methods (such as Dirichlet distribution partitioning) for non-IID partitioning.\n\nThese designs support the author's theoretical conclusions, and the data volume and process are sufficient to ensure the credibility of the results.", "relation_to_broader_scientific_literature": "This paper is closely related to recent research on the role of neural tangent kernel(NTK) in overparameterized networks, and also connects with various schemes to reduce non-IID impact in federated learning. It extends the NTK analysis paradigm to the commonly used federated learning algorithm FedAvg, and points out that network widening can bring dynamics similar to centralized training. Although these conclusions have some prototypes in existing research on overparameterized networks, this paper clearly applies them to federated learning under non-IID data scenarios, making it more complete in both theory and experiment.", "essential_references_not_discussed": "The author has cited the main research related to NTK and overparameterization in recent years, as well as common literature on federated learning's sensitivity to non-IID, without finding any critical missing citations. However, if there are more analytical studies on the convergence behavior of multi-layer NTK in real-world applications, the author could supplement some more examples, but this seems not to affect their main argument.", "other_strengths_and_weaknesses": "The advantages of the article lie in the provision of rigorous theoretical proofs, as well as relatively systematic experiments to verify the conclusions, with the analysis scope covering a variety of network architectures and multiple loss functions. The writing structure of the article is clear, the mathematical derivation is complete, and the illustrative examples are also relatively sufficient.\n\nThe potential shortcomings is:  Although there are experiments on CIFAR-10, they are still limited compared to industrial-scale data. Whether the same patterns can be maintained on larger-scale non-IID data requires more experimental evidence.", "other_comments_or_suggestions": "1) If more discussion on the impact in actual industrial scenarios or similar extension ideas could be provided in the text, it would enhance the paper's heuristic value for applicators.\n2) The author has explained the premise assumptions for gradient continuity in the appendix, but if a more intuitive overview could be provided in the main text, it would be more convenient for readers to understand.", "questions_for_authors": "When the non-IID degree is extremely high (for example, the categories assigned to users almost completely do not overlap), if the network width is limited, have additional bias phenomena been observed in the experiment? Will this affect the error upper bound in the theorem?", "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": "LpZtnj5mUI", "reviewer": "Reviewer_7NNj", "summary": "The paper provides theoretical analysis showcasing that increased neural network width provably reduces the effects of data heterogeneity, and approaches the convergence of centralized gradient descent. However, it is also proven that data heterogeneity slows convergence.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful comments. We will clarify your questions and concerns below.\n\n**1. Novelty compared with Lee (2019) and Shi (2024):**  The main novelty of our work compared to Lee (2019) and Shi (2024) lies in addressing *data heterogeneity*, which was not considered in those earlier analyses. Specifically:\n\n- Lee (2019) studied centralized learning using NTK theory where data heterogeneity does not arise because all data are pooled together. In contrast, our work focuses on a federated learning setting where clients possess diverse local data distributions, introducing heterogeneity that fundamentally alters the training dynamics.\n- Although Shi (2024) extended the NTK analysis to the federated learning scenario, their framework did not analyze the effect of data heterogeneity. Our work proves that the model divergence caused by data heterogeneity can be bounded by $\\mathcal{O}(n^{-\\frac{1}{2}})$, revealing that increasing the network width can reduce the impact of data heterogeneity. Additionally, we bridge centralized learning and FedAvg, proving that their training dynamics are equivalent in the infinite-width limit, which is another novel aspect that was not explored in (Shi 2024).\n\n**2. Weakness regarding testing on larger models to reflect the setting proposed in the paper:** We suspect that there may be a misunderstanding due to our naming convention. For example, in our CNN32, the number of neurons in each layer is up to $3840$ and the total number of model parameters is $\\mathcal{O}(10^7)$. The detailed configuration is provided in Appendix E and we would like to clarify that the models used in our experiment are already sufficiently large to validate our analysis. Nevertheless, to address your concerns, we increase FNN16 and ResNet16 to FNN512 and ResNet32, respectively, both with parameters about $\\mathcal{O}(10^7)$, which already exceed those considered in existing NTK literature (Park et al. 2019).\n\n>Park, Daniel, et al. The effect of network width on stochastic gradient descent and generalization: an empirical study. ICML, 2019.\n\nFor your convenience, we provide the following table demonstrating the width and number of parameters in each layer of our models, and we update experiment results  ([**click here**](https://anonymous.4open.science/r/ICML2025-Rebuttal/Fig%20A.md)), demonstrating that our conclusions still hold.\n\n\n| FNN512 | Width | Parameters/Layer |\n| ------ | ----- | ---------------- |\n| Layer1 | 4096  | 4100096          |\n| Layer2 | 4096  | 16781312         |\n| Layer3 | 10    | 40970            |\n\n| CNN32  | Width | Parameters/Layer |\n| ------ | ----- | ---------------- |\n| Layer1 | 192   | 4992             |\n| Layer2 | 512   | 2457600          |\n| Layer3 | 3840  | 49155840         |\n| Layer4 | 2688  | 10324608         |\n| Layer5 | 10    | 26890            |\n\n| ResNet32 | Width | Parameters/Layer |\n| -------- | ----- | ---------------- |\n| Layer1   | 16    | 4992             |\n| Layer2   | 512   | 2437048          |\n| Layer3   | 512   | 4723760          |\n| Layer4   | 512   | 4723760          |\n| Layer5   | 10    | 5130             |\n\n**3. Suggestion on describing FL:**  Thank you for pointing this out. In our revised version, we will rephrase\n\n> \"Federated Learning (FL) is a decentralized machine learning paradigm\"\n\nto\n\n> \"Federated Learning (FL) is a distributed machine learning paradigm\"\n\n**4. Question 1 — How easily can this analysis be generalized to SGD?** Thank you for the question. Extending the theoretical analysis from GD to SGD is non-trivial. The main challenge lies in the inherent mini-batch sampling noise in SGD, which introduces stochastic perturbations to parameter update trajectories. This differs from the deterministic dynamics assumed in NTK analysis. One possible way to address this is to derive probabilistic bounds on gradient variation and establish a stochastic differential equation framework to model the behavior of SGD. \n\n**5. Question 2 — How easily can the generalization performance be expanded to non-linear models?** Thank you for raising this point. While characterizing the generalization of practical nonlinear models remains a significant challenge in machine learning theory, our work specifically focuses on the infinite-width regime where non-linear neural networks in FedAvg become provably linearizable through the NTK framework. Our experiments further show that sufficiently wide networks in FedAvg can indeed be well approximated by linear models (Figure 4), which serves as the foundation for establishing equivalence of generalization between federated and centralized learning. This leads to the vanishing of heterogeneity in FedAvg, which is the central conclusion of our work. Nevertheless, we admit that in finite-width regime where the model cannot be well approximated by a linear model, the generalization performance analysis is still a challenge to be dealt with in future works.", "claims_and_evidence": "The theory clearly supports the claims made by the authors that data heterogeneity is reduced as network sizes increase.", "methods_and_evaluation_criteria": "The paper is predominantly theory, so the bulk of the novelty lies there. The experiment section was not super useful in evaluation due to the extremely small nature of the ML models. In practice, larger models perform better in Centralized FL and these results are not anything particularly new (especially due to the small nature of the models).", "theoretical_claims": "I did not closely verify the proof correctness.", "experimental_designs_or_analyses": "The experimental design seemed a bit lacking. I understand that the paper is theory-based, with an emphasis on linear and GD, so validation in typical FL settings is not straightforward. However, it would be nice for the authors to work on larger models (that have huge width). Testing on small networks with few parameters does not accurately reflect the setting the authors are evaluating.", "relation_to_broader_scientific_literature": "The paper seems to follow closely to NTK literature (Lee 2019 and Shi 2024). I am not familiar enough with those papers to ensure that the contributions in this paper are completely novel.", "essential_references_not_discussed": "Not to my knowledge.", "other_strengths_and_weaknesses": "## Strengths\n\n1. Clear and concise theory. Explanations and remarks are nice for readers.\n\n## Weaknesses\n\n1. Experiment section could be improved upon. Test on larger models to reflect the setting proposed in the paper.", "other_comments_or_suggestions": "> \"Federated Learning (FL) is a decentralized machine learning paradigm\"\n\nI would change this. FL can be, and often is, centralized, with a server leading training and receiving updates. FL is, however, a distributed ML paradigm.", "questions_for_authors": "1. How easily can this analysis be generalized to SGD? It is a great place to start when analyzing GD, but most ML models require SGD (since memory and computational costs are unfeasible due to large dataset sizes).\n2.  How easily can the generalization performance be expanded to non-linear models?", "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": "1gR0K7waCI", "reviewer": "Reviewer_aWck", "summary": "This paper provides an NTK based analysis of FedAvg for training deep neural networks under the MSE loss, showing that the effect of heterogeneity disappears in the infinite-width limit. Essentially, the model divergence (distance between local model parameters) is inversely proportional to the square root of the network width, so the usual \"client drift\" term becomes neglible in the infinite-width limit. The authors also provide experimental results showing that the gap between networks trained with IID and non-IID data shrinks as the network width increases.\n\n## Update after rebuttal\nAfter discussion with the authors, I still lean towards acceptance, but I will not increase my score due to remaining concerns about the rigor of the mathematical arguments. In particular, I am not totally satisfied with Assumption 1 and Theorem 3: stating that terms \"can be neglected\" is not a rigorous mathematical statement in general.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive comments. We will clarify your questions and concerns below. \n\n**1. Non-rigorous statement:** We will revise Assumption 1 and Theorem 3 as:\n\n> **Assumption 1.** The terms of order $\\mathcal{O}(n^{-1})$ and higher are neglected for $n \\to \\infty$.\n\n> **Theorem 3.** Under Assumptions 2 to 4,  when $n\\to \\infty$ and the terms of order $\\mathcal{O}\\left(\\eta_0^2\\tau^2\\right)$ and higher are neglected for $\\eta_0\\tau\\to 0$, ......\n\nRegarding using gradient flow to approximate GD, we confirm this is a standard way used in NTK literature (Jacot 2018 and Lee 2019). Also, we will remove the word *rigorously* in lines 80 and 432 as you suggested.   \n\n**2. Technical novelty:** We extend the NTK analysis to a federated setting where clients have diverse local data distributions. This heterogeneity requires more novel bounding techniques compared to the iid setup in prior NTK studies. Specifically, we bound the model divergence term based on mathematical induction, with a newly proposed induction hypothesis that underpins the proof. To further analyze how model divergence affects learning convergence, we characterize the relationship between the parameters before and after global aggregation via Taylor expansion. Crucially, the remainder terms of this expansion (which encapsulate the model divergence term) are bounded in a novel way by constructing two auxiliary functions (lines 931 and 937) and applying the mean value theorem. Relying on the above techniques, we are able to further prove the linearized model and the equivalence between FedAvg and GD in the infinite-width limit. \n\n**3. FedAvg vs GD:** We would like to first clarify that we do not use the equivalence of FedAvg and GD in the infinite-width limit (Theorem 3) to explain how increasing network width mitigates heterogeneity. That phenomenon is instead explained by Theorem 2 and Remarks 1 & 2. A key difference in the preconditions between Theorems 1 and 3 is that Theorem 1 omits only $\\mathcal{O}(n^{-1})$ and higher-order terms, making it an asymptotic result, whereas Theorem 3 omits all terms that vanish as $n\\to\\infty$, making it an exact limiting result. Therefore, Theorem 1 can indeed explain the phenomenon in the large-but-finite-width regime.  \n\nRegarding the different phenomenon reported in [1, 2] that *local SGD (a.k.a FedAvg) outperforms large mini-batch SGD*, those works focus on the iid setting. By contrast, we consider non-iid data,  where the client drift causes FedAvg to underperform compared to full-batch GD. Another difference is that we compare FedAvg and GD under the **same number of iterations**, whereas [1, 2] compare them over the **same number epoch**, resulting in more updates in local SGD (small batch) than for mini-batch SGD (large batch). \n\n**4. Experimental results:**\n\n4a: We only consider mini-MNIST with FNN for Figs. 3-4 because they require computing the NTK and the output of the linear model, both involving the Jacobian matrices, whose dimension is extremely high for larger models and datasets. Considering that Fig. 5 and the last subfigure of Fig. 3 do not involve computing the NTK, to address your comments, we add experiments ([**click here**](https://anonymous.4open.science/r/ICML2025-Rebuttal/Fig%20B&C.md)) on CIFAR-10 with CNN, showing that FedAvg and GD achieve almost the same loss and accuracy, and exhibit lazy training behavior. \n\n4b: The quantity in Fig. 4 is the *root-mean-square error* (RMSE) between the outputs of the neural network and linear model *across the entire dataset*. Hence, there is a division by the dataset size. While metrics like loss comparison could offer complementary insights, we selected RMSE as it directly quantifies the discrepancy in model outputs and is aligned with Theorem 2. Nevertheless, to address your concern, we add an experiment to compare the loss between the neural network and linearized model ([**click here**](https://anonymous.4open.science/r/ICML2025-Rebuttal/Fig%20D.md)), showing that the loss difference is marginal.\n\n4c: We note that in proving the equivalence between GD and FedAvg (Theorem 3), we make the precondition that $\\mathcal{O}(\\eta_0^2\\tau^2)$ and higher-order terms are neglected for $\\eta_0\\tau \\to 0$, where $\\eta_0$ is related to the learning rate and $\\tau$ is the number of local updates. Therefore, to maintain the equivalence, a smaller learning rate is required with the growth of $\\tau$. This explains why the gap grows with $\\tau$ in Fig. 5, where $\\eta_0=0.1$ is fixed. \n\n**5. Limitation of NTK theory:** We admit that the NTK framework requires certain assumptions, e.g., the infinite width and gradient flow. Nevertheless, we conducted experiments with finite-width models using SGD or GD, and while there is indeed a gap between theory and practice, we observed trends consistent with our theoretical predictions. This gap deserves further investigation, and extending the analysis to more realistic settings is an important future direction.", "claims_and_evidence": "The claims are moderately well supported. The standard of mathematical rigor in this paper is not as high as some might like, though this seems to be common among papers using NTK techniques. Examples include: the statement of Assumption 1, the statement of Theorem 3, and the use of gradient flow to approximate the trajectory of gradient descent in Equation (22). None of these are so non-rigorous as to completely invalidate the results, but I think that the mathematical statement of theorems and assumptions should not involve non-rigorous language. To be clear: saying that a quantity \"can be neglected\" is not a rigorous mathematical statement in general, and such language should only be used when the meaning is totally clear; this principle is violated in Assumption 1 and Theorem 3.", "methods_and_evaluation_criteria": "Yes, the experimental setup is reasonable.", "theoretical_claims": "I did not check the proofs thoroughly, although it is not hard to believe the result.  Just as in the centralized NTK, we can show that the distance from any iterate to its initialization is vanishingly small; therefore, client drift is also vanishingly small.  According to the local linear approximation of the objective around the initialization, client drift should not be an issue because the gradient does not change. However, this is just my hand-wavy understanding. I am not an expert in NTK techniques and I defer to other reviewers on the matter of the validity of the proofs.", "experimental_designs_or_analyses": "In general, the experimental setup is reasonable, but certain key results are omitted. See Weakness #4 for more details.", "relation_to_broader_scientific_literature": "his paper contributes to the theory of distributed deep learning, which itself is the intersection of deep learning theory and the practice of using distributed systems to leverage compute and data from many devices to train machine learning models. The authors investigate how data heterogeneity across devices affects the behavior of distributed optimization of deep neural networks, which is a central topic in the theory of distributed optimization. Most works on the subject characterize the heterogeneity in terms of a black-box assumption (i.e. bounded gradient dissimilarity), but this work examines the effect of heterogeneity by directly considering the model architecture.", "essential_references_not_discussed": "I do not know of any essential references that were not included.", "other_strengths_and_weaknesses": "Strengths:\n1. The question tackled by the authors is a very important one. For several years, the question of heterogeneous objectives in federated learning has been central to the theory of optimization in this setting, and a lot of papers were dedicated to new algorithms (e.g. SCAFFOLD) or new theoretical conditions (e.g. bounded Hessian dissimilarity) to theoretically handle the effect of heterogeneity on optimization.  According to the classical theory that uses bounded gradient dissimilarity, heterogeneity should have a significant negative effect on optimization. Importantly, this classical theory does not usually consider deep neural networks when examining the effect of heterogeneity on optimization, though the model choice (e.g. linear vs neural network) must surely play a role in how heterogeneity affects convergence. In that sense, I think this paper provides an important perspective on heterogeneity in distributed optimization.\n2. Some of the experimental results are compelling. In particular, Figures 1 and 2 paint a pretty convincing picture that model width as a strong effect on how heterogeneity affects convergence: the gap between networks trained with IID data vs non-IID data shrinks significantly for wide networks, but not thin networks. This is a significant observation and could inspire future work on the distributed optimization of deep networks.\n\nWeaknesses:\n1. The results of the paper are not totally rigorous. Nothing is so non-rigorous as to invalidate the paper, but there are several instances of non-rigorous statements and approximations (see Claims and Evidence section). Perhaps this is the standard of rigor for NTK papers, so I defer to other reviewers on the matter of rigor.\n2. It is unclear to me if there is significant technical novelty in this paper. I'm not an expert in NTK techniques, so it isn't clear to me when reading the proof whether there are novel techniques being used or if everything is standard for this area. Again, I defer to other reviewers on this point.\n3. While the experiments in Figures 1 and 2 suggest that network width plays an important role in the effect of heterogeneity, I am skeptical that the authors' explanation of this phenomenon is the correct one. The results of this paper say that GD and FedAvg are completely equivalent in the infinite-width limit, but this does not align with the behavior of FedAvg in practice [1, 2]. I would say that it is commonly believed that FedAvg has qualitatively different behavior than GD for finite-width networks, and this makes me doubt that the infinite-width regime actually explains what happens in practice.\n4. Some of the experimental results are lacking. Figures 3, 4, and 5 try to confirm certain implications of the theory (network parameters and NTK remain constant, network is well approximated by a linear model, FedAvg is equivalent to GD), but there are some issues.\n\n    4a. Figures 3-5 only apply for mini MNIST with fully connected networks. To me, it is important to see the same Figures for harder datasets and CNN architectures, in order to verify whether these theoretical implications actually happen in practical situations.\n\n    4b. The quantity graphed in Figure 4 is not very interpretable, since there is no frame of reference for this quantity. Why do we divide by the dataset size? What counts as \"big\" or \"small\" for this quantity? It would probably be more informative to look at the loss of the linear model vs. the loss of the neural network, or maybe the proportion of data samples for which the two networks make the same prediction; just some quantity that is interpretable.\n\n    4c. Figure 5 does not totally support the authors' claims. It shows that the gap between GD and FedAvg is small when the number of local steps is tiny ($\\tau = 2$) and grows with the number of local steps (i.e. $\\tau = 5$). In some ways, this goes against the theoretical prediction, which says that GD and FedAvg should be equivalent with wide networks; yet the gap between them grows with $\\tau$. I am not convinced that this aspect of the theory aligns with practice: by these experiments, it seems that FedAvg and GD are not really the same.\n\n[1] Lin, Tao, et al. \"Don't Use Large Mini-Batches, Use Local SGD.\" Proceedings of the 8th International Conference on Learning Representations. 2019.\n\n[2] Gu, Xinran, et al. \"Why (and When) does Local SGD Generalize Better than SGD?.\" The Eleventh International Conference on Learning Representations. 2022.", "other_comments_or_suggestions": "I would suggest that the authors provide an in-depth discussion about whether their results actually describe practice. There is already some degree of skepticism in the community about whether NTK really explains practice; if you do not acknowledge this alternative perspective then your own case is weakened. In my opinion, the paper would be strengthed if the authors acknowledge the possibility that finite-width networks in practice may not have the behavior predicted by NTK theory, and point to this as a direction of future work. Also, I suggest weakening the claims of rigor (e.g. line 432), since many will be skeptical about the rigor immediately upon reading e.g. Assumption 1 and Theorem 3.", "questions_for_authors": "1. Is there significant technical novelty in your proofs? If so, can you briefly summarize the novel techniques you used?", "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:

Meta-review: The paper studies the impact of overparameterization on statistical heterogeneity in FL. Notably, the widening of network's width impact is characterized by extending the theory of NTK and overparameterized DL models to the FL setting.

The paper seems to rely on assumptions deed restrictive in the FL setting, e.g., not handling partial participation. However, I believe the extension of NTK and overparameterized to the FL setting requires nontrivial steps and this work could serve as a reference to incorporated FL-specific considerations by future works.

Indeed, the reviewers hold a positive view and are supportive of accepting the paper despite the aforementioned shortcoming as well as other concerns, e.g., lack of detailed experiments as noted by all reviewers and occasional lack of rigorous arguments as noted by Reviewer aWck.
Decision: Accept (poster)