FedMix: A Simple and Communication-Efﬁcient Alternative to Local Methods in Federated Learning
Abstract
Federated Learning (FL) is an increasingly popular machine learning paradigm in which multiple nodes try to collaboratively learn under privacy, communication and multiple heterogeneity constraints.
A persistent problem in federated learning is that it is not clear what the optimization objective should be: the standard average risk minimization of supervised learning is inadequate in handling several major constraints speciﬁc to federated learning, such as communication adaptivity and personalization control.
We identify several key desiderata in frameworks for federated learning and introduce a new framework, FedMix, that takes into account the unique challenges brought by federated learning.
FedMix has a standard ﬁnitesum form, which enables practitioners to tap into the immense wealth of existing (potentially non-local) methods for distributed optimization.
Through a smart initialization that does not require any communication, FedMix does not require the use of local steps but is still provably capable of performing dissimilarity regularization on par with local methods.
We give several algorithms for solving the FedMix formulation efﬁciently under communication constraints.
Finally, we corroborate our theoretical results with extensive experimentation.
Introduction
Federated Learning (FL) aims to enable machine learning in the decentralized setting while respecting data privacy.
Application domains of federated learning include healthcare, learning language models for virtual keyboards and speech recognition (Kairouz et al., 2019).
The promise of federated learning is that by participating in a distributed training process, clients can learn better machine learning models than they can using only their own data.
The main cost in using federated learning over local training lies in the network bandwidth used for the distributed training process.
Hence, federated learning must be ﬂexible enough to provide a beneﬁt to users without a prohibitive communication cost.
The standard formulation of FL is to cast it as an optimization problem of the form min x ∈ R d (cid:20)
[Equation]
where f i is the loss function on client i .
Thus, the goal of classical FL is for the n clients to collaboratively learn a single model, x ∗ = arg min f , to be deployed on all clients.
Recent development shows that using a single model for all clients can be severely detrimental to individual performance on many clients (Yu et al., 2020), defeating the purpose of joining distributed training.
Furthermore, (ERM) offers no clear tunable knobs that can accommodate constraints on the network bandwidth.
[Figure]
Can we ﬁnd a formulation for federated learning that is ﬂexible enough to accommodate the needs of federated learning, yet also solvable using standard methods?
1.1 Key properties of the FedMixture framework
Our main contribution is FedMixture , a novel and ﬂexible formulation for federated learning: deﬁne α i > 0 to be the personalization parameter for node i , and let x i def = min x ∈ R d f i ( x ) be the local solution to the i -th objective– note that x i can be found by solely running a local optimizer, and hence computing it requires no communication at all.
The FedMixture problem is
[Equation]
Once we ﬁnd a solution x ∗ of sufﬁcient quality for (FedMix), we deploy T i ( x ∗ )
=
α
i x ∗ +(1 − α i ) x i on node i as its ﬁnal personalized model.
We now enumerate some of the key properties of (FedMix): • Efﬁciently solvable as a ﬁnite-sum problem.
(FedMix) preserves the standard ﬁnite-sum formulation of empirical risk minimization.
Moreover, it preserves problem structure: we show (in Section 3) that when the f i are smooth (and/or convex), ˜ f is also smooth (resp. convex).
• Adaptive to communication constraints.
Communication efﬁciency is an important concern in federated learning, as often bandwidth is valuable and limited (Kairouz et al., 2019; Kone ˇ cný et al., 2016; Li et al., 2019).
Computing x i , a precondition to solving (FedMix), requires no communication at all, and can be done purely locally on node i .
If α i = 0 , then no communication at all is needed to compute the personalized model
T i ( x ∗ ) .
By varying α i between 0 and 1 , we can control the amount of communication needed to compute T i ( x ∗ ) .
We show (in Section 3) that given a communication budget of R steps, we can ﬁnd parameters
α i that allows us to solve (FedMix) in no more than R communication steps. • Adaptivity to personalization.
Our end-goal in federated learning is to generalize well on each client : this means that the solution deployed on node i should be tailored to its local data distribution, which may differ from the data distributions on other nodes.
In FedMix, varying α i enables us to amplify or reduce the effect of other objectives on the solution deployed on node i . In situations where the data on all of the nodes is sufﬁciently heterogeneous, we set α i to be small and the effect of other data on node i will be neglibile.
On the other hand, when the data on the different nodes is related we may set α i to be closer to 1 .
We observe a beneﬁt to varying α in this manner in practice: Figure 1 shows the effect of varying the α i on real data (see Section 4 for the details and for other experiments). FedMix
ﬁlls a gap that is unsatisﬁed by existing methods.
To the best of our knowledge, there is no other method for federated learning that is efﬁciently solvable via standard algorithms and also adaptive to communication and personalization constraints, and indeed both constraints are important in practice (Li et al., 2020).
We believe the key properties we enumerate can also serve as natural desiderata in the development of new formulations and methods for federated learning.
1.2 Related work
Personalization has garnered signiﬁcant recent interest in federated learning as personalized models often perform well in practice compared to non-personalized models (Jiang et al., 2019; Yu et al., 2020).
FedMixture is a model mixture method: the personalized solution is a mixture of a global model and a local model.
In recent work, Deng et al.
and Mansour et al.
propose model mixture methods and prove their statistical beneﬁts, while Zec et al.
(2021) introduce a similar formulation based on the mixture of experts framework.
Unfortunately, we show in the supplementary material that, from the perspective of optimization and without additional data, the formulations in all three works are trivially minimized at the local minimizers x 1 , . . .
, x n .
An alternative to model mixing is mixing in function space, where we optimize a mixture of objectives rather than a model mixture.
This mixture is often constructed to control model variance: examples of this approach can be found in (Dinh et al., 2021a;b; Hanzely and Richtárik, 2020; Huang et al., 2021).
In FedMix, we take the model mixture approach as it allows us to use pretraining to better solve the problem while still regularizing model variance (see Section 2.1).
A parallel line of work applies meta-learning methods like MAML to federated learning (Fallah et al., 2020a; Jiang et al., 2019): in Section 2.2 we motivate FedMix by taking MAML as our starting point.
Chen et al.
(2021) discuss the statistical limits of personalization and show that either solving empirical risk minimization or local training is optimal, depending on certain problem parameters;
However, as of yet there is no single optimal adaptive algorithm (from the statistical perspective).
There are several other techniques in federated learning that can be combined with our approach for better results, such as clustering (Sattler et al., 2020) or robust optimization (Reisizadeh et al., 2020).
The FedMix formulation
In this section we reintroduce and motivate the FedMix formulation in detail.
We deﬁne the FedMix objective as
[Equation]
where α i ∈ (0 , 1) is the personalization coefﬁcient for node i and x i is the minimizer of f i , for all
i = 1 , 2 , . . . , n .
We will use ˜ f
( x ) to refer to the objective in (1) when the α i and x i are clear from the context.
The FedMixture problem is then
[Equation]
Let α =
[ α 1 , . .
, α n ] be the vector of the personalization coefﬁcients.
If x ∗ = x ∗ ( α ) is a solution of (2), we call T i ( x ; α i , x i )
=
α
i x ∗ + (1 − α i ) x i the deployed solution on node i .
Like with ˜ f , we will refer to the deployed solution on node i as T i ( x ) when x i and α i are clear from the context.
2.1 Motivation 1: from Local GD to FedMix
The most popular algorithm for solving federated learning problems is the Federated Averaging algorithm (Kairouz et al., 2019), also known as Local (Stochastic) Gradient Descent (Local GD/SGD).
Local GD alternates steps of local computation on each node with steps of communication and aggregation.
More concretely, the Local GD update is:
[Equation]
where H is the number of local steps.
Early papers on federated learning (such as e.g. (Koneˇcný et al., 2016)) motivated local methods as communication-efﬁcient ways of solving (ERM), but subsequent theoretical development reveals that local methods are, in fact, quite bad solvers for (ERM) whenever there is signiﬁcant statistical heterogeneity among the clients (Woodworth et al., 2020).
Moreover,
Pathak and Wainwright (2020) show that for the linear least-squares problem, Local GD converges to a different point than the minimizer of (ERM).
More generally, the ﬁxed points of Algorithm (3) can be very different from the minimizer of (ERM) whenever H > 1 (Malinovskiy et al., 2020).
Hanzely and Richtárik (2020) show that a mild variant of Local GD can be interpreted as SGD applied on the nd -dimensional regularized objective f λ deﬁned by
[Equation]
where ¯ y = 1 n (cid:80) ni
=
y i is the counterpart of x in (ERM), and where λ is a regularization parameter determined according to the number of local steps.
Objective f λ is the summation of two terms: the ﬁrst asks that each node i ﬁnds a solution y i that minimizes its local objective well, while the regularizer ψ ( y 1 , . . .
, y n ) = 12 n (cid:80) ni
=1 (cid:107) y
i − ¯ y (cid:107) 2 forces the solutions y 1 , y 2 , . . .
, y n to be close to their average ¯ y .
Hence, Local GD incentivizes ﬁnding personalized solutions y 1 , y 2 , . .
, y n that have small population variance.
Hanzely and Richtárik (2020) note that as the λ parameter varies between 0 and ∞ , the solutions found by Local GD interpolate between the pure local optimal models (i.e. x i = argmin
x f i ( x ) ) and the solution of the global problem x ∗ (the minimizer of (ERM)).
We observe that the solutions y 1 , . . .
, y n found by Local GD are an implicit mixture of the local minimizers x 1 , . . .
, x n and the global empirical risk minimizer x ∗ .
Rather than seeking an implicit mixture of the local and global optimal models, we instead propose to ﬁnd an explicit mixture of the local optimal models and a global model: given any global model x (not necessarily the empirical risk minimizer), we choose coefﬁcients α 1 , α 2 , . .
. , α n (all between 0 and 1 ) and then deploy on node i the mixture
[Equation]
) we may then choose x as the best such global model by explicitly solving the optimization problem min x ∈ R d 1 n n
[Equation]
i =
1 f i ( α i x + (1 − α i ) x i ) , and this is exactly the FedMix formulation.
Observe that coefﬁcients α 1 , α 2 , . .
, α n should regularize the population variance of the deployed solutions
T 1 ( x ) , T 2 ( x ) , . .
, T n ( x ) , as in local methods.
We show this rigorously for equal α i in Proposition 3.
Our development thus leads us to a natural framework that captures the strength of local methods while also satisfying the desiderata speciﬁed in Section 1.1.
2.2 Motivation 2: from model-agnostic meta-learning to FedMix
We now motivate FedMix differently by starting with personalization via ﬁne-tuning .
The ordinary formulation of the federated learning problem (ERM) asks for a single global model to be used on all clients.
If the clients are sufﬁciently heterogeneous, a single model may perform badly on many of them (Jiang et al., 2019).
Personalizing a global model to each of the users’ custom data is often beneﬁcial in practice; For example, Wang et al.
(2019) study the beneﬁts of personalizing language models for a virtual keyboard application used by tens of millions of users.
They observe that a sizeable fraction of the users beneﬁt from personalization.
Personalization is often done in two steps:
Step I: initial model training
Find a “good” global model x global
Step
II: ﬁne-tuning
Personalize the global model x global on each client to get the personalized local models x i
Methods that ﬁt this framework are known as ﬁnite-tuning approaches : they include the modelagnostic meta-learning (MAML) family of methods (Finn et al., 2017).
In addition to its practical popularity, recent theoretical investigations reveal that ﬁne-tuning approaches, such as MAML, are also beneﬁcal from a statistical perspective (Chua et al., 2021; Fallah et al., 2021).
In MAML, we ﬁnd x global by optimizing for the loss after a single step of gradient descent, i.e. the MAML objective is
Find
[Equation]
where γ is a given stepsize.
Once x global is found, we may then ﬁne-tune it by running gradient descent for a number of steps on each node i locally using its own objective f i (Finn et al., 2017).
To gain further insight into what ﬁne-tuning is doing, we now consider the case when each f i is a quadratic function.
Because this problem is amenable to analysis, several authors have used it to study the theoretical properties of MAML (Charles and Koneˇcný, 2021; Collins et al., 2020; Gao and
Sener, 2020), and we follow in their footsteps.
Assume that each f i is a quadratic function, suppose that we have some initial global model x 0 , and we ﬁne-tune it by running gradient descent for H steps on node i : the next proposition shows the ﬁnal iterate is a matrix-weighted average of the initial solution and the optimal local solution:
Proposition 1.
Suppose that we run gradient descent for H steps on the quadratic objective f i = 12 x T A
i x − b Ti x + c starting from x 0 with stepsize γ > 0 .
Suppose that the stepsize satisﬁes γ ≤ 1 L i , where L i = λ max ( A i ) .
Then the ﬁnal iterate x Hi can be written as
[Equation]
where x i minimizes f i and J
i ∈ R
d × d is a matrix with maximum eigenvalue smaller than 1 , i.e. λ max ( J ) < 1 .
The proof of Proposition 1 and all subsequent proofs are relegated to the supplementary.
Plugging the result of Proposition 1 into Equation (6), observe that in MAML we ﬁnd the initial model x 0 by solving the problem min x ∈ R d 1 n n (cid:80) i =1 f i (( I − J i ) x i + J i x ) .
Hence, MAML is optimizing for a speciﬁc weighted average of the initial model x 0 and the local solutions x 1 , x 2 , . . .
, x n .
We thus propose to dispense with the speciﬁc matrix J i and instead optimize an average weighted with an arbitrary constant α i : min x ∈ R d 1 n n (cid:80) i =
1 f i ( α i x + (1 − α i ) x i ) , and this is exactly the FedMix formulation.
Observe that by properly normalizing or whitening the data and tuning α i we may accomplish a similar effect to multiplying by J Hi for any H .
This gives FedMix a new interpretation as an approximate generalized MAML, where we optimize the global model for performance after potentially many gradient descent steps rather than just a single step.
Theory and algorithms
In this section we aim to develop algorithms to solve (FedMix) in a communication-efﬁcient manner.
Before discussing concrete algorithms, we study a few algorithm-independent properties of (FedMix) that will come in handy for understanding the formulation and proving convergence bounds.
The following proposition shows that the formulation preserves smoothness and convexity.
This is in contrast, for example, to MAML, where the objective may be nonsmooth (Fallah et al., 2020b).
Proposition 2.
Suppose that each objective f i is L i -smooth.
That is, for any x, y ∈ R d we have (cid:107)∇ f i ( x ) − ∇ f i ( y ) (cid:107) ≤ L i (cid:107) x − y (cid:107) .
Then the FedMix objective ˜ f deﬁned in (1) is L α -smooth for
L α def = 1 n (cid:80) ni
=
α 2 i L i .
If each f i is convex, then ˜ f is also convex.
If each f i is µ i -strongly convex, then ˜ f is µ α strongly convex for µ α def = 1 n (cid:80) ni =
α 2 i µ i .
Our next result offers some insight into the variance-regularizing effect of the α i : in particular, when all the α i are equal, increasing α in (FedMix) directly decreases the variance of the deployed local models from their mean.
As discussed in Section 2.1, this is a key property of local descent methods that the FedMix formulation captures.
Proposition 3.
Suppose that α 1 = α 2 = . . .
=
α n = β in the FedMix formulation (FedMix) .
Let
T 1 ( x ) , T 2 ( x ) , . .
, T n ( x ) be the deployed models deﬁned in (5) .
If y 1 , . .
, y n are vectors in R d and ¯ y is their mean, we deﬁne V ( y 1 , . . .
y n ) as the population variance V ( y 1 , . .
, y n ) def = 1 n n (cid:80) i =1 (cid:107)
y
i − ¯ y (cid:107) 2 .
Then, V ( T 1 ( x ) , T 2 ( x ) , . . .
, T n ( x )) =
(1 − β ) 2 V ( x 1 , x 2 , . . . , x n ) .
One-shot learning is a learning paradigm where we may use only a single round of communication to solve the federated learning problem (Guha et al., 2019; Salehkaleybar et al., 2019).
When the personalization parameters are small enough, we can provably solve the FedMix problem with a single round of communication by computing a certain weighted average of the local solutions x 1 , x 2 , . . .
, x n .
Theorem 1.
Suppose that each objective f i is L i -smooth, let ˆ L def = 1 n (cid:80) ni
=1 L i .
Given the pure local models x 1 , x 2 , . . .
, x n , deﬁne the weighted average x avg def = n (cid:80)
[Equation]
We further deﬁne the constants D def =
[Equation]
any (cid:15) > 0 .
Assume that either max i =1 ,...,n α i ≤ √ 2 (cid:15)/ (cid:112) ˆ LD , or α i = β for all i and β ≤ √ 2 (cid:15)/ (cid:112) ˆ LD .
Then x avg is an (cid:15) -approximate minimizer of (FedMix) .
For α i larger than this, we need more communication rounds.
In the next subsection, we describe how distributed gradient descent can be used to solve the problem.
3.1 Distributed gradient descent
The simplest approach to solving (FedMix) is via distributed gradient descent (DGD): given the local models x 1 , x 2 , . .
, x n (precomputed before starting the process) and an initial global model x 0 , we run the update x k +
= x
k − γn n (cid:80) i =
α i ∇ f i (cid:0)
α i x k + (1 − α i ) x i (cid:1) .
The next theorem shows that under smoothness and strong convexity, DGD converges linearly to the (FedMix) solution.
Theorem 2.
Suppose that each f i in (FedMix) is L i -smooth and µ i -strongly convex.
Deﬁne x avg , L α , ˆ L , V and D as in Theorem 1.
Suppose that we run DGD for K iterations starting from x 0 = x avg .
Then the following hold:
i)
If the α i are allowed to be arbitrary, then for α max
[Equation]
we have ˜ f ( x k ) − min x ∈ R d ˜ f ( x ) [Equation]
(cid:16)
[Equation]
µ
[Equation]
L
[Equation]
(cid:17) K α
[Equation]
ˆ LD
[Equation]
ii)
Let ˆ µ def = 1 n (cid:80) ni
=1 µ i .
If α i = β for all i , then
[Equation]
There are four ways of making the right hand side in (7) (the communication complexity) small:
• Communicate more.
Increase the number of communications K .
• Homogeneous data .
The variance V = (cid:80) ni
=
w i (cid:107) x
i − x avg (cid:107) 2 can be seen as a measure of data heterogeneity.
More homogeneous data means smaller V , which leads to better performance.
• Train simpler models.
Focusing attention on models with smaller ˆ L (adjust model design), or larger ˆ µ (e.g., add more regularization).
• Put more weight on local models.
If we prefer local models to the global model, then α i is small, and hence fewer communications are needed to achieve any given accuracy.
Armed with Theorem 2, we make good on our promise in Section 1.1 and show that FedMix can be solved using any communication budget.
Looking at (7) we see that for any ﬁxed (cid:15) >
we have ˜ f ( x k ) − min x ∈ R d ˜ f
( x ) ≤ (cid:15) as long as β ≤ Aq k , where A = (cid:113) 2 (cid:15)/ (ˆ LV ) and q = 1 / (cid:113) 1 − ˆ µ ˆ L .
Putting this together leads to the following observations:
• If β = 0 , the problem can be solved with 0 communications (i
e. , each device i independently computes the pure local model x i ).
• If 0 < β ≤ A , the problem can be solved with 1 communication (i
e., compute x avg ).
This follows from Theorem 1, and also from the more general result Theorem 2 by setting K = 0 .
If A < β ≤ Aq , the problem can be solved with 2 communications (1 communication to compute x 0 = x avg , followed by one iteration of distributed GD).
• If Aq k − 1 < β ≤ Aq k , the problem can be solved with k + 1 communications (1 communication to [-computex 
-]compute x
0 = x avg , followed by K iterations of distributed gradient descent).
• If β = 1 , we need 1 communication to compute x 0 = x avg , followed by k ≥ ¯ L ¯ µ log ¯ LV 2 ε iterations of distributed gradient descent.
This is recovers the standard communication complexity of gradient descent needed to ﬁnd the optimal solution of the average risk minimization problem (ERM). In the supplementary, we develop other algorithms for solving (FedMix) such as distributed gradient descent with compression (Alistarh et al., 2017) and DIANA (Mishchenko et al., 2019). We note that because (FedMix) has a standard ﬁnite-sum form, many more algorithms can be used to solve it, e.g. accelerated minibatch SGD (Cotter et al., 2011) or SARAH (Nguyen et al., 2017). [Figure]
L o ss mushrooms ® = 1 ® =
1e-1 ® =
1e-2 ® =
1e-3 ® = 1e0 25 50 75 100 125 150 175 Communication rounds 1010 - 1010 -
L o ss a6a ® = 1 ® =
1e-1 ® =
1e-2 ® =
1e-3 ® =
1e-
Experiments
Logistic regression with l 2 regularizer.
For our ﬁrst experiment, we consider a setup where each device runs regularized logistic regression: f i ( x ) :
=
1 k i (cid:80) k i j =1 (cid:2) log (1 + exp ( − a (cid:62) i,j x )) (cid:3) + λ/ 2 (cid:107) x (cid:107) 2 , where a i,j ∈ R d are given for all j = 1 . . .
k i , k
i is a number of data points associated with device i and λ is a regularization parameter.
The objective is smooth and strongly convex.
We use four datasets from LIBSVM (Chang and Lin, 2011) for this task: w6a, mushrooms, ijcnn1.bz2, a6a .
We set all α
i ’s to be equal.
We divide data equally between all machines while preserving the order of data points such that i -th machine owns data with indices (cid:98) ( i − 1) r / n (cid:99) +up to (cid:98) ir / n (cid:99) , where r is the total number of data points.
We set the regularization parameter λ to 0 .
To ﬁnd pure local models for each machine we run gradient descent with step size 1 / L i , where L i is f i ’s smoothness parameter until the norm of the gradient is below 10 − 6 .
We set the condition number to be κ = 1 n (cid:80) ni
=1 L i /λ .
In this experiment, we investigate the convergence of gradient descent and look into the dependence between convergence and value of α .
As expected, Figure 2 conﬁrms that smaller values of α lead to better convergence as we rely more on local solutions and thus start closer to the optimal solution.
Also note that the speed of the convergence appears to be constant among different values of α which is also predicted by our theory as the same α on each machine does not affect the conditioning of the global problem, see Proposition 2.
We also use the DIANA algorithm (Mishchenko et al., 2019) with the random sparsiﬁcation (also known as Rand- k compression), where set k coordinates to zero at random before communicating gradients to the server.
Figure 3 shows the effect of varying k on the convergence of the method in terms of communication rounds.
We consider 7 values of k linearly spaced between 1 and d .
Similar to Figure 2, we observe that smaller values of α lead to better convergence.
The rate at which the algorithm converges linearly is controlled by the compression constant ω + 1 = d / k or the effective conditioning κ ( ω +1) / n .
[-M S E o v e r c li e n t s 70vs 0.0 0.
-]
This is in line with the theory for DIANA (Mishchenko et al., 2019).
For completeness, we include convex experiments with unregularized logistic loss, extra experiments for all the combinations of 5 datasets and 3 algorithms – gradient descent, compressed gradient descent, and DIANA, as well as a detailed experimental description in the supplementary material.
Generalization experiment 1: Fitting Sine Functions.
Following Finn et al.
(2017) and Zhou et al.
(2019), we show the generalization advantages of FedMix on the following regression problem.
We deﬁne i -th client’s function f i ( x )
=
a i sin ( x + b i ) , where amplitude a i and phase b i lie in the intervals [0 . 1 , 0 . 5] and [0 , 2 π ] , respectively.
For each client, we ﬁx a i and b i and sample 50 points uniformly at random from the interval [ − 5 . 0 , 5 . 0] .
We measure regression ﬁt in terms of mean squared error (MSE) loss.
To train a model each client adopts a neural net with 2 hidden layers of size 40 with tanh activation.
Further technical details are deferred to the appendix.
For the experiment,
[Figure]
[Figure]
a g
[Figure] .5 0.6 0.7 0.8 0.9 1.0 alpha 104×10 -2 6×102×10A v e r a {+g e M S E o v e r c li e n t s 70vs 0.0 0.
 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 alpha 104×10 -2 6×102×10A+}
v e r a g e
M S E o v e r c li e n t s 90vs
we ﬁrst sample 2 pairs { a i , b i } and each of 200 clients is assigned one pair, we investigate different proportion– (30 , 170) , (50 , 150) , (70 , 130) , (90 , 110) .
We then train our FedMix formulation with α i = α = 0 .
For testing for each client generates a new dataset of size 2000.
Figureshows average MSE over clients against different values of α for different proportions.
As this ﬁgure indicates, optimal α for which test average MSE is minimal can dramatically outperform the edge cases of either global model for all tasks or personalized model trained only on the local dataset.
Generalization experiment 2: Comparison to FOMAML and Reptile.
Inspired by Reddi et al. (2021), we conduct a similar experiment to compare generalization capabilities, i.e., test accuracy, of FedMix and its two baselines FOMAML (Finn et al., 2017), and Reptile (Nichol et al., 2018).
For the ﬁrst experiment (see Figure 1a), we take 500 train data points of two clients (with client ids ‘00000267’ and ‘00000459’) from the Stack Overﬂow dataset (TensorFlow Developers, 2021) and divide them among 100 workers so that there are 50 workers with 10 train data points from the ﬁrst client and another 50 with 10 data points from the second client.
For the second experiment (see Figure 1b), a worker gets 90 train data points from a distinct client.
For both experiments, each objective component f i is a cross-entropy loss for multi-class logistic regression.
Further technical details and the hyperparameters tuning for a fair comparison can be found in the supplementary.
In the test phase, for each client, we used a hold-out testing dataset of size 300 (the same dataset has been used for workers related to the same client in the ﬁrst experiment).
It can be observed from
Figure 1a, that for wide range of α i , α i ∈ { 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 } FedMix exhibits a better generalization than its classical meta-learning competitors–FOMAML and Reptile, and it can lead to improvement of up to 11% in recall@5.
Figure 1b shows that in the more real-world scenario FedMix outperforms
FOMAML and Reptile while showing its best test accuracy in non-edge α .
7 Other algorithms
Strongly convex objectives . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
DCGD . .
DIANA . .
Convex objectives . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
DCGD . .
DIANA . .
Nonconvex objectives . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
DCGD . .
DIANA . .
Discussion of other model mixture
methods
11 Limitations[-
-]
Basic facts and notation
The convexity of the norm (cid:107)·(cid:107) and Jensen’s inequality implies the triangle inequality for norms:
[Equation]
We say that a function g is L g -smooth if for any x,
y ∈ R
d we have
[Equation]
Note that (9) implies
[Equation]
We say that a function g is µ g -strongly convex for µ g >
if for all x, y ∈ R
d we have
[Equation]
If (11) holds with µ g = 0 , we say that g is convex.
We say C ∈ B d ( ω ) is a compression operator if C is unbiased (i.e., E [ C ( x )]
= x for all x ∈ R d ) and if the second moment is bounded as
[Equation]
Note that if C ∈ B d ( ω ) , then
[Equation]
Let f be a convex function.
Then the Bregman divergence associated with f for points x, y ∈ R d is deﬁned in the following way:
[Equation]
Nesterov (2018) shows that if function f is convex and L -smooth, then for all x and y
[Equation]
Let a, b ∈ R n be arbitrary vectors.
Then, it holds that
[Equation]
Missing proofs for Sections 2 and
This section collects the proofs of all propositions and theorems mentioned in the paper.
6.1 Proof of Proposition
In Section 2.2 we provided the motivation for FedMix through the lens of ﬁne-tuning.
We used the following proposition on ﬁne-tuning quadratics:
Proposition 1.
Suppose that we run gradient descent for H steps on the quadratic objective f i = 12 x T A
i x − b Ti x + c starting from x 0 with stepsize γ > 0 .
Suppose that the stepsize satisﬁes γ ≤ 1 L i , where L i = λ max ( A i ) , and suppose that A i is positive deﬁnite.
Then the ﬁnal iterate x Hi can be written as
[Equation]
where x i minimizes f i and J
i ∈ R
d × d is a matrix with maximum eigenvalue smaller than 1 , i.e. λ max ( J ) < 1 .
Proof.
The gradient descent update is
[Equation]
Note that because x i minimizes f i , we have ∇ f i ( x i ) = 0 by ﬁrst-order optimality.
Hence,
[Equation]
Using this in (17), x
[Equation]
Subtracting x
i from both sides, x
[Equation]
Iterating the above equality for H steps we get x Hi −
[Equation]
− x
[Equation]
Rearranging the terms we get (16) with J i def =
( I − γA i ) .
Observe that when γ ≤ 1 λ max ( A i ) and λ min ( A i ) > 0 we have that λ max ( I − γA i ) < 1 .
6.2 Proofs for algorithm-independent results 6.2.1 Proof of Proposition
Proposition 2.
Suppose that each objective f i is L i -smooth.
That is, for any x, y ∈ R d we have
[Equation]
Then the FedMix objective ˜ f deﬁned in (1) is L α -smooth for L α def = 1 n (cid:80) ni
=
α 2 i L i .
If each f i is convex, then ˜ f is also convex.
If each f i is µ i -strongly convex, then ˜ f is µ α strongly convex for µ α def = 1 n (cid:80) ni
=
α 2 i µ i .
Proof.
We separate the proofs in two cases: when each f i is smooth, and when each f i is (strongly) convex.
(i) Suppose that each f i is L i -smooth, and let x, y ∈ R d .
Then by direct computation we have
[Equation]
hence ˜ f is L α -smooth.
(ii) Suppose that each f i is µ i -strongly convex for µ i ≥ 0 (where µ i = 0 corresponds to just convexity).
Then for x, h ∈ R d we have
[Equation]
hence ˜ f is µ α -strongly convex if all the µ i are positive (resp.
convex if they are equal to 0 ).
6.2.2 A proposition for
bounding the gradient norm
Proposition 2 provides us with a simple way to bound the gradient of ˜ f in terms of the distance to the optimum x α :
Proposition 4.
Deﬁne x α as the solution to the FedMix problem (FedMix) :
[Equation]
Deﬁne L α as in Proposition 2.
Then for any x ∈ R
d we have
[Equation]
(19) and
[Equation]
Proof.
By Proposition 2, we have for any x, y ∈ R d that
[Equation]
Putting y =
x α and using that ∇ ˜ f
( x α ) =
we get (19).
The L α -smoothness of ˜ f implies
[Equation]
Putting y
=
x α recovers (20).
6.2.3 Proof of Proposition
Proposition 3.
Suppose that α 1 = α 2 = . . .
= α n = β in the FedMix formulation (FedMix) .
Let T 1 ( x ) , T 2 ( x ) , . .
, T n ( x ) be the deployed models deﬁned in (5) .
If y 1 , . .
, y n are vectors in
R d and ¯ y is their mean, we deﬁne V ( y 1 , . . .
y n ) as the population variance V ( y 1 , . .
, y n ) def =n (cid:80) ni
=1 (cid:107) y
i − ¯ y (cid:107) 2 .
Then,
[Equation]
Proof.
By direct computation observe
V (
[Equation]
( x [Equation]
[Equation]
( x
[Equation]
[Equation]
[Equation]
n ( x
[Equation]
= 1 n n (cid:88) i = (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
T i ( x ) − 1 n n (cid:88)
j =
T j ( x ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) =
n n (cid:88) i = (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) βx + (1 − β ) x
i − 1 n n (cid:88)
j = ( βx + (1 − β ) x j ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) =
n n (cid:88) i = (cid:13) (cid:13) (cid:13) (cid:13)(cid:13)(cid:13) (
[Equation]
6.3 Result on one shot averaging
Before proving Theorem [Equation] we will need the following lemma which shows that the gradient of the
FedMix and the functional suboptimality can be bound using a weighted average of the iterate norms:
Lemma 1.
Suppose that each f i is L i -smooth.
Then for any x ∈ R d we have,
[Equation]
and,
[Equation]
Proof.
{+Using the fact that ∇ f i ( x i ) = 
+}
for all
i
[Equation]
and applying the triangle inequality (8) and the
L i -smoothness of f i [Equation] we get
[Equation]
α i x + (1 − α i ) x i ) − ∇ f i ( x i )] (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (8) ≤n n (cid:88) i =
| α i | (cid:107)∇ f i ( α i x + (1 − α i ) x i ) − ∇ f i ( x i ) (cid:107) (9) =n n (cid:88) i =
α 2 i L i (cid:107) x − x i (cid:107) .
Using L i -smoothness of f i and that ∇ f i ( x i ) = 0 , we get f i ( α i x + (1 − α i ) x i ) = f i ( x i + α i ( x − x 0 ))
(10) ≤ f i ( x i ) + (cid:104)∇ f i ( x i ) , α i ( x − x i ) (cid:105) + α 2 i L i 2 (cid:107) x − x i (cid:107)= f i ( x i ) + α 2 i L i 2 (cid:107) x − x i (cid:107) 2 .
Averaging the above inequality yields (22).
As communicated earlier in the paper, we can solve (FedMix) by taking a weighted average of the pure local models x i if the α i are small enough.
This is a consequence of preprocessing step where we compute x i , which is communication-free, and also a consequence of the new formulation, where for small α the FedMix objective is less responsive to argument change (see Proposition 2).
Theorem 1.
Suppose that each objective f i is L i -smooth, let ˆ L def = 1 n (cid:80) ni
=1 L i .
Given the pure local models x 1 , x 2 , . . .
, x n , deﬁne the weighted average
[Equation]
We further deﬁne the constants
[Equation]
and
[Equation] [Equation]
Fix any (cid:15) > 0 .
Assume that either max i =1 ,...,n α i ≤ √ 2 (cid:15)/ (cid:112) ˆ LD , or α i = β for all i and β ≤ √ 2 (cid:15)/ (cid:112) ˆ LD .
Then x avg is an (cid:15) -approximate minimizer of Problem (FedMix) .
That is,
[Equation]
Proof.
Lemma 1 provides the following bound on the suboptimality of any point x :
[Equation]
where
[Equation]
Note that since the vectors { x i } are known, the expression E α ( x ) can be minimized in x , leading to a weighted average of the pure local models:
[Equation]
Note that x avg , can be computed using a single communication round.
By plugging (27) into (26), we can evaluate the error of the average model:
[Equation]
where
[Equation]
can be interpreted as the variance of the local optimal models { x i } .
We now give two ways how
E ( x avg ) can be bounded:
• Assume the customization parameters { α i } are allowed to be arbitrary.
Since x avg is in the convex hull of the set { x 1 , . .
, x n } , we have (cid:107) x
avg − x i (cid:107) ≤ max i,j (cid:107) x
i − x j (cid:107) .
Therefore,
[Equation]
Moreover,
[Equation]
where α max def = max i α i and ¯ L def = 1 n (cid:80) ni
=1 L i .
By plugging (30) and (31) into (28), we get
[Equation]
Require: Number of communication rounds K , stepsize γ , initial point x1: for k = 0 , 1 , . . .
, K − 1 do 2: for i = 1 , . . . , n in parallel do 3: Compute ∇ ˜ f i ( x k
[Equation] = α [Equation] ∇ f i (cid:0)
[Equation]
(cid:1) and communicate it to server. Average and broadcast the new iterate
[Equation]
• Assume the customization parameters { α i } are all equal: α i = β for all i . Then w i = L i (cid:80) j L j ,
and hence x avg and V α = V are independent of β .
Since L α = β 2 ¯ L , by plugging these expressions into (28) we get
[Equation]
Note that V ≤ D .
Plugging (32) or (33) into (25) yields the following observation: ﬁx any (cid:15) > 0 and assume one of the following conditions holds
• α max ≤ (cid:113) 2 (cid:15)LD .
• α i =
β for all i and β ≤ (cid:113) 2 (cid:15)LV .
Then the weighted average of the local optimal models x avg satisﬁes
[Equation]
where x α = argmin x ∈ R d ˜ f
( x ) .
6.4 Distributed Gradient Descent
Gradient descent is a simple but yet informative way to solve many optimization problems.
Here we provide the proof for the convergence of DGD.
Theorem 2.
Suppose that each f i in (FedMix) is L i -smooth and µ i -strongly convex.
Deﬁne x avg , L α , ˆ L , V and D as in Theorem 1.
Suppose that we run DGD for K iterations starting from x 0 = x avg .
Then the following hold:
i)
If the α i are allowed to be arbitrary, then for α max def = max i =1 ,...,n α i we have
[Equation]
ii)
If α i = β for all i , then
[Equation]
where ˆ µ def = 1 n (cid:80) ni
=1 µ i .
Proof.
Recall the standard result that gradient descent for an L -smooth and µ -strongly convex objective g satisﬁes for any initial point x[-
-]
[Equation]
[Equation]
where g ∗ = min x ∈ R d g ( x ) .
For a proof, see (Nesterov, 2018).
Note that by Proposition 2 we have that ˜ f is L α -smooth and µ α -strongly convex.
Specializing (34) to this case yields
[Equation]
Note that our initialization is the same x avg from Theorem 1.
Plugging the bounds of that theorem into (35) yields the theorem’s claims.
[Table]
In this section
[Table]
we include additional algorithms to solve (FedMix) formulation, namely, DCGD (Gorbunov et al., 2019) and DIANA (Mishchenko et al., 2019). In Table
[Table] we display convergence guarantees for these algorithms.
One can see that similarly to our results in Section 3 both of the algorithms requires less iterations in terms of convergence in both local deploy iterates T i ( x k ) ’s and functional value ˜ f
( x k ) .
In addition, for the standard (ERM) problem, i.e., α = 1 , we recover the best known convergence guarantees.
Below, we provide a derivation of these claims.
7.1 Strongly convex objectives
1 DCGD
We ﬁrstly introduce the DCGD algorithm followed by general convergence results that we later exploit to obtain a convergence guarantee for DCGD applied to (FedMix).
Algorithm 2 Distributed Compressed Gradient Descent with different noise levels ω i
Require: x 0 ∈ R d , learning rate γ g
ki = C
i ( ∇ i ( x k ) : g
[Equation]
(cid:80)
[Equation] [Equation]
=
[Equation]
Lemma 2.
Suppose each f i is L i -smooth and convex, and f is L -smooth.
Let C i : R d → R d be randomized compression operators satisfying C i ∈ B d ( ω i ) .
Let g k = 1 n (cid:80) ni
=1 C i ( ∇ f i ( x k )) .
Then
[Equation]
) where σ DCGD = 2 n 2 (cid:80) ni =1 ω i (cid:107)∇ f i ( x ∗ ) (cid:107) 2 .
Proof.
By applying bias-variance decomposition to E (cid:107) g k − ∇ f
( x ∗ ) (cid:107) 2 , we get
[Equation]
Since all functions are convex holds, function f as a linear combination of convex functions is also convex.
That is why the ﬁrst term enjoys the ’classic’ bound for convex and smooth functions expressed in Bregman divergence between x k and x ∗ (see inequality 14):
[Equation] [Equation]
We start looking into the second term by expanding the quadratic:
[Equation]
Since C i ∈ B ( ω i ) ∀ i is drawn independently [Equation] the expected value of each scalar product in the second sum becomes the scalar product of expected values, each of which is zero due to unbiasedness of a compression operator (see Section 5):
[Equation]
Then, we apply 12 (the second property of compressed operators) to get the ﬁnal upper-bound on E (cid:107)
[Equation]
Curiously but not surprisingly, the result resembles the Law of large numbers (indeed, if each squared gradient norm and noise level are bounded above by values R and ω 0 respectively, then the right side is bounded by ω 0 Rn , which converges to zero as n converges to inﬁnity), what partially veriﬁes the correctness of the proof.
The error E (cid:107) g k − ∇ f ( x k ) ||
is now estimated by squared gradient norms at iterate x k , which is, in general, a random point generated by CGD.
Since we already have some dependence on Bregman divergence in 37, which perfectly ﬁts assumptions for uniﬁed theories from papers Gorbunov et al.
and Khaled et al.
(2020) required in later theorems, we express the term in the right side of the last inequality in Bregman divergence, too.
By subtracting and adding the same vector ∇ f i ( x ∗ ) inside the norm operator we get
[Equation]
where in the second line we used (15).
Applying again the bound from 14 to the ﬁrst term and using the linearity of Bregman divergence we have
[Equation]
Incorporation of the results above gives the statement.
Proposition 5.
Suppose each f i is L i -smooth and µ i -strong.
Assume α
i ≡ α ∈ R for all i .
For any x 0 ∈ R d , it holds that
[Equation]
Proof.
[Equation]
Proposition 6.
Suppose each f i is L i -smooth and µ i -strong.
Assume α
i ≡ α ∈ R for all i .
Then,
[Equation]
Proof.
According to Proposition 5, for any x ∗ l where l ∈ { 1 , . . . , n } it follows that
[Equation]
Taking the average of both sides ﬁnishes the proof.
Theorem 3.
Assume all conditions of Lemma 2 hold, and each function f
i is µ i -strongly convex.
Then, if γ ≤
[Equation]
, then
[Equation]
Proof.
The lemma is a direct corollary of Theorem 4.1 from Gorbunov et al.
and Propositionwith constants A =
[Equation]
i
[Equation]
i
[Equation]
Corollary 1.
Assume all conditions of Theorem 3 hold, and α i ≡ α ∀ i .
If γ =
[Equation]
γ 0 1 α 2 , then
[Equation]
and
[Equation]
where
C
=
L i 2 µ L i
[Equation]
x
[Equation]
− x i
[Equation]
2 and C 2 =
L i
[Equation]
Proof.
Notice that n (cid:88) i = ω
i (cid:107)∇ [ f
[Equation]
( α i x
[Equation]
+ (1 − α i ) x i )] (cid:107) 2 = n (cid:88)
[Equation]
This brings us to the following bound on the neigbourhood:
[Equation]
We further investigate the rate of convergence
[Equation]
As can be seen, the rate of convergence does not depend on α .
Combining last two results we obtain the following dependence on α for CGD convergence:
[Equation]
where constants (cid:101) C 1 = 1 µ L i [Equation] x
[Equation]
− x i [Equation] 2 and (cid:101)
[Equation] = [Equation] do not depend on α .
In terms of suboptimality convergence for CGD we have
[Equation]
where C 1 = (cid:101) C 1 L 2 and C 2 = (cid:101) C 2 L 2 .
7.1.2 DIANA
We follow the same procedure as for the previous subsection.
Below, we introduce DIANA algorithm, followed by general theorem, which is then applied to (FedMix).
Algorithm
Require: x 0 , h 01 , . . .
, h 0 n ∈ R d , h
= 1 n (cid:80) ni
=
h
i 1: for k = 0, 1, 2, . . .
do 2:
Broadcast x k to all workers 3: for i = 1 , . . .
, n in parallel do 4:
[Equation]
Sample [Equation]
[Equation]
7: ˆ
[Equation]
[Equation]
ki 9: x
[Equation]
10: h
[Equation]
Lemma 3.
Suppose each f i is L i -smooth and convex, and f is L -smooth.
Let C i : R d → R d be randomized compression operators satisfying C i ∈ B d ( ω i ) .
Let g k = 1 n (cid:80) ni =
[Equation] ki + C i [Equation]
Then
[Equation]
where σ
[Equation]
Proof.
We start with bias-variance decomposition:
[Equation]
where the ﬁrst term in RHS bounded due to convexity and smoothness of function f : (cid:107)∇
[Equation]
[Equation]
LD f ( x k , x ∗ ) .
Function f is, indeed, convex as a linear combination of convex functions f i .
When expanding the second term in RHS of equation 48, we encounter scalar products, each of which is zero in expectation due to independence of C i for all i and unbiasedness of a compression operator (see the deﬁnition in Section 5).
[Equation]
Then
[Equation]
we apply bounded variance property of compression operators:
[Equation]
We subtract and add ∇ f i ( x ∗ ) inside each norm operator to split ∇ f i ( x k ) and h ki from each other using (15):
[Equation]
While leaving the second term in the last line unchanged (it is basically 2 σ k /n ), we apply (14) to each norm of gradient differences in the ﬁrst term and use max-function over L
i ω i to take out Bregman divergences:
[Equation]
where in the last part we used linearity property of Bregman divergence.
It remains to incorporate all results to get the statement.
Lemma 4.
Assuming all conditions of Lemma 3 hold, let
h k
+1 i
=
h
ki + β i C i ( ∇ f i ( x k ) − h ki ) and β i ≤ 1 ω
i +1 , where i ∈ { 1 , . . . , n } .
Then,
[Equation]
Proof.
Gorbunov et al.
(2019) state for β i ≤
[Equation]
From this it follows that
[Equation]
Theorem 4.
Assume all conditions of Lemmas 3, 4 hold, and each function f
i is µ i -strongly convex.
If γ ≤
[Equation]
, then
[Equation]
where D k = (cid:107) x k − x ∗ (cid:107) 2 + 4 n min
β
i γ 2 σ 2 k and L α is deﬁned as in Proposition 2.
Proof.
The lemma is a direct corollary of Theorem 4.1 from Gorbunov et al.
(2019), 3 andand strong convexity with constants:
A = L α + 2max { L i
α 2 i[-} , D 2 = 0 , M = 
 n min β i-]
[Equation]
n
[Equation]
Corollary 2.
Assume all conditions of Theorem 4 hold and α i ≡ α ∀ i .
If γ =
[Equation]
2 , then
[Equation]
[Equation]
and
[Equation]
(53) where C =
L i 2 (cid:16) 1 +
[Equation]
Proof.
Let us investigate the quantity D 0 .
We assume memory h
i 0 in DIANA algorithm is equal to ∇ [ f i ( T i ( x 0 ))] for i ∈ { 1 , . . . , n } .
Plugging in deﬁnitions and initializations we obtain expanded deﬁnition of D 0 :
[Equation]
We further upper bound the sum in the last term using smoothness properties of functions f i :
[Equation]
Plugging this result back into the deﬁnition of D 0 and applying further bounds we get
[Equation]
what shows that D 0 can be upper bounded by value independent of α .
Since Mγ 2 σ 2 k ≥
D {+k = (cid:107) x k − x ∗ (cid:107) 
 + Mγ 2 σ 2 k ≥ [Equation]+}
x k − x ∗
[Equation]
2 , and that is why
[Equation]
what means that the convergence of iterates is completely independent from α .
To derive the result for the convergence of functional values we use L α -smoothness of function ˜ f
(see Proposition 2):
[Equation]
7.2 Convex objectives 7.2.1 DCGD
As for the previous section, we ﬁrst introduce the algorithm, followed by general theorem, which is then applied to (FedMix).
Theorem 5.
Assume all conditions of Lemma 2 hold.
Let 0 < [Equation]
≤
[Equation]
Then,
[Equation]
where x k = (1 /k )
k (cid:80) j
=1 x j .
Proof.
The theorem is a corollary of Corollary [Equation] 1 from Khaled et al [Equation]
(2020) with constants
[Equation]
Corollary 3.
Assume all conditions of Lemma 2 hold.
Let α i
= α , ω i = ω for all i , γ =
( L i + 2 ω max i Li n ) α 2 and sup
[Equation]
[Equation]
, 1] inf x
[Equation]
= R .
Then,
[Equation]
Proof.
Let us investigate ˜ f
( x 0 ) − ˜ f ( x ∗ ) .
Let α i ≡ α ∈ R for all i .
Since min 1 n (cid:80) ni =
f i ( T i ( x ))
≥n (cid:80) ni
=1 min f i ( T i ( x ))
= 1 n (cid:80) ni =1 min f i ( x ) and each f i is L
i -smooth, we get
[Equation]
This observation shows that the dependence between ˜ f ( x 0 ) − ˜ f
( x ∗ ) and α is quadratic.
˜ f
( x 0 ) − ˜ f ( x ∗ ) diminishes to zero as α goes to zero what proves our intuition: the closer α is to zero, the less steps the algorithm needs to make till convergence.
We verify that the stepsize in the corollary ﬁts the restriction on stepsizes in Theorem 5.
[Equation]
Further assuming that inf x ∗ ∈X (cid:107) x 0 − x ∗ (cid:107) is bounded above for all α ∈ [0 , 1] by value R ∈ R , we obtain the following result:
[Equation]
The neighborhood governing term is γσ DCGD .
Let us assume α i ≡ α and ω i ≡ ω for all i .
[Equation]
Since x ∗ i minimizes f i ( x ) , ∇ f i ( x ∗ i )
= 0 .
Thus,
[Equation]
Inequality (14) further bounds each norm of gradient differences.
[Equation]
In the rightmost term point x ∗ depends on α because it is an optimal point of a problem dependent on α .
But the observation below explicitly reveals the dependence.
Since each f i is convex,
[Equation]
Taking the minimum of both sides of inequality above we get to
[Equation]
We established that
[Equation]
Together with equality 55 it brings us to the following bound on the neighbourhood:
[Equation]
where C does not depend on α .
7.2.2 DIANA
As for the previous section, we ﬁrst introduce the algorithm, followed by general theorem, which is then applied to (FedMix)
[Equation] Theorem [Equation]
Assume all conditions of Lemmas 3 and 4 hold.
Let< γ ≤
[Equation]
Then,
[Equation]
Proof.
The theorem is a corollary of Corollary 4.1 from Khaled et al.
(2020) with constants
[Equation]
Corollary 4.
Assume all conditions of Lemmas 3 and 4 hold.
Let α i
= α , ω i = ω , h 0 i = [-
-]0 ∈ R d for all i , γ =
[Equation]
Then
[Equation]
where
[Equation]
and
[Equation]
Proof.
Let us investigate stepsize γ .
The statement of the corrolary requires the stepsize equals its maximum value
[Equation]
Since α i
=
α, ω
i = ω, β i = 1 1+ ω for all i ,
[Equation]
[Equation]
From deﬁnitions it follows that σ 20 = n 2 σ DCGD if h 0 i = 0 ∈ R d .
Thus,
[Equation]
That is why
[Equation]
Remaining terms are analyzed as for DCGD in convex case.
7.3 Nonconvex objectives
In this section we analyze algorithms in the nonconvex case of functions.
Throughout this section we assume that x ∗ i is a stationary point of function f i .
For the sake of convenience we transfer Lemmafrom Khaled and Richtárik (2020) to here.
Lemma 5 (Lemma 1 from Khaled and Richtárik (2020)) .
Suppose f ∗ i =
min
f i ( x ) exists and f i is
L i -smooth.
Then, for any x ∈ R d it holds that
[Equation]
Throughout this section we will use notation f ∗ def = 1 n (cid:80) ni
=1 f ∗ i .
7.3.1 DCGD
Theorem 7.
Suppose f ∗ i =
min f i ( x ) exists and f i is L
i -smooth for each i .
Then, if γ ≤ 1 L α ,
[Equation]
Proof.
The proof is a direct application of Theorem 2 in Khaled and Richtárik (2020) to our setting.
First, let us analyze E (cid:107) g
[Equation]
x
[Equation]
[Equation] To analyze the ﬁrst term we refer ourselves to Lemma 5:
[Equation]
Plugging this in the previous inequality we get
[Equation]
what means that Assumption 2 from Khaled and Richtárik (2020) holds with A
=
2max { L i ω i
α 2 i } n , B = 1 , C = 0 .
Then, as Theorem 2 from the same paper states
[Equation]
for stepsize γ ≤ 1 L α it holds that
[Equation]
Corollary 5.
Suppose assumptions of Theorem 7 hold, α i ≡ α , ω i ≡ ω .
Let γ = γ 0 1 α 2 ≤ 1 L i α 2 {+and ∆ 
+} = sup α ˜ f
( x 0 ) [Equation] f ∗ .
[Equation]
Proof.
We note that when α i ≡ α , L α =
Lα 2 , which follows from the deﬁnition of L α .
Plugging equations L α = Lα 2 , γ = γ 0
α 2 , α i ≡ α and ω i ≡ ω into the main result of Theorem 7 we get
[Equation]
Noting that ˜ f
( x 0 ) − f ∗ ≤ ∆
we ﬁnish the proof.
7.3.2 DIANA
Theorem 8.
Suppose f ∗ i =
min f i ( x ) exists and f i is L
i -smooth for each i .
Let β k ≡ β and ω k ≡ ω for all k ∈ { 1 , . . . , n } .
Suppose β ∈
[Equation]
[Equation]
and γ
[Equation]
, and h [-
-]0 i = ∇
[ f i ( T i ( x 0 ))] .
Let ˆ x be a point chosen uniformly at random among iterates x 0 , x 1 , . . .
, x k − 1 generated by DIANA.
Then,
[Equation]
Proof.
The proof is following the proof of Theorem 1 in Li and Richtárik (2020) with small modiﬁcations: we rewrite the main proposition in the different way and give slightly different bound on the stepsize.
First, Lemma 9 in Li and Richtárik (2020) says that
[Equation]
where
[Equation]
and r is an arbitrary positive number.
We choose r = 2(1 + w ) > 0 .
For the sake of readability, we deﬁne function ζ in the following way
[Equation]
That is why we can write that θ ∈ min { 1 − β 2 ω, ζ ( β ) } .
Second, let us analyze the lower bound for θ
[Equation]
Note that for β lying in the range
[Equation]
we have
[Equation]
and ζ
[Equation]
and
[Equation]
Since ζ is a quadratic function with the ends of the parabola pointed downwards, for all β in the range we have ζ ( β ) ≥ 1 2(1+ ω ) , which altogether means that θ ≥ [-
 [Equation]-]1 2(1+ ω ) > 0 .
If we enforce ω (1 + r ) n
L 2 α γ 2 ≤ 4(1 + ω ) or γ ≤ 2 L α (cid:114) n (1 + ω ) ω (1 + r ) =
L α (cid:114) n (1 + ω ) ω (1 + 2(1 + ω )) , then ρ = θ − ω (1+ r ) n
L 2 α γ 2 ≥ 1 4(1+ ω ) .
Now let us switch to the proof of the convergence.
We ﬁrst note that due to smoothness of f i s and inequality 62 we get
[Equation]
Let us ﬁx ξ > 0 .
Then according to 63 we have
[Equation]
+1 ]
[Equation]
Let us notate ∆ k = f
( x k ) − f ∗ + ξσ 2 k and set ξ =
L α γ 2 2 ρ 1+ ωn .
Then from previous inequality it follows that
[Equation]
) (cid:107)
[Equation]
Let us deﬁne γ (cid:48)
the coefﬁcient in front of (cid:107)∇ f ( x k ) (cid:107) 2 in the last inequality.
When is γ (cid:48) larger than γ 2 ?
[Equation]
if
[Equation]
Let us deﬁne y :
=
L α γ and a := ω (1+ r ) n .
Then according to 65 ρ = θ − ay 2 .
Plugging this back to the last inequality we get
[Equation]
The last inequality holds for all y lying between zero and the positive root of the quadratic equation.
The positive root is
[Equation]
That means that the last inequality holds if
[Equation]
or the tightest bound for γ achieved on smallest θ gives
[Equation]
Going back to 67 we ﬁnally get
[Equation]
That is why
[Equation]
The proof is written for general function f ( x ) = 1 n (cid:80) ni
=
f i ( x ) .
To apply the theorem for our setting, we replace f by ˜ f .
Corollary 6.
Suppose all conditions of Theorem 8 hold and α i ≡ α for all i from { 1 , . . . , n } .
[-Letγ = 
-]Let γ
= 1 α 2 γ 0 , where γ 0 = 1 L min (cid:110) 1 2 √ η 0 ,
[Equation]
1+8 η 0 + (cid:111) .
Let ∆
= sup
α ˜ f ( x 0 ) − f ∗ Then,
[Equation]
Proof.
According to the deﬁnition of L α = 1 n (cid:80) ni
=1 L i
α 2 i , in the case of equal α
i s we get L α
= Lα 2 .
It remains to notice that the stepsize in the corollary condition still satisﬁes the condition on the stepsize γ in Theorem 8 and ∆ 0 ≥ ˜ f
( x ) − f ∗ .
Discussion of other model mixture methods
Deng et al.
(2020) propose that each client solves the local problem
[Equation]
where w ∗ is the minimizer of (ERM).
This does not result in any personalization since the “personalized” solution on each node is a reparameterization of each local solution x i .
Furthermore, the convergence theory that Deng et al.
(2020) develop does not recover the linear convergence of gradient descent at α i = 1 .
Mansour et al.
(2020) introduce a similar method, MAPPER, where they propose to solve
[Equation]
Again, this objective is trivially minimized by setting α i = 1 , y i = min x ∈ R d f i ( x ) , and
z =
(i.e. with no personalization at all).
Zec et al.
(2021) also introduce a similar formulation based on the mixture of experts framework, where they propose to ﬁrst learn the minimizer x ∗ of (ERM), learn the optimal local models x 1 , x 2 , . . .
, x n , and then learn a mixture of both the global and local models (i.e. the α i ) on each client.
Unfortunately, this is also ill-deﬁned, as α i = 1 will always perform best on the local training set, and hence if there is no additional data the optimization process cannot improve over the local minimizers x 1 , . . .
, x n .
- -       6 stepsize 0 .0 .0 .0 .1 .a
l p h a -4.
Validation loss on Grid Search
Experimental Details
Generalization experiment 1: Fitting Sine Functions.
We train each local model i until the gradient norm (cid:107)∇ f i ( x ) (cid:107) is below 10 − 2 .
Stopping criteria for global models is the same but with respect to FedMix formulation: (cid:107) 1 n (cid:80) ni
=1 ∇ f i ( T i ( x )) (cid:107) < 10 − 2 .
Gradient descent with line search at each iteration was used as a local and a meta optimizer.
Generalization experiment 2: Comparison to FOMAML and Reptile. We preprocess the raw data in the same way as in Reddi et al. (2021) for Stackoverﬂow logistic regression task, i.e., each feature vector is a bag-of-words representation of a user’s sentence, each label vector is a binary vector showing if a sentence relates to a particular question tag or not. Word vocabulary for the feature dataset is restricted to the 10000 most common words.
We restrict the task to the 500 most used tags.
For preprocessing, we used Tensorﬂow computational procedures from Xu et al.
To select 50 clients for the Figure 1, we map ﬁrst 5000 clients from train dataset to vector space with BERT Devlin et al. and run k-means with 10 clusters. The ﬁrst 50 clients from the ﬁrst cluster have been selected for the experiment.
The hold-out validation dataset is of size 100, for the ﬁrst experiment, and of size 110, for the second one.
To compute pure local models we run gradient descent until the norm of the gradient is less than [10 − 1 , 10 − 2 , 10 − 3 , 10 − 4 , 10 − 5 ] , respectively.
We train pure local models until the gradient norm is less than 10 − 4 , for the ﬁrst experiments, and 10 − 5 , for the second, as this tolerance level was observed to have the lowest generalization error on the validation dataset in our experiments.
Then, we use a grid-search to ﬁnd optimal stepsizes for gradient descent used for training FedMix.
Our grid for step-sizes are [5 · 10 − 2 , 5 · 10 − 1 , . . .
The results are presented in Figuresand 6.
As table shows, the smaller α is, the higher step-size the task needs to achieve the best generalization, which is in line with our convergence results, see Theorem 2 and Table 1.
To train FOMAML and Reptile, we set the number of inner steps to ﬁve and grid search outer and inner loop step-sizes.
The explored outer step-sizes are the same as for FedMix. Inner step-sizes iterate over the set [10 − 3 , 10 − 2 , . .
- -       6 stepsize 0 .0 .0 .0 .1 .a
l p h a -4.
Validation loss on Grid Search
After grid search, we run gradient descent for each value of alpha of FedMix, FOMAML, and Reptile for 10 000 and 50000 iterations in the ﬁrst and seconds experiments accordigly and report obtained test accuracy.
Extra Experiments
Logistic regression with l 2 regularizer.
On ﬁgures 7, 8, 9, and 10 we present the plots for regularized logistic regression loss for all datasets mentioned in the previous section and for allalgorithms: GD, CGD and DIANA.
Number of machines n is set to 50 for DIANA and GD, and is set to 8 for CGD.
Plots showing the dependence between loss and communication cost show that although in terms of communication rounds DIANA and CGD lose to classical GD, they are better with respect to communication cost, which is of more practical importance.
We run all the algorithms with their best theoretical step-sizes.
[Figure]
problem l 2 regularizer for four datasets, k is a sparsiﬁcation parameter of Random- k compressor.
Unregularized logistic regression.
Unregularized logistic regression task has got the following form f i ( x ) :
=
1 k i (cid:80) k i j =1 log (1 + exp ( − a (cid:62) i,j x )) , which is well-known to be a convex problem.
Similarly to regularized case, we present four plots ( 11, 12, 13, and 14) exhibiting convergence of three algorithms in terms of communication rounds and overall communication cost for four LIBSVM datasets.
Number of machines n is set to 100 for DIANA and GD, and to 8 for CGD.
We run all the algorithms with their best theoretical step-sizes.
Generalization experiment 1: Fitting Sin Functions.
On ﬁgure 15 we present the rest of the plots of this subsection.
[Figure]
[Figure]
[Figure]
o ss a6a, ® = 0.
k = 1 k = 24 k = 49 k = 73 k =k =
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
r a g
[Figure]
M
[Figure]
v e r c li
[Figure]
v e r a g e M S E
o
v e r c li e n t s
[Figure]
Limitations
Our current (FedMix) formulation is dependent on a computation of the local optimal solutions x i ’s.
Arguably, this is not well-deﬁned beyond strongly convex models.
Therefore, applying (FedMix) formulation to weakly convex or non-convex models might lead to some problems because of the existence of potentially several local optima and selection of x i affects (FedMix) formulation.
One possibility to remedy such issue is to use the ensemble of global and local models with the weights α i and (1 − α i ) , respectively, instead of the linear combination of x and x i .
Another possible limitation is the availability of local data to compute an optimal local solution as the small number of local data points might lead to a noisy estimate and subsequently, it can potentially lead to bad generalization of (FedMix).
To mitigate this limitation, we propose to use some clustering method to group several clients and compute x i ’s on clusters rather than single clients.
Further investigation of both of these limitations is out of the scope of our current manuscript and we plan to investigate this in future work.