Initialization Matters: Unraveling the Impact of Pre-Training on Federated Learning
Authors: Divyansh Jhunjhunwala
Pranay Sharma
Zheng Xu
Gauri Joshi
Venue: Accepted by TMLR ()

Abstract: Initializing with pre-trained models when learning on downstream tasks is becoming standard practice in machine learning. Several recent works explore the benefits of pre-trained initialization in a federated learning (FL) setting, where the downstream training is performed at the edge clients with heterogeneous data distribution. These works show that starting from a pre-trained model can substantially reduce the adverse impact of data heterogeneity on the test performance of a model trained in a federated setting, with no changes to the standard FedAvg training algorithm. In this work, we provide a deeper theoretical understanding of this phenomenon. To do so, we study the class of two-layer convolutional neural networks (CNNs) and provide bounds on the training error convergence and test error of such a network trained with FedAvg. We introduce the notion of aligned and misaligned filters at initialization and show that the data heterogeneity only affects learning on misaligned filters. Starting with a pre-trained model typically results in fewer misaligned filters at initialization, thus producing a lower test error even when the model is trained in a federated setting with data heterogeneity. Experiments in synthetic settings and practical FL training on CNNs verify our theoretical findings.

Published in Transactions on Machine Learning Research (10/2025)
Initialization Matters: Unraveling the Impact of Pre- Pre-Training on Federated Learning
Divyansh Jhunjhunwala Cornegie Mellon University
@hunghus@nndreu.cmu..d
Pranay Sharma IIT Bombay
pronaysh
Zheng Xu Google
mucheng@@google.com
Gauri Joshi
gourij@andrew.com cww. edu
Carnegic Mellon University
Reviewed on OpenReview: https: // openreview. net/ forum? id- wwkkuu
//
Abstract
Initializing with pre-trained models when learning on downstream tasks becoming standard practice machine learning. Several recent works explore the benefits of pre-trained initialization in a federated learning (FL) setting, where the downstream training performed at the edge clients with heterogencous data distribution. These works show that starting from pre-trained model can substantially reduce the adverse impact of data heterogeneity the test performance of model trained in a federated setting, with no changes to the standard FedAvg training algorithm. In this work, we provide deeper theoretical understanding of this phenomenon. To do so, we study the class two-layer convolutional neural networks (CNNs) and provide bounds on the training error convergence and test error of such network trained with FedAvg, We introduce the notion of aligned and misaligned filters at initialization and show that the data heterogencity only affects learning on misaligned filters. Starting with pre-trained model typically results in fewer misaligned filters at initialization, thus producing a lower test error even when the model is trained in a federated setting with data heterogeneity. Experiments in synthetic settings and practical FL training on CNNs verify our thcoretical findings.
1 Introduction
Federated Learning (FL) (McMahan et al. 2017) has emerged as the de-facto paradigm for training a Machine Learning (ML) model over data distributed across multiple clients with privacy protection due to its no data-sharing philosophy. Ever since its inception, it has been observed that heterogeneity in client data can severely slow down FL training and lead model that has poorer generalization performance than model trained on Independent. and Identically Distributed (IID) data (Kairouz et. al.) 2021 Li et al.) 2020 Yang et al., 2021a). This has led works to propose several algorithmic modifications to the popular Federated Averaging (FedAvg) algorithm such varianco-reduuciion (Acar et al. 2021 Karimireddy et al. 2020), contrastive learning (Li et al.
2021 Tan et al. 2022 and sophisticated mnodel-aggregatinn techniques to combat. the challenge of data heterogeneity.
(Li et al. 2020 Wang et al. 2020)
A recent line of work (Chen et al. 2022
Nguyen et al.] 2022) has sought to understand the benefits of starting from pre-trained models instead of randomly initializing the global model when doing FL. This idea has
been popularized by results in the centralized setting Devlin al. 2019 Radford et al. 2019 He et al.

---

Published in Transactions on Machine Learning Research (10/2025)
Dosovitskiy et al. 2021). which show that starting from pre-trained model can lead to sate-offthheart accuracy and faster convergence on downstream tasks. Pre-training is usually done on internet-scale public data (Schuhraan et al. 2022 Thomee et al. 20166; Raffel et al.) 2020 Gao et al., 2020 in order for the model to learn fundamental data representations (Sun et al. 2017 Mahajan et al. 2018: Radford et al.) 2019), that can be easily applied for downstream lasks. Thus, while would not. be unexpected to see some gains of using pre-trained models even in FL, what is surprising is the sheer scale of improvement. In many cases Nguyen et al. (2022); ; Chen et al. (2022) show that just starting from a pre-trained model can significantly reduce the gap between the performance of model trained in federated setting with non- IID versus IID data partitioning with no algorithmic modifications. Figure D shows our own replication of this phenomenon, where starting from a pre-trained model can lead to almost 14% improvement in accuracy for FL with non-lID data (i.e., high data heterogeneity) compared to 4% for FL with IID data and 2% in the centralized setting. This observation leads us to ask the question:
2019
Why can pre-trained initialization drastically improve model performance in FL?
One reason suggested by Nguyen et al. 2022 is lower value
90
of the training loss at initialization when starting from pre-
FL Non-IID (a 0.1)
Centralized
trained models. However, this observation can only explain improvement in training convergence speed (see Theorem V
85-
FL HID (a 10)
Test Acouracy
80-
81.7
in Karimireddy et al. (202T)) and not the significantly im- proved generalization performance of the trained model. Also, a pre-trained initialization can have larger loss than random initialization while contimuing to have faster convergence and better generalization (sce Table in Nguyen et al. (2022) )
79.6
79.7
75-
6
70-
65
Chen et al. 2022); ; Nguyen et al. (0022), also observe some opúimiztiion-rllteed factors when starting from a pre-trained model including smaller distance to optimum, better condi- tioned loss surface (smaller value of the largest eigen value of Hessian) and more stable global aggregation. However, has not been formally proven that these factors can reduce the adverse effect of non-lID data, Thus, there is still lack of fundamental understanding of why pre- pre-trained initialization benefits generalization for non-IID FT..
63.6
60
Random
Pretrained
Figure 1: Test accuracy (%) SqueezeNet model landola et al. (0066) under random and pretrained initializations for FL and centralized training. Pre-training benefits FL more than centralized setting and signifi- cantly reduces the gap between IID and non- IID FL model performance.
CIFAR10 with
Our contributions. In this work we provide a deeper theoretical understanding of the importance of initialization for FedAvg by studying two-layer ReLU Convolutional Neural Networks (CNNs) for binary classification. This class of neural networks lends itself to tractable analysis while providing valuable insights
that empirically extend to training deeper CNNs as shown by several recent works [Cao et al. 2022 Du
et al. 2018 Kou et al. 2023 Zou et al. 2023 Jelassi & Li 2022 Bao et al. 2024 Oh & Yun 2024). Our
data generation model, also studied in Cao et al. (2022); Kou et al. (2023), allows us to utilize a sigmal-noise
decomposition result (sce Proposition I to perform a fine-grained analysis of the CNN filter weight updates
than can be done with general non-convex optimization. Some highlights of our results are as follows:
1. We introduce the notion of aligned and misaligned filters at initialization (Lemma and show that data heterogeneity affects signal learning only on misaligned filters while noise memorization is unaffected by data heterogeneity (see Lemma 2 A pre-trained model is expected to have fewer misaligned filters, which can explain the reduced effect of non-IID data.
2. We provide a test error upper bound for FedAvg that depends on the number of misaligned filters at initialization and data heterogencity. The effect of data heterogeneity on misalignod filters is exacerbated as clients perform more local steps, which explains why FL benefits more from pre-trained initialization than centralized training, To our knowledge, this is the first result where the test error for FedAvg explicitly depends on initialization conditions (Theorem 2
3. We prove the training error convergence of FedAvg by adopting two-stage analysis: a first stage where the local loss derivatives lower bounded by constant and second stage where the model is in the
2

---

Published in Transactions on Machine Learning Research (10/2025)
neighborhood of global minimizer with nearly convex loss landscape. Our analysis shows a provable benefit of using local steps in the first stage to reduce communication cost.
4. We experimentally verify our upper bound on the test error in a syuthetic data setting (see Section 3 as well as conduct experiments on practical FL tasks which show that our insights extend to deeper CNNs (soe Section
Related Work. The two-layer CNN model that we study in this work was originally introduced in Zou et al. (2023 for the purpose of analyzing the generalization error of the Adam optimizer in the centralized setting. Later Cao et al. (2022) study the same model to analyze the phenomenon of benign overfitting in two-layer CNN, i.c., give precise conditions under which the CNN can perfectly fit the data while also achieving small population loss. [hh & Yun (2024) use this model to prove the benefit. of patch-level data augmentation techniques such Cutout and CutMix. Kou et al. 2023) relaxes the the polynomial ReLU activation in Cao et al. (2022) to the standard ReLU activation and also introduces label-flipping noise when analyzing benign overfitting in the centralized setting. We do not consider label- fleppliipiing in our work for simplicity; however this can be easily incorporated as future work. To the best of our knowledge, we are only aware of two other works (Huang et al.] 2023 Bao et al.] 2024) that analyze the two-layer CNN in a FL setting. The focus in Huang et al. 2023) is on showing the benefit of collaboration in FL by considering signal heterogeneily across the data in clients while Bao al. (2024) considers signal heterogeneity to show the benefit of local steps. Both Huang et al. [2023] and[Bao et al. (2024) do not consider any label heterogencity and there is no emphasis on the importance of initialization, making their analysis quile different from ours. We defer more discussion on other related works to the Appendix.
2 Problem Setup
We begin by introducing the data generation model and the two-layer convolutional neural network, followed by our FL objective and a brief primer ou the FedAvg algorithm. We note that given integers a, b, we denote by [a : b] the set of integers {a, a + +,+1,..... b}. Also, [n] denotes (1.
Let D be the global data distribution. A datapoint (x, y) ~ D contains feature vector x = T E 222d with two components x(1), x(2) E TRd and label y E (+1,-1}, that are generated as follows:
Data- Generation Model.
Label y E 1-1,1} is generated as Ply = 1] = [y = - = 1/2.
2. One of x(1), x(2) is chosen at random and assigned as yu,, where μ E Rd is the signal vector that we are interested in learning. The other of x(1,, x(2) is set to be the noise vector & E IRd, which is generated from the Gaussian distribution 03
pepe |||111
This data generation model is inspired by image classification tasks Cao et al. (022)) where it has been observed that only some of the image patches (for example, the foreground) contain information (i.c. the signal) about the label. We would like the model to predict the label by focusing on such informative image patches and ignoring background patches that act as noise and are irrelevant to the classification. Note that by definition, the noise vector is orthogonal to the signal μr, i.e., E fu- 0. We assume orthogonality just for simplicity of analysis and can be casily relaxed as done in Kou et al. 2023 by assuming a stronger condition the dimensionality d (see Condition (C2) defined as part of Main Condition []. The key idea is that we can show that (5.m) (L/d) s 0 which implies that and μ are nearly orthogonal with high probability when d is sufficiently large. Consequently, terms involving (E. can be treated as negligible noise, absorbed into existing noise terms, and do not significantly affect the analysis.
[|||
Measure of Data Heterogeneity. We consider n datapoints drawn from the distribution D, and partitioned across K clients such that each client has N = n/K datapoints. The assumption of equal-sizod client datasets is made for simplicity of analysis and can be casily relaxed. The data partitioning determines the level of heterogeneity across clients.
3

---

Published in Transactions on Machine Learning Research (10/2025)
Let D+ and D denote the set of samples at client & with positive (y +1) and negative (v -1) labels respectively. Define
Ex= min (|D+,a) (|D+AA)|(-xx)) E [0, 1/2].
(1)
Note that a smaller h implies a higher data heterogeneity on average. In the IID setting, with uniform partitioning across clients, expect min([D+,sl, [D-,x)) "/2K for all & [K], and therefore h 1/2. In the extreme non-ITD setting where each client only has samples from one class, h = 0.
Two-Layer CNN. We now describe our two-layer CNN model. The first layer in our model consists of 2m filters (w5.-h=1.3 E {=1}, where each Wjr E I performs a 1-D convolution on the feature x with stride d followed by ReLU activation and average pooling Lin et al. (2013); Yu et al. (2014). The weights in the second layer then aggregate the outputs produced after pooling to get the final output and are fixed as 2/m for = +1 filters and -2/m for j = 1 filters. Formally, we have,
1(W,x) x) = lo ((w (.w,,.r,))) + a ((W+1,r, E))
la ((W1,,,8p)) + a
(2)
[W 1xx
1( 1.x)
Here W E 222d parameterizes all the weights of our neural network, W +1, W-1 E Rnd parameterize the
weights of the j = +1 filters and j = -1 filters respectively, and (z)) = max (0, 2) is the ReLU activation.
Intuitively
represents the 'logit score' that the model assigns to label j.
FL Training and Test Objectives. Let ( be the local dataset at client k. Then the global FL objective can be written follows:
min WER'd [(( = 1 Ek 1 LW().
L (w( N Ein Cyu, Monni::(WW) J(W, xx.i)),
(3)
where Lx(W) (W) is the local objective at client and 8(2, 2) = log(1 + exp(-2 . 2)) is the cross-entropy loss. We also define the lest-error L6-1 (W) as the probability that W will misclassify a point (x,y) ~ D:
T.g T (W) ((xx)) (y + sign(((Wx))).
(4)
The FedAvg Algorithm. The standard approach to minimizing objectives of the form in Equation (3) is the FedAvg algorithm. In each round of the algorithm, the central server sends the current global model (! to the clients. Clients initialize their local models to the current global model by setting w((..) W for all & E [K], and run local steps of gradient descent (GD) as follows
Local GD: (((+++1)
-(( (
(5)
for all E [0: 7-: 1] and for all k E [K]. After steps of Local GD, the clients send their local models (W((.)))
to the server, which aggregates them to get the global model for the next round: W(++1) = Ek w(t.)) KK. While we focus on FedAvg with local GD in this work, we note that several modifications such as stochastic gradients instead of full-batch GD, partial client participation Yang et al. (2021b) and server momentum Reddi al. (2021) are considered in both theory and practice. Studying these modifications is an interesting future research direction.
3
Main Results
In this section we first introduce our definition of filter aligument at initialization and a fundamental result regarding the signal-noise decomposition of the CNN filter weights. We then state our main result regarding the convergence of FedAvg with random initialization for the problem setup described in Section P and the impact of data heterogeneity and filter alignment at initialization on the test-error. Later we discuss why starting from pre-trained model can improve the test accuracy of FedAvg.
4

---

Published in Transactions on Machine Learning Research (10/2025)
3.1 Filter Alignment at Initialization
Given datapoint (x,y), y), for the CNN to correctly predict the label y and minimize the loss E(y, f(W, x));
from Equation 2 - (3,, we want wf(W, x) = (W,, x) F- -w,, > O. At an individual filter r E [m]] this can happen either with (Wy.ry' > 0 or (Wy.r, E > 0. However, we want the model to focus on the signal yu in x while making the prediction. Therefore, for filter (j.r) we want (W y/4) > 0 if j = y and (Wg.r. Upl) < 0f -y Depending the initialization of our CNN, we have the following definition of aligned and misaligned filters.
Definition 1. The ri,r)-t filter (with j E {+1), E [ml) is said to be aligned (with signal) at initialization if (w/ ju) 2 and misaligned otherwise.
We shall see Section 3.4 that the alignment of a filter at initialization plays a crucial role in how well it learns the signal and also the overall generalization performance of the CNN in Theorem 2
Extension to multi-class settings and other architectures. At a high level, the notion of alignment captures whether a filter consistently responds to the underlying signal in the data. Specifically, once a filter is aligned, the sign of its inner product with the signal remains unchanged over time, i.c., if (w(*), p4) > 0 then for all L' 2 T, sign(((((), px)) = sigu((w, 14) (see Lemma 27). Importantly, this definition not. restricted to the binary setup. The key property the preservation of sign consistency with respect to the signal - provides principled way to quantify misalignment (sce Equation ( While our current analysis formalizes this in the binary case, we believe the same notion can be naturally extended to multi-class classification. In particular, one could define aligmment with respect to multiple class-specific signals, and then study whether filters consistently align with the relevant signal subspaces. We also discuss how alignment can be extended to deeper CNNs and Transformer architectures in Section
3.2 Signal Noise Decomposition of CNN Filter Weights
One of the key insights in Cao al. (2022) is that when training the two-layer CNN with GD, the filter weights at each itcration can be expressed as linear combination of the initial filter weights, signal vector and noise vectors. Our first result below shows that this is true for FedAvg as well.
Proposition 1. Let (1)), for j E {+1} and r E [m], be the global CNN filter weights in round t. Then there ezist unique coefficients r(t)) j.r 2 0 and (P such that
w jor
( j.r ||u|| p E
p(t) tinks |E
En.is
(6)
jr
1= i=
Signall Term.
Noise Term
where [K] and e [N] denote the client and sample indez respectively.
This decomposition allows us to decouple the effect of the signal and noise components the CNN filter weights, and analyze them separately throughout training.
As we run more communication rounds (denoted by t), we expect the weights to learn the signal yp, hence it is desirable for T()) j.r increase with f. In addition, the filter weights also inevitably memorize noise & and overfit to it,, therefore the noise coefficients {p(t) j.r.k.i) will also grow with t. We are primarily interested
the growth of positive noise coofficients ((() p(t)
J.r.k, 2 0) since the negative noisc-cocfficients i.r. J.rkk. I 0) remain bounded (see Theorem in Appendix Ch and can show that
j.r.k,i = p(t)
J.r.k.e (P(()
Pirnk := p(t)
(p((0)
j.r and Σx, (tt) j.r.k.s as the signal learnóng and noise memorization coefficients of filter (ir)) respectively. As we see later in Theorem 2 the ratio of signal learning p(!)
Ex, p(t) j.rkk.i = O(E,, p(t) a. Henceforth, we refer to (())
to noise memorization jrr Ex, Jink, is fundamental to the generalization performance of the CNN.
5

---

Published in Transactions on Machine Learning Research (10/2025)
3.3 Training Loss Convergence and Test Error Guarantee
Next, we state our main result regarding the convergence of FedAvg with random initialization. We assume the CNN weights are initialized as jjr ~ (N o for all filters, where Ia the (d x d) identity matrix. We first state the following standard conditions used in our analysis.
Condition 1 (Main Condition). Let e be desired training error threshold and ò E (0,1) be some failure probability 4
(C1) The allowed number of commaniiatiio rounds t is bounded by T 7, ,, m, n, n.
(C2) Dimension d is sufficiently large: d Z max "f
[ .n²
(C3) Training set size n and neural network width m satisfy: m Z log(n/0), log(m/0). vn
}
(C4) Standard deviation of Gaussian initialization is sufficiently small: oo S min dr 111ell2 S
(C5) The norm of the signal satisfies: 1| 2 op.
(C6) Learning rale is sufficiently small: 1) S min
The above conditions are standard and have also been made in Cao et al. 2022); ; Kou et al. (0033 for the purpose of theoretical analysis. [CI]]s is a mild condition needed to ensure that the signal and noise coefficients remain bounded throughout the duration of training. Furthermore, we see in Theorem 1 that we only need T = (ranvn - -=-(=-))(((T rounds to reach training error of e, which is well within the admissible number of rounds. [C2) is used to bound the correlation between the noise vectors and also the correlation of the initial filter weights with the signal and noise. [C3) is needed to ensure that sufficient number of filters have non-zero activations at initialization so that the initial gradient is non-zcro. (C4) is needed to ensure that the initial weights of the CNN are not too large and that it has bounded loss for all datapoints. (C5) is needed to ensure that signal learning is not too slow compared to noise memorization. Finally, a small enough learning rate in (C6) ensures that Local GD does not diverge. Additional discussion on these With this assumption we are now state our main results.
assumptions is provided in Appendix C
Theorem 1 (Training Loss Convergence). For any > 0 under Condition
there erists a T =
mn
+
go
0 mn log(r/e) such that FedAug satisfies L(W(T)) S e with probability 2 1-6. nofdle
Proof Sketch. The proof divided into 3 parts. Tn the first part (Appendix C.2 we show that the magnitude of the signal and noise memorization coefficients for the global model is bounded for the entire duration of training (see Theorem 3 where [r((]] I 4 log(TT>>) and [p(t) S 4loog(T))) for all 0 < T 11 s.nk.a
Next, we divide our training into two stages. In the first stage (Appendix C.3 we show (see Lemma 21 that the noise (and also signal) memorization coefficients grow fast and are lower bounded by some constant
after T1 rounds i.c., pP (1). In the second stage (Appendix C.4), the growth of the noise and signal coefficients becomes relatively slower and the model reaches a neighborhood of a global minimizer where the loss landscape is nearly convex (see Lemma 25 Using this we can show that our objective is monotonically
decreasing in every round (see Lemma 26 which establishes convergence (in Appendix C.5
Note that our analysis does not require the condition x 1/7 as is common in many works analyzing FedAvg. Therefore, by setting 7 large enough we can make the number of rounds in the first stage as small as O (1), thereby redueing the communication cost of FL. However, in the second stage we do not see any continued benefit of local steps; in fact the number of rounds required grows as log(r). This suggests an optimal strategy would be to adapt throughout training: start with large and decrease after some rounds, which has also been found to work well empirically Wang & Joshi 2019).
Theorem 2 (Test Error Bound). Define signal-to-noiss ralio SNR =- Ilulla and A₂ = {r [m] (w...n)) v,vd 0} to be the set of aligned filters (Definition corresponding to label j. Then under the same conditions Theorem 1 our trained CNN achieves
We use 2 dennte inequalities that hide constants and logarithmin factors. See Appondix exact conditions.
6

---

Published in Transactions on Machine Learning Research (10/2025)
[ - 1/2)
m//2 alggus ins.
6.12
my22 fitters aligned
Nea-IID(h 0)
All filtess algued
All filters algnnd init.
Tise Ibror
0:10
Tist Frre
6.10
(022
0.0
0.1
0.2
C33
No. tniealigned initislisnsion
local stepe
Heterogenadiy
(a)
(b)
(c)
Figure Empirical results on synthetic dataset to verify the upper bound on test enror Theorem We fix the training error = 0.1. Figure 2a Test error increases as we increase the number of misaligned filters, with much larger
rate of increase in the non-lID setting. Figures 2b and 2c Test error increases with local steps and heterogeneity
when m/2 filters are misaligned at initialization, remains constant when all the filters are aligned.
1. When SNR" S 1/nn,, test error LB L (W(!)) 2 0.1.
2. When SNR 2 \/vnd, test error
1 Σ exp T (w(() S
(-a 11 [A d m
[-
SNR? + | ISNR m
h 2n+ 11) (1+ )²)
2
jE(21}
Proof Sketch. The upper bound on the test error Theorem relies on upper bounding the probability that Pr(y f(W()) x) S 0) for randomly sampled test data point (x, y)- Following the Gaussian concentration of Lipschitz function (see Lemma 35/. we can show that this comes down to lower bounding the ratio -yrrkks which can be intuitively interpreted filters to the sum of noise memorization across the = filters. As outlined later in Section 3.4. we see that noise memorization remains unaffected by initialization and data heterogeneity, i.c., we get Σrk, p(o) P = 0(03md) (sce Lemma 30), whereas signal learning depends on both the initial alignment and data heterogeneity ie., Em af ((w ypi) = 2(n Il (IA,) + (m - [Ay])(h + +1 h))) (see Lemma 33 Substituting these quantitics in our upper bound in Appendix D.1 completes the proof.
E ((w (9)
yor, wp)) En., Tot)
bounding the sum of signal
learning across the
For our lower bound, we use a similar argument done in Kou al.
2023
The first step involves showing
that the Pr(y- s(W()), x) < 0) 2 0.5 Pr(2) where Ω ΩZ { EE,, jo (10 j.r: <) 2 C maxx
E, r
j.r] where
C some positive constant. Now given SNR² < 1//vnd, we can show for any given & there exists vector v such that one of +v E must belong to Ω (sce Lemma 5.8 in Kou et al. (2023)). By union bound it follows that min (Pr(S,, Pr(- Pr(0- v), Pr(-Ω- v)} 2 0.25. Now using arguments based on symmetry and TV distance we can show that Pr(Ω)) 2 0.22, which completes the bound.
Impact of SNR on cermfuu/benngge overfitting. Intuitively, if the SNR is too low (SNR² S 1/Vnd), then there is simply not enough signal strength for the model to learn compared to the noise. Hence, we cannot expect the model to generalize well no matter how we train it. This generalizes the centralized training result in (Kou et al., 2023] Theorem 4.2) (with p = 0), which corresponds to 7 1 in FedAvg. In this case, the model is in the regime of harmful overfitting. However, if the SNR is sufficiently large (SNR² 2 2 1/vma), enter the regime of benign overfitting, where the model can fit the data and generalize well with the test error reducing exponentially with the global dataset size
Empirical Verification. We now provide empirical verification of the upper bound on the test error in Theorem 2 in the benign overfitting regime. We simulate a synthetic dataset following our denaratnnottion model in Section with n = 20 datapoints, = 2 clients and m 10 filters. Additional experimental details can be found in Appendix F. We fix training crror threshold of c 0.1 and then measure the test error of our CNN under various settings in Figure 2 Figure shows the test error as a function of the number of misaligned filters (m [Ayl in Theorem under different. data partitionings with the number of local steps
7

---

Published in Transactions on Machine Learning Research (10/2025)
fixed at 100. While the test error grows with the number of misaligned filters in both data settings, the rate of growth is much larger in the non-IID setting. Figure 2b shows the test error a function of local steps T under different initializations for fixed h = 0 while Figure 2 shows the test error as function of heterogencity under different initializations for fixed r = 100. As predicted by our theory, heterogencity and the number of local steps do not affect test error when all the filters are aligned al initialization. On the other hand, the test error grows with T and heterogeneity when the number of misaligned filters is non-zero (m/2 = 5) for each j E {+1}. Therefore, our empirical results strongly validate our theoretical results showing the effect of heterogeneity, number of local steps and number of misaligned filters on the test error.
3.4 Impact of Filter Alignment and Data Heterogeneity on Signal Learning and Noise Memorization.
The key results in our analysis are the following lemmas which bound the growth of the signal learning and noise coefficient during the first stage of training, that is 0 S S Ti (see discussion under Theorem D Using our definition of A, := fr E [m]] (w) ju) 2 0} as the set of aligned filters, we have the following lemma for growth of the signal learning coefficient in the first stage.
(n || ) if E A, and j.r
Lemma 1. Under Condition 1. for all 0 S t S Ti, we have [ 52
Ω by 1 (1 + M(T (07-11) 12) if r & Ag-
(6) This lemma shows that for aligned filters (r E Ay), j.r does not depend on heterogeneity and grows linearly with the number of local steps T. On the other hand, for misaligned filters (r & Aj), the growth depends on the heterogeneity parameter h. Furthermore, under extreme data heterogeneity (h = 0), for misaligned filters pl0) does not scale with the rember of local steps T. For the growth of noise coefficients we have the following corresponding lemma,
[tyToBd ( m
Lemma 2. Under Condition 1 for all allost we have Σn, P j.r.k, =
This lemma shows that noise memorization does not depend on dale-bheeeroyeneity or filler alignsmnt and always seales linearly with the number of local steps t. Intuitively, this can be expected because the noise vectors are independent of the label information y in a datapoint following our data generation model in Section [ and for any given filter we can show there are Ω (N) noise vectors that are aligned with the filter at initialization for every client with high probability (see Lemma 7
Using the above two lemmas, we have the following bound on the ratio of signal learning to noise memorization for filter (j,r) at the end of the first stage of training
SNR?, if E Aj, SNR + (1--)), if E [m]) A₁-
ΣN,
B( Jinkk
This ratio is key to bounding the generalization performance of the CNN model as we show later in the proof of Theorem 2 in Appendix [.1] For aligned filters (r E Ay), the ratio is unaffected by data heterogencity h and the number of local steps r. However, for misaligned filters (r E [m]] \ Aj), the ratio becomes smaller as heterogencity increases (h becomes smaller) or increases. Thus, for misaligned filters we see a corresponding dependence on heterogeneily and local steps in our upper bound on test error in Theorem 2 Note that in and thus we do not see any impact of heterogeneity
centralized training with 1, have & (1-A)) at misaligned filters. Therefore, we recover the bound L (DW(() S exp(-nSNR') /d) in (Kou et al. 2023 Theorem 4.2). It. is only in FL training with t > 1 local sleps that encounter the adverse effect of dala heterogeneity at the misaligned filters.
Empirical Verification. We empirically verify the results above in the IID (h = 1/2) and non-IID (h 0) setting following the same simulation setup as done in Figure 2 Figure 4 shows the aligmment of filters at initialization (note that only filter r = is aligned). Figure 3a shows that in the IID setting signal learning coelficients are similar for all the filters regardless of aligument and increases with the number of local step.
8

---

Published in Transactions on Machine Learning Research (10/2025)
1.5
0.1
60
0.
No. load step
Load
No. lood stego
(a) ITD Signal Learning
(b) ITD Noise Memorization
(c) ITD Sig. Learning/ /Noie Mem.
E
6.1
0.3
00.3
No. loaa steps
looa stepp
(d) Non-IID Signal Learning
Non-IID Noise Memorization
(f) Non-IID Sig. Learning/Noii Mem.
Figure 3:
Signal learning and noise memorization for our CNN model in the IID (h 1/2) and non-IID (h 0)
setting after round. Figures 3a
In the IID setting signal learning coefficients are similar for all the filters and increase with the number of local steps T but in the non-IID setting they saturate (Lemma D for misaligned filters
3d
(r= 2, 5). Figures 3b 3e Noise memorization is similar for all filters in both settings and grows with 7 Lemmna 2
Figures 31 in the IID setting, the ratio of signal learning to noise memorization remains independent of r. But
the non IID setting, the ratio decreases to zero as 7 increases for misaligned filters (r 4,5,
However, shown by Figure 3d in the non-lID setting signal learning saturates for misaligned filters. Figures 3b and 3e show that the growth of noise coefficients for all the filters is similar in the IID and non-IID case.
011
-01
In Figure Be we see that ratio of signal learning to noise memorization lower bounded by a constant for filters in the IID setting whereas in the non-IID setting it docays as 7 increases for misaligned filters (Figure 3f thus verifying theoretical analysis.
Figure 4: Initial alígnment the filters in Figure 3 Only filter r 3 is aligned.
3.5 Impact of Pre- Training Federated Learning
Given the result in Theorem we return to our question in Section [ about the effect of pre-trained initialization on improving generulization performance FL. We focus on centralized pre-training but our discussion here can be extended to federated pre-training well (see Lemmal 31 which states a federated counterpart of the lemma below).
Suppose we pre-train a CNN model in a centralized manner on dataset with signal μ (pre) generated according to the data model described in Section 2 Now if we train for sufficient number of iterations, then we can show that all filters will be correctly aligned with the pre-training signal.
Lemma 3 (All Filters Aligned After Sufficient Training). There exists Ti O
such that for all
nood
12 Tıjj E (+1), r E [m] we have (w (pre.)) (pplpre) 2 O.
j.r
Now suppose we pre-train for 2 Ti iterations to get model W(pre,s) and use this model to initialize for downstream foderated training (i.c., W(0) = (pre,, ")) with signal vector μ.. Then for all j.r filters, we have (pro.)
(w (0)
jp (w(rre,") (pro,") j.pp(pre)) + (w(nee.) .j( (pre))). We also know that (pre:") (jj((e)))) 2 using
Lemma 3 Therefore, if || (pre))] is small, all the filters (w (.r) correctly aligned with the signal ju.

---

Published in Transactions on Machine Learning Research (10/2025)
CIFAR-I0 10
CIFAR-10
40
% of Misalignod Filters Init.
80-
30
25.5
Test. Accursey
70-
20
60-
12.6
10-
50-
Prestrined
Randon
40
Pre-trained
Random
200
Round No.
(b)
ThyylmageNet
Tiyy TuylmngeNed
% of Misaligned Filters Init.
33.77
10-
30
Test Accuraey
40-
30-
20-
15.0
20-
10-
10
Pre-trained
Handon
Pre-trained
Random
200
300
Round No.
(d)
Figure 5: The percentage of misaligned filters (see Equation CIFAR-10 (Figure 5a and Figure 5b and TinylmageNet (Figure 5c and Figure 5d As the complexity of the signal information in the data grows from CIFAR-10 to TinyImageNett, filters for random initialization, explaining why pre-trained initialization offers larger improvements for TinyImageNet.
[ and test accuracy for different initializations
sharp increase in the ratio of misaligned
As result, in Theorem 2 A, im] for j E {+1} and in the benign overfitting regime (SNR² 2 1/vnd), we recover the centralized result. L (w(1)) S exp(-nSNR²/d) (Kou el al. 2023 Theorem 4.2). Hence, the adverse effects of cross-client heterogeneity are mitigated with pre-trained initialization.
4 Experiments
In this section we provide empirical results showing how insights from Section extend to practical FL training on real world datasets with deep CNN models. Unless specified otherwise, we use the ResNet18 model He et al. (2016) in all our experiments and split the data across 20 clients using the Dirichlet sampling scheme Hsu et al. (2019) with non-iid parameter a = 0.3. For pre-training, we use a ResNet18 pre-trained on ImageNet Russakovsky et al. (2015), available in PyTorch Paszke et al. (2019). We emphasize that these results are primarily qualitative and indicative, as several simplilying assumptions from our theoretical analysis such binary classification and a two-layer CNN do not strictly hold in this setting. Nevertheless, observe that the key insights derived from theory continue to hold empirically. For completenoss, also perform experiments small binary-class subset that directly mirrors the theoretical setup and provide additional details in Appendix F
Empirical Measure of Misalignment. Measuring filter aligument for deep CNNs is challenging since we cannot explicitly characterize the signal information present in real world datasets and furthermore different layers will learn the signal different levels of gramularity. Nonetheless, our thcoretical findings suggest: that given sufficient number of training rounds, filters will be aligned with the signal (see Section 3) and once filter aligned, the sign of the output produced by the filter with respect to the signal does not.
10

---

Published in Transactions on Machine Learning Research (10/2025)
CIFAR-10 0.03)
CIFAR-10/a 0.051
30
% of Misslgnedd Pilters at luit.
26.1
Pre-trained
Ranzdou
Test Accuracy
20-
60
15
11.5
10-
Pre-trained
Random
100
200
300
Round No.
(b)
CIFAR- 1 10)
CIFAR-101 10/a 10)
30
% Misadignod Filters at Init.
80-
25
24.1
20-
Test Accuracy
70
15
60-
10
9.2
50
Prestrained
Random
40-
Pre-trained
Random
100
200
300
Round No.
(c)
(d)
Figure 6: The percentage of misaligned filters (see Equation and test accuracy for different initializations CIFAR-10 with a = 0.05 heterogeneity (Figure 6a and Figure 6b and a = 10 heterogeneity (Figure 6c and Figure [d]). Although the percentage misaligned filters does not vary significantly across the two settings for both initializations (signal information is the same in both settings), pre-training offers more improvement in the higher heterogeneity setting (a 0.05), as suggested by our theoretical analysis.
change, i.c, if (w( jpe) > 0 then sign((w (t) use the sign the output produced by a filter at the end of training as a reference for alignment at any given round. Formally, let W(O), W()... (T) be the sequence of iterates produced by federated training and let x) = [[w, ((1), ((2) ... (w, x(p)i] E Iº be the feature map vector generated by filter for input x. For given batch of data B, we define the empirical mcasure of alignment of filter w(t) relative to w(T) follows:
14)) = sign((w j.rt p4),, for all 2 t. Therefore, we propose to
A(w(!)= = Σ (x)
x)).
(8)
TEBIEE
We say that the weight (()) at round is misaligned if A(w < O, because this implies that the sign of the output produced by the filter at round t eventually changed for majority of the inputs, hence indicating that the filter was misaligned at round f. We compute this measure over batch of data to account for signal information coming from different classes of data as well as reduce the impact of noise in the data.
Measuring Misalignment on Real World Datasets with Varying Signal Information. In this experiment our goal is to empirically demonstrate that (a) pro-trained initialization leads to much fewer number of misaligned filters than random initialization and (b) the number of misaligned filters for random initialization increases as we increase the complexity of the signal. To demonstrate this, we consider federated training on the CIFAR-10 Krizhevsky (2009) and 2. TinyImagoNos [e Yang (2015) datascts. Figure shows the test accuracy and percentage of misaligned filter across training rounds for both datasets with pre-trained and random initialization. Firstly, we that the percentage of misaligned filters is 2 - 3x
11

---

Published in Transactions on Machine Learning Research (10/2025)
smaller when starting from pre-trained initialization compared to random initialization. Furthermore, as the complexity of the signal information in the dataset increases (CTFAR-10 < TinyImageNet), we see a sharp increase in the percentage of misaligned filters (25% to 40%) for random initialization. In contrast, with pro-trained initialization, the percentage of misaligned filters remains less than 15% across datascts leading to larger improvement in lest accuracy for Tiny TinyeeageNet. These results align with our theoretical findings: as the ratio of misaligned filters increases, the benefits of pro-training become more pronounced.
Measuring Misalignment with Varying Heterogeneity Levels. We extend the experiment in Figure conducted on CIFAR-10 with a 0.3 Dirichlet heterogeneity to other levels of heterogeneity 1. a = 0.05 which is an extreme non-lID split and 2. a = 10 which can be thought of as close to IID split. Figure G shows the test accuracy and percentage of misaligned filters plots for these two heterogencity levels with pre-trained and random initialization. We observe that in both cases the percentage of misaligned filters remains approximately 25% with random initialization and 10% with pre-trained initialization, regardless of the level of heterogeneity. However, as heterogeneity increases, the improvement in test accuracy provided by pre-trained initialization becomes more pronounced. This trend is consistent with our theoretical analysis in Theorem 2, which suggests that the percentage of misaligned filters will have a greater impact on test. performance as data heterogencity increases.
5 Conclusion and Future Work
In this work we provide a deeper thcoretical explanation for why pre-training can drastically reduce the adverse effects of non-ITD data in FL. by studying the class of two layer CNN models under a signal-noise data model. Our analysis shows that the reduction in test accuracy seen in non-llD FL compared to IID FL is only caused by filters that are misaligned at initialization. When starting from pro-trained model we expect most of the filters to be already aligned with the signal thereby redueing the elfect of heterogeneity and leading to a higher ratio of signal learning to noise memorization. This is corroborated by experiments on synthetic setup as well as more practical FL training tasks. Our work also opens up several avenues for future work as we discuss below.
Extension to Deeper CNNs and Transformer Architectures. For deeper CNNs, we anticipate that will need hierarchical feature model (Li & Li, 2024 Allen-Zhu & Li 2023b Wang et al., 2023b) to capture alignment in deeper layers. Intuitively, while deeper architectures provide the capacity to learn more complex hicrarchical signals, aligning a filter with the signal will also progressively become harder across layers with the aligument of the L-th layer potentially depending on the alignment of the (L 1)-th layer. In non-lID FL, this suggests that feature learning deteriorates as we move from the base layer toward the final layer. This intuition is supported by empirical evidence (sce experiment in Table Yu et al. (2022b) and explains why pre-trained initialization offers greater benefits for deeper networks compared to shallower ones.
For extending to Transformers, recall that alignment of convolutional filter is defined via the sign of its inner product with input patch containing the signal (see Definition D. For Transformers, the convolution operation is replaced by the self-attention mechanism. However, the core principle remains: intermediate representations are computed as inner products between the previous-layer embeddings and the current-layer attention weights. Specifically, parameterizing a self-attention layer by (Q, (2 V) E gd and input 2 E R N we have:
g = 2Q, k =zK, =V (intermediate representations)
gli Vdl
8' sootmax
(self-attention aggregation).
(10)
Notice that cach column of Q, K, and V is involved in an inner product with the input 2, analogous to filters in CNNs. Thus, our notion of filter alignment and its empirical measurement (Equation (5) can naturally be extended to these matrices. This connection provides a promising foundation for generalizing alignment analysis to Transformer-based architectures. Lastly, we note that recent work has also begun to explain benign overfitting in Transformer models (Jiang et al. 2024 Frei & Vardi 2024), which could serve as a natural starting point for extending our analysis.
12

---

Published in Transactions on Machine Learning Research (10/2025)
References
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. In International Conference on Learning Representations, 2021.
Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. International Conference on Learning Representations, 2023a.
Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep (hier- archical) learning. In The Thirty Sirth Annual Conference on Learning Theory, pp.
PMLR,
2023b.
Yajie Bao, Michael Crawshaw, and Mingrui Liu. Provable benefits of local steps in heterogencous federated learning for neural networks: A feature learning perspective. In Forty-first International Conference Machine Learning, 2024.
Leighton Pate Barnes, Alex Dytso, and H Vincent Poor. Improved information theoretic generalization bounds for distributed and federated learning. In 2022 IEEE International Symposium on Information Theory (ISIT), pp. 1465 1470. IEEE, 2022.
Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convolutional neural networks. Advances in Neural Information Processing Systems, 35:25237 22252037-2525,, 2022.
Hong- You Chen, Cheng-Hao Tu, Ziwei Li, Han-Wei Shen, and Wei-Lun Chao. On the importance and applicability of pre-training for federated learning. International Conference on Learning Representations, 2022.
Shuxiao Chen, Qinqing Zheng, Qi Long, and Weijie J Su. A theorem of the alternative for personalized federated learning. XrXie preprint arXiv:2103.01900 2021.
Gary Cheng, Karan Chadha, and John Duchi. Fine-tuning fine in federated learning. arXie preprint arXiv:210S.07913, 3, 2021.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Fedavg with fine tuning: Local updates lead to representation learning. Advances in Neural Information Processing Systems, 35:10572- 2022.
Soham De, Leonard Berrada, Jamie Hayes, Samuel L Smith, and Borja Balle. Unlocking high-accuracy diffcrentially private image classification through scale. arXiv preprint arXiv:2204.13650 2022.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for langnage understanding. In North American Chapter of the Association for Computational Linguistics, 2019.
Lue Devroye, Abbas Mchrabian, and Tommy Roddad. The total variation distance between high-dimensional gaussians with the same mean. arXiv preprint XrXi::11000088 0.0689,, 2018
Alexey Dosovilskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations, 2021.
Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, and Barnabas Poczos. Gradient descent learns one- hidden-layer cnn: Don't be afraid of spurious local minima. Tn International Conference on Machine Learning, pp. 1339 1348. PMLR, 2018.
Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise to sensitivity in private data analysis. In Theory of Cryptography: Third Theory of Cryptography Conference, TCC 2006, New York, NY, USA, March 4-7, 2006. Proceedings 3, pp. 265 284. Springer, 2006.
13

---

Published in Transactions on Machine Learning Research (10/2025)
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Generalization of model-agnostic meta-learning algorithms: Recurring and unseen tasks. Advances in Neural Information. Processing Systems, 34:5469-544 -5480, 2021.
Eros Fani, Raffaello Camoriano, Barbara Caputo, and Marco Ciccone. Fed3r: Recursive ridge regression for fedcrated learning with strong pre-trained modcls. In International Workshop on Federated Learning in the Age of Poundation Models in Conjunction with NewrIPS 2023, 2023.
Spencer Frei and Gal Vardi. Trained transformer classifiers generalize and exhibit benign overfitting in-context. arXiv preprinl 2xXi:24400001 2024.
Arun Ganesh, Mahdi Haghifam, Milad Nasr, Sewoong Oh, Thomas Steinke, Om Thakkar, Abhradeep Guha Thakurta, and Lun Wang. Why is public pretraining necessary for private model training? International Conference on Machine Learning, pp. 10611-10627. PMLR, 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint orXiv:2101.00022 2020.
Peyman Gholami and Hulya Seferoglu. Improved generalization bounds for communication efficient federated learning. arXiv preprint arXiv:204..117 117544 2024.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Yee Whye Teh and D. Mike Titterington (eds.), Proceedings of the Thirteenth International Conference Artificial Intelligence and Statistics, AISTATS 2010, Chia Loguna Resort, Sardinia, Italy, May 13-15, 2010, volume 9 of JMLR Proceedings, pp. 249 256. JMLR.org, 2010. URL httt::/proccedding mlr gress/v9/glorot100 html
Samyak Gupta, Yangsibo Huang, Zexuan Zhong, Tianyu Gao, Kai Li, and Danqi Chen. Recovering private text in federated learning of language models. Advances in Neural Information Processing Systems, 35: 8130 8143, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human- level performance on imagenet classification. In 2015 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pp. 1026- 1034. TEFE Computer Society, 2015. doi: 10.1109/ICCV.2011 2015.123. URL https://doi. org/ 10. 1109/ICCV. 2015. 123
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp. 770-778. IEEE Computer Society, 2016. doi: 10.1109/CVPR.2D0 https: org/ 10. 1109/ CVPR. 2016.90
URL
Kaiming He, Ross Girshick, and Piotr Dollár. Rethinking imagenet. pre-training. In Proceedings of the IEFE/CVE /VFF International Conference on Computer Vision, pp. 4918 4927, 2019.
Charlie Hou, Akshat Shrivastava, Hongyuan Zhan, Rylan Conway, Trang Le, Adithya Sagar, Giulia Fanti, and Daniel Lazar. Pre-text: Training language models on private fodcrated data in the age of 1ms. arXio preprint arXiv:24060000 02958, 2024.
Tau-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the elfects of non-identical data distribution for federated visual classification. In International Workshop on Federated Learning for User Privacy and Data Comfidentiality in Conjunction with NewrIPS 2019 (FL-NeurIPS' 19)), Docember 2019.
Xiaolin Hu, Shaojie Li, and Yong Liu. Generalization bounds for federated learning: Fast rates, unparticipating clients and unbounded losses. In The Eleventh International Conference on Learning Representations, 2022.
Baihe Huang, Xiaoxiao Li, Zhao Song, and Xin Yang. Fl-ntk: A neural tangent kernel-based framework for foderated learning analysis. In International Conference on Machine Learning, pp. 4423 4434. PMLR, 2021.
14

---

Published in Transactions on Machine Learning Research (10/2025)
Wei Huang, Ye Shi, Zhongyi Cai, and Taiji Suzuki. Understanding convergence and generalization in federated learning through feature learning theory. In The Twelfth International Conference on Learning Representations, 2023.
Forrest N landola, Song Han, Matthew W Moskewicz, Khalid Ashraf, William J Dally, and Kurt Keutzer. Squeezenet: Alexnet-level accuracy with 50x fewer parameters and< 0.5 mb model size. arXiv preprinl arXiv.1602.079960 2016.
Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves genoralization in deep learning. In International Conference on Machine Learning, pp. 9965 10040. PMLR, 2022.
Jiarui Jiang, Wei Huang, Miao Zhang, Taiji Suzuki, and Liqiang Nie. Unveil benign overfitting for transformer vision: Training dynamics, convergence, and generalization. Advances in Newral Information Processing Systems, 37:135464 135625, 2024.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trends0 in Machine Learning, 14(1 2):1 210, 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In Inter- national Conference on Machine Learning, pp. 1132--1483 PMLR, 2020.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Breaking the centralized barrier for cross-device federated learning. Advances in Neural Information Processing Systems, 34:28663 226766, 2021.
Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting in two-layer relu convolutional neural networks. In International Conference on Machine Learning, pp. 17615 17659. PMLR, 2023.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. URL httpp://sww. http://wwww..ttoo.teeeeeeeete torooto. socoott...et///tttt//aanntttttte TRo9
Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma, and Percy Liang. Fine-tuning can dis- tort pretrained features and underperform out-of-distribution. International Conference Learning Representations, 2022.
Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015.
Yann LeCun, Léon Bottou, Genevieve B. Orr, and Klaus-Rober Müller. Efficient backprop. In Grégoire Montavon, Genevieve B. Orr, and Klaus-Robert Müller (eds.), Neural Nelworks: Tricks of the Trade Second Edition, volume 7700 of Lecture Notes in Computer Science, pp. 9 48. Springer, 2012. doi: URL http:::////o...//1111.110000 10tp//00.007//78---1--11 642- 1-
Gwen Legale, Nicolas Bernier, Lucas Page-Caccia, Edouard Oyallon, and Eugene Belilovsky. Guiding the last layer in federated learning with pre-trained models. Advonces in Neural Information Processing Systems, 36, 2024.
Binghui and Yuanzhi Li. Adversarial training can provably improve robustness: Theoretical analysis of feature learning process under structured data. orXiv preprint arXiv:2410.08550 0800000,, 2024.
Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10713 10722, 2021.
Tian Li, Anit Kumar Sahu, Amoet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEFE Signal Processing Magazine, 37(3):50 60, 2020.
15

---

Published in Transactions on Machine Learning Research (10/2025)
Xuechen Li, Daogao Liu, Tatsunori B Hashimoto, Huseyin A Inan, Janardhan Kulkarni, Yin-Tat Lee, and Abhradeep Guha Thakurta. When does differentially private learning not suffer in high dimensions? Advances in Neural Information Processing Systems, 35:28616 28630, 2022a.
Xuechen Li, Florian Tramer, Percy Liang, and Tatsunori Hashimoto. Large language models can be strong differentially private learners. International Conference on Learning Representations, 2022b.
Min Lin, Qiang Chen, and Shuicheng Yan. Network in network. arXiv preprint arXiv: 2013.
Tao Lin, Lingjing Kong, Sebastian Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2363, 2020.
Dianbo Liu and Tim Miller. Federated pretraining and fine luning of bert using clinical notes from multiple silos. arXiv preprint Xic:2202.088562 2020.
Dhruv Mahajan, Ross Girshick, Vignesh Ramanathan, Kaiming He, Manohar Paluri, Yixuan Li, Ashwin Bharambe, and Laurens Van Der Maaten. Exploring the limits of weakly supervised pretraining. In Proceedings of the Europeon Conference on Computer Vision (ECCV), pp. 181-196, 2018.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication- efficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pp. 1273-1282. PMLRR, 2017.
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In International Conference on Machine Learning, pp. 4615-4625. PMLR, 2019.
John Nguyen, Jianyu Wang, Kshitiz Malik, Maziar Sanjabi, and Michael Rabbat. Where to begin? on the impact of pre-training and initialization in foderated learning. International Conference on Learning Representations, 2022.
Junsoo Oh and Chulhee Yun. Provable benefit of cutout and cutmix for feature learning. In High-dimensional Learning Dynamics 2024: The Emergence of Structure and Reasoning, 2024.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32, 2019.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Dly Sutskever, el al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21 21140)):-667 2020.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. International Conference on Learning Representations, 2021.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenct large scale visual recognition challenge. International Journal of Computer Vision, 115:211 252, 2015.
Christoph Schuhmann, Romain Beaumont, Richard Veneu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al. Laion-5b: An open large-scale dataset for training next generation image-text models. Advances in Neural Information Processing Systems, 35:25278 25294, 2022.
Milad Sefidaaran, Romain Chor, and Abdellatif Zaidi. Rato-distortion theoretic bounds on generalization error for distributed learning. Advances in Neural Information Processing Systems, 35:19687-19977 2022.
16

---

Published in Transactions on Machine Learning Research (10/2025)
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In Proceedings of the IEEF International Conference on Computer Vision, pp- 843 852, 2017.
Zhenyu Sun and Ermin Wei. communication-efficintt algorithm with linear convergence for federated minimax learning. Advonces in Neural Information Processing Systems, 35:6060-600 6073, 2022.
Zhenyu Sun, Xiaochun Niu, and Ermin Wei. Understanding generalization of federated learning via stability: Heterogeneity matters. In International Conference on Artificial Intelligence and Statistics, pp. 676 684. PMLR, 2024.
Yue Tan, Guodong Long, Jie Ma, Lu Liu, Tianyi Zhou, and Jing Jiang. Federated learning from pre- trained models: A contrastive learning approach. Advances in Neural Information Processing Systems, 35: 19332 18332-19344, 2022.
Bart Thomee, David A Shamma, Gerald Friedland, Benjamin Elizalde, Karl Ni, Douglas Poland, Damian Borth, and Li-Jia Li. Yfec100m: The new data in multimedia research. Communications of the ACM, 59 (2):64 73, 2016.
Yuanyishu Tian, Yao Wan, Lingjuan Lyu, Dezhong Yao, Hai Jin, and Lichao Sun. Fedbert: When fodcratod learning meets pre-training. ACM Transoctions on Intelligent Systems and Technology (TIST), 13(4): 2022.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi- task benchmark and analysis platform for natural language understanding, arXiv preprint 1xx::11144.00. 2018.
Boxin Wang, Yibo Jacky Zhang, Yuan Cao, Bo Li, H Brendan McMahan, Sewoong Oh, Zheng Xu, and Manzil Zaheer. Can public large language models help private cross-device federated learning? arXiv preprint ar Xrv::23005.12 12132, 2023a.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris S. Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. In International Conference on Learning Representations, 2020.
Jianyu Wang and Gauri Joshi. Adaptive communicatio strategies to achieve the best error-runtime trade-off in local-update sgd. Proceedings of Machine Learning and Systems, 1:212-221 229, 2019.
Peng Wang, Xiao Li, Can Yaras, Zhihui Zhu, Laura Balzano, Wei Hu, and Qing Qu. Understanding deep rep- resentation learning via layerwise feature compression and discrimination. arxiv preprint arXiv22311.0000 2023b.
Shanshan Wu, Zheng Xu, Yanxiang Zhang, Yuanbo Zhang, and Daniel Ramage. Prompt public large language models synthesize data for private on-device applications. COLM, 2024.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018.
Zheng Xu, Maxwell Collins, Yuxiao Wang, Liviu Panait, Sewoong Oh, Sean Augenstein, Ting Liu, Florian SchrofT, and H Brendan McMahan. Learning to generate image embeddings with user-level differential privacy. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognétion, pp. 7969-7980, 2023a.
Zheng Xu, Yanxiang Zhang, Galen Andrew, Christopher Choquette, Peter Kairouz, Brendan Mcmahan, Jesse Rosenstock, and Yuanbo Zhang. Federated learning of gboard language models with differential privacy. In Sunayana Sitaram, Beata Beigman Klebanov, and Jason Williams (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 5: Industry Track), pp. 629 639, Toronto, Canada, July 2023b. Association for Computational Linguisties. doi: 10.18653/vv (6653/vl// v1122023,aa222 ael-aaclinnn id-industryy60. URL htt.atpss//acclmaoooooooooot
17

---

Published in Transactions on Machine Learning Research (10/2025)
Chengxu Yang, Qipeng Wang, Mengwei Xu, Zhenpeng Chen, Kaigui Bian, Yunxin Liu, and Xuanzhe Liu. Characterixing impacts of heterogeneity in federated learning upon large-scale smartphone data. In Proceedings of the Web Conference 2021, pp. 935 946, 2021a.
Haibo Yang, Minghong Fang, and Jia Liu. Achieving linear speedup with partial worker participation in non-iid federated learning. International Conference Learning Representations, 2021b.
Jiayuan Ye, Zhenyu Zhu, Fanghui Liu, Reza Shokri, and Volkan Cevher. Initialization matters: Privacy-utility analysis of overparametorized neural networks. In Alice Oh, Tristan Naunann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL htttp://papers nipa. wwtpapprrfillls/eeeeeeeeeeeeeeee www..ss/pppee//2//aaaaaaaa 1165ff60116336622204444660000 -aeess-t-ttttttttttnaaa Confeeenceeeeeeeceeeeee html
Da Yu, Saurabh Naik, Arturs Backurs, Sivakanth Gopi, Huseyin A. Inan, Gautam Kamath, Janardhan Kulkarni, Yin Tat Lee, Andre Manoel, Lukas Wutschitz, Sergey Yekhanin, and Huishuai Zhang. Differentially private fine-tuning of language models. In International Conference on Learning Representations (ICLR), 2022a.
Dingjun Yu, Hanli Wang, Peiqiu Chen, and Zhihua Wei. Mixed pooling for convolutional neural networks. In Rough Sets and Knowledge Technology: 9th International Conference, RSKT 2014, Shanghai, China, October 24-26, 2014, Proceedings 9, pp. 364-375. Springer, 2014.
Yaodong Yu, Alexander Wei, Sai Praneeth Karimireddy, Yi Ma, and Michael Jordan. Tet: Convexifying federated learning using bootstrapped neural tangent kernels. Advances in Neural Information Processing Systems, 15:30882-38991 2022b.
Honglin Yuan, Warren Morningstar, Lin Ning, and Karan Singhal. What do we mean by genoralization in federated learning? Xix preprint arXiv:2110. 10X1:2210..4216, 2021.
Tuo Zhang, Tiantian Feng, Samiul Alam, Dimitrios Dimitriadis, Mi Zhang, Shrikanth S Narayanan, and Salman Avestimehr. Gpt-fl: Generative pre-trained model-assisted federated learning. arXiv preprint arXim:2306.022) 2023.
Weiming Zhuang, Chen Chen, and Lingjuan Lyu. When foundation model meets federated learning: Motivations, challenges, and future directions. arXiv preprint arXiv:2306.1553 2023.
Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. Understanding the gencralization of adam learning neural networks with proper regularization. International Conference on Learning Representations, 2023.
18

---

Published in Transactions on Machine Learning Research (10/2025)
Appendix
A Additional Related Work
20
B Theory Notation and Preliminaries
21
B.1 Local Model Update
22
B.2 Proof of Proposition
22
B.3 Co-efficient Update Equations
23
C
Training Error Convergence of FedAvg with Random Initialization
23
Preliminary Lemmas
C.2
Bounding the Scale of Signal and Noise Memorization Coefficients
................................
26
C.3
First Stage of Training,
42
C.4
Sccond Stage of Training
44
Prooř of Theorem
48
Proof of Theorem
49
D.1 Test Error Upper Bound
D.2 Test Error Lower Bound
57
E Main Paper Lemma Proofs
58
Proof of Lemma
58
E.2 Proof of Lemma
58
Proof of Lemma
59
Additional Experimental Details and Results
59
19

---

Published in Transactions on Machine Learning Research (10/2025)
A Additional Related Work
Use of Pre-Trained Models in Federated Learning. Tan et al. (2022) 9 explore the benefit of using pre-trained models in FL by proposing to use multiple fixed pre-trained backbones as the encoder model at each client and using contrastive learning extract useful shared representations. [Zhuang al. (2023) discuss the opportunities and challenges of using large foundation models for FL including the high commuication and computation cost. One solution to this as proposed by Legate et al. (2024) is that instead of full fine-tuning as done im Chen et al. (002)) Nguyen et al. (2022). we can just fine-tune the last layer. Specifically Legate et al. (2024) proposes a two-stage approach to federated fine-tuning by first fine-tuning the head and then doing a full-finetuning This approach is inspired by results in the centralized setting Kumar et al. (2022 which show that in some case fine-tuning can distort the pre-trained features. Fani et al. (2023) also study the problem of fine-tuning just the last layer in federated setting by replacing the softmax classifier with a ridgo-regression classifier which enables them to compute a closed form expression for the last layer weights.
There has also been some recent. work on exploring the benefit of pre-training for federated natural language processing tasks including the use of Large Language Models (LLMs). [Wann et al. (2023a) discuss how to leverage the power of pre-trained LLMs for private on-device finc-tuning of language models. Specifically, Wang et al. (2023a) proposes a distribution malching approach to select public data that closest to privale data and then use this selected public data to train the on-device language model. Zhang et al. (2023) propose to first pre-train on synthetic data construet the initialization point followed by federated fine-tuning. Hou et al. (0244 propose that clients send DP information to the server which then uses this information to generate synthetic data and fine-tune centrally on this synthetic data. [iu & Miller (2020) discuss the challenges of pre-training and tunettnnng BERT in federated manner using clinical notes from multiple silos without data transfer. Tian et al. (2022) propose to pre-train a BERT model in a federated manner in a more general setting and show that their pre-trained model can retain accuracy on the GLUE (Wang et al.) 2018 dataset without sacrificing client privacy. Xu et al.] (2023b) pretrain production on-device language models on public web data before fine-tuning in federated learning with differential privacy, and Wu et al. (2024) later replace the pretraining data with data synthesized by LLMs. Gupta et al. (2022) propose defense using pre-trained models to prevent an attacker from recovering multiple sentences from gradients in the federated training of the language modeling task.
Importance of Initialization for Private Optimization. We note that an orthogonal line of work has explored the benefits of starting from a pre-trained model when doing differentially private optimization
Dwork et al. (2006) and seen similar striking improvement in accuracy De et al. 2002 ; Li et al. (2022b): Yu Xu et al. (2023a), as we see in the heterogeneous FL setting. Ganesh et al. (2023) study this phenomenon for a stylized mean estimation problem and show that public pre- pre-training can help the model start from a good loss basin which is otherwise hard to achieve with private noisy optimization. Li et al. (2022a) study differentially private convex optimization and show that starting from a pre-trained model can leads to dimension independent convergence guarantees. Specifically Li et al. (002aa) define the notion of restricted Lipschitz continuity and show that when gradients are low rank most of the restricted Lispchitz coefficients will be zero. Ye et al. (2023) studies the impact of different random initializations on the privacy bound when training overparameterii meterized neural networks and shows that for some initializations (LeCun LeCun et al. (2012), Xavier Glorot & BBnggoo (2010)) the privacy bound improves with increasing depth while for other initializations (Hc He et al. (2015). NTK Allen-Zhu & Li (023a it degrades with increasing depth.
et al. (002a)) :
Generalization performance in Federated Learning. Several existing works have studied the general- ization performance of FL in different settings Cheng et al. (2021) Gholami & Seferoglu (024h; Huang et all (2023): : Yuan et al. (2021). Some of the initial works either provide results independent of the algorithm being
used Mohri el al. (2019) Hu el al. (2022); Sun & Weil 2022) or only study convex losses Chen el al. 2021 )
Fallah et al. (2021). Barnes et al. ( Sefidgaran et al. (0022) derive iioomationnthhoreee bounds, but
these bounds require specific forms of loss functions and cannot capture effects of heterogeneity. [Huang ct al.
(2021) study the generalization of FodAvg on wide two-layer ReLU networks with homogeneous data. Collins
et al. (022) studies Fed under multi-task linear representation learning setting. Tn Sun et al. (2024).
20

---

Published in Transactions on Machine Learning Research (10/2025)
the authors have demonstrated the impact of data heterogencity on the generalization performance of some popular FL algorithms.
Theory Notation and Preliminaries
We follow a similar notation as Kou et al. (0023 in most of the analysis.
Table 1: Summary of notation
Symbol
Description
jE [-1,1}
Layer index
m d
Number of filters
Dimension of filter
r [m]]
Filter Index
K
Number of clients
& [K]
Client index
N
Number of datapoints at each client Datapoint index
E [N] n KN 1h,i E {1, -1}
Global dataset size Label of i-th datapoint at k-th client
Signal vector Variance of Gaussian noise
of
Ék,
Noisc vector for k-th client and i-th datapoint
"
Local learning rate Number of local steps
(12 2) log(1 1 exp( -2
Cross-entropy loss function ReLU function
(2)) max (0, 2)
(z) = (2 2 0) t
Derivative of ReL.U function Round index Iteration index
8
h
Heterogeneity parameter
SNR 1alll//a,v
Signal to Noisc Ratio
W
Parameterized weights of the k-th client
(j. r)-tt filter weight of the k-th client Local signal co-efficient for k-th client Local noise cocfficient for k-th client and i-th datapoint
Positive local noise coefficient. for k-th client and i-th datapoint
Binkk
Negative local noise coefficient for k-th client and i-th datapoint
Shorthand for -1/ exp(yx.:((!!." .x.))) which is the derivative of cross-entropy loss for i-th datapoint at k-th client Parameterized weight vector of the global model
w
j, r-th filter weight of the global model
Global signal co-efficient Global noise coelficient for (k, i)-th datapoint
pe
Jnka
Arrka
Positive global noise coefficient for (k, i)-th datapoint
pe Jinkt
Negative global noise coefficient for (k:, i)-th client datapoint
21

---

Published in Transactions on Machine Learning Research (10/2025)
B.1 Local Model Update
Using local GD updates in equation 5 to minimize the local loss function in equation 3 the local model update for the r))) filter at client round can be written as,
(17) Jonk
(1)
"
geltan) o (( Ekai)) Ekc kg
Nm
KE[N]
"
Σ Σ
(e(t.)) k2
o' ((w (1,3)
(inde sk.Ht)) ju
Nm
==0 iE(N!
w/e + j7j.n, 11
++ E
Plinks 11Ex.ill2 Ex,
(11)
iE[N]
where, use w j.r jir Further, we define
(2.0) A (0)
Birn
(1rr
n
((t.s) hi -a (1.0) M.ske) (o.s))) 11
(12)
Nm
iE[N]
((,T)
17
(r(e.))
d'((ww//ooa) ((w! Ex.ll) 11Ex,e122 jun.
(13)
Pinkk
Nm
which respectively, denote the local signal (1 and local noise (1 components of w (f,r) j.nkk" We also
define Pj.r,k, Pjr,k, (phre) 2 0) and Li..k. Pirrk (t.T) < 0), where 1 1--) denotes the indicator
((,TT) (t,r)
(1,7) (37)
function, and which can allernatively be written as
(17) Pirck
plt.s) a jir,i (aa))) 11Ex.il22 (yx. = 5);
(1.a)
(14)
Nm
17
(ttss) á' ((w/tra) En.i)) ||Ex, 122 1 1yMes -)).
(15)
Ljrkk
Nm
B.2 Proof of Proposition
The global model update round & 1 can be written as
K
(1+1)
(f,r) j.rkk
jr
K
W
h=1
j.r
(1)
(1.r) jink
||l1
Σ
Panki [=(N)
Σ
K
1 ((r)
2
E
(16)
Mimicking the signal-noise decomposition equation 11 model as follows.
we can define a similar decomposition for the global
(0) f.r
(1)
|4
p +
p(t) pinks ||Ex, EK.
(17)
J.r
VE(N]
22

---

Published in Transactions on Machine Learning Research (10/2025)
B.3 Co-efficient Update Equations
Comparing with equalion 16 using n = KNN.
we have the following recursive update for the global signal and noise coeflicients
K
[t+1)
(1)
(1,T)
3r
K
jrrk
1=
(1)
gett,e) k2
. -'' ((w
Jonko Wk, |111122
(18)
1r
k= [E[N]
==0
p(+1) p(()
+ ((.))
jink
J.rkk,
K Jrnkk
=p()) Jr,kk
=
Σ
plt.s) ki
o (IwSrrd Jirkk Ea,i)) ||En, jun,
(19)
8=0
Analogously, we can also define the positive and negative global noise coefficients,
p((+1) P(()
jirks
"
Σ ple,s)
(t.s) o ((w (ink) ||ER..122 (vu.s 5)
(20)
num
and,
Pink = kin
1
Σ
plt.s) hi
. ((w Sa,))) ||En, 112 (yu. = -j).
(21)
Jin,kr Lemma (Measuring local and global signal coefficient)
mm
From equation 11 it follows that
(wilia) (() jink Jart (k, 580, (t.a)
(22)
and from equation 17 il follows that
(0) Jor se) = =II(()) J.
(23)
Since [ are nen-tegative and non-decreasing in t, the global weights [w]f]r become increasing aligned with the actual signal ykii corresponding to the filters j = yk.i- Similarly, as fn are non-negative and non-decreasing in s for fixed I,, the local weights (t,e) Br become increasing aligned with the signal yk
corresponding to the filters j = yk.
C Training Error Convergence of FedAvg with Random Initialization
For the sake completeness, we state the conditions used in our analysis (Condition in full detail.
Assumptions. Let be desired training error threshold and & E (0, 1) be some failure probability. Let T" poly(e m,,, d) be the maximum admissible rounds.
Suppose there exists sufficiently large constant C, such that the following hold. Assumption 1. Dimension d is sufficiently large, i.e.,
2 C max n [1llll (T)(())) n² log(nm//0)log((">
o3
23

---

Published in Transactions on Machine Learning Research (10/2025)
Assumption 2. Training sample size and neural network width m satisfy
C'log(n/6), C'log(m/6).
Assumption 3. The norm of the signal satisfies,
| 2 Co log(1/6))).
Assumption 4. Standard deviation of Gaussian initialization is sufficiently small, i.e.,
1 min
vn Opdt'
S
c
[lg(mm/5) Ne
Assumption 5. Learning rate is sufficiently small, i.e.,
rem Vinyloun log(m/6)
c min
1
o3
The assumptions are primarily used to ensure that the model is sufficienlly overparameterized, i.e., training loss can be made arbitrarily small, and that we do not begin optimization from a point where the gradient is alrcady zero or unbounded. We provide a more intuitive reasoning behind each of the assumptions below:
Bounded number of communication rounds: This is needed to ensure that the magnitude of filter weights remains bounded throughout training since they grow logarithmically with the number of updates (see Theorem 3). We note that this quite mild condition since the max rounds can have polynomial dependence on 1/e where is our desired training error.
Dimension d is sufficiently large: This is needed to ensure that the model is sufficiently overparame- terized and the training loss can be made arbitrarily small. Recall that our input x consists of a signal component μ E Rd that is common across all datapoints and noise component & E IR that is independently drawn from NO, og I)- Having a sufficiently large d ensures that the correlation between any two noise vectors, i.e. (E, (') is not too large. Otherwise if the correlation between Iwo noise vectors is large and negative, then minimizing the loss on one data point could end up increasing the loss another training point which complicates convergence and prevents loss from becoming arbitrarily small.
Training set size and network width is sufficiently large: The condition ensures that a sufficient mumber of filters get activated at initialization with high probability (see Lemma 6 and Lemma. 7) and prevents cases where the initial gradient is zero. The condition training set size also ensures that there are sufficient. number of datapoints with negative and positive labels (see Lemma 8).
Standard deviation of Gaussian random initialization sufficiently small: This condition is needed to ensure that the magnitude of the initial correlation between the filter weights and the signal and noise components, ie kw!!) / p)], Kw (W I) is not too large. This simplifies the analysis and prevents cases where none of the filters get activated initialization (see Lemma 21). It also ensures that after some number of rounds all filters get aligned with the signal (see Lemma 30).
Norm of signal is larger than noise variance: This condition is needed to ensure that all misaligned filters at initialization event become aligued with the signal after some rounds (see Lemnna 30). This allows us derive meaningful bound on test performance that is not dominated by noise memorization.
Learning rate is sufficiently small: This a standard condition to ensure that gradient descent does not diverge. The conditions are derived from cnsuring that the signal and noise cocfficient remain bounded the first stage of training and that the loss decreases monotonically in every round in the second stage of training.
24

---

Published in Transactions on Machine Learning Research (10/2025)
For case of reference, restate Theorem below.
Theorem (Training Loss Convergence). Let Ti = initialization, for all Ti s T T have,
myz noger
With probability 1 8 over the random
|W(n] - W n(T T +1 +
T- Σ L(W
Therefore we can find on iterate with training error smaller than 2c within T= Ti + W(T)) w 1/12 12/(m)) =
1722 O
+ + mn log(t/e)
rounds.
norde
C.1 Preliminary Lemmas
Lemma 5. (Lemmna B.4 in Cao et al. (2022)) Suppose that 8 > 0 and d = 92 (log(4n/0))
Then with
probability at least 8,
o,d/2 S | S 30dd/2.
KE Ev,e) < 20 dlog(6n2/0))
for all k [K], i E [N], and (k, + i').
Lemma 6. (Lemma B.5 in Kou et al. (2023). ). Suppose that d = Q (log(mn/00), Ω (log(1/))). Then
with probability at least 1 - 6,
odd/2
(0)
S 30/d/2,
(w/0r (0) fom 14) S 22log oo || (w. Ensa)
I 2/ ooo, ooo,
for all E [m], jE [+1,, [K] and E [N].
Lemma 7. (Lemma B.6 in Kou et al. (023)) Let (0) 2 50 50lgg(2n/0). Then with probabilily al least 1-8, 8,
{o {re [m]
(w.... Ex,i) 2 0). Suppose 5 > 0 and
((0) k2
2 0.4m, vi E [n].
Lemma 8. (Lemma B.7 in Kou et al. (2023)) Let Sj.
((0)
{ [K] [N] Vk.s = (w Ex,) 2 20.
Suppose 8 > and n 2 32 log 4m/0). Then with probability at least -- 8,
((0) j.r
2 n/8, Vie [n].
Lemma 9. Let Dj- = {k E at least 8,
yk.s = j}- Suppose ô > 0 and n 2 8log(4/6). Then with probability
n Vj {+1) |D,
4
Proof. We have |Dyl = Ex., (yu.; = )) and therefore E|D, = Σx. P(y.. = j) = n/2. Applying Hoefdding's inequality have with probability 28,
D
log 4/8)
2n
Now if n 2 8log(4/8), by applying union bound, we have with probability at least. --5, 5,
|D 2 n
4 V (+1).
25

---

Published in Transactions on Machine Learning Research (10/2025)
C.2 Bounding the Scale of Signal and Noise Memorization Coefficients
Our first goal is to show that the coefficients of the global model, L..,, 1 p( j.r.k.i and p(t)
kr,k, are bounded
O Tl). To do so, look at virtual iteration index given by v = 0, 1, 2, ..... T T 1. For any v, we can define the filter weights at virtual iteration v in terms of the filter weights we have seen so far. In particular,
(v) A (L#J.0 mod r) j.rkk W
j.r.k
We also define the following virtual sequence of local coefficients which will be used in proof. Let (0) k = 0, ( J.r..r... = 0.P O. We have the following update equation for (v) p(l) j.nk jorkk, and p((t) kjnnk, for
u 2 L.
G(u-!))
E
e(v-1)
(w (1-1) jorkk se) |1 if (mod r) # 0,
Jr,k
Nm
((+)
iE(N
J.rck
T-1
(24)
((U-T)=
IL
ΣΣΣ
((= ((
((W ye, ||
else,
Jrnk
jinc
==0
NE[N]
where we slightly abuse notation, using g(v) k2 to denote er (())) mod
k,
F ((v-1)
y p(v-l)
(( (0-1). Ek.)) 112. 1j yk..), if (mod 7) # 0,
Nm ki
((()
P (uu-))
T
r-1
l (vv-T+s)
((W!rr
(v-((())) Ex.,)) 1l =y.a))
else.
(25)
Jrr
mmm
k,j
jrnk
==0
p((-1)
j.rkk Ek,i)) 11E,, 112 [ = -y..), i (mod 7) # 0, Jinki + J ((v-)+s) o ((wfo-rt)) En.i)) ||E 1jj yy..)
Junk + n pi(u-1) o' ((W ((w(orh))
Nm ki
((0)
T-1
(26)
kinkk
else.
ki
==0
Note that have the relation (/TT = (t) p((tr)
(()) j.t.ki k.r, jjn,k,, jon
pl(tr)
pl()
for all t = 0, 1.2...... T ---1 1. Intuitively, if we can bound the virtual sequence of coefficients, we can also bound the actual coefficients of the global model at every round.
C.2.1 Decomposition of Virtual Local Filter Weights
The purpose introducing the virtual sequence of coefficients is write the local filter weight at cach client as the following decomposition.
(0) + j.G j.nkk 11aa1122 2 #
Σ
[ j.rkk.
ppTlw/T])) + Ejnkk || Evr
j.nkk
Σ (p())
jinkkt + +l)))
(.nn.) ||ER,ill" En.i.
(27)
(E[M]
Note that (r(0/T]) denotes the last iteration at which communication happened. v (mod 7) = 0, then (()
jink is the same for all l: E [K].
26

---

Published in Transactions on Machine Learning Research (10/2025)
C.2.2 Theorem Scale of Coefficients
We will now state the theorem that bounds our virtual sequence of coelficients and give the proof below. We first define some quantities that will used throughout the proof.
{|
(ww j.r: u2) (w for ER,A) } ;9= n |
a == log r): B = max
kj
oRd
Theorem 3. Under assumptions, for all 0,1.2, 2............ T T- 1, we have that,
G Jrr
0, jinkk,, 0. j.rikes O,
(0)
0 S pp()
jinks a,
(28)
-3 8V log(6n2/6)
na 2 aa,
(29)
0 I Glv) jrn S C'Na,
(30)
for all E [m],j [#1}, l e [K] [N], where C' is some positive constant.
We will use induction prove this theorem. The statement clearly true at Now assuming the statement holds al of we will show that it holds o' 1. We first state and prove some intermediate lemmas that we will use in our proof.
C.2.3 Intermediate Steps to Prove the Induction in Theorem 3
Lemma 10.
log(6n2 /8)
1
max B, 4
1 nac
d
12"
Proof. From Lemma 6 we have 8 40 max v log( (12mn/8) op sd Vlog (log(12m/6) [|alll]. Now from
Assumptions D and by choosing C large enough, the inequality is satisfied.
Lemma 11. Suppose, equation 28 equation 29 and equation 30 holds for all iterations 0 S S v. Then for all r E [m], j E (+1), [K] N] we have,
(w()) w (0) ( jG((')
(31)
j.nk
(w)) w(O)) i Ev,e) -p(()
Tog(Gnz (5)
S 4V
na.j yk,
(32)
jirsk
(w)
w/ jor EM.t)
(0) -B!")
S <AV
log 10)
na.j j # yk.i.
(33)
j.rck
Proof of equation l [K] [N].
31 It follows directly from equation 27 by using our assumption that (1, Ev.)) 0 for all
Proof of equation
32
Note that
27

---

Published in Transactions on Machine Learning Research (10/2025)
for yk = j we have p((") jir,hi O. Now using equation 27 for yk,i we have,
w Wirr En.i) PanAd
(0)
Jink
Σ E
E
(Ex (e. Se
1|5 112
K.K
"E(NN
E
E
(I
Sinder + (pp [essoker))
A
log(6n
K 4 ME[N]
MEIN]
(6)
log log(6n2 (8)
na.
where (a) follows from triangle inequality and Lemma 5 (b) follows from the induction hypothesis.
Proof of equalion 33 Note that for
p(()) Jirkk O. Using equation 27 for j + yk. we have,
+ Ves have
(wolo') jink
(0) j.r Ek.)) - P S
Σ
(P (
(er TE 112
Σ (PP)))
+poo')
(En
+
|E 11
ME(N),
Σ
( + 1
Σ
(D
+ /es
log(Gn 18) 4.
W.N & FEIN
N'IN
(6)
1on( (8)
1/
no.,
d
where (a) follows from triangle incquality and Lemma
(b) follows from the induction hypothesis.
This concludes the proof of Lemma 10
Lemma 12. Suppose equation and i [N],
28 equation
29
and equation 30 hold at ileration v. Then for all ls [K]
1. For j + yh.s, Fj( (W(('))
J.k xx.) S 0.5.
2. For j = Yk.t F2 J,k x 2 1 E (1')
j.rkk, - 0.25.
3. yk.if (WW (v')
m Em pl()
x
,,, 0.75.
Proof of 1 First note that for j # yk. from Lemma we have,
(0) Jinke 11) S (w j.rt 14).
(34)
since G (v') by the induction hypothesis. Also from Lemma 11 for j + yk.t we have,
|log(6n2 (6)
jor.or Ex,e) S (0) Ex.a) + P.it + 4v
d na.
(a) (w/ j.r: Ek,s) +4V
log(6n2/0)
d na.
(35)
28

---

Published in Transactions on Machine Learning Research (10/2025)
where (a) follows from p p(()) we have,
S
(induction hypothesis). Now using the definition of (W, x) for j + yk.
Jr,ki
P (W()) Jik
1
(wile') Jin,kr mk.s + a (wwiind Sk.1i]
xk.)
sindk
m
[
(0)
/w ( Emxi .
log(6n2
3 max rE(m)
j,r
d
na
(6)
log(6n2/0) 1a
3 max
B, 3.41
d
0.5.
(36)
Here (a) follows from equation 34 and equation
35
(b) follows from the definition of B3 (c) follows from
Lemma 10
Proof of 2 For j = Uk. we have,
Fj( (W(W
xk,s) = Σ
(ww jirkkr Me.st)) + a (woffrr) dirckr SR.e)) )
a
m
r=1
IVB
1
172
(w) Wirde (k + (ort Enas)
m
IV
(0) (w/0)) 3k,iHH (w Ex.e) + P IF('))
log(6n2/0 d na
Jinki 4
((')
J.r,k,s 233 - -4V log(6n2²/1 /8)
71
711
((e)
Jrrkk 0.25.
(37)
Here (a) follows from (2) 2 (b) follows from Lemma 11 and that (x)) Gjrck 0: (c) follows from the definition
of B; (d) follows from Lemma 10
Proof of 3 Combining the results in equation
36 and equation 37 we have,
y /(W xx,)) = Fyn.: (W
F-n.: (Wlo'))
ke Xx.s)
Be Xk.i)
(a)
Fge, (W(o')
Whisk xk.)) 0.5
IV@
Mea.rr 0.75.
where (a) follows from equation 36
(b) follows from equation 37
This concludes the proof of Lemma 12
Lemma 13. Suppose equation
28 equation
29 and equation
30
hold at iteration v. Then for all j E (+1},
k E [K] and E [N], pole))
< exp
P Ke x:) + 0.5
29

---

Published in Transactions on Machine Learning Research (10/2025)
Proof. We have,
F+i (W()) +1,6 Xx,i) - F- ( +1, XX,1)
1 exp
Yhii
(2)
Uk,t F41(W +1. xk.i) F- F-I(W ((v)
exp
+1,6 XL.
F Fyu (Wl)) Ve.ck Xb,i) + F-ve.. (W.. (W(o)
exp
(0)
Fe (W Un xx,i) + 0.5).
-
exp
where (a) uses 1/(1 + exp(2)) v exp(-2) (b) uses part 1 of Lemma 12
Lemma Let g(z) = (1))) - 1/(11 + exp(=)). Purther suppose 22 21 S e where e 0. Then,
g(21)
S exp(c).
(38)
((22
Proof. We have,
(21)
1 + exp(22)
< (1, (1xx(2e 21} a exp(c),
1 exp(21)
where (a) follows from 2 0.
Lemma 15. Suppose equation and i [N],
28
equation
20
and equation 30
hold at iteration v'. Then for all k E [K]
(w(())
Ih Ex,a) 2 -0.25,
(39)
Ba de Ev.s) S (wwlo' Bhinkk Ex.i))
(w()') Be Ex.s) + 0.25.
(40)
Proof of equation
30 From Lemma 11 we have,
( (v')
Wrr EniE) 2 (w (0)
srrks Ek.i) + p(o")
log(6n2/0) /8) mn.ikni 1V
d nax
log 1/)
8 AV
d 110
-..55.
Here (a) follows from the definition of and ((v)
Ve..niki 2 0 for all o' 2 0; (b) follows from Lemma 10
Proof of equation 40 The first inequality of equation 40 follows naturally since ((2) for all z R. For the second incquality have,
(w/rr)
Ve. Ea,i) S (wle') Da. Ex.a) + 0.25, if (w)) Mesrk Ek,s) 2 0 (w))
(( Eni))
o
(a)
(o')
S Whnorn Ex,i) + 0.25,
Wasid Ex,a) < 0,
where (a) follows from (w
wr Ex,a) -0.25. This completes the proof.
This concludes the proof of Lemma 15
30

---

Published in Transactions on Machine Learning Research (10/2025)
Lemma 16. Suppose equation 28 and i, [M],
equation
29
and equation 30
hold at iteration v'. Then for all k. k' [K]
Iplo')
Buiisr,ki pitt>
V [WW Xx,i) ,e J(W
S 1.75.
Proof. We can write,
yk ( xx.i)
(v')
- ,k - ffff/WW ((W()) (v')
Fun.: (W.((W(e)) (Wre')
Mo. Xx.i) F- F-ye. (Wie')
De Xk.i)
Pev (Wil'")
Bese MN xx, + F Mer Wilo')
=F-ww, (W(r')
-Ve Xx xx.)) F- Fnn "u..k.e Xk.s)
Wk..k xk. Fo (W = F- (W()))
+ Fy (W(6')
Nal of XH.
i M XW xW..) F- wa. (W(v')
Me Xk.
(1)) win wn.s
m
1
m
0
(w')
Ve.ii Skir))
(wwle)
murp Exr
m
13
Next bound 4, In and Is as follows.
|4 F- (Wlo)
(a) te xg.s) S 1, Me xw.) + F-... -..(W (W(r)
where (a) follows from part 1 of Lemma
12 For [Z] we have the following bound,
1
Σ9 (woll) we (h..n--
1
Σo (mmf)) Ber We s
[22 < max
m
(e)
{ (w wessrr 10) (mm We (ww. ,(10) C(V') Mesk (el
max
rE(m]
2 max rE[m]
f3, C²
0.25.
(b) follows from the definition of and the induction hypothesis, (c) follows from Lemma 10 and Assumption
Here (a) follows Lemma
31

---

Published in Transactions on Machine Learning Research (10/2025)
Next derive upper bound on I3 follows.
(wler)
Deaein Skee) - (Im Werresn Sk.,)))
13
IVG
(wwlo)
(w Vel an Ew,ee 0.25
IA3
Iple')
p(o) Metemman
log(6n2/0)
+ + VV
0.25
m
d
IA
(0')
F(v) "e
+ +.5.
Here (a) follows from Lemmna
15
(b) follows from Lemia 11 follows from Lemma 10
Similarly, we can get a lower bound for 13 as follows,
13
m
(W (0))
Dn Ek.i)) a (mm())
aa Ex.+))
m
IV
(wol))
Se. Enii) (wol))
(Weee (ii 0.25
3
Bhiist,k,i MUrerrt - 22 - 8V
d na 0.25
IV
(m))
Ver Mer
0.5.
Here (a) follows from Lemma 15 (b) follows from Lemma
11
(c) follows from Lemma 10
Combining the above results, we have
i (W())
Xx.. Wk a f f(W!!) xxM.)
S 11 142 13
1
p(o')
+ 1.75,
m
and,
Wk.i. (WW ls Xx.i) - 3 f(W!
x
2 -111 [/2 + la
m
p (v')
(o))
Ve..rrkk. Ber 1.75.
m71
This implies,
Mk.i (W/)) xxx) M xx.)
f(W!))
m
F (p')
F(() Uererr,k. B.
I 1.75.
m71
We will now state and prove a version Lemma C.7 that appears in Cao et al. (002). Note that Cao et al, (2022) only considers the heterogeneity arising due to different datapoints for the same model. Interestingly, we show that the lemma can be extended the case with different local models and different datapoints long the local models start from the same initialization.
32

---

Published in Transactions on Machine Learning Research (10/2025)
Lemma 17. Suppose equation holds for all 0 S v v..
28
equation
29
and equation
30
hold for all 0 v. Then the following
p(0)
m 2r=1
P
S n. for all k, k' E iK), E (N].
2. yk ((W()
(v) xk.i)
3 (W xwxx)) S Cr for all k k' E [K] and i [N].
pla)
3. p
I C2 exp(Ci) for all k, k E [K] and i, [N].
4. Sk, S $ where s r E [m] (W o En,i) 2 of. and hence | 2 0.4m for all e [K], i e [M].
5. // S ((t) where 5 = { [K],EE [M] ykut = j. (w Ex,a) 20}. and hence | = g
Here we take K = 5 and C₁ = 6.75.
1 We will use proof by induction. For == 0, it is simple to verify that 1 holds since F = for all j E (+1),r r e [m]. k [K] E [N] by definition. Now suppose [ holdd for all 0 I v S T < v. Then we will show that also holds at v = 0 + 1. We have the following cases.
Proof of
Case 1: (0 + 1) (mod r) + 0
In this case, from equation
25
(++1)
P (0)
pe10) Nm
Ko ((w (8) brrkk Ex.i)) 1
Thus,
m
(6+1)
p (0+11)
p())
po(0) Ver (rrkeee
r=1
+ 1 Nm
(())))
(-e(g)
k.Y) 1|Ex.el12 - ((8) - (-e |1 113]
(41)
where ((() s(0)
are defined in
We bound equation 41 in two cases, depending the value of Em m
F (0)
((()
m If L
P b
((0)
S 0.9k. From equation 41 we have,
p (0+1)
1(++1)
S 0.9k +
(- (-e()) 1|ER.
Nm
m77
(g)
W
0.9k
+
||E
Nm
I&G
K.
(a) follows from $ < m, - --(-) < 1;(b) follows from Lemma 5 and Assumption 5
33

---

Published in Transactions on Machine Learning Research (10/2025)
ii If T E 1=
F
Fl0)
> 0.9k. From Lemma 16
we know that,
yk (W h xx,))
WK ((W l x 2
o 1.75
m
0.9k 0.35x
0.55k.
(42)
(0)
F (0)
where (a) follows from y. Also note that since
m Dr=1
2 m
Ex=1
Me P 0.9k
0.9k 4.5, we have from Lemma 12 that
yk Xk.)) 2 3.75.
(43)
Now from the definition of f(-) we have,
1 + +xp(yye v/ (( xK,z))
+ exp(yk, (W xk,i))
(0) 1 + expLuR, / XR,i) - 0.55x)
exp( xx,i))
(6)
1/7.5.
(44)
Here (a) follows from equation
42
(b) follows from equation 43
Thus,
(B)
||
(()
(a)
11E {
(-e
(1)))
k.i
(
(c)
S 2.5
S
2.5
S
|| 11 (- ((0)
k(
Here (a) follows from
k S m,
( k
0.4m using our induction hypothesis; (b) follows from
Lemma (c) follows from equation
44
This implies s(0) ||Ex
((0)
(-(4))))
Now from equation 41 have,
1 m
Σ
(6+1) Meaaor
((+1)) ww.
m I
T0()) Da Srr S K.
m
where the last inequality follows from our induction hypothesis.
Case 2: (T + 1) (mod r) =0
In this case, using equation 25
write our update equation follows:
(o+1)
F (6+1)
m
P (
Uk...ikke - pp (8+1--7)
1-))
Bier on
g(0+1-7+0) ki
(- kid
((++1-++s)
|
[s(++1-++0)
(-10
(-d(811-r+))
m
==
16-1-1
Po (0+1-7)
J
(45)
34

---

Published in Transactions on Machine Learning Research (10/2025)
From induction hypothesis
we know that
(B)
p Mar
S k.
(46)
Now unrolling the LHS expression in equation 46
using equation 25
we see that this implics
p plot1-r) Bh.i
p(0-1-1)
S
(47)
m
N
Case 2a): 1 2
In this case it directly follows equation 45 and equation 47 that N <
Em=
P
((+11
P 1
S since
Case 2b): If 1 < 0.
In this case from equation 45 we have,
Σ
11-1) ve..kkkke
Ip!
(0+1)
1
p
(8+1-T)
]
S k.
where the last inequality follows from our induction hypothesis.
Proof of 2 For any 0 S v S o' we have,
yk. WW.
(W
xk.) yk 1 f(W! xp.)
1
P
(1)
P (0)
1.75
m
< K 1.75 = C1.
Here (a) follows from Lemma 16 (b) follows from
Proof of 3 For any 0 S v S o' we have,
(v) (a)
(6) exp exp ( xx.i) ye NfW. 1 xxx))) exp(C1).
max 1,
(n))
k.i
Here (a) follows from Lemma 14 (b) follows from
7
Proof of To prove 4 we will use the result in 3 and show that (w( for all 1 < 1 v'. We use proof by induction. Assuming (w(o) BRaAK Ex.s) > 0 implies (w (1)
pea.rk Ev.s) > 0
Bn Ev.s) > 0 for all 0 s Sese V, will
show that (0+1)
Dereirke Ex,e) > We have the following cases.
Case 1: (8 + 1) (mod r) # O.
Using the fact that (w(") (8)
W...r,k Ex,i) > 0 have,
35

---

Published in Transactions on Machine Learning Research (10/2025)
(( Ck.I)
(w (0)
ye Sk.a)
1
(- Ekill1
Nm
Nm Σ
(- k.tt) )o' ( Vearr, SKF) (Ex, Expr)
PE
UA... Ek,I) + no
(-(
11 20 dllog(nn2/0
Σ
(-e
2Nm
kit
Nm
PE(N],'A
n0
Ba Ek.)) +
2Nm
(-(
kit
m
11 20 v dllog(4nn//0))
4n² (
Waark Ex,s)
Here (a) follows from Lemma 5 (b) follows from large d.
(c) follows from Assumption 1 by choosing sufficiently
Case 2: (0 1) (mod r) = O.
From induction hypothesis we know that (w (w (0tl-rt)))
Werre Ek,i) for all 0 S 1. Then,
1o Ek.i)) = (w (0+)) Ev,i) +
(((+1- ||
==0
(-enet
(-e10t1-rta)) you
(w
(0+1))+))
EA.r)) (Ek. Ek.e)
Ki
a=0 J'E'N
EE Σ TumI
(ww
(w( (1+1-7+0) .Ex.s)) (EN Exi
(48)
- K MAAAEE
13
Using Lemma 5 we can lower bound I as follows:
no Σ
I 2
2nm
ki
==0
5 where the inequality follows from Lemma
For [Z2] we have,
Lemma 5 as follows:
(a) <
n220z v d log 4n 2/2/8)
[T
Σ Σ
Rid
nm
V'E
(b) n(N 1)02 / d log 4n /0)
(0+1-1++))
nm
36

---

Published in Transactions on Machine Learning Research (10/2025)
Here (a) follows from Lemma
5 (b) follows from 3
Similarly can bound [Is] follows,
1220p log 4n /6)
[aa
Σ
num
K
P'ENN
(b) 77n - 22 log 4n²
nm
0
Here (a) follows from Lemma 5 (b) follows from
Substituting the bounds for 11,1221, (4s in equation
48 we
have,
(w
Weanni Ex,i)
2 (w 1-) Ek.i) + +1 1 -
[131
2 (w (f-1-7) Ek,s) + no Σ
((041-7)+8)
2mm
k,t
8=0
7/2 20 / /dlog((422//
k,i
==0
(6+1-7) Mh.in Ex,e)
0.
Here (a) follows from Assumption I by choosing a sufficiently large d. Thus we have shown that (wlo)
te Ex.i)
for all 0 S y' and such that (w(l) Baank Ex.i) 2 0. This implies s S
for all 0 < Susa v. Furthermore we know that s(0)]
kkt 2 0.4m for all & E [K] i E [N] from Lemma and thus
(( k,
2 0.4m for all k E [K] E N,,0suset
Proof of 5 Note that as part of the proof of 4 we have already shown that (w(e)
jock Ex.s) 2 0 for all 0 S v S of jink Ek.e) 2 0. This implics 5 S 5 for all 0 S S v.. Furthermore
and k, é such that yk,s j and (w(0)
know that 5/0) 3rr 2 n/8 for all j E (+13, E [m]] from Lemma and thus ((()) 2 n/8 for all j e (+1), [m].
This concludes the proof of Lemma
We are ready to prove Theorem 8
C.2.4 Proof of Theorem 3
We will again use a proof by induction to prove this theorem.
Proof of equation 29 For yk, we know from equation 26 that [[ +1)
0 and hence we look at the case
where j + Wk.i.
Case 1: (n' + 1) (mod+ 7) + O.
a) If (((()) kin -0.53 4v log(Gn2/0 /8)
na, then from equation 33 in Lemma 11 know that,
(W (0')
log 6n /8) Jinbbr Ex.e) S (W. En.o) + Iplot + +4
na
d
log(6n2
/8)
0.5,3 + plo') + +
d na.
(b)
< 0.
37

---

Published in Transactions on Machine Learning Research (10/2025)
Here (a) follows from definition of 3 in Theorem 3 (b) follows from -j.nk,
-0.5 -41 /oc(66n2 /01 na. Sinder Sk,E)) = O, which implies plo't1)
Now using the lact that (werr Jinkks Sk,z) < O we have o' (wofind
=P(1)) -j.r, 2
-j,r,
-B 2--g 8V log(6n2/ /8) na using the induction hypothesis.
b). If (PP) j,r,k,s 2 -0.58 - AV
log(6n2/0 (8)
na; then from equation 26 we have,
d
=p(1)
((
(n")
(k,I)) [|5R,1112 11122166 1(j = -Vk..)
Jir,ki
Nm
j.rkkt
1og(Gn2 (n2//// /0)
3no
-0.53 4V
na 2.
d
log(6n' 18)
- 3 8V
(49)
Here (a) follows from [f'(-) < 1 and Lemma 5
(b) follows from
3n0
I 4
log(6n"/6)
na using Assumption 5
2Nm
d
Case 2: (v' + 1) (mod 7) = 0.
In this case, from equation 26 we have,
p((++)
p(
( +1-r+s) o ((Wfo'=) Ev.i)) 1Ex.:11:(( 1(j -vk.))
Jonks
Jirkk
ki
jink
Tum
1=12
p(1++1-7)
12.
(50)
Sjink,
Now suppose instead of doing the update in equation 50 we performed the following hypothetical update:
p(o'+1)
=p(v')
+ n pe(r) o' ((W)
(v')
Ex.i)) E 12 1j = -Vk..)
jonkk = J.r,k,
Nm k.
jrrk
(0) +(14+1-))
ploo 1-+s) ((w(o')) En.il) 11Eu.il2 1 = -kh,i)
+ Nm ==0
p(v++1-)) Jir,k,
7 Iz.
+
Nm
Here (a) uses equation 26 for v = ju' + 1 +1--:001 From the argument in Case 1 we know that p
3 - -8 log(6n2 /6)
no. Observe that plo'+1)
jrr 2 p(o'=1) since I2 S and s n and thus p(o'+1)
IV
jinkit
3 8 log(Bn? (8)
na..
Proof of equation 28 We know from equation 25 that for j # yk.s. F = 0 for all 0 d' S TT* - and hence we focus on the case where = yk,i-
Case 1: (V' + 1) (mod 7) + O.
Let Uy be the last iteration such that Vy (mod 7) = and pl (hrnn)) jrrhi S 0.5c and let a be the maximum
value in {0, 1, ..... 1} such that T ts)
0.5a.
Define Wjnkki U
We that for all
jakk
38

---

Published in Transactions on Machine Learning Research (10/2025)
U kg we have 10 jinkki 0.5a. Furthermore,
p (a)
pl jr,ki
finkk " "e ellgeine Nm Rit
o ((w (lsiinn) Ekii) 1|E,,:112 1 - wu.))
Jrnks -
Σ
1 edlo
kidd (Ww
Jon Sk,N)) 11ER.il22 112 (j (j= = yk.s)
(51)
Nm
Here (a) uses the fact that we are avoiding the scaling down by a factor of K which occurs at every v
(mod r) = 0 (see equation 25) for V/ynk, < U.SSA v'.
We know p(t)eka)) 0.5a. We can bound L1 and Le follows:
(a) L S
(b) 3no d (c) (d)
Nm E S
2Nm S 1 0.25a.
Here (a) uses (((1) & (.)) S (b) uses Lemma (c) uses Assumption
(d) uses T 2
Now note that for Ujrnki < v S of since p(0)
Jinis 2 0.5a we have,
(w())
(w (0)
Jrrkk Ek.I) + Jiniit - 41 log(6n2/0) /0)
p(o)
J..kt Sk,))
max
v Iv3
-0.53 + 0.5a - aV log [oog(Bn2// 15)
na
0.25a.
(52)
Here (a) follows from Lemma 11 (b) follows from the definition of 3 (sec Theorem 3 and ((v)
jarkk, 0.5a, (c)
follows from 3 < T2 I 0.la and 4v log(Gn na < 0.2a using Assumption
Substituting the bound above in Le we have,
L2
Σ
7 exp (-(w(o) (Wi...ee Ex,e) + +0.5) (WW Ex,a)) I|Ex,el22 112 1 = yk.s)
Nm
Uj.nkii SuS
2n
exp
(-(w(1)
jink Ex.T)) Ex,i)
|E 12
(53)
Nm
WV
Σ
2n
exp( -0.20)))
30dd
Nm
2n(v' Ujrnki - exp( log T 7)
30
Nm
(T)
exp(- - [ol-log11 T) 30dd
<
Nm
2
3
Nm
0.25a.
For (a) we use Lemma 13 for (b) we use exp(0.5) < 2 and (v)
from Lemma and equation 52 (d) follows from Assumption Ankk Ek.i) 2 from equation (c)
follows
39

---

Published in Transactions on Machine Learning Research (10/2025)
Thus substituting the bounds for L1 and Lg we have,
pll't1) j.nkk Sa,
which completes our proof.
Case 2: (v' + 1) (mod 7) = 0.
Suppose instead of doing the update in equation 25 performed the following hypothetical update
(0'+1) 1p00 5. = jonke Nm k,s
1 Lllo o' (w
(W(o')
jink Ex,i)) |IEx (jj = sn.s).
(54)
((0411) From the argument in Case 1 we know that P S
(1'+1)
Observe that jirck
P Pink and thus
p(o4+1)
Jr,ki S a.
Proof of equation 30 This part bounds G(((+1) To do so we show that the growth of (( jink
is upper
jink
bounded by the growth of p(o(+1)
Ma. 3,1 for any r E s that is,
(v'+1)
Jank
CC'.
We will again use a proof by induction. We first argue the base case of our induction. Since E $ S (
so,
p
(0)
ello o' ((w pe ER.1) )) ||Ex
(0)
Nm
=I(:rres
|E 11
(a) nond
7)
2Nm ((0)
Nm
where (a) follows from Lemma
5
On the other hand,
1111
(())
G jink
Σ
((w 1 H1) 11
Ar,k
Nm NE[N]
m
Therefore,
(() j.nk
2N |
C²G
o7d
Ve
k,1
if C" 2 2. Now assuming equation 55 holds at of we have the following cases for (v' 1).
jank S C²T.
p (0)
Case 1: (0' + +1) (mod 7) * 0. From equation 24 we have,
G (1/+1)
(v')
7 Nm
e((o')) ((w(((")))
jonker (h. |1re1
jr,k =
jrnk
EE[M]
(a)
(ee)
nC2
(-(((e)) k.
||14112
j.r.k
(55)
40

---

Published in Transactions on Machine Learning Research (10/2025)
where (a) follows from part (3) in Lemma 17 At the same time since (w (v) and for all 0 S s - 1, we have from equation 25
10... A Ek.1) 2 for any r* S
p(r't1)
F
(-
||Ex, l12
Me... k.
R,1
Nm k,1
(((n)) ozd
1
(-0
Wk.i k,1 +
Nm
k,
where (a) follows from Lemma
Thus,
((0++1) +1) (11
(o')
Jink
202N
Hull
(a)
2,,9, _ C²R.
(
max
ogd
3e.1. 1
Here (a) follows from the definition of 7: (b) follows from setting C' === =222.
Case 2: (v' + 1) (mod+ 7) = 0.
We have from equation 24
G( jink
G (v'+1-r)
7
(( (1-(++))
jink (3.,)()) 11
j.nk
s=0
k''
kE|N]
(0)
(1441-))
1
(v)+1-
|11112.
Jrnck
k,1
where (a) follows from part (3) Lemma 17 At the same time since (w v do Cx.1) 2 for any r* Sk
and for all S T 1, we have from equation 25
p p(('+1) pp
(0
1|
shrr R11
k2
(v' +1 +1-++)) ozd
p
plo' +1-
7
10,1,
k
nm
where (a) follows from Lemma 5 Thus,
(1)+1)
(1'+1-r)
jack
202n
|11112
(a)
(6) max (C' 2C27) 2 C²R.
S max
p
0zd
Here (a) follows from the definition of 93 (b) follows from setting C' 2C2-
Thus we have shown (v'(1) j.r.k
Coppl't MII C2aa where the last incquality follows from p(1'+1)
M A.
Now that have proved Theorem g that is, equation 28 equation 20 and equation 0 0Sus T - state a simple proposition that exlends the result in Lenma 17 for all S T'r-1
30 hold for all
Proposition 2. Under assumplions, for all 0 S Sest STSTT 11 we have
p(0)
-P(")
1.
m r=1
w. S for all k, k [K], E [N].
2. y (w) xk.i)
M x < C₁ for all k, k' E [K] and i, [N].
((e)
3. V
< C₂ = exp(C1) for all k, k E [K] and i e [N].
41

---

Published in Transactions on Machine Learning Research (10/2025)
4. S/. S (h where s = r E [m]] : (w Cx,s) 2 o). and hence |
2 0.4m for all l E [K], E [M].
5. 5000 C jir where 5/00 = { [K] E [M] : yest i. (wih Ex.) 2 0). and hence | 2 8
Here we take k = 5 and C₁ = 6.75.
C.3 First Stage of Training.
Define,
Cynm gozdT
Ti
(56)
P.T where C3 = (1) is some large constant. In this stago, our goal is to show that ye,,,r ,,, 2 2 for all r* such that r" E 5000 := r E [m] (w!")) (0)
yr... Ex,e) 2 =0) To do so, we first introduce the following lemmas.
Lemma 18. For all 0 s t Ti- and 0 S S T 1 we have
{1100 (t:s) 1
Czn || (1). ozd
max j.rkk binkk
Proof. We have,
t-1
((0) +
(1,0)
1
(l ki o ((w o serse)) 1ul1
Pjrnk
K=0
KE[N] ==0
"
pltis') á' ((wtted (ksH))) ||m11/ kii
Nm
s'=N
1/ ΣΣΣΣ ( 1M Nm [c[]
p(tas')
mmm
1E[N] ==0
(6) 1/(1 + 1)T 111112
m 7/Tir ||aa|
m Can 1
azd
O (1).
Here (a) follows from d(() E {0, 1}, (b) follows from [((-)] S 1, (c) follows from Assumption 1
Lemma 19. For all 0 < t < Ti - and 0 SSST - 1 we have,
(10) Jorki + s/tee))
max j.r.k,i
Piickk,e = 0
42

---

Published in Transactions on Machine Learning Research (10/2025)
Proof. We have from equation 14 and equation 20
s, + Pinkks
(nnde I1Ex,ill2 (yu. - j)
rum
v= ==
elltea "2d 0 ((w!rrd) Ek.))) [|Ex. 122 1 (UK. = j)
Nm
1
((tt)s)) b 11Ex.1112
7 Σ e(tae 115u.1112
Nm
y'm0 0==0=0
s'm0
(t + 1)T Elx.
Nm
2N
30nn
2N
O (1).
Here (a) follows from o'()) S 1, (b) follows from / S Ti and Lemma 5
]
Lemma 20. For any k E [K] and i E [N], we have F2 (w (1,2) xx..) = (1) for all j (+1), 0 0-1ST-1
and 0 S SSSTT- 1.
Proof. We have,
Fy Fy (t.s) j.k Xx,s)
(w (t.s) + (ww En.e)) ]
ve,s De,
Sor,kr
(w (t,s) jin, sk.
+
(w (t,o)
jonkk J.rkk Sk.i)
(6)
(w (0)
+ p(t)
(1,5)
(w (0) Ea.t) + sirck, + + + 4/
p(1))
log Log(6n2/0
J.r 14)
+
+
na
d
log(6n2
max rE[m]
(w j.r p4)
(t)
jor + 7g.nkk
j.r Ex.i)
Jorkke + (1.0) Piiris.s: 4V
d na
(c)
x 5 max rE[m]
0, 1 jir +
jin,ee j.r.k, + Plank VV
(1.0) p(()
log(6n2//) d na
4 ((1)
Here (a) follows from (22) < |z]. (6) follows from Lemma 11 from Lemma 10 Lemma 18 and Lemma 19
(c) follows from the definition of B, (d) follows
Lemma 21. For all t 2 Ti and 0 < S we have,
p(() Vhe Rss + p (t.0)
P k 2 (Ti)
"k..T", k. 2 2.
(57)
where r" s(0) = {r E [m]] : (w
B Ex,A) > 0}.
(0)
Proof. First note that from Lemma 20 we have for any &s [K], E [N], F+i(W(ta) +1.6: F- (W((ss) Xx,))
-1,k
(1) for all & (0, 10, Ti 1}. e (0, 1, ,...,, 1}. Thus there exists positive constant C' such that for
43

---

Published in Transactions on Machine Learning Research (10/2025)
all 0 0 Ti -1 and 0 T we have,
ple .0)
C.
(58)
Next we know from Proposition 2 part that,
(wlfis) Vnner do Ex.s) > for all - -
where r" E E = {rr [m]] (w me. Ek,s) > 0). This implies that for t t T1
plt)
Mhi Ri + ((e))
Pmncirt 2 2 p2(T))
T
T-1
"
|15x,112
k,
Tmm
==0
IV
2nm
IV
(59)
Here (a) follows from equation 20 (b) follows from equation 58 and Lemma 5
(b) follows from the definition
of Ti in equation 56 and setting C3- 4/C.
C.4 Second Stage of Training
In the first stage we have shown that for any l [K] and E [N], P(o) p RRs + (1.a)
(t,s)
P 2 2 for all t Tj
and 8 E T - 1]. Our goal in the second stage is to show that for every round in Ti S - the loss of the global model is decreasing. To do so, we will show that our objective satisfies the following property
L (W((e))) W(()) W 2 La (W((
2T
where W is defined as follows.
W j.r = w jrr + +log((27/c)
Σ
Σ KE[N]
[ = yk..))
Ex. | 11
(60)
Using this we can easily show that the loss of the global model is decreasing in every round leading to convergence. We now state and prove some intermediate lemmas.
Lemma 22. Under Condition we have
W - W w -0 (
| ozd
log(T/0)
44

---

Published in Transactions on Machine Learning Research (10/2025)
Proof.
( - w $ -
- -w
+ W - w
2
0
m 1/2 |111 mmax
max ((Ti)
+0
12
1/2 a ld-1/2 1/2
max jinkk,
jinki p(T.) J.nkke))
+
1
1/2 no d 'd 3/4 + W W--W" -w
0
(m 12 ||1|| op )
+ +0
(2 1/2 p
id d 1/2) 1-12) +
122 ( a
9
0
12 d 1/2 0
( 1/2 log(t (o)g(T/)00//00 1-12)
0
(m
log (og(T//)) (ogl///) /
Here (a) follows from the following argument:
w
Σ
(1000 -11e1112 -1
2
11ER.ill2 Ex.
Sonke
1r
j.
2m
((0) j.r
|1ee|12 ,r,
((t) j.r.k, ||Em.il12
ER,s
KE[N]
0
m
max(
mn
ankk 32
max (P(1)
+0
mn² max K,
|111 fir
[|Ex.el2 mn
0
m
max(T(()
1l5x,1122 j..k,t Brnkk,))
max (P(t)
mn?
| 2
o3/
j.
where the last equality follows from Lemma 5
Getting back to our proof, we see that (b) follows from
Lemma 18 Lemma 19 and definition of w in equation 60 (c) follows from Assumption !
Lemma 23. For any k E [K], i E [N] we have for all t E (Ti,Ti + 1,..., T 1}, sE (0, 7- 1}.
Uk. (V (G.)) xx.)), W > log(2r /c).
45

---

Published in Transactions on Machine Learning Research (10/2025)
Proof.
(VWW & xk.)), W") Uk. 3
1
Σ o' (ww (inkk sasee) (u, jw 1 o' (1,)))) Ekr)) (yR ER,,,
Σ
sir
of (n.)) 5log(2/e)1(( - yy (y) (Kr.i jEw
m
11EM
ΣΣ
(ww wa.))) blog(2/))1 um(ee) jEw,er)
||Ex 112
for W
Σd jr
o'
(nn (w m
Eo' 0' (Iws En.e)) Eui jw
Σ o' ((w)in) 5log(27/c)
j Dhri.r
Σ
E
o' ((wliik 5l0g(2r/e)
KEN Ekr,e)
||E 112
Jrr ((,,))))(()))
ΣΣ o' (Iwfledd finke We 5 log(27 (2//E)
K1,, Swtel)
Dor Wis
Σ Ji
d' (ww wkale) jun s00))
1 Σo'
o'
(wwfook (inner En.i)) (er Kee. (0))
3.
Jo
Now noting that dz) < and (1s, Ex.s) wk E [K], e [N] we have the following bounds for Iz, Iz, Ia, Is
using Lemma
Lemma and Lemma 10
12 (og(27/c)00 (nV/og(n2//))// (m/logn/// , -I=
44 0 (2.ovvolog (Voggmm66) oo [1ull)) . IIs = 0 (Vlog(mmn/6) :1==0(Vog((mn//)
GOo aoo, vd).
For I1 we know that, (w)))))) - Ex,s) 2 0 V [0: 11 E [0: 5 (Lemma 21 and r" such that 20
r E (0) = {r e [m] (w (0)
me En,i) 2 0) Thus,
4 4 m
[ 5 log((2/))) 2 log(2T/e).
where the last inequality follows from Lemma 7 Applying triangle inequality we have,
VV. (t,s) xx,,), W*)
1/2 [4/] /4|
[4s log((2/ee) (2T/e),
where the last inequality follows from Assumption and Assumption
Lemma 24. (Lemma D.4 in Kou et al, 20223 Under assumptions, for 0 S and 0 SSsST-b 1, the following result holds,
LL W (:)))) O max (ull2oozoo La LA (W((.)))
Lemma 25. For all & E [K], Ti STST" 0 ST- we have,
LV W((a))
W > La (W( ((( C 2T
46

---

Published in Transactions on Machine Learning Research (10/2025)
Proof.
(VL (W
w( W
ell ke (U V (6.0) x.,)), W
Wh (1.0) W
glt.s)) (W.f( x) - V... xx.,)) xk.)), W*)
N
EE|N
glle.) kii
yuaf(WW .x.,)) - 1og(27/E)
N
Σ
((W. xxx.)) (xx.)) - e/2T]
N
AE[N]
2r
Here (a) follows from the property that (VS(W, x), w) - x(W,x) for our two-layer CNN model; (b) follows
from equation 23 (note that gr((.e) S 0), (c) follows from (= z') 2 {(z) - E(z²) since f(-) is convex and
hi
log(1 + 2) $ 2.
Lemma 26. (Local Model Convergence) Under assumptions, for all t Ti we have,
w(( -..0---l1. w-1l s | E (W((i))) + ny..
Proof.
(((.+tl)) w -w
w
wl(,0) W.
1/3
|
2n(( Tu (())))
- - +,² VL (W(())))))) / + n? (W( /
(a)
w(( W
2
2mLx (w (1.0))
+ DE
+
w
W
nL (W(()))
+ nE
where (a) follows from Lemma 25 (b) follows from Lemma 24 and Assumption and unrolling the recursion we have,
Now starting from =--1
T-1
w w
(1.0) W
- La (W(!..)) + ne.
<
47

---

Published in Transactions on Machine Learning Research (10/2025)
C.5 Proof of Theorem 1
For any / 2 Ti we have,
K
2
W W
W((.7
W
K
b
w((r)
W.
K
k=1
K
W(
W
1
La (W(h))) + me
17
K
K
W
W
2
"
1
La (W(l)) + m/6
K
1=
W W
(L(w(() + n/e,
(61)
where (a) follows from Jensen's inequality,
(b) follows from Lemma 26 (c) follows from E.e I& (W!!:))
La w (((5)) = Lx((())) From equation 61 we get,
/Wittt W 12 + 1/c.
nL((())
w W 2
Summing over t t= Ti,7; and dividing by n7T 71 + 1) we have,
| W n(T TT +e,
1
)(W(1))
(62)
T
t=
for all 7i STST" 1. Now equation 62
implies that can find an iterate with training error less than
within,
T Ti
W
mn no?d
O
mn log(T/e)
T/E
nozde
rounds where the last equality follows from the definition of Ti equation 56 and Lemma our proof of Theorem
22
This completes
48

---

Published in Transactions on Machine Learning Research (10/2025)
D Proof of Theorem 2
We first state some intermediate lemmas that will be used in the proof.
Lemma 27. Suppose (w]rr .jp)) 2 0 for some 2 0. Then for all 2 1', 10: T 1], k E [K], we have (wilaa
jarkjj) 2 0.
Proof. We will use proof by induction. We will show that our claim holds for t= ,', E [0: - 1] and also t= (t' + +),, = 0. Using this fact we can argue that the claim holds for all t 2 t' and s E [0 7 1].
Case 1: First let us look at the local iterations s [0 [ 1| for t'. From Lemma
wc have,
(woft (( Jink jpm
(w Jor jpm + +./%/02
(w(r) jjm)
0,
where (a) uses 1vr) 2 O by definition; (b) uses (w (w j14) 2 0.
Case 2: Now let us look at the round update = 1' + 1,8 = O. We have,
(w 1,0)
J.r,k .14) (w//(+)), .jpm) (
K
-(w (1) .jm) + K
Grnk
1=1
(t) .jme)
0,
(-)) where (a) uses Tank 2 O by definition; (b) uses (w (1') j.r jpm} 2 0.
Lemma 28. Under Condition 1 for any 0 < t T 1 we have
(() 2 (t-1)
min 10)) ki
if (w!!
(e-1) .jm1) 2 0,
(63)
4m ==
and,
17 |1
((0) > (t-11 +
4m
min k.
10() min
ki
+ h
Ri
k,j
if (w < 0.
(64)
Proof.
From equation
18 we have the following update equation for j,.
P(0) 1)
(1-1,s) (tt)1,,)) o (fw jirkk (kk)))) ||11112
(65)
Proof of equation 63 this case we know from Lemma 27 that (w!)) (w.... jmm 2 0, then
(t,s)
jinkjpp ju) 2 0 for all & e [K] E 10:7-11 -1].
(66)
49

---

Published in Transactions on Machine Learning Research (10/2025)
Using this observation we have from equation 65
r j.r 2
(0)
p(t-1) + m/ 1
Σ
(x(t-1,))
min
b
mmm
(R.
(b) 1(t=1) +
77 1|ا
Σ
min
(((-1
(67)
4m
==0
where (a) follows from the definition of D, == {kr E [K] i i [N] yk.s = j}: (b) follows from Lemma and [ell 2,s
(j,je, ((t))))
b 2 ming,
Proof of equation
64 First let us look al the ileration = In this case we know that (w0-110))
ji,k jjp)
(w]r .jp) < O and thus (w) (t-1) yk.H)) > 0 for yk. = =-j. Using this observation we have,
7 m|D_a|lled|
"
1,0) o' ((w/t-1 (a, |1
(r((-1,0)
min (R.))D
k,s
7um
nu
k,
(a) 1) |1
min
RR
4m
R
where (a) follows from Lemma and min(x.)eD, (e(()))
2 mink, ee(e'1))
k,i
Now let look at the case 1 <SST 1. In this case if (wirk (w((-)))) .jm) then,
. ((w(t (), ||ell22 n mID-jaul || ki
11)))1,))
min
k2
(68)
num
nm
(k,))
and if (w) (t-1,>)
j.nkk jjm) 2 0 then,
n n1 |
gr(t-11 1,8) . ((wll-ii
b. jink
(Bs,)))) |1a112 2
min
nm
nm
Thus,
(iirk sk)s 11 > mmnn 1 111 |111
" Σ 1 nm
min (2
e(t-11s)
(69)
nm
Using the results in equation
68
and equation 69 we have,
(0) fr 2
((-1)
1)
77 |1
lettt
1.0)
17
min min((D++all +...ll1D--... #B r-1
+
min k,
min (k,3) kef
4m
m
(0)
11-1)
7 ||
(10))
1,0)
(((t=1ee)
2
+
4m
min
k.
+ h
min
k,i
where (a) follows from our definition of h in equation
Lemma 29. Let A₂ = (r E [m]
(w)) (0)
(un 2 0}. For any 0 0515T we have,
1. For any j E {+1), [m] r(f
7lllll m
Et= EEEd maxk. (er(e
2. For any r E A₁ (t)
fr
>
4m
Et En min(k.s)
(e((
k,i
50

---

Published in Transactions on Machine Learning Research (10/2025)
For any r $ Ay: : 2
فللاا 4m
Ev=
ming,
+ Ex ming, (2( (0((..0) ki
Proof.
Unrolling the iterative update equation 18 we have,
p&t)
MI t==
(WW (l(er))) -0' ((w) s (r.s)) ||
(70)
J.r
Tum
Proof of equation
1
Using equation
70
we get an upper bound
p as follows.
n 3.r S
11
ΣΣ
f-1
(e(tt
max
ki
8,8
1'=0
where the inequality follows from d') < 1.
Proof of equation From Lemma 27 we know that (w).rsjp) 2 0 then (w 2 for all 1' 2 0. Thus
using equation 63 repeatedly for all <t I we get,
[(] " 1
((((.s))
min R.S
4m
Note that the bound in equation equation 64 repeatedly for all 0 < t' <t- we get,
Proof of equation 9
64 holds even if (w .jp) 2 0. Thus applying
(1)
01
min
min
4m
Lemma 30. Under assumptions, for any 0 S STST we have,
Σn, p(0)
Jinkk I
3 2m
Ei= EED
maxk, Lee((
2. Σx, p(()
ya 16m
(e(( ki
Jirck 2
Eiz
EED min (M
where 3(( [K]
E e [N] (w(()) En,a) 2 0).
Proof.
From equation 20 we have the following update equation for (t))
p()) Σ
Σ
(t--1)
Σ
pr(t-1,s)
o ((w Ex,i)) 1|Ex.el1
nm
b
(1-11)
(1-1,)))
11Ex.
(71)
mm
(tt
where the last equality follows from the definition of ((t,)) 3.r
51

---

Published in Transactions on Machine Learning Research (10/2025)
Proof of equation
Now using equation 71 we have,
E
1) jonk
(e) S
P
1)-11 Jr,ki
3nd
Σ max R
r-1
(e- k
1,8)
2m
where (a) follows from Lemma
Unrolling the recursion above have the following upper bound,
3 (-1
Σ
P
ΣΣ
((() k,
Junkh
S
2m
max
8=0
Proof of equation 2
From equation 71 have,
Σ pl ( Jinks
(0)
noge 16m
(
min ((
k,
where (a) follows from Lemma and Proposition recursion above have,
part
which implics J.r
((t-1,0)
2 2//8. Unrolling the
2 t-1
Σ
Epir jinki 2
min (f
16m
t'=D
Lemma 31. For all 2 Tı, we have (w mp) > 0.
Proof. We have,
(w()) bim y14) (w) (a)
yor: y/a)) D
- (v /loogmm/8) oo [Ilella) +
r(tt
j,r
1111112
2 Ti
0
log Vlogmm///// (m/ 000 [1alll))
min
(( k,s
4m
kid
#=0
( θ
log (og(mm/)) on [|le|l2
+1
n 1|
12
(d)
Vn ||lll|12 0 ( (Voo(m/o))
n |||
2
2
+ ft
opdr
of dr
0.
(72)
Here (a) follows from Lemma
(b) follows from Lemma 29
(c) follows from the definition of Tj in
Equation 56 (d) follows from Assumption (e) follows from Assumption 3 and Assumption 2
Lemma 32. Under Condition for any Ti tt T 1 we have,
1.
01)
0 d 1/2 1/2)
((0) 11111
2. 0 1d11/22-1/2)
(1)
En,
p(o)
52

---

Published in Transactions on Machine Learning Research (10/2025)
Proof of equation 1 Note from our proof of Lemma 21 we know that for all Ti tt T 1 have ((() jrr
p(t) ",,,"" 2 2 for all (k","*) E
= k [K], E [N] yh J. (w (0)
Jin EnLi) 2 0). Thus,
Σ R,
pl j.nkk, 2 2 ((0) (+) (n),
(73)
where (a) follows from Lemma
S This implies,
(0) w
(oovd)
jir
2
(op
d 1d11/2-1122 En,e
P
0p d 1/1/1/1111/ E P00
jinka
0o0 1/2)
(1).
Here (a) follows from Lemma 6
(b) follows from equation 73
(c) follows from Assumption
Proof of equation
From Lemma 28 and Lemma 30 we have
pr(t(ss) kg
(())
16 |4
2
max a
16C₂ |el1
S
(()) ΣR, j.rkk,
oRd
Er= P==
-==u min
(KiDESft
(e(((
ki
07d
(0((
(er(t- 11,))
where (a) follows from Proposition
part 3 which implies maxk, / k,s
I C2 min
(k.))s
for all 0 st T 1,0 I. Thus,
((t) j.r
||
00 |Al|2
0
(a)
O (1).
(op d 1/2n 1/2) En,e P 11) j.rcks
opdl/2
where (a) follows from Assumption i
Lemma 33. For any Ti 1LST 1 we have,
Σ, o (WWrrr
C4 1111112 ormd
[Anl + (m - AAyl) (n h- (1 (1-)))
m))
Enpii P
where C4 is some constant.
Proof.
We can write,
Eo (wr()) yi (u:
Σ
(w100h8144) +
Σ
(74)
r:(wooP-aa) 20
ys(W (AA) <0
12
First note that if (w ypu 2 O then from Lemma 27 we know that,
(1.)) yo w/44) 2 0 for all & E -
(75)
53

---

Published in Transactions on Machine Learning Research (10/2025)
We can bound I follows:
1
Σ
((w))) mu1)
rs n/ p.
Σ
(w
Mir: mue)
r:(w v.r)))22000 20
3."
w: 20
4-1
Σ
min Rat
(76)
t' s=0
Here (a) follows from equation have the following bound:
75 (b) follows from Lemma
(c) follows from Lemma 29 part
For T2, we
1
Σ
((w ((wg)r-u1)))
r:(w (0)
We
(w/0)
6.r 014) + (1)
r:(w (0)
BH) <0
(m 1ny110 (Vlog(m/8) oo [leell) +
Σ
(+)
3,r
r:(w
Ω
Σ
r dar
r:(w bomm3443 <0
(0)
T1-17-1
(m - |Ayly
min R,
+
min
kit
P==0
min (r((1()))
(m A 11
(77)
p=
0=0
Here (a) follows from (22) 2; (b) follows from Lemma 6 and Assumption
(c) follows from Lemma 31
(d)
follows from Lemma 29 Substituting equation 76 and equation 77 in equation 74 we have,
Σo a ( (wo (0rr 2 12 (Ayli 1
(e((t s)) k,
min Rid
1' r'=0
==
+ (m |A,l) 2
min
larte .0] k,
+
min
(((e
ki
#==0 aml
(m -1 1|
Σ
Σ
min
(r(ttss) kid
(78)
Now using equation 78 and Lemma 30 we have,
54

---

Published in Transactions on Machine Learning Research (10/2025)
Σx σ
(t)
Endii P!.)
pale .0)
(0)
||eal12
=
ming,
k,
Ω
Ayy
oBmdd
2V==0 Ex=0 $
(r(t's)) 22
Σ Ii -1 =0=
mink, e(le200)
+ E
ming, (e(e. "i Ex= Σi
mings (e(e k,
+ (m - (Ayl)
Ti -1 Ei=d maxk, ((((111)) + =i= Ei maxg, (1(()s) k,
=1'=0
(b) Ω
||ee|12
1 (1 11) |A,l (m (m- (Ayl) h +
opmd
where (a) follows from Lemma 30
(b) follows from Proposition part 3 and Equation
58
Lemma 34. Under assumptions, for all Ti STST 1 we have
(t) Σ, o ((whin ww)
1/2 |1
1 (1 h)) )) [Ag + (m - -14y)) h++ 7
20
m pp -- -g.r
(1)
o2mdl//
|w W.r = o (o-td-1/2n 1123) Ex, P(O)
Proof.
To prove this, we first show that
Jonk for all j E (+1}.
We first bound the norm of the noise components as follows.
Σ
p(t)
[|Ex, Ex.
2
Σ
(p(t) E(p(t))))) 3irdsss)
2
||ER, 1122 2
Σ
p(t) p(t)
Jackh J.r,y ||Ex.ill? -1E 1122 (En,
40
2d
((t)
jiridess)
+
Σ
p(t)
Sirck dirkkr, p((t)
p(t)
(160,'d /dlogg/
dlog(6n2/8))
ki K
2
40,
2
((t)
J.rckee)
+ +320, + -21
Σ
p(t)
Birdess
-
(((t)
Jirckie)
)2
bi
R
2 (o 2d d
Σ
(ptt) )2 + (072)-312)
p(t)
jirskes)
kid
2
(
0 ( - 2 1 + -
1 2322)
Σ
((1)
Σ
((t)
+
Rit
R
2
1a d '' 1)
Σ E
(79)
55

---

Published in Transactions on Machine Learning Research (10/2025)
Here for (a) uses Lemma 5 (6) uses ,a
p(0 j.r.
S
3 + / (nn /6) na O (1) from Theorem 3 and
Σ
p(() Ji
O Er., P Birdsss Now from equation
17 know that,
w (c) w
SI ||eel12
plo Janki 11Ex,il12 Ex.s.
Jrr
+
1= [E
Using triangle incquality and equation
79
we have,
(0) j.r
(r() 3.r |,e1
+ + -0 1/ 1/2
(()) Jink
j.r
2
w
+
k
(a)
1-d 1/2 112)
Σ
pit)
jirkk,
k,
where (a) follows from Lemia Thus,
Σ, o ( (wrhh (t) wir 810)
Σ, ((w 8/4e)
pp -x=1
θ 1211222" -1/2)
Ex.t p(t)
1/2 || |ell
(Aul) (1,l) 21 h))
|Ayl + (m
ozmdl1/2
where (a) follows from Lemma
Lemma 35. (sub-result in Theorem E.1 in Cao et al. (022).) Denote g(E) = E, o (ww Ei) Then for any x 2 it holds that
cx
Pr(g(E) - Eg({)) > 2) exp
w
(t) y.r
2
where is constant and Eg(E) V2x
-
w (t)
yor
D.1 Test Error Upper Bound
We now prove the upper bound on our test error in the benign overfitting regime as stated in Theorem 2
First note that for some given (x, v) we have,
P(y sign = -P(y((W()) x) s 20).
We can write,
y vs((() x) = F,(w(() x) F-, (W(())
m
m
(w (w m111) + ((w/t) 4)
((w!!) (1) (/r + a ((w"")
(80)
Now nole that since / 2 Ti we know that a (w = 0 for all r E [m]] from Lemma 31 Thus,
56

---

Published in Transactions on Machine Learning Research (10/2025)
m
m
P(yS(W() x) S 0) SP
Σ 0
((w"))
-y.r: E)
Σ a r=1
((w)))
r=1
(4)
g(E) Ey(E)
m
0 (ww!)) (01))
Op
Σ
m
w
-
/T
- 0
((worr v11))
2
(1)
b)
2
22n
-y-""
5 exp
o
E
(()))
Σ r=
0
(w (W yp4
exp
(1)
Op
E1=1
W
-yir
Σm= 0
c
exp
2π
2
pp >m
[w u.
n ||ll (|Ayl + (m (m-1Ap) (h + +1 h)))²
c
exp
2T
Csorm2d
n ا|ااا1 (14y) + (m - (A (h + [11 -h))) 2
exp
203o3m²d
Here (a) follows from the definition of g(E) in Lemmma 35 (b) follows from the result in Lemma 35
uses (a b)² 2 a²/2 b2, Va, b 2 0; (d) uses Lemma 34 (e) follows from the benign overfitting condition
n |||| = Qt (o/d) and choosing sufficiently large Cg- Now note that,
16 (w(T)
Σ )E(+1)
P(y = j)P(v + sign (i x))
n n1l14ll2 (lA,1+ + (m - (A,1) (h (1 - 1))) =))))²
Σ jE(#1}
exp
2C3orm2d
This completes our proof for the upper bound the test error in the benign overfitting regime.
D.2 Test Error Lower Bound
We first state some intermediate lemmas that we use in proof.
Lemma 36. (Lemma 5.8 in Kou et al. (2023)) Let g(E) Eg,jo ((w (W .E)) If Elallll o (ofd)
(harmful overfilingg condition) then there exists fixed veclor v with 1|v12 S 0.060p such that
Σ løti't v) - -(5'E)) 2 4Co max
Σn 3or
jE(Ll)
for all ddd. Lemma 37. (Proposition 2.1 Deuroye al. 2018)) The TV distance between N(0,ofTa) and N(V,
is less than 1Ivl22 /2op.
Proof.
We have,
57

---

Published in Transactions on Machine Learning Research (10/2025)
13 (W(T))
P(xN))D (U + (f(,,))
P(XNNDDD ( x) 0)
(4) P(xxy)ND
Σ
(w (T)
- v.r: 6) - -o - ( E) 2 E
(ww!?7 (.r. (811))
- -oo
(ww uor muie
(8)
((xDD)DDD
Σ
((w2))
(2) 5) - -E ((w!7),)) 5)) z Co max
Σ
p()) 1,+
(2)
1,5
2 0.5I (x.))>D
Σ [
((w
(T) {) - Σ a (w {)) 2 Co max
Σ
Ir
Σ
(T)) 1)
(c) 0.5P 0.5P (x.y)>D
(9(E)] 2 Cs max Eri.. C1 -
(4) 0.5P()).
(81)
Here (a) follows from equation
80
Ply
+ sign (s(Wlo, x)) PPys(W()) x) S 0); (b) fol-
lows from ((w").. (11)
0 (Lemma
31 and (W))r.B))
e (T (c) follows
from defining g(E) Σ,
(
(w
E)
Σ, o (ww"))rr -1,r 6)) (d) follows from defining Ω =
(T)
{ (g(()) 2 Co max
I (T) 1,r E,
Now know from Lemma Lemma 36 that Ej Kgr v) (iggg(ii) 4Cs max, m (T)). This implies that one one of the E, c. EE 6 must belong to Ωt Therefore,
min P(-)), P(2 P(-2 v)) 2 0.25
(82)
Also note that by symmetry IP()) = P(-?). Furthermore,
PP (2) (2 v) = P PE-NNO0.03WLE E .2) - PE 2)]
TV (,) ov))))
||v
20p
0.03.
(83)
Here (a) follows from the definition of TV distance; (b) follows from Lomma Lemma 37 Thus we see that
equation 83 along with equation 82 implies that P(2) = 0.22. Substituting this equation we get. r--1 ( (T) = 0.1 as claimed.
E Main Paper Lemma Proofs
E.1 Proof of Lemma
This lemma follows from directly from Lemma 29 and the constant lower bound on cross-entropy loss derivatives, i.c., Equation (58.
E.2 Proof of Lemma
This lemma follows from directly from Lemma 30 and the constant lower bound on cross-entropy loss derivatives, i.e. Equation 588
58

---

Published in Transactions on Machine Learning Research (10/2025)
E.3 Proof of Lemma 3
Using our result in Lemma 28 and r E [m]]
with T 1 and h = 0, we have after Ti =0
m2 no
ilerations for all j E {+1)
TT)
"
(pre)) 1/
T
min ((pre.)) (a) (llle(ere)) 1/
CTi
111u (puppr))1
Ω
4m
4m
ozdd
Here (a) follows from equation
58
Now for any 2 Ti we have from Lemma
4
(w/pre (p)e,t) (p((re))
[w]] w./rroo
(pro.0)
.jp((rre) + r (pre.t)
IV2
(w (pre.0)
-j(u(1e)) + (()re.2))
IVE IV
(Vl vnl Anlla(ngg (Dre))
| + S op 2
|elonos13
0.
where (a) follows from the fact that [ is non-dccreasing with respect to t, (b) follows from Assumption 4
and Lemma 6 (c) follows from Assumption 3
F Additional Experimental Details and Results
Implementation. We use PyTorrch Paszke et al. (019 to run all our algorithms and also simulate our synthetic data setting. For experiments on neural network training we use one H100 GPU with 2 cores and 20GB memory. For synthetic data experiments our experiments neural networks is about 36 hours. The approximate total run-time for all experiments the synthetic data setting is about 1 hour.
use one T4 GPU. The approximate total run-time for all
Details for Figure I We simulate FL setup with K 10 clients on the CIFARIO data partitioned using Dirichlet(a) with a 0.1 for the non-IID sctting and a 10 for the IID sctting. For pre-training, we consider Squeezenet model pre-trained on ImageNet, Russakovsky et al. (2015) which available in PyTorch. Following Nguyen et al. (2022) we replace the BatchNorm layers in the model with GroupNorm Wu & He (2018). For FL optimization
use the vanilla FedAvg optimizer with server step size Tg = 1 and train the model for 500 rounds and 1 local epoch at each client. For centralized optimization we use SGD optimizer and run the optimization for 200 epochs. Learning rates were tuned using grid search with the grid {0.1, 0.01, 0.001). Final accuracies were reported after averaging across 3 random seeds.
Details for and Figure 2 and Figure 3. For these experiments simulate synthetic data setup following data model in Section 2 We set the dimension d 200, 20 datapoints (we keep small to ensure we are in the paramearaetrri regime), m 10 filters, K = 2 clients, N = 10 local datapoints. The signal strength is ||11 3, noise variance is o² 0.1 and variance of Gaussian initialization is do 0.01. The global dataset has 10 datapoints with positive labels and 10 datapoints with negative labels. We also create a test dataset of 1000 datapoints following the same setup to evaluate our test error.
Details for Figure and Figure 6 We simulate a FL setup with K 20 clients using Dirichlet(a) Hsu et al. (2019). For pre-tnaining, we consider ResNet18 model pre-trained on ImageNet [Russakovsyy al, (2015] which is available in PyTorch. Following Nguyen al. (0022) we replace the BatchNorm layers in the model with GroupNorm Wu & He (2018). For FL optimization we use the FedAvg optimizer with server step size f2 =1 and 1 local epoch at each client. In the case of random initiation, for local optimization we use SGD optimizer with learning rate of 0.01 and 0.9 momentum. In the case of pre-trained initiation, for local optimization we use SGD optimizer with a learning rate of 0.001 and 0.9 momentum. The learning rate is decayed by a factor 0.998 every round in the case for both initializations. Each experiment is repeated with 3 different random seeds.
59

---

Published in Transactions on Machine Learning Research (10/2025)
25-
20-
% Misal ligned Fiteer
(%)
70-
15-
65
60
58
50
Random
Protrained (Tiny ImageNet)
Pretraized (CIPAR-10)
Random
Prelesined (Tiay ImageNNt))
Pretrained (CIFAR-10)
40
10
20
Round No.
(b)
Figure 7: The percentage misaligned filters (Figure 7a and test accuracy (Figure 7b for different initializations on CIFAR-10 with 0.05 heterogeneity when training 2-layer CNN.
Additional Experiment. To simulate an experimental setting aligned with our theoretical analysis, we conduct the following study. We partitioned two classes of CIFAR-10 (airplane and car) across = 20 clients with a 0.05 heterogeneity. We then considered FL training of a 2-layer CNN model with m = 64 filters and evaluated three initialization strategics:
1. Pre-training on centralized CIFAR-10: The model is pre-trained centrally for 5 epochs on the same CTFAR-10 dataset later used for federated training.
2. Pre-training Tinylmagenet: We pre-trained the CNN on two semantically related classes in TinyImagenet albatross and racecar, which resemble the airplane and car categories in CIFAR-10.
3. Random initialization
As expected, pre-training on the same dataset yields the strongest improvement, with 0% misaligned filters at initialization and the highest test accuracy at the end of FL training. Importantly, pre-training on Tiny Imagene,, while not perfectly aligned, still substantially reduces the proportion of misaligned filters (3.1% vs. 25.7%) and leads to better test accuracy than random initialization. These findings strongly support our theoretical insights: the reduction in misaligned filters at initialization directly correlates with improved FL performance, thereby reinforcing the connection between our theory and empirical results.
60


Reviews:
[{"review_id": "5nnHexeKN1", "reviewer": "Reviewer_p1Zg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThis paper is well-written in general (pending a few minor points)\nThe research angle is interesting and can potentially lead to new insights on using pre-trained model in federated learning.\n\nWeaknesses:\n\nFirst, I am concerned that the theoretical setting here is too detached from existing practice in federated fine-tuning. For example, the CNN construct is restricted to a binary classification setting with 1 convolution layer followed by a linear layer. This is disconnected from most federated learning work. \n\nI am aware that this setting has been adopted in previous works for different forms of analysis in centralized and/or federated learning w/o initializing with pre-trained models. I am however not sure if this simplified setting will fit with the broad research question raised by the authors (2nd page):\n\n\"Why can pre-trained initialization drastically improve model performance in FL?\"\n\nDrastic performance improvement requires the pre-trained model to be expressive and well-trained with large volumes of relevant data. How is this going to be the case with such simplified CNNs? \n\nI do not see how this is demonstrated in the empirical verification. It is also unclear to me how the pre-trained model generalizability (i.e., its quality) does not come into the picture? \n\nSecond, the experiment was designed with a much more generic pre-trained ResNet which does reiterate known observations that initializing with pre-trained model helps but it does not support the theoretical model here. To connect better between experiment & theory, the authors can consider re-partition the 10-class CIFAR-10 dataset into a 2-class dataset, pre-train the CNN & use it as initialization for FedAvg. We expect to see that (1) the pre-trained CNN has good generalization and (2) using it as initialization, the FL performance improves significantly. \n\nThird, the overall analysis seems to prove that with a sufficient no. of training iterations, the FL model can align the learning signal with that of the pre-trained model which is perhaps implicitly assumed to align with the data generation model. This raises another question: suppose the pre-trained model was not created using the same (centralized) dataset, how will it affect alignment theoretically? This is important to address because most of the time, FL algorithms were not initialized with models pre-trained on the exact target data.\n\n--\n\nMinor points:\n\nI feel that the concept of aligned and misaligned filters need additional elaboration on the high level. It is not clear what is being aligned when these concepts were mentioned the first time in the introduction. Furthermore, going by the inline definition in Lemma 1, I have the impression that such directional alignment might not be sufficient when we generalize to multi-class classification. Can the author elaborate more on this?\n\nAt multiple points, the authors mentioned that some assumptions (e.g., orthogonality) are meant to ease the technical exposition and relaxing them is not difficult. Can you please elaborate on high level how can these be easily relaxed?\n\n--\n\nOverall, I feel that the paper might have over-claimed their theoretical contribution. The key research question probably needs to be rephrased to reflect more accurately the theoretical setting. Otherwise, there is a risk of over-generalization.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper provides a theoretical analysis of the impact of pre-training on federated learning with non-IID data. This is achieved via studying the class of two-layer convolutional neural networks (CNNs) for binary classification. The studies provide bounds on the training error convergence and generalized test error of the network trained with FedAvg. At the core of this analysis is the notion of aligned and misaligned filters at initialization and a demonstration that data heterogeneity only affects learning on misaligned filters. The study goes on and show that starting with a pre-trained model results in fewer misaligned filters at initialization, reducing the generalized test error.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I suggest the following changes:\n\n1. To connect better between experiment & theory, the authors can consider re-partition the 10-class CIFAR-10 dataset into a 2-class dataset, pre-train the CNN & use it as initialization for FedAvg. We expect to see that (1) the pre-trained CNN has good generalization and (2) using it as initialization, the FL performance improves significantly. \n\n2. Repeat the above experiment with the pre-trained CNN on a different dataset, e.g., TinyImageNet, and see if we also see a drastic improvement on the target dataset with FL.\n\n3. Elaborate on the alignment notion. It is currently mentioned in Lemma 3 and very briefly even though this is a central concept. Particularly, I would like to see discussion on how this notion can or can be generalized to accommodate multi-class classification.\n\n4. Elaborate on how the pre-trained model's generalizability comes into the picture. \n\n5. It'd also be good to discuss how the current model can be generalized to deep CNN settings (e.g., ResNet-18) with multiple CNN layers. There might be a difference between filter alignment at raw input level and alignment at higher feature level. In other words, it is unclear from the current theory whether filter alignment at input level alone is sufficient to characterize the entire benefit of pre-training on federated learning.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "5PMkMum9cF", "reviewer": "Reviewer_5iQi", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths \n\n- The goals of the paper are clear and the theoretical analysis seems sound.  \n- The theoretical results (especially Theorem 2) tie all the important quantities with the test error and motivate the study of misalignment. \n- Empirical results on CIFAR-10 and TinyImageNet at the very least show correlation with what is expected from theory as well as the correlation between complexity of data and expected misalignment when starting from a random initialization. \n\n## Weaknesses\n\n- The analysis, while strong, is limited to a binary classification task with a two-layer CNN and the standard FedAvg algorithm (the authors acknowledge this in the conclusion as future work). I think the paper would be strengthened by a brief discussion on the challenges of extending this \"filter alignment\" concept to other architectures (e.g., Transformers, where the notion of a filter is less direct) or other federated algorithms.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work studies the effect of model initialization for federated learning, specifically focusing on the FedAvg algorithm. Studies have shown that there are advantages to using a pre-trained instead of a randomly-initialized model when wanting to fine-tune on a downstream task, especially when it comes to generalization. The authors look into this phenomenon from the point of view of CNNs and a particular learning algorithm, FedAvg. \n\nIn this, mostly theoretical work, the authors consider a data model and two-layer CNN that is quite similar to the one from Cao, et al, 2022[^1]. In this setup, data points contain features $x=[x(1)^T, x(2)^T]\\in \\mathbb{R}^{2d}$ and each data point is assigned one of two labels, $y\\in \\{-1,1\\}$. At random, one component is set to a signal vector $y\\mu$ or to a Gaussian noise vector.  \n\nWith this setup, the authors introduce the notion of aligned filters – a filter is aligned with a signal vector if their inner product is positive. \n\nThe main contributions after this setup are: \n1. Based on the data model and assumptions on the SNR, dimensionality of input, training set size, etc., the authors show the training loss convergence of FedAvg as well as lower and upper bounds on the test error according to a respective assumption on the value of SNR. The upper bound is especially interesting as it ties together the contribution of the aligned and misaligned filters to the worst test loss we can get. \n2. The authors also show that there is a time during pre-training after which all filters are aligned and so we get the benefits of Theorem 2. without any misaligned parts (and thus no impact from heterogeneity). \n3. An empirical measure of misalignment is defined and used with ResNet18, both a randomly-initialised version and a pre-trained one, and image datasets to show some of the predictions of the theory for more realistic examples. The authors demonstrate that the pre-trained version shows smaller misalignment and better test accuracy.\n\nOverall I enjoyed the notion of misalignment enough that I'm wondering if parallel notions to that are applicable for other setups. \n\n[^1]: Cao, Y., Chen, Z., Belkin, M. and Gu, Q., 2022. Benign overfitting in two-layer convolutional neural networks. _Advances in neural information processing systems_, _35_, pp.25237-25250.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Typos:\n- Page 3, missing citation?!, “can be easily relaxed as done in ”.\n-  Page 4, towards the end, please capitalize proper nouns, “equation 2 - equation 3”. \n- Page 8, Empirical Verification part, “Figure Figure 3a” (delete first “Figure). Please proofread that paragraph for typos. \n\n## Not necessary for recommendation \n\n- Equation (8), the misalignment measure: was that perhaps meant to be an average over the batch?", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "3HCvIEFCdy", "reviewer": "Reviewer_8fwx", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**S1.**  Providing an analysis and efficient rigorous computational results is always welcome in machine learning.\n\n**S2.**  This is especially true for federated learning that is much less studied than batch learning.\n\n**S3.**  The authors define meaningful quantities and obtain reasonable and intuitive results.\n\n**S4.**  The authors provide detailed analysis for their results and the appendix goes even beyond that and provides additional information on related work and also has a table with notation.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates the idea of using pre-trained models as starting points for models that are used in a federated learning setting.  In particular, the authors provide a justification as to why the use of such pre-trained models can provide faster and more stable convergence to a good enough learnt model.  In this direction the authors study convolutional neural networks (CNNs) that have two layers and are trained using the FedAvg algorithm for a federated learning setting.  One justification for the choice for this class is that they have been investigated recently in the context of batch learning.  Another justification (this is my interpretation) is that they allow for the proof of some intuitive theoretical results that the authors provide in their work.  In particular, the authors provide rigorous results on the convergence rate and on the test error bound for such models.  In addition, the authors also provide an empirical evaluation for the behavior of the tools that they have defined (or developed) in this context using the CIFAR-10 and TinyImageNet datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**W1.**  While the simplicity of the class of models that the authors investigate (two-layer CNNs) allows them to provide some rigorous guarantees, nevertheless, it is unclear how many people will work with such a class of models in the real-world, especially outside of standard benchmarks of computer vision.\n\n**W2.**  While Strength S4 is welcome, it should be stated that the paper is longer than usual papers.  Personally, I would prefer in some cases sketches of the proofs of some theorems in the main text (i.e., first 12 pages), but given the length of the proofs in the appendix, I understand that this may not be a feasible after all.  So, I would try to urge the authors to consider providing 1-2 paragraphs giving a sketch of the proofs and the important pieces that when put together can justify intuitively even more the results that the authors obtain.\n\n**W3.**  Missing discussion on class-imbalance?\n\nThe parameter h for heterogeneity is like an \"average\" metric for the class-imbalance that you have among the various clients.  In other words, while you use some global distribution D in page 3 (paragraph \"Measure of Data Heterogeneity\") the fact that you are partitioning the generated dataset arbitrarily, leads to different class ratios among the clients. Which now begs the question if you are missing discussion on class imbalance. \n\n\n**W4.**  Minor issues in the text.\n\n- In page 3, before the paragraph \"Measure of Data Heterogeneity\" you mention Condition (C2).  At that point in time (C2) is undefined.  First of all, I think you need to write \"Condition (C2)\" instead of just \"(C2)\".  Then you also need some surrounding text to refer to \"Condition 1\" in page 5, which has these six \"sub-conditions\".  Of course you will lose a line in the text there and this may push back some of your content, but sentences should be complete.  Related to that, I acknowledge that it is a little bit cumbersome to have \"Condition 1\" and then the conditions (C1)-(C6) listed there, because now in the actual text if one refers to \"Condition (C1)\" this can be confused with the environment \"Condition 1\".  I think something needs to change there but there is no easy way out.  Especially since you are also using the terms \"Assumption 1\", etc. in the appendix!  Anyway, I do not have very strong feelings, but it seems a little bit cumbersome there.\n\n- Defining the loss function $\\ell$ to be taking only one argument is odd.\n\n- I have a feeling that Figure 4 in page 8 is never mentioned in the actual text.  Also, while we are at it, I believe you should introduce it it before you start discussing Figure 3 and its results because it really allows one to understand the values of the different filters and why the results of Figure 3 make sense as they are.  In the current flow of presentation they did not make much sense to me initially, until I actually \"ignored\" my issue in the paragraph of page 8 and then later focused on Figure 4 that really elucidated the results of Figure 3.", "broader_impact_concerns": "**Q1.**  Will you make your code publicly available upon acceptance?\n\n**Q2.**  Condition (C2) states that we have much fewer total examples in all the clients compared to the dimensionality of the data.  Is this a reasonable assumption?  I believe it is the other way around (e.g., n is number of patients in a federated healthcare system).  In that light, the results that you obtain are now weaker because one, in principle, would care about the situation where our data points far exceed their dimensionality.  Even in traditional batch learning, interesting model spaces have VC-dimension that is linear or superlunar in the dimensionality of the data and therefore the number of examples that we need for generating \"good\" hypotheses is linear/super-linear in the dimensionality of the data.  Having said all these, the question is to what extent your results may hold for much smaller d?  If you do not have any rigorous results in this direction, do you have any experimental results?\n\n**Q3.**  Do you have any minor comments on the class-imbalance issue that I raised above?", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)