FedAvg with Fine Tuning: Local Updates Lead to Representation Learning
Anonymous Author(s)
AfﬁliationAddress email
Abstract
The Federated Averaging (FedAvg) algorithm, which consists of alternating between a few local stochastic gradient updates at client nodes, followed by a modelaveraging update at the server, is perhaps the most commonly used method in
Federated Learning.
Notwithstanding its simplicity, several empirical studies haveillustrated that the output model of FedAvg, after a few ﬁne-tuning steps, leads to amodel that generalizes well to new unseen tasks.
This surprising performance ofsuch a simple method, however, is not fully understood from a theoretical pointof view.
In this paper, we formally investigate this phenomenon in the multi-tasklinear representation setting.
We show that the reason behind generalizability of the
FedAvg’s output is its power in learning the common data representation among theclients’ tasks, by leveraging the diversity among client data distributions via localupdates.
We formally establish the iteration complexity required by the clientsfor proving such result in the setting where the underlying shared representationis a linear map.
To the best of our knowledge, this is the ﬁrst such result for anysetting.
We also provide empirical evidence demonstrating FedAvg’s representationlearning ability in federated image classiﬁcation with heterogeneous data.
Introduction
Federated Learning (FL) [1] provides a communication-efﬁcient and privacy preserving means tolearn from data distributed across clients such as cell phones, autonomous vehicles, and hospitals.
FL aims for each client to beneﬁt from collaborating in the learning process without sacriﬁcingdata privacy or paying a substantial communication cost.
Federated Averaging (FedAvg) [1] isthe predominant FL algorithm.
In FedAvg, also known as Local SGD [2–4], the clients achievecommunication efﬁciency by making multiple local updates of a shared global model before sendingthe result to the server which averages the locally updated models to compute the next global model.
FedAvg is motivated by settings with homogeneous data across clients, since multiple local updatesshould improve model performance on all other clients’ data when their data is similar.
In contrast,
FedAvg faces two major challenges in more realistic heterogeneous data settings: learning a singleglobal model may not necessarily yield good performance for each individual client, and, multiplelocal updates may cause the FedAvg updates to drift away from solutions of the global objective[5–9].
Despite these challenges, several empirical studies [10–12] have observed that this sharedglobal model trained by FedAvg with several local updates per round when further ﬁne-tuned forindividual clients is surprisingly effective in heterogeneous FL settings.
These studies motivate us toexplore the impact of local updates on post-ﬁne-tuning performance.
Meanwhile, a large number of recent works have shown that representation learning is a powerfulparadigm for attaining high performance in multi-task 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
[Figure]
[Figure]
as a whole are heterogeneous.
Consider, for example, heterogeneous federated image classiﬁcationin which each client (task) may have images of different types of animals.
It is safe to assume theimages share a small number of features, such as body shape and color, which admit a simple andaccurate mapping from feature space to label space.
Since the number of important features is muchsmaller than the dimension of the data, knowing these features greatly simpliﬁes each client’s task.
To explore the connection between local updates and representation learning, we ﬁrst study multi-tasklinear regression sharing a common ground-truth representation (Figure 1).
We observe that FedAvgconverges (exponentially fast) to the ground-truth representation in principal angle distance, while
Distributed-GD (D-GD), which is effectively FedAvg with one local gradient update, fails to learnthe shared representation.
A similar concept can be shown in the nonlinear setting.
We study amulti-layer CNN on a heterogeneous partition of CIFAR-100 (Figure 2).
Since there is not necessarilya ground-truth model here, we evaluate representation learning as follows.
We ﬁrst train the modelswith FedAvg and Distributed-SGD (D-SGD) then ﬁne-tune the pre-trained models on clients from anew dataset, CIFAR-10.
Finally we evaluate the quality of the learned representation by measuringthe amount that each model layer changes during ﬁne-tuning using CKA similarity [13].
Observe thatthe early layers of FedAvg’s pre-trained model (corresponding to the representation) change muchless than those of D-SGD.
More details for both experiments are in Section 5 and Appendix C. Theseobservations suggest that FedAvg learns a shared representation that generalizes to new clients, evenwhen trained in a heterogeneous setting.
Hence, a 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-truth representationin the case of multi-task linear regression.
Critically, we show that FedAvg’s local updates leveragethe diversity among client data distributions to learn their common representation.
This is surprisingbecause FedAvg is a general-purpose algorithm not designed for representation learning.
Our analysisthus yields new insights on how FedAvg ﬁnds 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 R d ,where k (cid:28) d .
Our results show that FedAvg with τ ≥ 2 local updates learns the representationat a linear rate when each client accesses population gradients.
To the best of our knowledge,this is the ﬁrst 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 k directions in the linear setting.
Incontrast, we prove that D-GD (FedAvg with one local update) 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 real-worlddata sets.
This conﬁrms the main massage of our theoretical results in the bilinear setting andsuggests that our results can generalize to more complex scenarios.
Related work.
Recently there has been a surge of interest motivated by FL in analyzing FedAvg/Local SGD in heterogeneous settings.
Multiple works have shown that FedAvg converges to a globaloptimum (resp. stationary point) of the global objective in convex (resp.
nonconvex) settings butwith decaying learning rate [14–17, 5], leading to sublinear rates and communication complexitysometimes dominated by Distributed-SGD [18].
These results are tight in the sense that FedAvg withﬁxed learning rate may not converge to a stationary point of the global objective in the presence ofdata heterogeneity, as its multiple local updates cause it to optimize a distinct, unknown objective[6, 19, 7–9, 20, 16].
Several methods have tried to correct this objective inconsistency via gradienttracking [5, 19, 21–25, 25], local regularization [26, 20, 27, 28], operator splitting [7], and strategicclient sampling [29–31].
In contrast, we show that local updates with ﬁxed learning rate beneﬁtlearning in heterogeneous settings by resulting in linear convergence to generalizable models.
Several papers have also studied FedAvg from a generalization perspective.
It was shown in [32]that in a setting with strongly convex losses, either local training or FedAvg with ﬁne-tuning (butnot both) achieves minimax risk, depending on the level of data heterogeneity.
Similarly, [33]argued that FedAvg with ﬁne-tuning generalizes as well as more sophisticated methods, includingmodel-agnostic meta-learning (MAML) [34, 35], in a strongly convex regularized linear regressionsetting.
Additional work has studied the generalization of FedAvg in kernel regression, but for convexobjectives that do not allow for representation learning [36], and the generalization of a 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 inrecent years as it is an expressive but tractable nonconvex setting for studying the sample-complexitybeneﬁts of learning representations and the representation learning abilities of popular algorithms indata heterogeneous settings [40–47].
Remarkably, our study of FedAvg reveals that it can learn aneffective representation even though it was not designed for this goal, unlike a variety of personalized
FL methods speciﬁcally tailored for representation learning [40, 48–52].
Notations.
We use N ( u , Σ ) to signify the multivariate Gaussian distribution with mean u andcovariance Σ .
O d × k denotes the set of matrices in R d × k with orthonormal columns.
The notationcol( B ) represents the column space of the matrix B , and col( B ) ⊥ is the orthogonal complement tothis space.
The norm (cid:107) · (cid:107) is the spectral norm and I d is the identity matrix in R d × d .
We use [ m ] toindicate the set of natural numbers up to and including m .
Problem Formulation
Consider a federated setting with a central server and M clients.
Each client i ∈
[ M ] has a trainingdataset ˆ D i of n i labeled samples drawn from a distribution D i over X × Y , where X is the inputspace and Y is the label space.
The learning model is given by h θ : X → Y for model parametersθ ∈ R D .
The loss of the model on a sample ( x , y ) ∈
X × Y is given by (cid:96) ( h θ ( x ) , y ) , which may be,for example, the squared or cross entropy loss.
The loss of model parameters θ on the i -th client isthe average loss of the model h θ on the samples in ˆ D i , namely f i ( θ ) :
= 1 n i (cid:80) n i j =1 (cid:96) ( h θ ( x i,j ) , y i,j ) ,where ( x i,j , y i,j ) is the j -th sample in ˆ D i .
The server aims to leverage all of the data across clientsto ﬁnd models that achieve small loss
f i ( θ ) for each client.
To do so, the standard approach is to ﬁnda single model θ that minimizes the average of the client losses weighted by number of samples:
[Equation]
where N = (cid:80) Mi =1 n i .
Due to communication and privacy constraints, the clients cannot share their 116 local data ˆ D i , so (1) must be solved in a federated manner.
FedAvg.
The most common FL method is FedAvg.
On each round t of FedAvg, the server uniformlysamples a set I t of m ≤ M clients.
Each selected client receives the current global parameters θ t ,executes multiple SGD steps on its local data starting from θ t , then sends the result back to the server.
The server then computes θ t +1 as the weighted average of the updates.
Speciﬁcally, upon receivingthe global model θ t , client i computes
[Equation]
for s = 1 , . . .
, τ − 1 , where τ is the number of local steps, θ t,i,
= θ t and g t,i,s ( θ t,i,s ) is a stochasticgradient of f i evaluated at θ t,i,s using b samples from ˆ D i .
The client then sends θ t,i,τ back to the server, which computes the next global iterate as:
[Equation]
where N t := (cid:80) i ∈I t n i .
Note that τ = 1 corresponds to D-SGD, also known as mini-batch SGDwhose convergence properties are well-understood [18, 53–55].
FedAvg improves the communicationefﬁciency of D-SGD by making τ ≥ 2 local updates between communication rounds.
Fine-tuning.
After training for T communication rounds, the global parameters θ T learned by
FedAvg are typically ﬁne-tuned on each client before testing.
In particular, starting from θ T , client iexecutes τ (cid:48) steps of SGD on its local data as follows:
[Equation]
for s ∈ [ τ (cid:48) − 1] .
The ﬁne-tuned model ultimately used for testing is θ T,i,τ (cid:48) .
Note that a new client,indexed by M + 1 , entering the system after FedAvg training has completed can also ﬁne-tune θ Tusing the same procedure to obtain a personalized solution θ T,M +1 ,τ (cid:48) .
Representation learning.
We aim to answer why the ﬁne-tuned models { θ T,i,τ (cid:48) }
M +1 i =1 perform wellin practice by taking a representation learning perspective.
We show that the output of FedAvg, i.e.,θ T , has learned the common data representation among clients.
To formalize this result, we considera class of models that can be written as the composition of a representation h rep and a predictionmodule, i.e. head, denoted as h head .
Let the model parameters be split as θ :=
[ φ , ψ ] , where φcontains the representation parameters and ψ contains the head parameters.
Then, for any x ∈ X , theprediction of the learning model is h θ ( x ) =
( h head ψ ◦ h rep φ )( x ) =
h head ψ ( h rep φ ( x )) .
For instance, if h θ isa neural network with weights θ , then h rep φ is the ﬁrst many layers of the network with weights φ ,and h head ψ is the network last few layers with weights ψ .
A standard assumption in multi-task settingsis existence of a common representation h rep φ ∗ that admits an easily learnable head h rep ψ ∗ ,i such thath head ψ ∗ ,i ◦ h rep φ ∗ performs well for task i .
It is thus of interest to all the clients to learn h rep φ ∗ .
Main Results
We employ the standard setting used for algorithmic representation learning analysis: multi-tasklinear regression [42, 56, 40, 57, 47].
In this setting, samples ( x i,j , y i,j ) for each client i are drawnindependently from a distribution D i on R d × R such that [Equation] ζ
[Equation]
j where ζ i,j i.i.d. ∼ p ζ for an unobserved ground-truth regressor β ∗ ,i ∈ R d and label noise ζ i,j .
We assume the distributionsp x and p ζ are such that E [ x i,j ]= 0 , E
[ x i,j x (cid:62) i,j ]= I d , E [ ζ i,j ]=0 .
To incentivize representation learning, each β ∗ ,i belongs to the same k -dimensional subspace of R d ,where k (cid:28) d .
Let B ∗ ∈
O d × k have columns that form an orthogonal basis for the shared subspace,so that β ∗ ,i =
B ∗ w ∗ ,
i for some w ∗ ,i ∈
R k for each i .
In other words, there exists a low-dimensionalset of parameters known as the “head” that can specify the ground-truth model for client i once theshared representation, i.e., col( B ∗ ) , is known.
It is advantageous to learn col( B ∗ ) because once it isknown, all clients (including potentially new clients entering the system) have sample complexity
O ( k ) (cid:28) d as they only need to learn the parameters of their head [42, 41].
Each client i ultimately aims to learn a model
ˆ β i that approximates β ∗ ,i in order to achieve goodgeneralization on its local distribution.
To eventually achieve this for each client, FedAvg withﬁne-tuning ﬁrst aims to learn a global model consisting of a representation B ∈ R d × k and ahead w ∈ R k that minimizes the average loss across clients.
The loss for client i is f i ( B , w ) :
=1 2 n i (cid:80) n i {+j =
+}
( y i,j − (cid:104)
Bw , x i,j (cid:105) ) 2 , i.e. the average squared loss on the local data, so FedAvg tries tolearn a global model that solves the nonconvex problem:
[Equation]
where N = (cid:80) Mi =1 n i . To solve (5) in a distributed manner, FedAvg dictates that each client makesa series of local updates of the current global model before returning the models to the server foraveraging (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 on the ground-truth heads.
Assumption 1 (Client normalization) .
There exists L max
< ∞ s.t. ∀ i ∈
[ M ] , (cid:107) w ∗ ,i (cid:107) 2 ≤ L max .
Assumption 2 (Client diversity) .
There exists µ> 0 s.t.
σ min ( 1 M (cid:80) Mi
=1 ( w ∗ ,i − ¯w ∗ )( w ∗ ,i − ¯w ∗ ) (cid:62) ) ≥µ 2 , where ¯w ∗ := 1 M (cid:80)
Mi =1 w ∗ ,i .
Deﬁne κ max :=
L max / µ .
Assumption 2 is very similar to typical task diversity assumptions except that it quantiﬁes the diversityof the centered rather than un-centered tasks [41, 42].
Intuitively, task diversity is required so that allof the directions in col( B ∗ ) are observed.
Next, to obtain convergence results we must deﬁne thevariance of the ground-truth heads and the principal angle distance between representations.
Deﬁnition 1 (Client variance) .
For γ > 0 , deﬁne:
γ 2 := 1 M (cid:80) Mi =1 (cid:107) w ∗ ,i − ¯w ∗ (cid:107) 2 , where ¯w ∗ isdeﬁned in Assumption 2.
For H > 0 , deﬁne H 4 := 1 M (cid:80)
Mi =1 (cid:107) w ∗ ,i w (cid:62)∗ ,i − 1 M (cid:80) Mi (cid:48)
=1 w ∗ ,i w (cid:62)∗ ,i (cid:107) 2 .
Deﬁnition 2 (Principal angle distance) .
For two matrices B 1 , B 2 ∈ R d × k , the principal angledistance between B 1 and B 2 is deﬁned as dist( B 1 , B 2 ) :
= (cid:107) ¯B (cid:62) 1 , ⊥ ¯B 2 (cid:107) 2 , where the columns of¯ B 1 ,
⊥ ∈
O d × ( d − k ) and ¯ B 2 ∈ O d × k form orthonormal bases for col( B 1 ) ⊥ and col( B 2 ) , respectively.
Intuitively, the principal angle distance between B 1 and B 2 is the sine of the largest angle betweenthe subspaces spanned by their columns.
Now we are ready to state our main result.
We suppose eachclient has n i = ∞ samples, i.e. it accesses the gradients of the population loss on its local distribution.
Theorem 1.
Consider the case that each client takes gradient steps with respect to their population loss
f i ( B , w ) :
= 12 (cid:107)
Bw − B ∗ w ∗ ,i (cid:107) 2 and all losses are weighed equally in the globalobjective.
Suppose Assumptions 1-2 hold, the number of clients participating each round satisﬁesm ≥ min( M, 20(( γ / L max )
+ ( H / L max ) 4 )( αL max ) − 4 log( kT )) , and the initial parameters satisfy (i)δ 0 :
= dist( B 0 , B ∗ ) ≤ √ 1 − E 0 for any E 0 ∈ (0 , 1] , (ii) (cid:107) I − α B (cid:62) [-
-]0 B 0 (cid:107)
= O ( α 2 τL 2max κ 2max ) and(iii) (cid:107) w 0 (cid:107)
= O ( α 2 . 5 τL 3max ) .
Choose step size α
= O ( 1 − δ 0 √ τL max κ 2max ) .
Then for any (cid:15) ∈ (0 , 1) , thedistance of the representation learned by FedAvg with τ ≥ 2 local updates satisﬁes dist( B T , B ∗ ) < (cid:15)after at most T =
O (cid:0) log( 1 / (cid:15) )
α 2 τµ 2 E 0 (cid:1) communication rounds with probability at least 1 − 4( kT ) − 99 .
Theorem 1 shows that FedAvg converges exponentially fast to the ground-truth representation whenexecuted on the clients’ population losses.
We provide intuition for the proof in Section 4 and the fullproof in Appendix B. First, some comments are in order.
Mild initial conditions.
Theorem 1 holds under benign initial conditions.
In particular, condition( i ) requires that the initial distance is only a constant smaller than 1.
Condition ( ii ) ensures that theinitial representation is well-conditioned with appropriate scaling, and ( iii ) guarantees the initialhead is not too large.
The last two conditions can be easily achieved by normalizing the inputs.
Generalization without convergence in terms of the global loss.
When each client accesses itspopulation loss as in Theorem 1, the global objective (5) becomes:
[Equation]
However, Theorem 1 does not imply that FedAvg solves (6).
In fact, our simulations in Section 5show that it does not even reach a stationary point of (6).
This is consistent with prior works that havenoticed the “objective inconsistency” phenomenon of FedAvg: it solves an unknown objective distinctfrom the global objective due to the fact that after multiple local updates, local gradients are no longerunbiased estimates of gradients of (6) [9].
Nevertheless, our results show that FedAvg is able to learna generalizable model even when it does not optimize the global loss in data heterogeneous settings .
Multiple local updates critically harness diversity, whereas Distributed GD (D-GD) does notlearn the representation.
Key to the proof of Theorem 1 is that the locally-updated heads becomediverse , meaning that they cover all directions in R k , with greater diversity corresponding to moreevenly covering in all directions.
We will show in Section 4 that the locally-updated heads becomeroughly as diverse as the ground-truth heads, and this causes the representation to move towards the ground-truth at rate depending on the diversity level.
Theorem 1 reﬂects this: the convergence rateimproves with the diversity metric µ / L max .
In this way FedAvg exploits data heterogeneity to learn therepresentation, as more diverse { w ∗ ,i } i ∈
[ M ] implies more heterogeneous data.
Moreover, Theorem 1shows that additional local updates improve the convergence rate up to τ
=
O ( α − 2 ) .
Importantly,head diversity only beneﬁts the global representation update if τ ≥ 2 .
We formally prove that D-GD(equivalent to FedAvg with τ =1 and m = M ) cannot recover col( B ∗ ) in the following result.
Proposition 1 (Distributed GD lower bound) .
Suppose we are in the setting described in Section 3and d>k> 1 .
Then for any set of ground-truth heads { w ∗ ,i } Mi =1 , full-rank initialization B 0 ∈ R d × k ,initial distance δ 0 ∈ (0 , 1 / 2] , step size α > 0 , and number of rounds T , there exists B ∗ ∈
O d × ksuch that dist( B 0 , B ∗ ) =
δ 0 , B ∗ ¯w ∗ ∈ col( B 0 ) and dist( B D-GD T , B ∗ ) ≥ 0 .
7 δ 0 , where B D-GD T ≡
B D-GD T ( B 0 , B ∗ , { w ∗ ,i } Mi =1 , α ) is the result of D-GD with step size α and initialization B 0 on thesystem with ground-truth representation B ∗ and ground-truth heads { w ∗ ,i } Mi =1 .
Proposition 1 shows that for any choice of δ 0 ∈ (0 , 1 / 2] , non-degenerate initialization B 0 , and groundtruth heads, there exists a B ∗ whose column space is δ 0 -close to col( B 0 ) , yet is at least 0 .
7 δ 0 -far fromthe representation learned by D-GD in the setting with B ∗ as ground-truth.
Therefore, even allowingfor a strong initialization, D-GD cannot guarantee to recover the ground-truth representation.
Thisnegative result combined with our previous results suggest even if we had an inﬁnite communicationbudget, it would still be advantageous to execute multiple local updates between communicationrounds 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 learnscol( B ∗ ) , while D-GD fails to achieve this goal.
Global update B t +1 .
To build intuition for why FedAvg can learn col ( B ∗ ) , we examine the globalupdate of the representation in the full participation case ( m = M ):
[Equation]
t
[Equation]
Notice that B t +1 is a mixture of B t and B ∗ with weight matrices in R k × k .
We aim to show that(I) the ‘prior weight’ on B t has spectral norm strictly less than 1, and(II) the ‘signal weight’ on B ∗ adds energy from col( B ∗ ) to B t +
so that σ min ( B t +1 ) ≈ σ min ( B t ) .
These imply the contribution from col( B t ) in col( B t +1 ) contracts, while energy from col( B ∗ )replaces the lost energy from col( B t ) .
Hence, col( B t +1 ) moves to col( B ∗ ) in all k directions.
The role of head diversity and multiple local updates.
To show (I) and (II), it is imperative touse the diversity of the locally-updated heads when τ ≥ 2 .
First consider (I).
Notice that for eachi , (cid:81) τ
− 1 s =
( I k − α w t,i,s w (cid:62) t,i,s ) has singular values at most 1, and strictly less than 1 correspondingto directions spanned by { w t,i,s } s ∈ [ τ − 1] .
Thus, the maximum singular value of the average ofthese matrices should be strictly less than 1 as long as { w t,i,s } s ∈ [ τ − 1] ,i ∈
[ M ] spans R k , i.e. the 245 locally-updated heads are diverse.
Similarly, the signal weight is rank- k if the locally-updated headsspan R k , which leads to (II) as discussed below.
In contrast, if τ = 1 , then the global update ofthe representation does not leverage head diversity, as it is only a function of the global head andthe average ground-truth head: B t +1 =
B t ( I k − α w t w (cid:62) t ) + α B ∗ ¯w ∗ w (cid:62) t in this case.
As a result,col( B t +1 ) can only improve in one direction, so D-GD ultimately fails to learn col( B ∗ ) (Prop. 1).
Achieving head diversity: the necessity of controlling
I k − α B (cid:62) t B t .
We have discussed theintuition for why head diversity implies (I) and (II) for FedAvg.
Next, we investigate why the headsbecome diverse.
Let us examine client
i ’s ﬁrst local update for the head at round t :
[Equation]
i
From this equation we see that if ∆ t :=
I k − α B (cid:62) t B t ≈ 0 and (cid:107) w t (cid:107) is bounded, w t
i ≈ then { w t, i
} i ∈ [ the diversity of w ∗ i } i ∈ which is indeed diverse due to Assumption 2, meaning that the local heads are diverse after just onelocal update.
Moreover, it can be shown that if ∆ t ≈ 0 and the heads become diverse after onelocal update, then they remain diverse for all local updates due to the observation that each B t,i,schanges slowly over s .
Note that in addition to implying local head diversity, ∆ t ≈ 0 for all t impliesσ min ( B t ) ≈ σ min ( B t +1 )
≈ 1 √ α , which directly ensures (II).
Thus we aim to show ∆ t ≈ 0 for allcommunication rounds, i.e. B t remains close to a scaled orthonormal matrix.
However, it is surprising why (cid:107) ∆ t (cid:107) remains small: B t +1 is the average of nonlinearly locally-updatedrepresentations, and the local updates could ‘overﬁt’ by adding more energy to some columns thanothers, and/or lead to cancellation when summed, so it is not intuitive why √ α B t remains almostorthonormal.
Nor does the expression above for B t +1 provide any clarity on this.
Nevertheless,through a careful induction we show that ∆ t indeed stays close to zero since the local heads convergequickly and the projection of the local representation gradient onto col( B t ) is exponentially decaying.
Inductive argument.
While the above intuitions seem to simplify the behavior of FedAvg, showingthat they all hold simultaneously is not at all obvious.
To study this, we are inspired by recent work[45] that developed an inductive argument for representation learning in the context of gradientbased-meta-learning.
To formalize our intuition discussed previously, in our proof we need to showthat (i) the learned representation does not overﬁt to each client’s loss despite many local updates andsimultaneously the heads quickly become diverse, and (ii) the update at the global server preserves thelearned representation despite averaging many nonlinearly perturbed representations gathered fromclients after local updates.
To address these challenges, we construct a pair of intertwined inductionhypotheses over time, one for tracking the effect of local updates, and another for tracking the globalaveraging.
Each induction hypothesis (local and global) itself consists of several hypotheses (in effect,a nested induction) that evolve both within and across communication rounds.
Local induction.
The proof leverages the following local inductive hypotheses for every t, i :
[Equation]
The local induction tracks the effect of updates at each client node:
At the end of τ local updates,
A 1 ,t,i ( τ ) captures the diversity of the local heads, A 2 ,t,i ( τ ) ensures that the heads remain uniformlybounded, A 3 ,t,i ( τ ) shows that the locally adapted representations stay close to a scaled orthonormalmatrix, and A 4 ,t,i ( τ ) shows that the locally adapted representations do not diverge too quickly fromthe ground-truth.
The second set of inductions below controls the global behavior.
Global induction.
The global induction utilizes a similar set of inductive hypotheses.
[Equation]
[Equation]
Hypotheses A 1 ( t ) , A 2 ( t ) and A 3 ( t ) are analogous to A 1 ,t,i ( s ) , A 2 ,t,i ( s ) and A 3 ,t,i ( s ) , respectively.
A 4 ( t ) shows that the energy of col( B t ) that is orthogonal to the ground-truth subspace is contracting,and A 5 ( t ) ﬁnally shows that the principal angle distance between the learned and ground-truthrepresentations is exponentially decreasing.
Our main claim follows from A 5 ( T ) .
However, provingthis result requires showing that all the above local and global hypotheses hold for all times t ≥ 1 ,as these hypotheses are heavily coupled.
As mentioned previously, the most difﬁcult challenge iscontrolling (cid:107) ∆ t (cid:107) ( A 3 ( t ) )
despite many local updates, and doing so requires leveraging both localand global properties.
The details of this local-global induction argument are in Appendix B.
Experiments
In this section, we conduct experiments to (I) verify our theoretical results in the linear settingand (II) determine whether our established insights generalize to deep neural networks.
Notably, [Figure] [Figure] demonstrating the competitive performance of FedAvg plus ﬁne-tuning for personalized FL is nota goal of this section, as this is evident from prior experiments [10, 12, 40, 33].
Rather, to achieve(II)
we test whether FedAvg learns effective representations when trained with neural networksin heterogeneous data settings via three popular benchmarks for evaluating the quality of learnedrepresentations.
Since our main claim is that local updates are key to representation learning, we use D-(S)GD as our baseline in all experiments [Figure]
5.1 Multi-task linear regression
We ﬁrst experiment with the regression setting from our theory.
We randomly generate B ∗ ∈ R d × kand { w ∗ ,i } i ∈ [ M ] by sampling each element i.i.d.
from the normal distribution, where d =100 , k =5and M = 40 , and then orthogonalizing B ∗ .
Then we run FedAvg with τ = 2 local updates and
D-GD, both sampling m =
M clients per round.
We have seen in Figure 1 that the principalangle distance between the representation learned by FedAvg and the ground-truth representationlinearly converges to 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 tostationary point of (6), while FedAvg does not converge to one at all.
Although D-GD optimizesthe global loss, it does not generalize as well as FedAvg to new clients as demonstrated by Figure 3(right).
Here, we ﬁne-tune the models learned by FedAvg and D-GD on a new client with n samplesgenerated by x M +1 ∼ N ( 0 , I d ) , ζ M +1 ∼ N (0 , 0 . 01) , and y M +1 =
(cid:104) B ∗ w ∗ ,M +1 , x M +1 (cid:105) + ζ M +1 .
We ﬁne-tune using GD for τ (cid:48)
= 200 iterations with batch size b = n , and plot the ﬁnal error(cid:107) B T,M +1 ,τ (cid:48)
w T,M +1 ,τ (cid:48) − B ∗ w ∗ ,M +1 (cid:107) 2 .
Both plots are generated by averaging 10 runs.
5.2 Image classiﬁcation with neural networks
Next we evaluate FedAvg’s representation learning ability on nonlinear neural networks.
For faircomparison, in every experiment all methods make the same total amount of local updates during thecourse of training (e.g. D-SGD is trained for 50 × more rounds than FedAvg with τ = 50 ).
Datasets and models.
We use the image classiﬁcation datasets CIFAR-10 and CIFAR-100 [58],which consist of 10 and 100 classes of RGB images, respectively.
We use a convolutional neuralnetwork (CNN) with three convolutional blocks followed by a three-layer multi-layer perceptron,with each convolutional block consisting of two convolutional layers and a max pooling layer. [Figure]
Cosine similarity of features.
A desirable property of representations for downstream classiﬁcation tasks is that features of examples from the same class are similar to each other, while features ofexamples from different classes are dissimilar [59].
In Figure 4 we examine whether the representations learned by FedAvg satisfy this property.
Here we have trained FedAvg with varying τ and
D-SGD (FedAvg with τ =1 ) on CIFAR-10.
Image classes are heterogeneously allocated to M = 100clients according to the Dirichlet distribution with parameter 0.6 as in [60].
Each subplot is a 10x10matrix whose ( i, j ) -th
element gives the average cosine similarity between features of images
fromthe i -th and j -th classes learned by the corresponding model.
Ideally, diagonal elements are closeto 1 (high similarity) and off-diagonal elements are close to 0 (low similarity).
Figure 4 shows that
FedAvg indeed learns features with high intra-class similarity and low inter-class similarity, withrepresentation 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 learnedby FedAvg to new classes and also new datasets.
An effective representation identiﬁes universallyimportant features, so it should generalize to new data, with perhaps a small amount of ﬁne-tuningneeded to learn a new mapping from feature space to label space.
The transfer learning performanceof ﬁne-tuned models is a popular metric for evaluating the quality of learned representations [61, 62].
We ﬁrst study how models trained by FedAvg and D-SGD generalize to unseen classes from the samedataset.
To do so, we train models on heterogeneous partitions of CIFAR-100 using both FedAvgwith τ = 50 as well as D-SGD.
In the left plot of Figure 5, we illustrate the case that models aretrained 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 ﬁne-tune the trained models on the new clients with 10 epochs of SGD, with varying numbersof samples per epoch as listed on the x-axis, before testing.
Next, we investigate how well modelstrained by FedAvg and D-SGD generalize to an unseen dataset.
In the right plot of Figure 5, we trainmodels with C classes/client from CIFAR-100, then test on new clients with samples drawn from CIFAR-10 (a different dataset, but with presumably similar “basic” features).
Speciﬁcally, for thesenew clients, we ﬁne-tune for 10 epochs as previously, then test the post-ﬁne-tuned models on the testdata for each client.
In both left and right plots, we observe that FedAvg signiﬁcantly outperforms
D-SGD, indicating that FedAvg has learned a representation that generalizes better to new classes.
Conclusion
We showed that FedAvg learns the ground-truth representation in the multi-task linear regressionsetting. To our knowledge, this is the ﬁrst theoretical study showing FedAvg learns an effectiverepresentation in any setting. Our analysis reveals that multiple local updates are critical to FedAvg’srepresentation learning ability, which is supported empirically on both linear and nonlinear models. These experimental results suggest future work can extend our ﬁndings to more complex settings.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] Our theoretical results only holdfor the multi-task linear setting.
This point is explicitly mentioned in the abstract,
Introduction (Section 1), 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]
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
B.
If you ran experiments...
(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 seeAppendix C and 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.
(c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
[Yes] Please see Figures 2, 3 and 5.
(d) Did you include the total amount of compute and the type of resources used (e g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] Please see Appendix C.
If you are using existing assets (e g., code, data, models) or curating/releasing 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 as a URL?
[No d)
Did you
discuss whether
and
how consent
was obtained
from people
whose data you’re
using/ curating?
N/A
](e
) Did you discuss whether
the
data
you are
using/
curating contains personally identiﬁable information or offensive
content?
N/
A]5.
If you used crowdsourcing or conducted research with human subjects...(
a)
Did you include the full text of instructions given to participants and screenshots, ifapplicable
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 amountspent on participant compensation?
N /A]