M IXED F EDERATED L EARNING : J OINT D ECENTRALIZED AND C ENTRALIZED L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Federated learning (FL) enables learning from decentralized privacy-sensitive data, with computations on raw data conﬁned to take place at edge clients.
This paper introduces mixed FL , which incorporates an additional loss term calculated at the coordinating server (while maintaining FL’s private data restrictions).
For example, additional datacenter data can be leveraged to jointly learn from centralized (datacenter) and decentralized (federated) training data and better match an expected inference data distribution.
Mixed FL also enables ofﬂoading some intensive computations (e.g., embedding regularization) to the server, greatly reducing communication and client computation load.
For these and other mixed FL use cases, we present three algorithms: P ARALLEL T RAINING ,
1- WAY G RADIENT
T RANSFER , and
2- WAY G RADIENT T RANSFER .
We perform extensive experiments of the algorithms on three tasks, demonstrating that mixed FL can blend training data to achieve an oracle’s accuracy on an inference distribution, and can reduce communication and computation overhead by more than 90%.
Finally, we state convergence bounds for all algorithms, and give intuition on the mixed FL problems best suited to each.
The theory conﬁrms our empirical observations of how the algorithms perform under different mixed FL problem settings.
I NTRODUCTION
Federated learning (FL) (McMahan et al., 2017) is a machine learning setting where multiple ‘clients’ (e.g., mobile phones) collaborate to train a model under coordination of a central server.
Clients’ raw data are never transferred.
Instead, focused updates intended for immediate aggregation are used to achieve the learning objective (Kairouz et al., 2019).
FL typically delivers model quality improvements because training examples gathered in situ by clients reﬂect actual inference serving requests.
For example, a mobile keyboard next-word prediction model can be trained from actual
SMS messages, yielding higher accuracy than a model trained on a proxy document corpus.
Because of the beneﬁts, FL has been used to train production models for many applications (Hard et al., 2018;
Ramaswamy et al., 2019; Apple, 2019; Ramaswamy et al., 2020; Hartmann, 2021; Hard et al., 2022).
Building on FL, we can gain signiﬁcant beneﬁts from ‘mixed FL’: jointly 1 training with an additional centralized objective in conjunction with the decentralized objective of FL.
Let x be model parameters to be optimized.
Let f denote a mixed loss, a sum 2 of a federated loss f f and a centralized loss f c :
[Equation]
Mixed loss f might be a more useful training objective than f f for many reasons, including:
Mitigating Distribution Shift by Adding Centralized Data to FL While FL helps with reducing train vs. inference distribution skew, it may not remove it completely.
Examples include: training device populations that are subsets of inference device populations (e.g., training on high-end phones, for eventual use also on low-end phones), label-biased example retention on edge clients (e.g., only retaining positive examples of a binary classiﬁcation task), and infrequent safety-critical example events with outsized importance (e.g., automotive hard-braking events needed to train a self-driving
AI) (Anonymous, a).
The beneﬁts of FL can be achieved while overcoming remaining distribution skew by incorporating data from an additional datacenter dataset, via mixed FL.
This affords a composite set of training data that better matches the inference distribution.
Reducing Client Computation and Communication In representation learning, negative examples are used to push dissimilar items apart in a latent embedding space while keeping positive examples closer together (Oord et al., 2018).
In federated settings, clients’ caches may have limited local negative examples, and recent work (Anonymous, b) showed this signiﬁcantly degrades performance compared to centralized learning.
This work also showed that using an additional loss (a regularization) to push representations apart, instead of negative examples, can resolve this performance gap.
However, if done naively this requires communicating and computing over a large embedding table, introducing massive overhead for large-scale tasks.
Applying mixed FL, where federated loss f f is the primary ‘afﬁnity’ loss and centralized loss f c is the ‘spreadout’ regularization, avoids communicating the entire embedding table to clients and greatly reduces client computation.
Though mixed FL can clearly be useful, an actual process to minimize f is not trivial.
FL requires that clients’ data stay on device, as they contain private information that possibly reveals personal identity.
Moreover, centralized loss/data is expected to differ signiﬁcantly 3 from client loss/data.
Contributions
• We motivate the mixed FL problem and present three algorithms for addressing it: P ARALLEL T RAINING (PT),
1- WAY G RADIENT T RANSFER (1- W GT), and
2- WAY G RADIENT T RANSFER (2- W GT).
These algorithms maintain the data privacy protections inherent in FL.
[Section 2] • We experiment with facial attribute classiﬁcation and language modeling, demonstrating that our algorithms overcome distribution shift.
We match the accuracy of hypothetical ‘oracle’ scenarios where the entire inference distribution was colocated for training.
[Section 4] • We experiment with user-embedding based movie recommendation, reducing communication overhead by 93.9% and client computation by 99.9% with no degradation in quality.
[Section 4] •
We state convergence bounds for each algorithm (in strongly, general, and non-convex settings), providing theoretical explanations for convergence behaviors we observe in the experiments.
This indicates how the algorithms will perform on new mixed FL tasks.
[Section 5]
A LGORITHMS
In FL, the loss function f f is an average of client loss functions
f i .
The client loss f
i is an expectation over batches of data examples
B i on client i .
[Equation]
F ED A VG (McMahan et al., 2017) is a ubiquitous, heuristic FL method designed to minimize Equation 2 w.r.t.
model x in a manner that allows all client data ( B i ) to remain at respective clients i .
Providing strong privacy protection is a major motivation for FL.
Storing raw data locally on clients rather than replicating it on servers decreases the attack surface of the system.
Also, using focused ephemeral updates and early aggregation follows principles of data minimization (White House Report, 2013).
While training with loss f f via F ED A VG can yield an effective model x , this paper shows there are scenarios where ‘mixing’ in an additional ‘centralized’ loss f c proves beneﬁcial to the training of x .
Such a loss term can make use of batches of centralized data examples B c , from a datacenter dataset:
[Equation]
The centralized loss f c will differ from the federated loss f f (else it would not be useful). This can be a difference in the distributions that B c and B i are drawn from, and/or in the functional forms of f c and f i .
We will present an expression that quantiﬁes the difference between f c and f f in Section 5.
We now state our mixed FL algorithms (Algorithms 1 and 2).
Appendix A has a few practical details.
• P ARALLEL T RAINING performs a round of F ED A VG (minimizing f f ) in parallel with steps of centralized training (minimizing f c ), merges (e
g., averages) the respective updates, and repeats.
Green in Algorithm 1 indicates added steps beyond F ED A VG for P ARALLEL T RAINING .
• 1- WAY G RADIENT T RANSFER starts a round by calculating a gradient of f c .
It is sent to participating clients and summed with clients’ gradients of f i during client optimization.
Blue in Algorithm 2 indicates added steps beyond F ED A VG for 1- WAY G RADIENT T RANSFER .
2- WAY G RADIENT T RANSFER is P ARALLEL T RAINING with gradient sharing.
Two gradients are now used, one based on f c and sent to clients (like 1- W GT), one based on f f and applied centrally.
Purple in Algorithm 1 is added steps beyond PT for 2- WAY G RADIENT T RANSFER .
[Table]
1 } do Initialize central model x ( t, 0) c = x ( t ) for central step k = 0 ,...,K c − 1 do
Sample centralized batch B ( k ) c ; compute stochastic gradient g c ( x ( t,k ) c ; B ( k ) c ) of f c ( x ( t,k ) c )
Perform central update x ( t,k +1) c
= C ENTRAL O PTIMIZER ( x ( t,k ) c ,g c ( x ( t,k ) c ; B ( k ) c ) + ˜ g ( t ) f ,η c ,t )
Compute central model delta ∆ ( t ) c
= x ( t,Kc ) c − x ( t )
Sample a subset S ( t ) of clients; for client i ∈ S ( t ) in parallel do ∆ ( t ) i , p i =
C LIENT U PDATE ( x ( t ) , ˜ g ( t ) c , C LIENT O PTIMIZER , η , t )
Aggregate client changes ∆ ( t ) =
(cid:80) i ∈S ( t ) p i ∆ ( t ) i / (cid:80) i ∈S ( t ) p
i
Compute federated model x ( t ) f
= S ERVER
O PTIMIZER ( x ( t ) , − ∆ ( t ) ,η s ,t )
Compute federated model delta ∆ ( t ) f
= x ( t ) f − x ( t )
Aggregate central model and federated model deltas ∆ ( t ) =
∆ ( t ) c
+ ∆ ( t ) f
Update global model x
( t +1)
=
M ERGE O PTIMIZER ( x ( t ) , − ∆ ( t ) ,η m ,t )
Update augmenting centralized gradient ˜ g ( t +1) c = − ∆ ( t ) c / ( η c K c ) − ˜ g ( t ) f
Update augmenting federated gradient ˜ g ( t +1) f
= − (cid:80) i ∈S ( t ) ∆ ( t ) i / ( η (cid:80) i ∈S ( t ) K i )
− ˜ g ( t ) c ( see Appendix A )
C LIENT U PDATE :
Input: Initial client model x ( t, 0) i ; (possible) augmenting gradient ˜ g ( t ) c ; C LIENT O PTIMIZER , learning rate η ; round t ; client loss f i (2) Initialize client weight p
i =for client step k = 0 ,...,K i − 1 do Sample batch B ( k ) i ;
compute stochastic gradient
g i ( x ( t,k ) i ; B ( k ) i ) of f i ( x ( t,k ) i ) ; update client weight p i
= p
i + |B ( k ) i |
Perform client update x ( t,k +1)
i =
C LIENT O PTIMIZER ( x ( t,k ) i ,g i
( x ( t,k ) i ; B ( k ) i )
+ ˜ g ( t ) c ,η,t )
Compute client model changes ∆ ( t ) i
= x ( t,Ki )
i − x ( t, 0) i and return ∆ ( t ) i , p i
Input: Initial model x (0) ; C LIENT O PTIMIZER , S ERVER O PTIMIZER with respective learning rates η,η s ; central loss function f c (3) for t ∈ { 0 , 1 ,...,T − 1 } do
Sample batch B ( t ) c ; compute stochastic gradient g c ( x ( t ) ; B ( t ) c ) of f c ( x ( t ) ) ; set augmenting gradient ˜ g ( t ) c
= g c ( x ( t ) ; B ( t ) c )
Sample a subset S ( t ) of clients; for client i ∈ S ( t ) in parallel do ∆ ( t ) i , p i =
C LIENT U PDATE ( x ( t ) , ˜ g ( t ) c , C LIENT O PTIMIZER , η ) (C LIENT U PDATE function deﬁned in Algorithm 1 )
[Table]
R ELATED W ORK
There are parallels between G RADIENT T RANSFER and algorithms addressing inter-client data heterogeneity in FL, like SCAFFOLD (Karimireddy et al., 2020b) or Mime (Karimireddy et al., 2020a).
Those algorithms calculate a gradient reﬂective of the entire federated population and transmit it to clients to reduce update variance, improving optimization on non-IID client data.
G RADIENT
T RANSFER calculates a gradient reﬂective of centralized data/loss, to augment client computations of decentralized data/loss (in 2- W GT, also the converse).
However, SCAFFOLD requires keeping state at the server (control variates) for each participating client, impractical in real large-scale FL systems (Kairouz et al., 2019).
2- W GT only requires state (augmenting gradients) for two entities, and so is easily implemented.
Another instance of a non-IID client approach (Zhao et al., 2018) inﬂuencing mixed FL is the
E XAMPLE T RANSFER algorithm (Anonymous, a).
Here, centralized examples are sent to federated clients, instead of gradients.
This is typically precluded in real FL applications.
The data volume needed to transfer may be excessive, and in general E XAMPLE T RANSFER does not offer solutions for one of the main motivations of this work, reducing client computation and communication.
‘Transfer learning’ also involves two different training datasets, but has a different purpose.
Transfer learning pretrains a model on a distribution (e.g., centralized data in a datacenter), then ﬁne-tunes it on the actual distribution of interest (e.g., decentralized data via FL).
It is desirable as a way to quickly train on the latter distribution (e.g., as in Ro et al.
But its sequential approach results in catastrophic forgetting (McCloskey & Cohen, 1989; Ratcliff, 1990; French, 1999); accuracy on pretraining data is lost as the model learns to ﬁt ﬁne-tuning data instead.
In mixed FL, we desire good inference performance on all data distributions trained on.
See Appendix B.4.5 for more.
In differentially private (DP) optimization, a line of work has aimed to improve privacy/utility tradeoffs by utilizing additional non-private data.
One way is to use non-private data to pretrain (Abadi et al., 2016).
Another avenue is to use non-private data to learn the gradient geometry (Zhou et al., 2020;
Amid et al., 2021; Asi et al., 2021; Kairouz et al., 2021a; Li et al., 2022), improving accuracy by enabling tighter, non-isotropic gradient noise during DP optimization.
Amid et al.
(2021) and Li et al.
(2022) consider the FL use case 5 .
As in transfer learning, additional data is used only to improve performance on a single distribution, and retaining accuracy on other distributions is a nongoal (in contrast to mixed FL).
Also, the non-private data used is generally matching (in distribution) to the private data, whereas in mixed FL we typically explicitly leverage distinct distributions.
A few recent works (Yu et al., 2020; Elbir et al., 2021; Jeong et al., 2021) present instances of mixed
FL problems; they shift computations to the server that are intensive or impossible at the clients.
These works do not address distribution shift or present more general mixed FL algorithms.
E XPERIMENTS
We now present experiments on three tasks (Table 1), showing a range of problems for which mixed
FL is useful.
They also show the comparative performance of each algorithm described in Section
[Table]
and see
[Table]
4. A DDRESS L ABEL I MBALANCE IN T RAIN D ATA (S MILE C LASSIFIER ; C ELEB A)
Earlier work (Anonymous, a) motivated mixed FL with the example problem of training a ‘smiling’- vs.-‘unsmiling’ classiﬁer via FL with mobile phones, with the challenge that the phones’ camera application (by the nature of its usage) tends to only persist images of smiling faces.
A solution
[Figure] [Figure]
for this severe label imbalance is to apply mixed FL, utilizing an additional datacenter dataset of unsmiling faces to train a capable classiﬁer.
To experiment, CelebA data 6 ( et al
Caldas et al., 2018) is split into a federated ‘smiling’ dataset and centralized ‘unsmiling’
Figures 1a and 2a show the AUC and loss convergence of our three algorithms when applied to this problem.
Empirically, we observe G RADIENT T RANSFER converge faster than P ARALLEL
T RAINING .
Section 5 will provide a theoretical explanation for this behavior.
4. M ITIGATE B IAS IN T RAIN D ATA (L ANGUAGE M ODEL ; S TACK O VERFLOW , W IKIPEDIA )
Consider the problem of learning a language model like a RNN-based next character prediction model, used to make typing suggestions to a user in a mobile keyboard application.
Because the ultimate inference application is on mobile phones, it is natural to train this model via FL, leveraging cached SMS text content highly reﬂective of inference time usage (at least for some users).
However, the mobile phones participating in the federated learning of the model might be only a subset of the mobile phones for which we desire to deploy for inference.
Higher-end mobile phones can disproportionately participate in FL, as their larger memory and faster processors allow them to complete client training faster.
But to do well at inference, a model should make accurate predictions for users of lower-end phones as well.
A purely FL approach can do an inadequate job of learning these users’ usage patterns.
(See Kairouz et al.
(2019) for more on aspects of fairness and bias in FL.)
Mixed FL overcomes this problem, by training a model jointly on federated data (representative of users of higher-end phones) and a datacenter dataset (representative of users of lower-end phones).
We simulate this scenario using two large public datasets: the Stack Overﬂow dataset 7 (Kaggle) for federated data, and the Wikipedia dataset 8 (Wikimedia Foundation) for datacenter data.
Figure 1b shows results.
The ‘only FL’ scenario learns Stack Overﬂow (but not Wikipedia) patterns of character usage, and so has reduced accuracy ( ∼ 0 .
60 ) when evaluated on examples from both datasets.
The
We see that mixed FL results in similar loss as the more expensive baseline scenario (see Table 2).
[Table]
mixed FL algorithms demonstrate learning both: they all achieve an evaluation accuracy ( ∼ 0 . 66 ) comparable to an imagined ‘oracle’ that could centrally train on the combination of datasets.
Unlike the smile classiﬁcation experiment, here P ARALLEL T RAINING converges about as fast as
G RADIENT T RANSFER .
Section 5 will provide a theoretical explanation for this behavior.
4. R EGULARIZE E MBEDDINGS AT S ERVER (M OVIE R ECOMMENDATION ; M OVIE L ENS )
The third task we study, movie recommendation, is one with an embedding regularization term as described in Section 1.
As Table 1 shows, a key difference from the previous two mixed FL experimental scenarios is that here we are mixing losses with different functional forms (instead of mixing different datacenter and federated data distributions).
We study this scenario by training a dual encoder representation learning model (Covington et al., 2016) for next movie prediction on the MovieLens dataset (Harper & Konstan, 2015; GroupLens).
As described in Section 1, limited negative examples can degrade representation learning performance.
Previous work (Anonymous, b) proposed using losses insensitive to local client negatives to improve federated model performance.
They observed signiﬁcantly improved performance by using a two-part loss: (1) a hinge loss to pull embeddings for similar movies together, and (2) a spreadout regularization (Zhang et al., 2017) to push embeddings for unrelated movies apart.
For clients to calculate (2), the server must communicate all movie embeddings to each client, and clients must perform a matrix multiplication over the entire embedding table.
This introduces enormous computation and communication overhead when the number of movies is large.
Mixed FL can alleviate this communication and computation overhead.
Instead of computing both loss terms on clients, clients calculate only the hinge loss and the server calculates the expensive regularization, avoiding costly computation on each client.
Also, computing the hinge loss only requires the embeddings of movies in a client’s local dataset.
Federated select (Charles et al., 2022) enables only those embeddings to be sent to that client, saving communication and on-client memory.
Experiments show all mixed FL algorithms achieve model performance (around 0.1 for recall@10) comparable to the baseline scenario where everything is computed on the clients.
Moreover, mixed
FL eliminates more than 99.9% of client computation and more than 93.9% of communication (see
Table 2).
For computation and communication analysis, see Appendix B. Note that real-world models can be much larger than this movie recommendation model 9 .
Without mixed FL, communicating such large models to clients and computing regularization would be impractical in large-scale settings.
P ARALLEL T RAINING converges slightly slower than either G RADIENT T RANSFER algorithm
(Figure 3) but reaches the same evaluation loss at around 1500 rounds.
C ONVERGENCE 5. P RELIMINARIES
We now describe convergence properties for each mixed FL algorithm from Section 2.
We assume mixed loss f has a ﬁnite minimizer (i.e., ∃ x ∗ s.t.
f ( x ) ≥ f ( x ∗ ) ∀ x ).
We assume client losses f i and centralized loss f c are β -smooth.
If f i are β -smooth, federated loss f f is also 10 .
For some results, we assume f i and f c are µ -convex (possibly strongly convex, µ > 0 ).
If f i are µ -convex, f f is also 11 .
For a parameter vector x , we use ∇ f i ( x ) to denote the full gradient of f i (i.e., over all data on client i ).
Similarly, ∇ f f ( x ) and ∇ f c ( x ) denote full gradients 12 of f f and f c at x .
We use g i ( x ) to denote an unbiased stochastic gradient of f i , calculated on a random batch
B i of examples on client i .
We focus on the impact to convergence when differences exist between the federated and centralized losses/data.
As such, we make the following homogeneity assumption about the federated data, which simpliﬁes the analysis and brings out the key differences.
Our analysis can be easily extended to heterogeneous clients by assuming a bound on variance of the client gradients.
Assumption 5.1.
The federated clients have homogeneous data distributions (i.e., with examples that are drawn IID from a common data distribution), and their stochastic gradients have bounded variance.
Speciﬁcally, for some σ > 0 , we have for all clients i and parameter vectors x :
[Equation]
Under such conditions, F ED A VG convergence can match SGD (Table 2, Karimireddy et al.
(2020b)).
Let g f denote an unbiased stochastic gradient of federated loss f f , formed by randomly sampling a cohort of S (out of N total) clients, randomly sampling a batch B i of client data examples, and averaging the respective client stochastic gradients.
Given Assumption 5.1 we bound variance of g f :
[Equation]
Let g c ( x ) denote a stochastic gradient of the centralized loss f c at x , calculated on a randomly sampled batch B c of centralized examples (from a datacenter dataset), with variance bounded by σ 2 c :
[Equation]
Summarizing Equations 4-6, a client’s stochastic gradient g i ( x ) has variance bounded by σ 2 , the federated cohort’s stochastic gradient g f ( x ) has variance bounded by σ 2 / S , and the centralized stochastic gradient g c ( x ) has variance bounded by σ 2 c .
Increasing client batch size |B i | reduces variance of g i ( x ) and g f ( x ) , increasing cohort size S reduces variance of g f ( x ) , and increasing central batch size |B c | reduces variance of g c ( x ) .
Note that σ 2 / S only bounds variance within the federated data distribution, and σ 2 c only bounds variance within the central data distribution.
To say something about variance across the two data distributions, we adapt the notion of ‘bounded gradient dissimilarity’ (or ‘BGD‘) introduced in
Karimireddy et al.
(2020b) (Deﬁnition A1), and apply it to the mixed FL scenario here.
Deﬁnition 5.2 (mixed FL ( G, B )-BGD) .
There exist constants G ≥ 0 and B ≥ 1 such that ∀ x :
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
In the deﬁnition, w f and w c are proportions of inﬂuence ( w f + w c = 1 ) of the federated and centralized objectives on the overall mixed optimization.
(The simplest setting is w f = w c = 1 / 2 .)
5. B OUNDS
We can now state upper bounds on convergence (to an error in mixed loss smaller than (cid:15) ) for the respective mixed FL algorithms.
For ease of comparison, the convergence bounds are summarized in Table 3.
The Theorems and Proofs of these convergence bounds are given in Appendix C. As mentioned
[Table]
previously, the analysis extends in a straightforward manner to the setting of heterogeneous clients assuming a bound on the variance of client gradients: for all x , 1 N (cid:80) Ni =1 (cid:107)∇ f i ( x ) − ∇ f f ( x ) (cid:107) 2 ≤ σ 2 f for some σ f ≥ 0 .
Under this assumption, the bounds in Table 3 change by an additional Kσ 2 f term in the expression involving σ 2 and σ 2 c in the parenthesis on the numerator of the leading term.
The derivation of these more general bounds follows on the same lines, so we omit it for brevity.
Analyzing Table 3, there are several implications to be drawn.
Signiﬁcant ( G , B )-BGD impedes P ARALLEL T RAINING
The convergence bounds for P ARALLEL T RAINING show a dependence on the G and B parameters from Deﬁnition 5.2.
If a mixed FL problem involves a large amount of dissimilarity between the federated and centralized gradients (i.e., if G (cid:29) 0 or B (cid:29) 1 ), then P ARALLEL T RAINING will be slower to converge than alternatives.
Signiﬁcant σ 2 c impedes 1-
WAY G RADIENT
T
RANSFER
1- WAY G RADIENT T RANSFER is more sensitive to central variance σ 2 c .
Unlike the other algorithms, the impact of σ 2 c on convergence scales with the number of steps K .
1- WAY G RADIENT T RANSFER requires a central batch size |B c | that is K times larger to achieve the same impact on convergence.
Intuitively, this makes sense; in a round, P ARALLEL T RAINING and 2- WAY G RADIENT T RANSFER sample K fresh batches during centralized optimization, while 1- WAY G RADIENT T RANSFER only samples a single central batch.
2- WAY G RADIENT T RANSFER should always converge at least as well as others The convergence bound for 2- WAY G RADIENT T RANSFER is unaffected by gradient dissimilarity (i.e., G (cid:29)or B (cid:29) 1 ), unlike P ARALLEL T RAINING .
Also, the bound for 2- WAY G RADIENT T RANSFER is less sensitive to σ 2 c than the bound for 1- WAY G RADIENT T RANSFER (as described above).
5. A NALYSIS OF E XPERIMENTS
P ARALLEL T RAINING has a convergence bound substantially different than the G RADIENT T RANSFER algorithms; the dependence on the BGD parameters G and B indicates there are mixed FL problems where P ARALLEL T RAINING is slower to converge than G RADIENT T RANSFER (in either form).
How can we know if a particular problem is one where P ARALLEL T RAINING will have slower convergence?
It would be useful to know G and B , but they cannot be exactly measured.
G and B (Deﬁnition 5.2) are upper bounds holding ∀ x , and the entire space of x cannot realistically be checked.
Instead, we introduce sampled approximations to empirically estimate these upper bounds.
Let x ( t ) be the global model at start of round t .
Let ˜ ∇ f f t , ˜ ∇ f c t , ˜ ∇ f t be approximations of federated, centralized, total gradients at round t .
Considering Deﬁnition 5.2, we deﬁne ˜ G t as a sampled approximation of G assuming B = 1 , and ˜ B t
as a sampled approximation of B assuming G = 0 :
[Equation]
[Table] [Table]
• Smile Classiﬁcation (4.1) As discussed, P ARALLEL T RAINING is at a disadvantage when G (cid:29) 0 or B (cid:29) 1 , and Table 4 (left column) shows that ˜ G t is signiﬁcantly large in smile classiﬁcation.
This aligns with our empirical observation in Figure 2a of P ARALLEL T RAINING converging slower than either G RADIENT T RANSFER algorithm.
• Language Modeling (4.2) Table 4 (center column) shows language modeling with more trivial gradient dissimilarity:
˜ G t ≈ 0 and ˜ B t ≈ 1 .
P ARALLEL T RAINING should be competitive.
Figure 2b empirically conﬁrms this to be true; all algorithms converge roughly equivalently.
C ONCLUSION
This paper has introduced mixed FL, including motivation, algorithms and their convergence properties, and intuition for when a given algorithm will be useful for a given problem.
Our experiments indicate mixed FL can improve accuracy and reduce communication and computation across tasks.
This work focused on jointly learning from a single decentralized client population and a centralized entity, as it illuminates the key aspects of the mixed FL problem.
Note that mixed FL and the associated properties we deﬁne in this paper (like mixed FL ( G, B )-BGD) are easily expanded to work with multiple ( > 1 ) distinct client populations participating.
E.g., a population of mobile phones and a separate population of smart speakers, or mobile phones separated into populations with distinct capabilities/usage (high-end vs. low-end, or by country/language).
Also, there need not be a centralized entity; mixing can be solely between distinct federated datasets.
It is interesting to reﬂect on the bounds of Table 3, and what they indicate about the beneﬁts of separating a single decentralized client population into multiple populations for mixed FL purposes.
The bounds are in terms of σ 2 (representing within population ‘variability’) and G and B (representing cross-population ‘variability’).
Splitting a population based on traits will likely decrease σ 2 (each population is now more homogeneous) but introduce or increase G and B (populations are now distinctive).
This might indicate scenarios where G RADIENT T RANSFER methods (only bounded by σ 2 ) become more useful and P ARALLEL T RAINING (also bound by G and B ) becomes less useful.
The limits of our convergence bounds should be noted.
First, they are ‘loose’; practical performance in particular algorithmic scenarios could be better, and thus comparisons between algorithms could differ.
Second, our bounds assume IID federated data, which is invalid in practice; convergence properties differ on non-IID data.
While our analysis, extended to handle non-IID data, shows that the bounds do not materially change, it is still a place where theory and practice slightly diverge.
It is important to note that mixed FL is orthogonal to the choice of algorithm for federated training.
For simplicity, this paper described mixed FL with F ED A VG used for federated training. However, one could use DP-FedAvg (McMahan et al. or DP-FTRL (Kairouz et al.
(2021b)) with mixed FL to ensure that DP protections are applied to the federated data. Similarly, an adaptive S ERVER O PTIMIZER like F ED ADAM (Reddi et al., 2020) could be used in the federated portion of training (preliminary mixed FL results with F ED ADAM are given in Appendix B.4.4).
In principle, mixed FL techniques are expected to have positive societal impacts insofar as they further develop the toolkit for FL (which has security and privacy beneﬁts to users) and improve accuracy on ﬁnal inference distributions.
Also, we’ve shown (Section 4.2) how mixed FL can address participation biases that arise in FL.
However, the addition of server-based data to federated optimization raises the possibility that biases in large public corpora ﬁnd their way into more applications of FL.
A P RACTICAL I MPLEMENTATION D ETAILS A. D OWNLOAD S IZE
G RADIENT T RANSFER (either 1- WAY or 2- WAY ) requires sending additional data as part of the communication from server to clients at the start of a federated round.
Apart from the usual model checkpoint weights, with G RADIENT T RANSFER we must now also transmit ‘augmenting’ gradients of the model weights w.r.t centralized data as well.
Naively, this doubles the download size as the gradient is the same size as the model.
However, the augmenting gradients should be amenable to compression, e.g. using an approach such as Mitchell et al.
A. U PLOAD S IZE
With P ARALLEL T RAINING and 1- WAY G RADIENT T RANSFER , no client gradient information is used outside of the clients themselves, so there is no additional information (apart from the model deltas and aggregation weights) to upload to the server.
With 2- WAY G RADIENT T RANSFER , client gradient information is used in centralized training, so needs to be conveyed to the server somehow.
When the FL client optimization is SGD, the average client gradient in a round (over all clients participating, over all steps) can be determined from the model deltas and aggregation weights that are already being sent back to the server, meaning no additional upload bandwidth is necessary.
Apart from bandwidth considerations, this also means there is no additional vector for private data leakage.
The algorithm to do this is as follows.
Each client i transmits to the server a local model change ∆ ( t ) i and an aggregation weight p i that is related to number of steps taken K i .
The average total gradient applied at client i during round t is:
[Equation]
The average client gradient (i.e., w.r.t.
just client data) at client i is:
[Equation]
(9) where ˜ g ( t ) c is the augmenting centralized gradient that was calculated from centralized data and used in round t .
The average (across the cohort) of average client gradients, weighted by K i , is:
[Equation]
This average client gradient ¯ g ( t ) f is in the spirit of SCAFFOLD (Karimireddy et al., 2020b) Equation 4, Option II.
It will be used as the augmenting federated gradient ˜ g ( t +1)
f in the subsequent round t + 1 , to augment centralized optimization.
See Algorithm 1.
A. D EBUGGING AND H YPERPARAMETER I NTUITION VIA K =
As these algorithms each involve different hyperparameters, validating that software implementations 13 are behaving as expected is non-trivial.
Something that proved useful for debugging purposes, as well as provided practical experience in understanding equivalences between the algorithms, was to perform test cases with the number of local steps K set to 1.
In this setting, the three mixed FL algorithms are effectively identical and should make equivalent progress during training.
Note that the convergence bounds of Table 3 hold for K ≥ 2 , so this takes us outside the operating regime where the bounds predict performance.
It also takes us outside an operating regime that is typically useful (FL use cases generally ﬁnd multiple steps per round to be beneﬁcial).
But it does serve a purpose when debugging.
[Figure] [Figure]
Table 4 is an informative comparison of the mixed FL optimization landscape of the three respective experiments conducted in Section 4.
It includes maximum values for the metrics ˜ G t and ˜ B t (the sampled approximations of the parameters deﬁned in ( G, B )-BGD
(Deﬁnition 5.2).
Here we provide some additional information.
Figure 4 plots these sampled approximation metrics over the ﬁrst 100 rounds of training.
We ransimulations per experiment and took the maximum at each round across simulations.
We used the same hyperparameters as described below (in Subsection B.2), except taking only a single step per round ( K = 1 ).
B. A DDITIONAL D ETAILS FOR E XPERIMENTS IN S ECTION
General Notes on Hyperparameter Selection For the various experiments in Section 4, we empirically determined good hyperparameter settings (as documented in Tables 5-10).
Our general approach for each task was to leave server learning rate η s at 1, select a number of steps K that made the most use of the examples in each client’s cache, and then do a sweep of client learning rates η to determine a setting that was fast but didn’t diverge.
For P ARALLEL T RAINING and 2- WAY
G RADIENT T RANSFER , which involve central optimization and merging, we set the merging learning rate η m to be 1, and set the central learning rate η c as the product of client and server learning rates: η c = ηη s (and since η s = 1 , this meant client and central learning rates were equal).
General Notes on Comparing Algorithms We generally kept hyperparameters equivalent when comparing the algorithms.
For example, we aimed to set batch sizes for all algorithms such that central and client gradient variances σ 2 and σ 2 c have equivalent impact on convergence (meaning |B c | =
S |B i | for PT and 2- W GT, and |B c | = KS |B
i | for 1- W GT).
In the case of language model training with 1- WAY G RADIENT T RANSFER , following this rubric would have meant a central batch size |B c | of 12800; we reduced this in half for practical computation reasons.
For a given task, we also generally kept learning rates the same for all algorithms.
Interestingly, we observed that as η (and η c , if applicable) is increased for a given task, the 2- WAY G RADIENT T RANSFER algorithm is the ﬁrst of the three to diverge, and so we had to adjust, e.g., in the language modeling experiment we used a lower η for 2- W GT than for PT and 1- W GT.
B.2. C ELEB A S MILE C LASSIFICATION
Datasets The CelebA federated dataset consists of 9,343 raw clients, which can be broken into train/evaluation splits of 8,408/935 clients, respectively (TFF CelebA documentation, 2022).
The raw clients have average cache size of ∼ 21 face images.
The images are about equally split between smiling and unsmiling faces.
In order to enlarge cache size, we group three raw clients together into one composite client, so our federated training data involves 2,802 clients with caches of (on average) ∼ 63 face images (and about half that when we limit the clients to only have smiling faces).
[Table] [Table]
[Table] [Table]
[Table] [Table]
Our evaluation data consists of both smiling and unsmiling faces
[Table]
and is meant to stand in for the inference distribution (where accurate classiﬁcation of both smiling and unsmiling inputs is necessary).
Note that as CelebA contains smiling and unsmiling faces in nearly equal amounts, a high evaluation accuracy cannot come at the expense of one particular label being poorly classiﬁed.
Model Architecture
The architecture used is a very basic fully-connected neural network 14 with a single hidden layer of 64 neurons with ReLU activations.
Hyperparameter Settings
The settings used in mixed FL training are shown in Tables 5 and 6.
SGD was used for all optimizers: C LIENT O PTIMIZER and S ERVER O PTIMIZER , and (if PT or 2- W
GT)
C ENTRAL O PTIMIZER and M ERGE O PTIMIZER .
B.2. S TACK O VERFLOW /W IKIPEDIA L ANGUAGE M ODELING
Datasets The Stack Overﬂow dataset is a large-scale federated dataset, consisting of 342, training clients and 204,088 evaluation clients (TFF StackOverﬂow documentation, 2022).
The training clients have average cache size of ∼ 400 examples, and evaluation clients have average cache size of ∼ 80 examples.
The Wikipedia dataset ( wikipedia/20201201.en ) consists of 6,210,110 examples (TFDS Wikipedia (20201201.en) documentation, 2022), which we partition by designating the ﬁrst 5,000,000 examples for training and the remainder for evaluation.
All raw text data is processed into sequences of 100 characters.
Our evaluation data is a combined dataset consisting of randomly shufﬂed examples drawn from the
Stack Overﬂow evaluation clients and the Wikipedia dataset.
After accounting for raw examples discarded during processing, the evaluation dataset is comprised of approximately 75% Stack Overﬂow and 25% Wikipedia.
Hence, we used this weighting proportion when performing mixed FL training in the language modeling problem (Table 8). [Figure]
Model Architecture
The architecture used is a recurrent neural network (RNN) 15 with an embedding dimension of 256 and 1024 GRU units.
Hyperparameter Settings
The settings used in mixed FL training are shown in Tables 7 and 8.
Note that in this experiment we also found clipping the norm of gradients (both federated and centralized) to be beneﬁcial, for all mixed FL algorithms.
We used adaptive clipping for the federated gradients (via the method described in Andrew et al.
(2021), without noise addition) and a ﬁxed clip value of 0 .
2 for the centralized gradients.
SGD was used for all optimizers: C LIENT O PTIMIZER and
S ERVER O PTIMIZER , and (if PT or 2- W GT)
C ENTRAL O PTIMIZER and M ERGE O PTIMIZER .
Evaluation Accuracy on Individual Data Splits Figure 5 shows the accuracy of models (trained either via mixed FL or pure FL) when evaluated on only federated data (Stack Overﬂow) or only centralized data (Wikipedia) individually.
Conﬁrming what we expect, the various mixed FL algorithms do a good job of achieving accuracy on both datasets.
But if we train only via FL (without mixing), then we do a good job of learning the federated data (Stack Overﬂow) character sequences, but aren’t nearly as accurate at predicting the next character in centralized data (Wikipedia) sequences.
B.2. M OVIE L ENS M OVIE R ECOMMENDATION
Dataset The MovieLens 1M dataset (GroupLens) contains approximately 1 million ratings from 6,040 users on 3,952 movies.
Examples are grouped by user, forming a natural data partitioning across clients.
For all mixed FL algorithms we study, we keep examples from 20% of users (randomly selecting 20% of users and shufﬂing the examples of these users) as the datacenter data, and use examples from the remaining 80% users as client data.
With this data splitting strategy, the server data won’t include the same individual client distributions but it will still be sampled from the same meta distribution of clients.
We then split the clients data into the train and test sets, resulting in 3, train, 483 validation, and 483 test users.
The average cache size of each client is ∼ 160 examples.
Model Architecture
The architecture used is the same as Anonymous (b), a dual encoder representation learning model with a bag-of-word encoder for the left tower (which takes in a list of movies a user has seen) and a simple embedding lookup encoder for the right tower (which takes in the next movie a user sees).
Hyperparameter Settings
The settings used in mixed FL training are shown in Tables 9 and 10.
In Subsection 4.3, SGD was used for all optimizers: C LIENT O PTIMIZER and S ERVER O PTIMIZER , and (if PT or 2- W GT)
C ENTRAL O PTIMIZER and M ERGE O PTIMIZER .
See Appendix B.4.4 for an additional experiment where ADAM is used as the S ERVER O PTIMIZER in 1- W GT.
Recall@10 As mentioned in Subsection 4.3, all mixed FL algorithms achieved similar global recall@10 compared to the baseline.
Figure 6 shows evaluation recall@10 over 2000 training rounds
[Table]
Comp.
( K · N 2 · d ) /( N 2 · d ) /( N 2 · d ) /( N 2 · d ) /
Comm.
2 · N · d 2 · n · d 3 · n · d 3 · n · d
[Table]
C OMPUTATION AND C OMMUNICATION S AVINGS FOR M OVIE R ECOMMENDATION
This section provides a detailed analysis of the computation and communication savings brought by mixed FL in the movie recommendation task.
For movie recommendation, both the input feature and the label are movie IDs with a vocabulary size of N .
They share the same embedding table, with an embedding dimension of d .
The input features and label embedding layers account for most of the model parameters in a dual encoder.
Therefore, we use the total size of the feature and label embeddings to approximate the model size:
M = ( N + N ) · d . Batch size is B i and local steps per round is K .
Let the averaged number of movies in each client’s local dataset for each training round be n , smaller than B i · K .
Computation As shown in the second row of Table 11, the amount of computation for regularization term is ( K · N 2 · d ) / 2 if calculating on-device (baseline).
When computing the regularization term on the server (mixed FL), the complexity is ( N 2 · d ) / 2 .
The total computation saving with mixed FL is (( K − 1) · N 2 · d ) / 2 .
We use ( N 2 · d ) / 2 instead of N 2 · d for regularization term computation which is more accurate for an optimized implementation.
The total computation complexity of the forward pass is O
( B i d + B i d
+ B 2 i d ) , where the three items are for the bag-of-word encoder, the context hidden layer, and similarity calculation.
The hinge loss and spreadout computation is O
( B i )
+ O (0 .
5 N 2 d ) .
The gradient computation is O (2 B i
d 2 + 2 B 2 i d ) for network backward pass and O ( B i )
+ O ( Nd ) for hinge and spreadout.
Therefore, when computing the regularization term on the server with mixed FL, the computation savings for each client is 1 − ( B i d + 3 B i d
+ 3 B 2 i d + 2 B i ) /
( B i d + 3 B i d 2 + 3 B 2 i d + 2 B i + 0 .
5 N 2 d + Nd ) , which is 99.98% for all mixed FL algorithms.
Communication The communication overheads of each algorithm are presented in the last row of
Table 11.
For the baseline, the server and each client need to communicate the full embedding table and the gradients, so the communication overhead is 2 · N · d or 494KB.
With P ARALLEL T RAINING , the server and each client only communicate movie embeddings and the gradients corresponding to movies in that client’s local datasets.
Thus the communication trafﬁc is reduced to 2 · n · d or 20KB.
G RADIENT T RANSFER requires the server to send both the movie embeddings and gradients to each client.
The communication overhead then becomes 3 · n · d or 30KB.
Overall, mixed FL can save more than 93.9% communication overhead than the baseline. [Figure] [Figure]
A DDITIONAL O BSERVATIONS AND E XPERIMENTS B.4. E FFECT OF σ 2 C ON CONVERGENCE
Table 3 shows that the theoretical bounds on rounds to convergence are directly proportional to the client variance bound σ 2 and central variance bound σ 2 c .
Also, as discussed in Subsection 5.2, 1- WAY
G RADIENT T RANSFER is more sensitive to high central variance than the other two algorithms.
Whereas in the other algorithms the impact of σ 2 c on convergence scales with cohort size S , in 1- WAY
G RADIENT T RANSFER it scales with cohort size S and steps taken per round K .
To observe the effect of σ 2 c in practice, and compare its effect on 1- WAY G RADIENT T RANSFER vs.
2- WAY G RADIENT T RANSFER , we ran sweeps of CelebA smile classiﬁcation training, varying the central batch size |B c | .
The plots of evaluation loss and evaluation AUC of ROC are shown in
Figures 7 (1- W GT) and 8 (2- W GT).
For each central batch size setting, we ran 10 trials; the plots show the means of each setting’s trials, with corresponding 95% conﬁdence bounds.
Figure 7 conﬁrms the sensitivity of 1- WAY G RADIENT T RANSFER to central variance, with experiments using larger central batches B c converging faster than experiments using smaller central batches.
However, at least in the case of this task, the beneﬁts of lower variance disappear quickly.
The convergence of AUC of ROC did not appreciably improve for central batch sizes larger than 25.
Presumably there is little effect at these larger central batch sizes because in these cases the convergence is now dominated by client variance (i.e., further convergence improvements would come from increasing client batch size |B i | ).
Comparing Figure 8 with Figure 7, we empirically observe that 2- WAY G RADIENT T RANSFER has lower sensitivity than 1- WAY G RADIENT T RANSFER to central batch size/central variance.
B.4. T RADING η FOR K
The convergence bounds of Table 3 have an additional implication, in regards to the trade off between client learning rate η (and central learning rate η c ) and number of local steps taken K .
It’s better to reduce η and η c and increase K , but there are limits The convergence bounds are not related to client or central learning rate ( η or η c ), but are inversely related to local steps K .
In general, it’s best to take as many steps as possible, and if necessary reduce learning rates accordingly.
But there are limits to how large K can be.
First, clients have ﬁnite caches of data, and K will always be limited by cache size divided by batch size.
Second, in the case of 1- WAY G RADIENT T RANSFER , [Figure] any increase in K means that central variance σ 2 c must be proportionally reduced (as mentioned above), necessitating even larger central batch sizes (which at some point is infeasible).
We empirically observed this relationship by running smile classiﬁcation (Figure 9) and language modeling (Figure 10) experiments where client learning rate η
(and central learning rate η c ) are inversely proportionally varied with K .
For each hyperparameter conﬁguration we ran 5 trials; the ﬁgures include 95% conﬁdence intervals.
The results conﬁrm that reducing these learning rates, and making a corresponding increase in the number of steps, is beneﬁcial.
It never hurts convergence, and often helps.
B.4. D IFFERENCES IN EFFECTIVE STEP SIZE
Table 12 in Appendix C shows that in order to yield the convergence bounds stated in this paper, each algorithm makes different assumptions of maximum effective step size.
From this we draw one ﬁnal implication in regards to comparing the mixed FL algorithms.
For given η , maximum K varies by algorithm, or , for given K , maximum η varies by algorithm
Consider just effective federated step size ˜ η
= ηη s K for the moment.
Assume that server learning rate η s is held constant.
Then each mixed FL algorithm has a different theoretical upper bound on the product of client learning rate η and local steps per round K .
If using a common η , the theoretical upper limit on K varies by mixed FL algorithm.
Alternatively if using a common K , the theoretical upper limit on η varies by mixed FL algorithm.
The maximum effective step sizes of Table 12 imply that 2- WAY G RADIENT T RANSFER has narrower limits than 1- WAY G RADIENT T RANSFER on the allowable ranges of η and K .
It also indicates that for P ARALLEL T RAINING the allowable range of η , η c , and K depends on the B parameter from mixed FL ( G, B )-BGD
(Deﬁnition 5.2).
Some of this behavior has been observed empirically, when hyperparameter tuning our experiments (discussed in Subsection B.2).
For example, for the language modeling experiment, assuming a [Figure] [Figure]
The increased learning rate boosts progress on evaluation loss and accuracy early in optimization, but does not change the number of rounds ultimately required for convergence.
All three algorithms have reached similar loss and accuracy by round 3000, and are still converging.
constant number of steps of K = 16 ,
2- WAY G RADIENT T RANSFER tends to diverge when learning rate η was increased beyond 1.0, whereas 1- WAY G RADIENT T RANSFER is observed to converge even with learning rate η of 5.0.
(P ARALLEL T RAINING is in-between; it still converges with learning rate η of 3.0, but diverges when learning rate η is 5.0.)
An interesting characteristic to note is that using different η in different algorithms does not really impact comparative convergence.
Figureshows convergence in the language modeling experiment, when 2- WAY G RADIENT T RANSFER uses η = 1 .
and 1- WAY G RADIENT T RANSFER and P ARALLEL T RAINING both use η = 3 .
(in all cases, with K = 16 ).
The higher learning rate of 1- W GT and PT helps a little early, but does not impact the number of rounds to convergence.
This holds with the theoretical convergence bounds of
Table 3, which show a relationship with steps K but not learning rates (as also discussed above).
B.4. 1- W GT WITH ADAPTIVE OPTIMIZATION
We brieﬂy studied the performance of 1- WAY G RADIENT T RANSFER when using ADAM in place of SGD as the server optimizer, i.e., F ED ADAM (Reddi et al., 2020).
Note that the server adaptive optimizer requires a smaller learning rate to perform well.
Figure 12 reports the results of using
ADAM as the server optimizer with a server learning rate of 0.01.
All the other hyperparameters are the same as in Tables 9 and 10.
We observe that (1) ADAM works better than SGD, leading to better convergence, and (2) 1- WAY G RADIENT T RANSFER performs almost the same as the baseline when using ADAM.
We will extend our investigation of mixed FL with adaptive optimization in the future.
This will include studying methods for applying adaptive optimization to P ARALLEL T RAINING and 2- WAY G RADIENT T RANSFER ; these algorithms are more complicated since they involve additional optimizers (C ENTRAL O PTIMIZER and M ERGE O PTIMIZER ).
B.4. C OMPARISON WITH T RANSFER L EARNING
Section 3 discussed how transfer learning is ill-suited for mixed FL.
Here we show why.
Recall the bias mitigation mixing task of Section 4.2.
The end goal is a single trained language model that performs well both for high-end and low-end mobile phone users.
We use Stack Overﬂow and
Wikipedia datasets as respective experimental stand-ins, as they model the problem well.
They are in the same domain (Latin alphabet) and same language (English) but with distinct usage distributions.
Transfer learning sequentially trains the language model, ﬁrst ‘pretraining’ on one dataset and then switching to ‘ﬁne-tune’ on the other.
One could ﬁrst centrally pretrain (with Wikipedia), then switch to federated ﬁne-tuning (with Stack Overﬂow), or do the converse.
Figure 13 shows what happens.
The switches take place at 250 rounds.
We also show P ARALLEL T RAINING as a point of comparison.
Figure 13 shows that in transfer learning, catastrophic forgetting of the pretraining task occurs after switching to the ﬁne-tuning task.
The centralized pretraining case precipitously drops its evaluation accuracy on datacenter data (Wikipedia).
The federated pretraining case does the same with federated data (Stack Overﬂow).
Figure 13a shows that P ARALLEL T RAINING achieves superior evaluation accuracy on the stand-in for the desired inference distribution (i.e. the combined Stack Overﬂow and
Wikipedia datasets).
[Figure] To simplify Figure 13 we don’t show G RADIENT T RANSFER results, but they would match P ARALLEL
T RAINING (e.g., as shown in Figures 1b and 2b).
Because the P ARALLEL T RAINING and G RADIENT
T RANSFER algorithms presented in this paper jointly train on both data distributions in parallel (as opposed to sequentially), they avoid catastrophic forgetting.
[Table]
[Table]
C C ONVERGENCE T HEOREMS
The three subsections that follow state theorems for convergence (to an error smaller than (cid:15) ) for the respective mixed FL algorithms.
The convergence bounds are summarized in Table 3 in Section 5.
Tables 12 and 13 convey some supporting aspects of the convergence bounds, about limits on effective step size ( ˜ η
= ηη s K ) and assumptions on learning rates.
C. P ARALLEL T RAINING
Given Assumption 5.1, one can view P ARALLEL T RAINING as a ‘meta-F ED A VG ’ involving two ‘meta-clients’.
One meta-client is the population of IID federated clients (collectively having loss f f ), and the other meta-client is the centralized data at the datacenter (having loss f c ).
As such, we can take the convergence theorem for F ED A VG derived in Karimireddy et al.
(2020b)
(Section 3, Theorem
I) and observe that it applies to the number of rounds T to reach convergence in the P ARALLEL
T RAINING scenario.
Theorem C.1.
For P ARALLEL T RAINING , where the federated data is IID (Assumption 5.1), for β -smooth functions f f and f c which satisfy Deﬁnition 5.2, the number of rounds T to reach an expected error smaller than (cid:15) is:
[Equation]
convex
[Equation]
where F = f ( x (0) ) − f ( x ∗ ) , D 2 = (cid:13)(cid:13)(cid:13) x (0) − x ∗ (cid:13)(cid:13)(cid:13) 2 .
Conditions for above: η m ≥ 1; η c , ηη s ≤ 1 6 ( 1+ B 2 ) βKη m .
Proof.
The analysis is exactly along the lines of the analysis in Karimireddy et al.
(2020b)
, Appendix
D.2, in the context of F ED A VG .
Effectively, the analysis applies to the ‘meta-F ED A VG ’ problem of P ARALLEL T RAINING , with two ‘meta-clients’, one being the central loss/data (with stochastic gradients with variance of σ 2 c ) and the other being the federated loss/data.
The homogeneity of the clients and the averaging over the sampled clients effectively reduces the variance of the stochastic gradients to σ 2 / S .
The analysis follows in a straightforward manner by accounting for the variance in appropriate places.
We omit the details for brevity.
C. 1- WAY G RADIENT T RANSFER
We now provide convergence bounds for the 1- WAY G RADIENT T RANSFER scenario.
Unlike
P ARALLEL T RAINING , which could be thought of as a ‘meta’ version of an existing FL algorithm (F ED A VG ),
1- WAY G RADIENT T RANSFER is an entirely new FL algorithm.
As such, we must formulate a novel proof (Appendix D) of its convergence bounds.
Given Assumption 5.1, the following Theorem gives the number of rounds to reach a given expected error.
Theorem C.2.
For 1- WAY G RADIENT T RANSFER , where the federated data is IID (Assumption 5.1), for β -smooth functions f i and f c , the number of rounds T to reach an expected error smaller than (cid:15) is: µ -Strongly convex:
T =
˜ O (cid:16) ( σ 2 + KSσ 2 c ) KSµ(cid:15)
+ βµ log( 1 (cid:15) ) (cid:17) when η s > (cid:113) 58 S, η ≤ 1
βKη s
General convex:
T
=
O
[Equation]
( σ 2 + KSσ 2 c ) D 2 KS(cid:15)+ βD 2 (cid:15) (cid:17) when η s > (cid:113) 58 S, η ≤ 1 8 βKη s
Non-convex:
T
=
O
[Equation]
( σ 2 + KSσ 2 c ) βF KS(cid:15)+ βF(cid:15)
(cid:17) when η s ≥ √ S, η ≤ 1
βKη s where F = f
( x (0) )
− f ( x ∗ ) , D 2 = (cid:13)(cid:13)(cid:13) x (0) − x ∗ (cid:13)(cid:13)(cid:13) 2 .
Proof.
Detailed proof given in Appendix D.
C. 2- WAY G RADIENT T RANSFER
Given Assumption 5.1, one can view 2- WAY G RADIENT T RANSFER as a ‘meta-SCAFFOLD’ involving two ‘meta-clients’ (analogous to the view of P ARALLEL T RAINING as ‘meta-F ED A VG ’ in Subsection C.1).
As such, we can take the convergence theorem for SCAFFOLD derived in
Karimireddy et al.
(2020b)
(Section 5, Theorem III) and observe that it applies to the number of rounds T to reach convergence in the 2- WAY G RADIENT T RANSFER scenario.
Theorem C.3.
For 2- WAY G RADIENT T RANSFER , where the federated data is IID (Assumption 5.1), for β -smooth functions f f and f c , the number of rounds T to reach an expected error smaller than (cid:15) is:
[Equation]
when η m ≥ 1; η c , ηη s ≤ min (cid:16) 1 81 βKη
[Equation] [Equation]
[Equation] [Equation]
when η m ≥ 1; η c , ηη s ≤
[Equation]
m
Non-convex:
[Equation]
when
[Equation]
m
[Equation] [Equation]
where F = f ( x (0) ) − f ( x ∗ ) , D 2 = (cid:13) (cid:13)(cid:13) x (0) − x ∗ (cid:13) (cid:13)(cid:13) 2 .
Proof.
The analysis is exactly along the lines of the analysis in Karimireddy et al.
(2020b)
, Appendix
E, in the context of SCAFFOLD.
Effectively, the analysis applies to the ‘meta-SCAFFOLD’ problem of 2- WAY G RADIENT T RANSFER , with two ‘meta-clients’, one being the central loss/data (with stochastic gradients with variance of σ 2 c ) and the other being the federated loss/data.
The homogeneity of the clients and the averaging over the sampled clients effectively reduces the variance of the stochastic gradients to σ 2 / S .
The analysis follows in a straightforward manner by accounting for the variance in appropriate places.
We omit the details for brevity.
D C ONVERGENCE P ROOFS FOR 1- WAY G RADIENT T RANSFER
We will prove the convergence rate of 1- WAY G RADIENT T RANSFER for 3 different cases: strongly convex, general convex, and non-convex.
While these proofs are inﬂuenced by those for SCAFFOLD in Karimireddy et al. (2020b), we note some additional technical challenges we face with 1- WAY
G RADIENT T RANSFER .
First, as the name indicates, gradient information only ﬂows one way (in
SCAFFOLD, gradient information ﬂows from and to all clients).
Second, we only sample a batch of centralized data once per round (whereas SCAFFOLD can draw multiple batches during a round).
The result is that (as described in Subsection 5.2),
1- WAY G RADIENT T RANSFER works out to be more sensitive to central variance σ 2 c and number of steps K than 2- WAY G RADIENT T RANSFER .
We will ﬁrst state a number of deﬁnitions and lemmas in Subsection D.1 that are needed in proving convergence rate of 1- WAY G RADIENT T RANSFER , before proceeding to the actual proofs in
Subsection D.2.
D. A DDITIONAL D EFINITIONS AND L EMMAS
Note that some of the lemmas below are restatements of lemmas given in Karimireddy et al. (2020b).
We opt to restate here (versus referencing the relevant lemma in Karimireddy et al.
(2020b) each time)
due to the volume of usage of the lemmas, to ease the burden on the reader.
We will ﬁrst present the subset of deﬁnitions and lemmas which don’t make any assumptions of convexity (Subsection D.1.1), followed by the subset that assume convexity (Subsection D.1.2)
D.
. G ENERAL D EFINITIONS AND L EMMAS
Deﬁnition D.1 ( β -Smoothness) .
A function h is β -smooth if it satisﬁes:
[Equation]
h
[Equation]
y
[Equation]
x − y
[Equation]
, for any x , y
This implies the following quadratic upper bound on h :
[Equation]
Lemma D.2 (Relaxed triangle inequality) .
Let { v 1 , . .
. , v τ } be τ vectors in R d .
Then for any a > 0 :
[Equation]
Also:
[Equation]
Proof.
The ﬁrst statement for any a > 0 follows from the identity:
[Equation]
The second statement follows from the convexity of v → (cid:107) v (cid:107) 2 and Jensen’s inequality:
[Equation]
[Equation]
v i (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
τ τ (cid:88) i =
(cid:107) v i (cid:107)
Lemma D.3 (Separating mean and variance) .
Let { Ξ 1 , . . .
, Ξ τ } be τ random variables in R d which are not necessarily independent.
First suppose that their mean is E [Ξ i ] = ξ i and variance is bounded as E [ (cid:107) Ξ i − ξ i (cid:107) 2 ] ≤ σ 2 .
Then:
[Equation]
Now instead suppose that their conditional mean is E [Ξ i | Ξ i − 1 , . . .
Ξ 1 ] = ξ i , i.e. the variables { Ξ i − ξ i } form a martingale difference sequence, and the variance is bounded same as above.
Then we can show the tighter bound:
[Equation]
Proof.
For any random variable X , E [ X 2 ] =
( E [ X − E
[ X ]])
+ ( E [ X ]) 2 implying:
[Equation]

Expanding the last term of the above expression using relaxed triangle inequality (Lemma D.2) proves the ﬁrst claim:
[Equation]
For the second statement, ξ i is not deterministic and depends on Ξ i − 1 , . . .
Ξ 1 .
Hence we have to resort to the cruder relaxed triangle inequality to claim:
[Equation]
Then we use the tighter expansion of the second term:
[Equation]
The cross terms in the above expression have zero mean since { Ξ i − ξ i } form a martingale difference sequence.
D.1. D EFINITIONS AND L EMMAS A SSUMING C ONVEXITY
Deﬁnition D.4 ( µ -Convexity) .
A function h is µ -convex for µ ≥
if it satisﬁes:
[Equation]
for any
[Equation]
When µ > 0 , we have strong convexity, a quadratic lower bound on h .
Proposition D.5 (Convexity and smoothness) .
If client losses f i and centralized loss f
c are each β -smooth
(Deﬁnition D.1), and x ∗ is an optimum of the overall loss f (as deﬁned in Equation 1), then the following holds true:
[Equation]
Proof.
Deﬁne the functions ˜ f
i ( x ) :
=
f i ( x ) − (cid:104)∇ f i ( x ∗ ) , x (cid:105) , for all clients i , and the function ˜ f
c ( x ) :
= f c ( x ) − (cid:104)∇ f c ( x ∗ ) , x (cid:105) .
Since f i and f c are convex and β -smooth, so are ˜ f
i and ˜ f c , and furthermore their gradients vanish at x ∗ ; hence, x ∗ is a common minimizer for ˜ f
i , ˜ f c and f .
Using the β -smoothness of ˜ f
i and f c , we have
[Equation]
Note that 1 N (cid:80) Ni =1 ˜ f
i + ˜ f c = f
since 1 N (cid:80) Ni =1 ∇ f i ( x ∗ ) + ∇ f c ( x ∗ ) = ∇ f
( x ∗ )
= 0 .
The claimed bound then follows from the above two facts.
Proposition D.6 (Convex bound on gradient of overall loss) .
If client losses f i and centralized loss f
c are each µ -convex
(Deﬁnition D.4) and β -smooth (Deﬁnition D.1), and x ∗ is an optimum of the overall loss f (as deﬁned in Equation 1), then the expected norm of the gradient of overall loss is bounded as:
[Equation]
Proof.
[Equation]
Applying the relaxed triangle inequality (Lemma D
[Equation]
2) twice:
[Equation]
Applying Proposition D.5:
[Equation]
Lemma D.7 (Perturbed strong convexity) .
The following holds for any β -smooth and µ -stronglyconvex function h , and any x , y , z in the domain of h :
[Equation]
Proof.
Given any x , y , and z , we get the following two inequalities using smoothness (Deﬁnition
D.1) and strong convexity (Deﬁnition D.4) of h :
[Equation]
Further, applying the relaxed triangle inequality (Lemma D.2) gives:
[Equation]
Combining all the inequalities together we have:
[Equation]
The lemma follows since β ≥ µ .
Lemma D.8 (Contractive mapping) .
For any β -smooth and µ -strongly convex function h , values x and y in the domain of h , and step-size (learning rate) η ≤ 1 β , the following holds true:
[Equation]
Proof.
Expanding terms, and applying smoothness (Deﬁnition D.1):
[Equation]
If step-size is such that η ≤ 1 β , then:
[Equation]
Finally, for µ -strong convexity (Deﬁnition D.4) of h we have:
[Equation]
D. P ROOFS OF T HEOREM C.
We will now prove the rates of convergence stated in Theorem C.2 for 1- WAY G RADIENT T RANSFER .
Subsection D.2.1 proves the convergence rates for strongly convex and general convex cases, and
Subsection D.2.2 proves the convergence rates for the non-convex case.
Let S be the cardinality of the cohort of clients S participating in a round of training.
Let the server and client optimizers be SGD.
Let the clients all take an equal number of steps K , and let ˜ η be the ‘effective step-size’, equal to Kη s η .
With 1- WAY G RADIENT T RANSFER , the server update of the global model at round t can be written as:
[Equation]
Henceforth, let E | t [ · ] denote expectation conditioned on x ( t ) .
As in Karimireddy et al.
(2020b), we’ll deﬁne a client local ‘drift’ term in round t as:
[Equation]
Lemma D.9 (Bound on variance of server update) .
The variance of the server update is bounded as:
[Equation]
Proof.
Let S denote the set of clients sampled in round t .
For brevity, we will use ∆ x to refer to x ( t +1) − x ( t ) .
[Equation]
x (cid:107) 2 = E | t (cid:20)(cid:13)(cid:13)(cid:13)
x ( t +1) − x ( t ) (cid:13)(cid:13)(cid:13) 2 (cid:21) =
E | t   (cid:13) (cid:13)(cid:13)(cid:13)(cid:13)
˜
η
KS (cid:88) i ∈S
K (cid:88) k
= (cid:16) g
i ( x ( t,k ) i ) + g c ( x ( t ) ) (cid:17) (cid:13) (cid:13)(cid:13)(cid:13)(cid:13) 2   ≤ E | t   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ˜
η
KS (cid:88) i ∈S
K (cid:88) k
= (cid:16) g
i ( x ( t,k ) i ) − ∇ f f ( x ( t ) ) (cid:17) + (cid:16) ∇ f f ( x ( t ) )
+ g c ( x ( t ) ) (cid:17)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2  
We separate terms by applying the relaxed triangle inequality (Lemma D.2):
[Equation]
In term A
[Equation]
we separate mean and variance for the client stochastic gradients
g i
[Equation]
using Lemma D
[Equation]
and Equation 4:
[Equation]
We apply the relaxed triangle inequality (Lemma D.2) followed by smoothness (Deﬁnition D.1), to convert it to an expression in terms of drift E ( t ) :
[Equation]
KS
[Equation]
KS
In term B we have a full gradient of the federated loss ∇ f f and a stochastic gradient of the centralized loss g c .
We use Lemma D.3 to separate the stochastic gradient into a full gradient of the centralized loss ∇ f c and a variance term, allowing us to express in terms of full gradient of the overall loss ∇ f .
[Equation]
Combining A and B back together:
[Equation]
[Table]
D.2. C ONVEX C ASES
We will state two lemmas, one (Lemma D.10) related to the progress in round t towards reaching x ∗ , and the other (Lemma D.11) bounding the federated clients ‘drift’ in round t , E ( t ) .
We then combine the two lemmas together to give the proofs of convergence rate for the strongly convex ( µ > 0 ) and general convex ( µ = 0 ) cases.
Lemma D.10 (One round progress) .
Suppose our functions satisfy bounded variance σ 2 , µ -convexity (Deﬁnition D.4), and β -smoothness (Deﬁnition D.1).
If ˜ η < 18 β , the updates of 1- WAY G RADIENT
T RANSFER satisfy:
[Equation]
Proof.
The expected server update, with N total clients in the federated population, is:
[Equation]
The distance from optimal x ∗ in parameter space at round t is (cid:13)(cid:13) x
( t ) − x ∗ (cid:13)(cid:13) 2 .
The expected distance from optimal at round t + 1 , conditioned on x ( t ) and earlier rounds, is:
[Equation]
For clarity, we now focus on individual terms, beginning with C :
[Equation]
We can use convexity (Deﬁnition D.4) to bound C 1 , with x = x ( t ) , and y
= x ∗ :
[Equation]
We apply perturbed convexity (Lemma D.7) to bound C 2 , with x = x ( t,k ) i , y = x ∗ , and
z = x ( t ) :
[Equation]
[Table]
Combining C 1 and C 2 back together:
[Equation]
Now we turn to term D , which is the variance of the server update (from Lemma D.9):
[Equation]
We can leverage Proposition D.6 to replace the norm squared of the gradient of the overall loss:
[Equation]
Returning to our equation for the expected distance from optimal x ∗ in parameter space, and making use of the bounds we established for C and D :
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Assuming that ˜ η ≤ 18 β
[Equation]
Lemma D.11 (Bounded drift) .
Suppose our functions satisfy bounded variance, µ -convexity (Deﬁnition D.4)
[Equation]
and β -smoothness (Deﬁnition D
[Equation]
Then the drift is bounded as:
[Equation]
Proof.
We begin with the summand of the drift term, looking at the drift of a particular client i at local step k .
Expanding this summand out:
[Equation]
Separating mean and variance of the client gradient, then using the relaxed triangle inequality (Lemma
D.2) to further separate out terms:
[Equation]
Term F is bounded via the contractive mapping lemma (Lemma D.8)
[Equation]
provided that η ≤ 1 β :
[Equation]
t
[Equation]
Putting back into the bound on drift on client i at local step k
[Equation]
and letting a = K :
[Equation]
Unrolling the recursion:
[Equation]
The second inequality above uses the following bound:
[Equation]
Now separating mean and variance of the central gradient:
[Equation]
Finally
[Equation]
we apply Proposition D
[Equation] [Equation]
Assuming that ˜ η ≤
[Equation]
β :
[Equation]
Proofs of Theorem C.2 for Convex Cases Adding the statements of Lemmas D.10 and D
[Equation]
[Equation]
and assuming that η s > (cid:113) 58 S , η = 1 8 βKη s so that ˜ η = 18 β , we get:
[Equation]
We can now remove the conditioning over x ( t ) by taking an expectation on both sides over x ( t ) , to get a recurrence relation of the same form.
For the case of strong convexity ( µ > 0 ), we can use lemmas (e.g., Lemma 1 in Karimireddy et al.
(2020b), Lemma 2 in Stich (2019)) which establish a linear convergence rate for such recursions.
This results in the following bound 16 for T ≥ 8 β 3 µ :
[Equation]
[Equation]
16 β (cid:19)(cid:19) , where ¯ x ( T ) is a weighted average of x (1) , x (2) , . . .
, x ( T +1) with geometrically decreasing weights (1 − 3 µ ˜ η 2 ) 1 − r for x ( r ) , r = 1 , 2 , . . .
, T + 1 .
This yields an expression for the number of rounds T to reach an error (cid:15) :
[Equation]
For the case of general convexity ( µ = 0 ), we can use lemmas (e.g., Lemma 2 in Karimireddy et al.
(2020b), Lemma 4 in Stich (2019)) which establish a sublinear convergence rate for such recursions.
In this case we get the following bound:
[Equation]
where ¯ x ( T )
= 1 T +1 (cid:80) T +1 t =1 x ( t ) .
This yields an expression for the number of rounds T to reach an error (cid:15) :
[Equation]
In the above expression, D 2 is a distance in parameter space at initialization,
[Equation]
− x ∗
[Equation]
D.
N ON -C ONVEX C ASE
We will now prove the rate of convergence stated in Theorem C.2 for the non-convex case for 1- WAY
G RADIENT T RANSFER .
We will state two lemmas, one (Lemma D.12) establishing the progress made in each round, and one (Lemma D.13) bounding how much the federated clients ‘drift’ in a round during the course of local training.
We then combine the two lemmas together give the proof of convergence rate for the non-convex case.
Lemma D.12 (Non-convex one round progress) .
The progress made in a round can be bounded as:
[Equation]
Proof.
We begin by using the smoothness of f to get the following bound on the expectation of f ( x ( t +1) ) conditioned on x ( t ) :
[Equation]
Substituting in the deﬁnition of the 1- WAY G RADIENT T RANSFER server update (Equation 11)
[Equation]
and using Assumption 5.1 for the expectation of the client stochastic gradient:
[Equation]
Next, we make use of the fact that − ab = 1 2 (( b − a ) 2 − a 2 − b 2 ) ≤ − 1 2 a 2 + 1 2 ( b − a ) 2 :
[Table]
[Equation]
Next, we use smoothness (Deﬁnition D.1)
[Equation]
and the deﬁnition of client drift (Equation 12):
[Equation]
The last term is the variance of the server update, for which we can substitute the bound from Lemma
D.9:
[Equation]
Assuming a bound on effective step-size ˜ η ≤ 118 β :
[Equation]
Lemma D.13 (Non-convex bounded drift) .
Suppose our functions satisfy bounded variance and β -smoothness
(Deﬁnition D.1).
Then the drift is bounded as:
[Equation]
Proof.
We begin with the summand of the drift term
[Equation]
looking at the drift of a particular client i at local step k .
Expanding this summand out:
[Equation]
Separating mean and variance of the client gradient:
[Equation]
Next we use relaxed triangle inequality (Lemma D.2) to further separate terms:
[Equation]
In the above inequality
[Equation]
term H can be converted via smoothness (Deﬁnition D.1)
[Equation]
and term J is the variance of the centralized stochastic gradient (Equation 6).
Letting a = K
[Equation]
we have:
[Equation]
Unrolling the above recurrence, we get:
[Equation]
Assuming η s ≥ 1 , and ˜ η ≤ 118 β , we have K +1 K + 2˜
η 2 β 2 Kη 2 s ≤ 1 + 163162 K , and hence k −
[Equation]
Adding back the summation terms over i and k
[Equation]
the bound on client drift is:
[Equation]
Proofs of Theorem C.2 for Non-Convex Case Adding the statements of Lemmas D.12 and D.13, and assuming η s ≥ √ S , we get:
[Equation]
With the above, we have a recursive bound on the loss after round t + 1 .
We can use lemmas (e.g.,
Lemma 2 in Karimireddy et al.
(2020b), Lemma 4 in Stich (2019)) which establish a sub-linear convergence rate for such recursions.
Assuming ˜ η ≤ 118 β and η s ≥ √ S , we get:
[Equation]
In the above expressions, F is the error at initialization, f ( x (0) )
− f ( x ∗ ) .
This yields an expression for the number of rounds T to reach an error (cid:15) :
[Equation]