FedAvg with Fine Tuning: Local Updates Lead to Representation Learning
Authors: Liam Collins
Hamed Hassani
Aryan Mokhtari
Sanjay Shakkottai
Venue: NeurIPS 2022 Accept (2022)
Keywords: Federated learning
Representation learning

Abstract: The Federated Averaging (FedAvg) algorithm, which consists of alternating between a few local stochastic gradient updates at client nodes, followed by a model averaging update at the server, is perhaps the most commonly used method in Federated Learning. Notwithstanding its simplicity, several empirical studies have illustrated that the model output by FedAvg leads to a model that generalizes well to new unseen tasks after a few fine-tuning steps. This surprising performance of such a simple method, however, is not fully understood from a theoretical point of view. In this paper, we formally investigate this phenomenon in the multi-task linear regression setting. We show that the reason behind the generalizability of the FedAvg output is FedAvg’s power in learning the common data representation among the clients’ tasks, by leveraging the diversity among client data distributions via multiple local updates between communication rounds. We formally establish the iteration complexity required by the clients for proving such result in the setting where the underlying shared representation is a linear map. To the best of our knowledge, this is the first result showing that FedAvg learns an expressive representation in any setting. Moreover, we show that multiple local updates between communication rounds are necessary for representation learning, as distributed gradient methods that make only one local update between rounds provably cannot recover the ground-truth representation in the linear setting, and empirically yield neural network representations that generalize drastically worse to new clients than those learned by FedAvg trained on heterogeneous image classification datasets.

FedAvg with Fine Tuning: Local Updates Lead to Representation Learning
Liam Collins
Hamed Hassani
ECE Department
ESE Department University of Pennsylvania lassani@seas upenn edu
The University of Texas at Austin liamcQutexas. edu
Aryan Mokhtari
Sanjay Shakkottai
ECE Department
ECE Department
The University of Texas at Austin mokhtari Secaarriuuutto tin. utexas edu
The University of Texas Austin sanjay. shakkottai @uteras.s.u.s... edu
Abstract
The Federated Averaging (FedAvg) algorithm, which consists of alternating be- tween a few local stochastic gradient updates at client nodes, followed by model averaging update at the server, is perhaps the most commonly used method Federated Learning. Notwithstanding its simplicity, several empirical studies have illustrated that the model output by FedAvg leads a model that generalizes well to new unseen tasks after few fine-tuning steps. This surprising performance of such simple method, however, is not fully understood from a theoretical point of view. In this paper, we formally investigate this phenomenon in the multi-task show that the reason behind the generalizability the FedAvg output is FedAvg's power in learning the common data representation among the clients' tasks, by leveraging the diversity among client data distributions via multiple local updates between communication rounds. We formally establish the iteration complexity required by the clients for proving such result in the setting where the underlying shared representation is linear map. To the best of our knowledge, this is the first result showing that FedAvg learns an expressive repre- sentation in any setting. Moreover, communication rounds are necessary for representation learning. as distributed gradient methods that make only local update between rounds provably cannot recover the ground-truth representation in the linear setting, and empirically yield neural network representations that generalize drastically worse new clients than those learned by FedAvg trained on heterogeneous image classification datasets.
linear regression setting.
show that multiple local updates between
1 Introduction
Federated Learning (FL) [1] provides a communicaiion-efficeen and privacy preserving means to learn from data distributed across clients such cell phones, autonomous vehicles, and hospitals. FL. aims for each client to benefit from collaborating the learning process without sacrificing data privacy or paying substantial communication cost. Federated Averaging (FedAvg) [T] is the predominant FL algorithm. In FedAvg, also known as Local SGD [2H4]. the clients achieve communication efficiency by making multiple local updates of shared global model before sending the result to the server, which averages the locally updated models to compute the next global model.
FedAvg motivated by settings with homogeneous data across clients, since multiple local updates should improve model performance on all other clients' data when their data similar. contrast,
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

CKA sm belween basee and FT model
FadAvg
Reg Representatio learning ertor
10
D-3G0
0.9
10"
10
0.8
10"
FodAug D-GD
0.4
50
100
150
200
250
900
350
400
0.0
25
5.0
7.5
10,0
12.5
15.0
17.5
Communication rounds
Model layer
Figure 1: In multi-task linear regression with population losses, FedAvg linearly converges to the ground-truth representation, while D-GD (FedAvg with one local update) fails to learn it.
Figure 2: The NN representation learned by Fe- dAvg on CIFAR- 100 with 5 classes/client does not change significantly when fine-tuned on new dataset (CIFAR-10), unlike D-SGD.
FedAvg faces two major challenges in more realistic heterogeneous data settings: learning a single global model may not necessarily yield good performance for each individual client, and, multiple local updates may cause the FedAvg updates to drift away from solutions of the global objective Despite these challenges, several empirical studies [0 have observed that this shared global model trained by FedAvg with several local updates per round when further fine-tuned for individual clients is surprisingly effective in heterogeneous FL. settings. These studies motivate us to explore the impact of local updates on post-fine- tinentuning performance.
[5
Meanwhile, large number of recent works have shown that representation learning a powerful paradigm for attaining high performance in multi-takk settings, including FL. This is because the tasks' data often share a small set of features which are useful for downstream tasks, even if the datasets as whole are heterogeneous. Consider, for example, heterogeneous federated image classification in which each client (task) may have images of different types of animals. images share a small number of features, such as body shape and color, which admit simple and accurate mapping from feature space label space. Since the number of important features much smaller than the dimension of the data, knowing these features greatly simplifies each client's 's task.
safe to assume the
To explore the connection between local updates and representation learning, we first study multi-task linear regression sharing a common ground-truth representation (Figure []. We observe that FedAvg converges (exponentially fast) the ground-truth representation in principal angle distance, while Distributed-O0 (D-GD). which is effectively FedAvg with one local gradient update, fails to learn the shared representation. multi-layer CNN on a heterogeneous partition of CIFAR- 100 (Figure Since there is not necessarily ground-truth model here, we evaluate representation learning as follows. We first train the models with FedAvg and Distributed-SCii (D-SGD) then fine-tune the pre-trained models on clients from new dataset, CIFAR-10 10. Finally we evaluate the quality the learned representation by measuring the amount that each model layer changes during fine-tuning using CKA similarity [13] Observe that
similar concept can be shown in the nonlinear setting. We study
the early layers FedAvg's pre-trained model (corresponding to the representation) change much less than those of D-SGD. More details for both experiments are in Section 5 and Appendix These observations suggest that FedAvg learns a shared representation that generalizes to new clients, even when trained in a heterogeneous setting. Hence, natural question that arises is:
Does FedAvg provably learn effective representations of heterogeneous data?
We answer "yes" to this question by proving that FedAvg recovers the ground trunchh-ruthhhu representation the case of multi-task linear regression. Critically, we show that FedAvg's 's local updates leverage the diversity among client data distributions to learn their common representation. This is surprising because FedAvg is general-purposs algorithm not designed for representation learning. Our analysis thus yields new insights on how FedAvg finds generalizable models. Our contributions are:
Representation learning guarantees. We study the behavior of FedAvg in multi-task linear regression with common representation. Here, each client aims to solve a d-dimensional regression with ground-truth solution that belongs to a shared k-dimensional subspace of Rd where k < d. Our results show that FedAvg with > 2 local updates learss the representation

---

at a linear rate when each client accesses population gradients. To the best of our knowledge, this is the first result showing that FedAvg learns an effective representation in any setting.
Insights on the importance of local updates. Our analysis reveals that executing more than one local update between communication rounds exploits the diversity of the clients' ground truth regressors to improve the learned representation in all & directions in the linear setting. In contrast, we prove that D-GD fails to learn the representation.
Empirical evidence of representation learning. We provide experimental results showing Fedavg learns a generalizable representation when we use deep neural networks on image classification datasets. In contrast, the representations learned by D-SGD generalize drastically worse to data from new clients. This suggests that the main message of our theoretical results that local updates facilitate representation learning can generalize to more complex scenarios beyond the bilinear setting.
Related work. Recently there has been surge of interest motivated by FL in analyzing FedAvg/Loca SGD in heterogeneous settings, Multiple works have shown that FedAvg converges global optimum (resp. stationary point) of the global objective in convex (resp. nonconvex) settings but with decaying learning rate [7] leading sublinear rates and communication complexity sometimes dominated by Distributed SGD [8] These results are tight in the sense that FedAvg with fixed learning rate may not converge to a stationary point of the global objective in the presence of data heterogeneity, as its multiple local updates cause it to optimize distinct, unknown objective 6.9 16 19 20). Several methods have tried to correct this objective inconsistency via gradient tracking [5 19 21 25]. local regularization [20] 28) operator splitting [7]. and strategic client sampling (1 In contrast, we show that local updates with constant learning rate benefit learning in heterogeneous settings by resulting in linear convergence to generalizable models.
Several papers have also studied FedAvg from generalization perspective. It was shown in 32 that in setting with strongly convex losses, either local training or FedAvg with fine-tuning (but not both) achieves minimax risk. depending on the level of data heterogeneity. Similarly, (33 argued that FedAvg with fine-tuning generalizes as well as more sophisticated methods, including Model-Agnostic Meta-Learming (MAML. 34 35 in a strongly convex regularized linear regression setting. Additional work has studied the generalization of FedAvg in kernel regression, but for convex objectives that do not allow for representation learning [6),. and the generalization of variant of FedAvg, known as Reptile (37], on wide two-layer ReLU networks with homogeneous data (38). We focus on the multi-task linear representation learning setting [39], which has become popular in recent years as it is an expressive but tractable nonconvex setting for studying the sample-complexin benefits of learning representations and the representation learning abilities of popular algorithms in data heterogeneous settings [11 46). Remarkably, our study of FedAvg reveals that it learn an effective representation even though it was not designed for this goal. unlike a variety of personalized FL methods specifically tailored for representation learning [1] (7-51.
Notations. We use (u, Σ) to signify the multivariate Gaussian distribution with mean u and covariance Σ. Odxk denotes the set of matrices in d with orthonormal columns. The notation (B) represents the column space of the matrix B, and cB) is the orthogonal complement to this space. The norm ||- || is the spectral norm and L is the identity matrix in Rdd We use [m] to indicate the set of natural numbers up to and including m.
2 Problem Formulation
Consider a federated setting with a central server and M clients. Each client i E [M] has a training dataset D. of nı labeled samples drawn from distribution D; over X x V, where X' is the input space and y is the label space. The learning model is given by he : X + y for model parameters 0 "R The loss of the model on sample (x, y) E X y is given by E(ho ( y), which may be, for example, the squared or cross entropy loss. The loss of model parameters 0 on the i-th client is the average loss of the model he on the samples in D., namely (0) = 1 Σソ= ((ho(x.))) Niy)) where (Xx.j, Vi.) is the j-th sample in D,. The server aims to leverage all of the data across clients to find models that achieve small loss f0) for each client. To do so, the standard approach is to find
3

---

single model 0 that minimizes the average of the client losses weighted by number of samples:
M m.
M Σ Σ
min 0 N
Σ
((ho(x,)):
N
==1 JED,
where N = Ei Ri: Due to communication and privacy constraints, the clients cannot share their local data D., ①T must be solved in federated manner.
FedAvg. The most common FL method is FedAvg. On each round of FedAvg, the server uniformly samples set Ie of m S M clients. Each selected client receives the current global parameters O,. executes multiple SGD steps on its local data starting from 0,. then sends the result back to the server. The server then computes 0z+1 the weighted average of the updates. Specifically, upon receiving the global model 0,, client i computes
01.i,11 = 0z,i, agt,i, (0t...)))......
(2)
for O,.... .,...,,--............ 1, where 7 is the number of local steps, θr,i,0 O, and B0.i..(0,.1.) is stochastic gradient of fa evaluated or,i,s using b samples from D. The client then sends Ot,ir back to the server, which computes the next global iterate as:
0++1
n,Or.i,rr
(3)
N VEI.
where N; = EiEI, Nh. Note that == corresponds to D-SGD, also known as mini-batch SGD whose convergence properties are well-understood [S efficiency of D-SGD by making 2 local updates between communication rounds.
54 FedAvg improves the communication
Fine-tuning. After training for T' communication rounds, the global parameters Or learned by FedAvg are typically fine-tuned on each client before testing. In particular, starting from 0r, client i executes 7' steps of SGD on its local data as follows:
0r,i = 0r,i. agr.i,. (01,5.))
(4)
for s = O.... 1. The fine-tuned model ultimately used for testing is fy,i,' Note that new client, indexed by M + 1, entering the system after FedAvg training has completed can also fine-tune Or using the same procedure to obtain personalized solution 0r.M-1.r'
Representation learning. We aim to answer why the fine-tuned models {OT Jiml perform well practice by taking representation learning perspective. We show that the output of FedAvg. i.e.. 0r. has learned the common data representation among clients assuming that such representation exists. To formalize this result, we consider a class of models that can be written as the composition of representation TP and prediction module, i.e. head, denoted as 1 head Let the model parameters be split as 0 == (Ø) v/.. where & contains the representation parameters and s/b contains the head parameters. Then, for any E X, the prediction of the learning model is (x) = (hhiadd h (x) =
M+1
h (h 1 (x)). For instance, if he neural network with weights 6, then hep is the first many layers of the network with weights Φ, and A is the network last few layers with weights æb. A standard assumption in multi-task settings, including the settings we consider, is the existence of a common representation 1P that admits easily learnable head h,5,, such that /.baad o tep performs P.,,
well for task i. As result, in these settings it is of interest to all the clients to learn top.
3
Main Results
To rigorously study the representation learning abilities of FedAvg, we employ the standard setting used for algorithmic representation learning analysis: multi-task linear regression [1] 40 41 46 55 (66. In this setting, samples (x..) Ji.)) for each client are drawn independently from a distribution
D. on d x IR such that
Xi.j lid Vi.j = (B,, Xi.)) + Ca.j. where Gi.5
Lid. ~ PC
for an unobserved ground-truth regressor B_, E Rdd and label noise Cijj. We assume the distributions Px and Pc are such that [[x,,]] 0, E(x,, Ia. [[(G,]]=0 =0.

---

To incentivize representation learning, each B.,; belongs to the same k-dimensional subspace of R", where k < d. Let B. Odx have columns that form orthogonal basis for the shared subspace, so that B... ; = B.w, w.,, for some w.,i E IR for each i. In other words, there exists a low-dimensional set of parameters known as the "head" that can specify the ground-truth model for client once the shared representation, i.e., (B.( is known. It is advantageous to learn col(B.) because once it is known, all clients (including potentially new clients entering the system) have sample complexity O(k) < d as they only need to learn the parameters of their head (40, 41.
Each client i ultimately aims to learn a model 3r that approximates B=,; in order to achieve good generalization on its local distribution. To eventually achieve this for each client, FedAvg with fine-tuning first aims to learn a global model consisting of representation B E Rdxk and a head w I that minimizes the average loss across clients. The loss for client is f w) == 1.
2nc 2==1 (8k.) (Bw, x 2 i.e. the average squared loss on the local data, so FedAvg tries to learn a global model that solves the nonconvex problem:
min BERAXA WER* N
1
M
n n4 f.
R2
Ues (Bw, (Bw, [.,,)))))))
(5)
2m
where = E n2 To solve (5) in a distributed manner, FedAvg dictates that each client makes series of local updates of the current global model before returning the models to the server for averaging (see Section 2). Our aim is to show that FedAvg training learns the column space of B.. First, we make standard diversity and normalization assumptions the ground-truth heads.
Assumption (Client normalization). There exists < oo S.I. Vi E [M]. ||w.,all S LVk.
Assumption 2 (Client diversity). There exists μ 0 s.t. minn ( 37 Σ (w.. -...-w,, (w..)--.)) 1², where w. 1 M 1 W.,s. Define k. = l/p.
EM
Assumption 2 is very similar to typical task diversity assumptions except that it quantifies the diversity of the centered rather than un-centered tasks [40] 41). Intuitively, task diversity is required that all of the directions in (B.( are observed. The smaller K, the more evenly spread the ground-trun -touh heads are, and the larger the task (i.e. client) diversity. Next, to obtain convergence results must define the variance of the ground-truth heads and the principal angle distance between representations.
Definition 1 (Client variance). For O, define: 72
Ei=1 M W.,i w.1l? where w, is T 37 EM 2,
EM Giml
defined in Assumption For H > 0, define H = R.M
E
[.,, W
Definition 2 (Principal angle distance). For two matrices Bi, Bg E Raxk, the principal angle distance between B1 and B is defined as dist(B B2) = [ B2ll2, where the columns of BI,1 E ddx(d-k) and B₂ E ddxk form orthonormal bases for (B, and col(B2), respectively.
Intuitively, the principal angle distance between B and B₂ the sine of the largest angle between the subspaces spanned by their columns. Now we are ready to state our main result. We suppose each client has n,, 00 samples, i.e. it accesses the gradients of the population loss its local distribution.
Theorem 1. Consider the case that each client takes gradient steps with respect to their pop- B.Wa.ill and all losses are weighed equally in the global objective. Suppose Assumptions Z hold, the number of clients participating each round satis- fies m 2 min(M, (innM))((7//)))
ulation loss Ja(B, w) == ||
and the initial parameters satisfy (i) 5u = dist B.) < /1- Eo for any Eo E (0. (ii) |I- oB. Bollz O(a²7 L k2) and (iii) ||wll22 = 2L 1.55. Choose step size O 1-60 ( Then for any (0, 1). the
(H/L)
distance of the representation learned by FedAvg with 5 2 2 local updates satisfies dist(Br,BB B.) < after at most T = 0 ( (( log('/+)
communication rounds with probability at least 1 - 4(KT)) -99
(a'Tp-E))
Theorem D shows that FedAvg converges exponentially fast to the ground-truth representation when executed on the clients' population losses. We provide intuition for the proof in Section Aanioon/anndd and the full proof in Appendix B First, some comments are in order.
holds under benign initial conditions. particular, condition (i) requires that the initial distance is only constant smaller than 1. Condition (ii) ensures that the
Mild initial conditions. Theorem |
5

---

initial representation is well-conditioned with appropriate scaling, and (tii) guaranteess the initial head is not too large. The last two conditions can be easily achieved by normalizing the imputs.
Generalization without convergence in terms of the global loss. When each client accesses its population loss in Theorem T
the global objective (5) becomes:
M Σ
min BER WERK M
Bw B.
(6)
However, Theorem I does not imply that FedAvg solves (6). In fact, our simulations in Section show that it does not even reach a stationary point of (6). This is consistent with prior works that have noticed the "objective inconsistency" phenomenon of FedAvg: it solves an unknown objective distinct from the global objective due to the fact that after multiple local updates, local gradients are no longer unbiased estimates of gradients of (G) (9) Nevertheless, our results show that FedAvg is able to learn a generalizable model even when does not optimize the global loss data heterogeneous settings.
Multiple local updates critically harness client diversity, whereas Distributed GD (D-GD) does not learn the representation. Key to the proof of Theorem D is that the locally-updated heads become diverse, meaning that they cover all directions in R* with greater diversity corresponding to more even covering in all directions. We will show in Section that the Docall--uddated heads become roughly as diverse as the ground -ruth heads, and this causes the representation to move towards the ground-truth at rate depending on the diversity level. Theorem reffeecs this: the convergence rate improves with the diversity metric s/r., In this way FedAvg exploits data heterogeneity to learn the M implies more heterogeneous data. Moreover, since also appears in the denominator of the communication round complexity, additional local updates improve the convergence rate up to T (a -2), which is the limit imposed due the upper bound on a.
representation, as more diverse {w..) Ne|M]
Importantly, head diversity only benefits the global representation update if 2. We formally prove that D-GD (equivalent to FedAvg with r= and m = M) cannot recover col in the following result.
Proposition (Distributed GD lower bound). Suppose we are in the setting described in Sec- IM full-rank initializa- tion Bo E Rdxkk initial distance do E (0, 1/2]. step size a > 0, and number of rounds T, there exists B. E Odxk satisfying dist (B,, B.) = 5g and dist(B .B.) 2 0.78p. where BIGD = ( B., () a) is the result of D-GD with step size a and initialization Bo in the setting with ground-truth representation B, and ground-truth heads {wo,)
tion and d >> k >> 1. Then for any set of ground-truth heads (wr,
Proposition shows that for any choice of d E (0, 1/2], non-degenerate initialization Bo, and ground-truth heads, there exists B. whose column space is ôp-close to col(Bo). yet is at least 0.76g- far from the representation learned by D-GD in the setting with B. as ground-truth. Therefore, even allowing for a strong initialization, D-GD cannot guarantee to recover the ground-truth representation. This negative result combined with our previous results suggest that even if we had an infinite communication budget, it would still be advantageous to execute multiple local updates between communication rounds in order to achieve better generalization through representation learning.
Intuitions and Proof Sketch
Next we highlight the key ideas behind the importance of local updates and why FedAvg learns col(B.), while D-GD fails to achieve this goal.
Global update B To build intuition for why FedAvg can learn col (B.), we examine the global update of the representation in the full participation case (m = M):
B = B
C o.W.i, (.
+B.
II
(
awa, WT (
M
prior weight
signal weight
Notice that B,41 is a mixture of B, and B, with weight matrices in 1x We aim to show that

---

(I) the 'prior weight' on B, has spectral norm strietly less than 1, and
(II) the 'signal weight' on B, adds energy from col( to Bu+1 so that miin (B141) miin (m))
These two conditions imply that the contribution from (B,( in col(Br+1) contracts, while energy from (B. replaces the lost energy from (B,( Hence, col col(Br+1) moves to col(B.) in all directions.
The role of head diversity and multiple local updates. To show (I) and (II), it is imperative to use the diversity of the locally-updated heads when т 2 2. First consider (I), Notice that for each i, IT: (I. - -..,,,.W..... (...)) has singular values at most 1, and strictly less than 1 corresponding to directions spanned by (weas. Thus, the maximum singular value of the average of these matrices should be strictly less than as long as (eer []]] spans Rt, i.e. the locally-updated heads are diverse. Similarly, the signal weight is rank-k if the locally-updated heads span R which leads to (II) as discussed below. In contrast, if of the representation does not leverage head diversity, as it is only a function of the global head and the average ground-truth head: B = (T, aw,w,) ) result, (B can only improve in one direction, so D-GD ultimately fails to learn col( (see Proposition]
1, then the global update
aB. w. w, in this case. As
Achieving head diversity: the necessity of controlling - aB, BB B., We have discussed the intuition for why head diversity implies (I) and (II) for FedAvg. Next, we investigate why the heads become diverse. Let us examine client i's first local update for the head at round
W1,111 1 (I, aB, B.) + a.B.B...
From this equation we see that if A = - aB, B, Ro O and ||well is bounded, then We,1.1 R aB, W.,i. If this approximation holds, then (wr,..1 hee [ inherits the diversity of (ee which is indeed diverse due to Assumption 2 meaning that the local heads are diverse after just one local update. Moreover, it can be shown that if A, ~ O and the heads become diverse after one local update, then they remain diverse for all local updates due to the observation that each B,,,,. changes slowly over s. Note that in addition to implying local head diversity, A, = O for all t implies amin (B.) R miin (B++1) 1) Rs 7a which directly ensures (II). Thus we aim to show A, 0 for all communication rounds, i.e. B, remains close to scaled orthonormal matrix.
However, it is surprising why A remains small: B++11 is the average of nonlinearly locally-updated representations, and the local updates could 'overfit' by adding more energy to some columns than others. and/or lead to cancellation when summed, so it is not intuitive why vaB, remains almost orthonormal. Nor does the expression above for B.+1 provide any clarity on this. Nevertheless, through careful induction we show that A, indeed stays close to zero since the local heads converge quickly and the projection of the local representation gradient onto col( is exponentially decaying.
Inductive argument. While the above intuitions seem to simplify the behavior of FedAvg, showing that they all hold simultaneously not all obvious. To study this, we are inspired by recent work [4] that developed an inductive argument for representation learning in the context of gradient mesed-mnta-learming. To formalize our intuition discussed previously, in our proof we need to show that (i) the learned representation does not overfit to each client's 's loss despite many local updates and simultaneously the heads quickly become diverse, and (ii) the update at the global server preserves the learned representation despite averaging many nonlinearly perturbed representations gathered from clients after local updates. To address these challenges, we construct a pair of intertwined inductive hypotheses over time, one for tracking the effect of local updates, and another for tracking the global averaging. Each inductive hypothesis (local and global) itself consists of several hypotheses (in effect, nested induction) that evolve within communication rounds.
Local induction. The proof leverages the following local inductive hypotheses for every i,
1. (s) = [l aB. B B. B. 111 = 7L ?
En
As'
E (1.
2. (s) = {|[Wr,i.e 112 S 22 cavaaaax Vs' E (L,.....
3. (s) = [114 -aB, B.,,BB,,,,,, Br, c3a Li 22 Eo
As
4. (s) = (dist(Br,..o B.) es dist B.) Vs' (I....,s})
7

---

global loss (?)
10
B. me--
Norm gradiemt giobal
10"
10"
10
10""
FedAvg
FadAug D-SGD
D-GD
10
11'
15
Fine-tuning
Figure 3: (Left) D-GD converges to stationary point of the global objective O unlike FedAvg, yet (Right) FedAvg achieves smaller error after fine-tuning with various numbers of samples.
The local induction tracks the effect of updates each client node: the end of T local updates, (T) captures the diversity of the local heads, (T),,(T) ensures that the heads remain uniformly bounded, (T) shows that the locally adapted representations stay close a scaled orthonormal matrix, and (4,,,;(T) shows that the locally adapted representations do not diverge too quickly from the ground-truti -tuth. The second set inductions below controls the global behavior.
Global induction. The global induction utilizes similar set of inductive hypotheses.
1. (t) = wl aI A B B. w. 2. TL3
Vt' {I....,t}}
2. A2(t) = {l/ vaValmax At' {l....,t}}
Vt' E {I,...,t}} 3. Aa(t) == {/) A ( = c Lmax Rmax Eo
4. A4(() == B B 112 < (1 4 Eo B Br-ille vt' (I.....,J
5. As(() == (dist((., B.) < (1 - Ga2 En) WI' (1,...,..
Hypotheses (t), (t) and At(t) are analogous to (s,,,(s)) (2,,,(s) and (s),,((s,, respectively.
(t) shows that the energy of col(B,) that is orthogonal the ground-truth subspace is contracting, and As(() finally shows that the principal angle distance between the learned and ground-truth representations exponentially decreasing. Our main claim follows from As(T). However, proving this result requires showing that all the above local and global hypotheses hold for all times / 2 as these hypotheses are heavily coupled. As mentioned previously, the most difficult challenge is controlling | (A)(t)) despite many local updates, and doing so requires leveraging both local and global properties. The details of this local-global induction argument are Appendix
5 Experiments
In this section, we conduct experiments to (I) verify our theoretical results in the linear setting and (II) determine whether our established insights generalize to deep neural networks. Notably, demonstrating the competitive performance of FedAvg plus fine-tuning for personalized FL is not a goal of this section, as this evident from prior experiments (10 33 Rather, to achieve (II) we test whether FedAvg learns effective representations when trained with neural networks in heterogeneous data settings via three popular benchmarks for evaluating the quality learned representations. Since main claim is that local updates are key to representation learning, use D-(S)GD as our baseline all experiments.
5.1 Multi-task linear regression
We first experiment with the regression setting from our theory. We randomly generate B, gdck and (we.) by sampling each element i.i.d. from the normal distribution, where d= 100, à = 55 and M = 40, and then orthogonalizing B.. Then we run FedAvg with
local updates and M clients per round. We have seen in Figure [ that the principal angle distance between the representation learned by FedAvg and the ground-truth representation linearly converges zero, whereas D-GD does not learn the ground-truth representation. Conversely, Figure 3 (left) tracks the gradient of the global loss (6) and shows that D-GD linearly converges to stationary point of (6), while FedAvg does not converge to at all. Although D-GD optimizes
D-GD, both sampling

---

D-5GD [t=
FedAvg == 10
FedAvg == 30
FedAvg == 50
1.0
CIFAN-10 clsos indox
0.8
0.6
0.4
0.2
0.0
CI AR-1-10 class index
CIFAR-10 class index
CIFAR110 10 class index
CIFAR-10 dass index
Figure 4: Average cosine similarity for features learned by D-SGD and FedAvg with varying numbers of local updates on a heterogeneous partition of CIFAR-10.
the global loss, it does not generalize as well as FedAvg to new clients as demonstrated by Figure (right). Here, we fine-tune the models learned by FedAvg and D-GD on new client with samples generated by XM+1 ~ (O0 Ig), SMS+1 ~ ~V 0.01), and 9Mr+1 (B, w. w,,M+ XM+1) + CM+11 We fine-tune using GD for
(0,1)
200 iterations with batch size b n, and plot the final error WT, B, w.w MI Both plots are generated by averaging 10 runs.
5.2 Image classification with neural networks
Next we evaluate FedAvg's representation learning ability on nonlinear neural networks. For fair comparison, in every experiment all methods make the same total amount of local updates during the course of training (e.g. D-SGD trained for 50x more rounds than FedAvg with 50).
Datasets and models. We use the image classification datasets CIFAR-10 and CIFAR-100 100 [57], which consist of 10 and 100 classes of RGB images, respectively. We use convolutional neural network (CNN) with three convolutional blocks followed by a three-layer multi-layer perceptron, with each convolutional block consisting of two convolutional layers and max pooling layer.
Cosine similarity of features. A desirable property of representations for downstream classification tasks is that features of examples from the same class similar each other, while features of examples from different classes are dissimilar [58]. In Figure 4 we examine whether the represen- tations learned by FedAvg satisfy this property. Here have trained FedAvg with varying and D-SGD (FedAvg with r =1 ) on CIFAR-10. Image classes are heterogeneously allocated to M 100 clients according the Dirichlet distribution with parameter 0.6 as [59]. Each subplot is 10x10 matrix whose (i, element gives the average cosine similarity between features of images from the i-th and j-th classes learned by the corresponding model. Ideally, diagonal elements are close to (high similarity) and off-diagonal elements are close to O (low similarity). F shows that FedAvg indeed learns features with high intra-class similarity and low inter-class similarity, with representation quality improving with more local updates between communications Meanwhile, D-SGD does not learn such features. The leftmost subplot shows that all of the features learned by D-SGD are dissimilar, regardless of whether two images belong to the same class.
Fine-tuning performance. We evaluate the generalization ability of the representations learned by FedAvg to new classes and also new datasets. An effective representation identifies universally important features, should generalize new data, with perhaps small amount of fine-tuning needed to learn new mapping from feature space to label space. The transfer learning performance of fine-tuned models a popular metric for evaluating the quality of learned representations We first study how models trained by FedAvg and D-SGD generalize to unseen classes from the same dataset. To do so, train models on heterogencous partitions of CIFAR-
(60) 61).
using both FedAvg with r = 50 as well as D-SGD. In the left plot of Figure 5 we illustrate the case that models are trained on 80 clients each with 500 total images from C classes sub-selected from 80 classes of CIFAR- 100, and tested on new clients with images from the remaining 20 classes of CIFAR- 100. We fine-tune the trained models on the new clients with 10 epochs of SGD, with varying numbers of samples per epoch as listed on the x-axis, before testing. Next, we investigate how well models trained by FedAvg and D-SGD generalize to an unseen dataset. In the right plot of Figure we train models with C' classes/client from CIFAR- 100, then test on new clients with samples drawn from CIFAR- (a different dataset, but with presumably similar "basic" features). Specifically, for these new clients, we fine-tune for 10 epochs previously, then test the post-fina post-fine-tuned models on the test

---

36
Finestuiing:
Fedkug
Fackogg
FedRvg 15
Faskvg C 15
FasAvg
D-8G0 C
D 8G0 ==
36
D-8G0 C=
D-50D C 15
960 25
1030
Fino-tuning samples
Fine-tunlng samples
Figure 5: Average fine-tuning accuracies
clients for models trained by FedAvg and D-SGD.
(Left) Models trained on 80 classes from CIFAR-100 (with C' classes/client) and fine-tuned on new clients from 20 new classes from CIFAR- 100. (Right) Models trained CIFAR-100 with C classes/client and fine-tuned on new clients from CIFAR-I0 (10 classes/client), For FedAvg, T = 50 in all cases, and error bars give standard deviations over five trials with five new clients tested per trial.
data for each client. both left and right plots, we observe that FedAvg significantly outperforms D-SGD, indicating that FedAvg has learned representation that generalizes better to new classes.
6 Conclusion
We showed that FedAvg learns the ground -ruth representation in the multi-task linear regression setting. To our knowledge, this is the first theoretical study showing FedAvg learns an effective representation in any setting. Our analysis reveals that multiple local updates are critical FedAvg's representation learning ability, which is supported empirically on both linear and nonlinear models. These experimental results suggest future work can extend our findings to more complex settings.
Acknowledgements
This research supported in part by NSF Grants 2127697, 2019844, 2107037, and 21 2112471, ARO Grant W91 INF21 10226, ONR Grant N00014-19-1-2566 the Machine Learning Lab (MLL) at UT Austin, and the Wireless Networking and Communications Group (WNCG) Industrial Affiliates Program.
References
[1] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. Arcas, "Coomunication--ffiiieen learning of deep networks from decentralized data," in Artificial Intelligence and Statistics, pp. 1273-1282, PMLR, 2017.
[2] U. Stich, "Local sgd converges fast and communicates little," in International Conference on Learning Representations, 2018.
[3] Stich and Karimireddy, "The freeb-eee framework: Better rates for sgd with delayed gradients and compressed communication," arXiv preprint (at9::9999003... 2019.
[4] Wang and G. Joshi, "Cooperative sgd: A unified framework for the design and analysis of cramunication-effccieee sgd algorithms." in ICML Workshop Coding Theory for Machine Learning, 2019.
[5] P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh, "Scaffold: Stochastic controlled averaging for federated learning," in International Conference on Machine Learning, pp- 5132-511433 PMLR, 2020.
[6] G. Malinovskiy, D. Kovalev, E. Gasanov, L. Condat, and P. Richtarik, "From local sgd to local fixed-point methods for federated learning," in International Conference on Machine Learning, pp. 6692-6701, PMLR, 2020.
10

---

[7] R. Pathak and M. J. Wainwright, "Fedsplit: An algorithmic framework for fast federated optimization," Advances Neural Information Processing Systems, vol. 33, pp. 7057-7066, 2020.
[8] Z. Charles and J. Konečný, "Convergence and accuracy trade-offs in federated learning and meta- learning." in International Conference on Artificial Intelligence and Statistics, pp. PMLR, 2021.
[9] J. Wang, Q. Liu, H. Liang, G. Joshi, and H. V. Poor, "Tackling the objective inconsistency problem in heterogeneous federated optimization," " arXiv preprint vrXiv:2007.07455 2020.
[10] Yu, F. Bagdasaryan, and V. Shmatikov, "Salvaging federated learning by local adaptation," arXiv preprint arXiv:2002.04755 2020.
[11] L. Collins, H. Hassani, A. Mokhtari, and S. Shakkottai, "Exploiting shared representations for
personalized federated learning," International Conference on Machine Learning, pp. 2089- 2099, PMLR, 2021.
[12] T. Li, S. Hu, A. Beirami, and V. Smith, "Ditto: Fair and robust federated learning through personalization," arXiv: 2012.04221, 2020.
[13] Kornblith, M. Norouzi, Lee, and G. Hinton, "Similarity of neural network representations revisited," in International Conference on Machine Learning, pp. 3519-3529, PMLR, 2019.
[14] A. Khaled, K. Mishchenko, and P. Richtárik, "Tighter theory for local sgd on identical and het- erogeneous data," in International Conference on Artificial Intelligence and Statistics, pp. 4519- 4529, PMLR, 2020.
[15] A. Koloskova, N. Loizou, S. Boreiri, M. Jaggi, and S. Stich, "A unified theory of decentralized sgd with changing topology and local updates," in International Conference on Machine Learning, pp. 5381-5393, PMLR, 2020.
[16] X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang, "On the convergence of fedavg on non-iid data," " in International Conference on Learning Representations, 2019.
[17] Z.. Qu, K. Lin, Kalagnanam, Li, Zhou, and Z., Zhou, "Federated learning's blessing: Fedavg has linear speedup," arXiv preprint arXiv:2007.05692 2020.
[18] B. E. Woodworth, K. K. Patel, and N. Srebro, "Minibatch vs local sgd for heterogeneous distributed learning," " Advances in Neural Information Processing Systems, vol. 33, pp. 6281 6292, 2020.
[19] A. Mitra, R. Jaafar, G. J. Pappas, and H. Hassani, "Linear convergence in federated learning: Tackling client heterogeneity and sparse gradients," Advances in Neural Information Processing Systems, vol. 34, pp. 14606-14619, 2021.
[20] K. Wang, R. Mathews, C. Kiddon, H. Eichner, F. Beaufays, and D. Ramage, "Federated evaluation of on-device personalization." arXiv preprint arXiv:1910.1 10252, 2019.
[21] Haddadpour, M. M. Kamani, A. Mokhtari, and M. Mahdavi, "Federated learning with compression: Unified analysis and sharp guarantees," in International Conference on Artificial Intelligence and Statistics, pp. 2350-2358, PMLR, 2021.
[22] X. Liang, S. Shen, J. Liu, Pan, E. Chen, and Y. Cheng, "Variance reduced local sgd with lower communication complexity," arXiv preprint arXiv: 1912. 12844, 2019.
[23] T. Murata and T. Suzuki, "Bias-variance reduced local sgd for less heterogeneous federated learning," " in International Conference on Machine Learning, pp. 7872-7881, PMLR, 2021.
[24] E. Gorbunov, F. Hanzely, and Richtárík, "Local sgd: Unified theory and new efficient methods." in International Conference Artificial Intelligence and Statistics, pp. 3556-3564, PMLR, 2021.
[25] C. Xi and U. A. Khan, "Dextra: A fast algorithm for optimization over directed graphs," IEEE Transactions on Automatic Control, vol. 62, no. 10, pp. 4980-4993, 2017.

---

[26] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith, "Federated optimization in heterogeneous networks," arXiv preprint arXiv:1812.06127 2018.
[27] F. Zhang, K. Kuang, Z. You, Shen, J. Xiao, Y. Zhang, C. Wu, Y. Zhuang, and X. Li, "Federated unsupervised representation learning," " arXiv preprint arXiv:2010.08982, 2020.
[28] C. T Dinh, N. Tran, and J. Nguyen, "Personalized federated learning with moreau envelopes," Advances in Neural Information Processing Systems, vol. 33, pp. 21394-21405, 2020.
[29] M. Ribero and H. Vikalo, "Communication-efiee federated learning via optimal client sampling," arXiv preprint arXiv:2007.15197, 2020.
[30] W. Chen, S. Horvath, and P. Richtarik, "Optimal client sampling for federated learning," arXiv preprint 2XXi:20010013223 2020.
[31] Y. J. Cho, J. Wang, and G. Joshi, "Client selection federated learning: Convergence analysis and power-of-choice selection strategies," arXiv preprint arX1v:2010.01243, 2020.
[32] S. Chen, Q. Zheng, Q. Long, and W. J. Su, "A theorem of the alternative for personalized federated learning." arXiv preprint arXiv:2103.01901, 2021.
[33] G. Cheng, K. Chadha, and J. Duchi, "Fine-tuning is fine in federated learning," arXiv preprint arXiv:2108.07313. 2021.
[34] C. Finn, P. Abbeel, and S. Levine, "Model-agnostic meta-learning for fast adaptation of deep networks," in Proceedings of the 34th International Conference on Machine Learning-Voil Volume 70, pp. 1126-1135, JMLR. org, 2017.
[35] A. Fallah, A. Mokhtari, and A. Ozdaglar, "Personalized federated learning: A meta-learning approach," 2020.
[36] L. Su, J. Xu, and P. Yang. "Achieving statistical optimality of federated learning: Beyond stationary points," arXiv preprint 1xXiv:2106.152/6 2021.
[37] A. Nichol and J. Schulman, "Reptile: a scalable metalearning algorithm," " arXiv preprint 18X.:91003.02999, vol. 2. 2, 2018.
[38] B. Huang, X. Li, 7. Song, and X. Yang. "Fl-ntk: neural tangent kernel-based framework for federated learning convergence analysis," arXiv preprint arXiv:2 105.05000, 2021.
[39] A. Maurer, M. Pontil, and B. Romera-Paredes, "The benefit of multitask representation learning." " Journal of Machine Learning Research, vol. 17, no. 81, pp. 1-32, 2016.
[40] S. S. Du. W. Hu, S. M. Kakude, J. D. Lee, and Q. Lei, "Few-shot learning via learning the representation, provably," in International Conference on Learning Representations, 2020.
[41] N. Tripuraneni, C. Jin, and M. Jordan, "Provable meta-learning of linear representations," in International Conference on Machine Learning, pp. 10434-10443, PMI.R, 2021.
[42] W. Kong, R. Somani, S. Kakade, and S. Oh, "Robust meta-learning for mixed linear regression with small batches," Advances in neural information processing systems, vol. 33, pp. 4683-4696, 2020.
[43] K. K. Thekumparampil, P. Jain, P. Netrapalli, and S. Oh, "Statistically and computationally efficient linear meta- rete-representttionn learning," Advances in Neural Information Processing Systems, vol. 34, 2021.
[44] L. Collins, A. Mokhtari, S. Oh, and S. Shakkottai, "Maml and anil provably learn representa- tions." arXiv preprint arXiv:2202.03488 2022.
[45] Y. Sun, A. Narang, I. Gulluk, S. Oymak, and M. Fazel, "Towards sample-efficient overpa- rameterized meta-learning." Advances in Neural Information Processing Systems, vol. 34, 2021.
12

---

[46] P. Jain, J. Rush, A. Smith, S. Song, and A. G. Thakurta, "Differentially private model personal- ization," 2021.
[47] P. Liang, T. Liu, L. Ziyin, R. Salakhutdinov, and L.-P. Morency, "Think locally, act globally: Federated learning with local and global representations." arXiv preprint arXiv:2001.011 01523, 2020.
[48] M. G. Arivazhagan, V. Aggarwal, A. K. Singh, and S. Choudhary, "Federated learning with personalization layers." arXiv preprint arXiv: 1912.00818, 2019.
[49] Oh, S. Kim, and S.-Y. Yun, "Fedbabu: Towards enhanced representation for federated image classification," arXiv preprint arXiv:2106.0600 2021.
[50] S.-J. Hahn, M. Jeong, and J. Lee, "Subspace learning for personalized federated optimization," arXiv preprint arXiv: 109.07629. 2021.
[51] W. Zhuang. X. Gan, Y. Wen, S. Zhang, and S. Yi, "Collaborative unsupervised visual repre- sentation learning from decentralized data," in Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4912-49211 2021.
[52] L. Nguyen, P. H. Nguyen, M. Dijk, P. Richtárik, K. Scheinberg, and M. Takác, "Sgd and hogwild! convergence without the bounded gradients assumption,"" in International Conference on Machine Learning, pp. 3750-3758, PMLR, 2018.
[53] O. Shamir and N. Srebro, "Distributed stochastic optimization and learning." in 2014 52nd Annual Allerton Conference on Communication, Control, and Computing (Allerton), pp. 850- 857, IEEE, 2014.
[54] R. M. Gower, N. Loizou, X. Qian, A. Sailanbayev, E. Shulgin, and P. Richtárik, "Sgd: General analysis and improved rates," in International Conference on Machine Learning, pp. 5200-5209, PMLR, 2019.
[55] K. K. Thekumparampil, P. Jain, P. Netrapalli, and S. Oh, "Sample efficient linear meta-learning by alternating minimization," arXiv preprint crXiv:2105.08366 2021.
[56] K. Chua, Q. Lei, and J. D. Lee, "How fine-tuning allows for effective meta-learning," Advances Neural Information Processing Systems, vol. 34, 2021.
[57] A. Krizhevsky, G. Hinton, el al., "Learning multiple layers of features from tiny images." Cileseer, 2009.
[58] R. Hadsell, S. Chopra, and Y. LeCun, "Dimensionality reduction by learning an invariant map- ping." " in 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'06), vol. 2. pp. 1735-1742, IEEE, 2006.
[59] D. A. E. Acar, Y. Zhao, Matas, M. Mattina, P. Whatmough, and V. Saligrama, "Federated learning based on dynamic regularization," in International Conference on Learning Represen- tations, 2021.
[60] W. F. Whitney, M. J. Song, D. Brandfonbrener, J. Altosaar, and K. Cho, "Evaluating represen- tations by the complexity of learning low-loss predictors," arXiv preprint arXiv:2009.07368, 2020.
[61] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton, "A simple framework for contrastive learning of visual representations," in International conference on machine learning, pp. 1597-1607, PMLR, 2020.
[62] F. Haddadpour, M. M. Kamani, M. Mahdavi, and V. Cadambe, "Local sgd with periodic averaging: Tighter analysis and adaptive synchronization." Advances Neural Information Processing Systems, vol. 32, 2019.
[63] H. Yu, S. Yang, and S. Zhu, "Parallel restarted sgd with faster convergence and less communica- tion: Demystifying why model averaging works for deep learning." " in Proceedings of the AAAI Conference Artificial Intelligence, vol. 33, pp. 5693-5700, 2019.
13

---

[64] A. Spiridonoff, A. Olshevsky, and I. Paschalidis, CCommuniation-efficcee sgd: From local sgd one-shot averaging." Advances Neural Information Processing Systems, vol. 34, 2021.
[65] B. Woodworth, K. K. Patel, S. Stich, Z. Dai, B. Bullins, B. Mcmahan, O. Shamir, and N. Srebro, "Is local sgd better than minibatch sgd?," in International Conference on Machine Learning. pp. 10334-10343, PMLR, 2020.
[66] Zhou and G. Cong, "On the convergence properties of a k-step averaging stochastic gradient descent algorithm for nonconvex optimization." arXiv preprint arXiv:/708.01012. 2017.
[67] Z. Charles and Konečný, "On the outsized importance of learning rates in local update methods," arXiv preprint www.::000700078,, 2020.
[68] Wang, Z. Xu, Z. Garrett, Z. Charles, Liu, and Joshi, "Local adaptivity in federated learning: Convergence and consistency," arXiv preprint arXiv:2106.02305, 2021.
[69] A. Nedié and Olshevsky, "Distributed optimization over uine-varying directed graphs," IEEE Transactions on Automatic Control, vol. 60, no. 3, pp. 601-615, 2014.
[70] M. Assran, N. Loizou, N. Ballas, and M. Rabbat, "Stochastic gradient push for distributed deep learning," in International Conference on Machine Learning, pp. 344-353, PMLR, 2019.
[71] M. Assran and M. G. Rabbat, "Asynchronous gradient push." IEEE Transactions on Automatic Control, vol. 66, no. 1, pp. 168-183, 2020.
[72] Q. Li, B. He, and D. Song, "Model-contrastive federated learning." in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10713-10722, 2021.
[73] N. Tripuraneni, M. Jordan, and C. Jin, "On the theory of transfer learning: The importance of task diversity," Advances in Neural Information Processing Systems, vol. 33, pp. 7852-7862, 2020.
[74] Z. Xu and A. Tewari, "Representation learning beyond linear prediction functions," Advances in Neural Information Processing Systems, vol. 34, 2021.
[75] D. Gross and V. Nesme, "Note on sampling without replacing from finite collection of matrices," arXiv preprint arXiv: 1001.2738, 2010.
[76] A. Subramanian, "torch-cka."
2021.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] Our theoretical results only hold for the multi-task linear setting. This point is explicitly mentioned in the abstract, Introduction (Section []. and Main Results (Section 3
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] Please see the Main Results (Section 3)
(b) Did you include complete proofs of all theoretical results? [Yes] Please see Appendix
3. If you experiments...
14

---

(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or as a URL)? [Yes] Please see Appendix [add the supplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Please see Appendix [
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes] Please see Figures 2
3 and 5
(d) Did you include the total amount compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Please see Appendix [
4. If you are using existing assets (e.g., code, data, models) or curating/releasnng new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or URL? [No]
(d) Did you discuss whether and how consent obtained from people whose data you're re using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsoureing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "gGsHugzF0Bk", "reviewer": "Reviewer_1GSs", "summary": "This paper analyzes both theoretically and empirically the performances of FedAvg with local fine-tuning as a post-processing step. More precisely, they consider the scenario where the model parameters are split into two distinct parts: a shared backbone model and a personalized head model. This scenario encompasses for instance convolutional neural networks where the backbone model extract high-level features and the head performs prediction. In order to derive theoretical results, the authors restrict their analysis to the linear case where the backbone parameter is a matrix $B \\in \\mathbb{R}^{d \\times k}$ with $k < d$ and $\\omega$ is the head.\nTheir main result is a Theorem showing the generalization power of FedAvg with at least 2 local updates via convergence towards the ground-truth backbone parameter $B_\\star$. They validate their theory with experiments.", "questions": "See the previous remarks and concerns.", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, the paper is well-written and clear. The related work section acknowledging previous works is sufficient.\n\nMy main concerns / remarks are:\n\n* The work by Collins et al. (ICML 2021) on shared representations should be compared more deeply with the proposed analysis since the two setting bear some similarity. Some comparison has been done in the supplement but should appear in the main paper in my opinion.\n* How difficult the analysis beyond the linear case is?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Yc4TRwxd5WA", "reviewer": "Reviewer_rvkq", "summary": "This paper provides theoretical analysis of FedAvg with Fine Tuning. When the local model weights for clients lie on a linear subspace, theorem 1 suggests FedAvg gets linear convergence, while proposition 1 suggests Distributed GD may diverge for representation learning. Some empirical results on a toy problem are used to verify the theorem, and results on CIFAR-10 and CIFAR-100 shows FedAvg+FineTune outperforms D-SGD+FineTune.", "questions": "Like I mentioned,  I find it difficult to connect Theorem 1 to the FedAvg algorithm described in Section 2 and the objective eq. (5).  Specifically, what is the connection between (5) and (7), why do we care about  \\|Bw − B∗w∗\\|^2 in (7)? \n\nCould the authors clarify what assumptions of FedAvg, and what assumptions of representation learning are made? \n\nRelated to previous questions, it is not quite clear to me what is the difference between a FL setting and a multi-task/transfer learning setting.  \n\nCould the authors discuss and compare to Adaptive Gradient-Based Meta-Learning Methods https://arxiv.org/pdf/1906.02717.pdf ? I am not an expert on meta-learning/representation learning theory, and this is only one of the papers that also discuss FedAvg. \n", "limitations": "NA", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n+ FedAvg is one of the most popular algorithms in FL. Though the connection between FedAvg and MAML/Reptile is known, the theoretical justification of the advantage of FedAvg for representation is (to my knowledge) novel. \n\nWeakness\n- The paper can probably improve on clarity. I find it difficult to connect Theorem 1 to the FedAvg algorithm described in Section 2 and the objective eq. (5). \n- As the authors acknowledged, previous work has drawn connections between FedAvg and MAML, and empirically shown FedAvg+FineTune is better than SGD+FineTune. The main contribution of this paper is theoretical, while the assumptions are very strong. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "QzL7bgaBdTj", "reviewer": "Reviewer_Qkio", "summary": "This work shows that FedAvg learns better representations than Distributed SGD; with a largely theoretical contribution, and some well executed small-scale empirical results. This paper:\n* considers multi-task linear regression in a non-convex setting (i.e., prediction is the composition of a learnable subspace matrix, and a linear head), where the shared component across tasks lives in a low-dimensional subspace. The works shows that FedAvg can regress to the correct concept (given a good initialization), because local client-updates in FedAvg encourage diverse high-rank representations, whereas Distributed SGD finds low-rank representations.\n* provide empirical experiments (both linear regression and CNN-based image classification), showing that FedAvg discovers better representations, that generalize better than Distributed SGD.\n", "questions": "* In related work, please discuss the relation to work in the gossip literature. In general, gossip-algorithm can be seen as a a generalization of Local-GD; e.g., with diminishing step-sizes, Local-SGD with an arbitrary number of local steps converges to a stationary point of the global objective [1,2]. If clients perform a different, and potentially time-varying, number of local-steps, then from [3] we can also see that they will minimize a weighted average of the global objective, and that this can be corrected with knowledge of agents’ local update rates. Moreover, gradient tracking was actually first proposed in the gossip literature (e.g., see [4])\n\n* Error in line 177; wbar_{*,i) with a summation over i^\\prime... I believe here you should just have wbar_* without a summation\n\n* On lines 205: if you analyze the sequence produced immediately after each synchronization step, and consider a diminishing step-size, this will converge (e.g., this is can be seen as a special case of Subgradient-Push, with a particular mixing matrix over a \\tau-strongly connected graph). The intuition is that the mixing converges at an R-linear rate (much faster than the sublinear rate of optimization with a diminishing step-size).\n\n* It is not clear to me how the method can discover the correct subspace, but not converge to a stationary point of the loss… e.g., after finding the correct subspace, fix B and optimize for w… this is a convex problem. I believe figure 3 supports this intuition, (i.e., the evolution of the loss with and without fine-tuning the head); please clarify this in the exposition.\n\n* From the analysis, one implicit requirement for the diversity of the discovered representation in the proof is that the clients do not overfit to their local objectives, as this would remove energy from other elements in the column-space, not aligned with the local objective, what is the role of \\tau in this case? is there an optimal range for the number of local steps?\n\n[1] Nedic and Olshevsky, Distributed optimization over time-varying directed graphs, IEEE TAC 2014.\n[2] Assran et al., Stochastic Gradient Push for Distributed Deep Learning, ICML 2019.\n[3] Assran and Rabbat, Asynchronous Gradient Push, IEEE TAC 2021.\n[4] Xi and Khan, DEXTRA: A fast algorithm for optimization over directed graphs, IEEE TAC 2017.", "limitations": "Some minor limitations on the analysis should be stated; e.g.,\n* The second term in the min condition in theorem 1 can easily be vacuous for any reasonable number of iterations, and thus the theorem requires all clients to participate in each round… this should be made clear in a remark following the theorem. Numerical results in section 5.1 also have all clients participate in each round.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality: Despite the growing literature of work on the role of federated learning in representation learning, this work is very interesting and quite timely. Rather than proposing a new method, the authors provide an analysis explaining why FedAvg can produce more generalizable representations than Distributed SGD.\n\nQuality: Despite minor errors (probably typos in the math), this work is sound.\n\nClarity: The exposition is clear, although the proof sketches are a little obscure; and cannot be as readily comprehended without going through the proofs themselves.\n\nI am happy to recommend this work for acceptance.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "G-_7VFsMQP3", "reviewer": "Reviewer_EBRg", "summary": "Motivated by the empirical success of Federated Averaging to generalize to new tasks (after fine-tuning), this paper seeks to connect this empirical success with a theoretical argument for how well FedAvg is capable of learning good representations, which is not an explicit objective of the optimization. Theoretical guarantees are presented for a multi-task linear regression setting and further empirical results demonstrate the effectiveness of learning representations with image classification tasks.", "questions": "While much of the empirical success of FedAvg was discussed in the paper, I would appreciate a discussion of related work that specifically addresses the representation learning aspect [1,2] or extensions that aim to improve the personalization aspect [3] and how they fit with the assumptions and insights provided by this paper's theoretical analysis.\n\n[1]: Zhuang, W., Gan, X., Wen, Y., Zhang, S., & Yi, S. (2021). Collaborative unsupervised visual representation learning from decentralized data. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 4912-4921).\n\n[2]: Li, Q., He, B., & Song, D. (2021). Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 10713-10722).\n\n[3]: Hahn, S. J., Jeong, M., & Lee, J. (2021). Connecting Low-Loss Subspace for Personalized Federated Learning. arXiv preprint arXiv:2109.07628.", "limitations": "The paper adequately addresses its limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- The paper is extremely well written and well organized. The mathematical notation is clearly described and easy to follow. \n- The theoretical results are well presented, with clear assumptions and, at each steps, the assumptions and statements are described in an intuitive way.\n- While the result under study (effectiveness of FedAvg for Representation Learning) has been shown (empirically) before, the formalization is an original contribution\n- The experiments are well described and effective at supporting the main theoretical claims\n\n### Weaknesses\n- The paper ends abruptly. It would be nice to have a conclusion section to tie up the paper. Perhaps some of the problem setting and notation could be abbreviated or moved to the appendix.\n- The main result is proved in simplified setting (linear regression tasks), which is far from the usual application of FedAvg for deep learning", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept