On Sample Optimality in Personalized Collaborative and Federated Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
In personalized federated learning, each member of a potentially large set of agentsaims to train a model minimizing its loss function averaged over its local datadistribution.
We study this problem under the lens of stochastic optimization,focusing on a scenario with a large number of agents, that each possess veryfew data samples from their local data distribution.
Specifically, we prove novelmatching lower and upper bounds on the number of samples required from allagents to approximately minimize the generalization error of a fixed agent.
Weprovide strategies matching these lower bounds, based on a gradient filteringapproach: given prior knowledge on some notion of distance between local datadistributions, agents filter and aggregate stochastic gradients received from otheragents, in order to achieve an optimal bias-variance trade-off.
Finally, we quantifythe impact of using rough estimations of the distances between local distributionsof agents, based on a very small number of local samples.
Introduction
A central task in federated learning [27, 34] is the training of a common model from local data setsheld by individual agents.
A typical application is when users ( e.g. mobile phones, hospitals) want tomake predictions ( e.g. next-word prediction, treatment prescriptions), but each has access to very fewdata samples, hence the need for collaboration.
As highlighted by many recent works ( e.g. Hanzelyet al.
[24], Mansour et al.
[33]), while training a global model yields better statistical efficiency on thecombined datasets of all agents by increasing the number of samples linearly in the number of agents,this approach can suffer from a dramatically poor generalization error on local datasets.
A solution tothis generalization issue is the training of personalized models, a midway between a shared modelbetween agents and models trained locally without any coordination.
An ideal approach would take the best of both worlds: increased statistical efficiency by using moresamples, while keeping local generalization errors low.
This raises the fundamental question: whatis the optimal bias/variance tradeoff between personalization and coordination, and how can it beachieved?
We formulate the personalized federated learning problem as follows, studying it under the lens ofstochastic optimization [5].
Consider N ∈ N ∗ agents denoted by integers
⩽
i ⩽ N , each desiringto minimize its own local function f
i : R d → R , while sharing their stochastic gradients.
Sinceonly a limited number of samples are locally available, we focus on stochastic gradient descent -likealgorithms, where agents each sequentially compute stochastic gradients g ki such that E (cid:2) g ki (cid:3) =
∇ f i .
In order to reduce the sample complexity, i.e. the number of samples or stochastic gradients requiredto reach small generalization error, agents thus need to use stochastic gradients from other agents, thatare biased since in general E (cid:2) g ki (cid:3) ̸ = ∇ f j .
Our algorithms are based on a gradient filtering approach: upon reception of stochastic gradients ( g kj ) j , agent
i filters these gradients and aggregates them usingsome weights λ j into (cid:80)
j λ j g kj , in order to achieve some bias/variance trade-off.
Contributions and outline of the paper In this paper, we consider an oracle model where at eachstep k
= 1 , 2 , . .
. , all agents may draw a sample according to their local distribution.
We aim atcomputing the number of stochastic gradients sampled from all agents, required to reach a smallgeneralization error, in terms of biases (distances between functions or distributions), regularity, andnoise assumptions.
The oracle model, main assumptions and problem formulations are given in
Section 2.
Our main contributions are then as follows.(i)
In Section 3 we prove information theoretic lower bounds: to reach a target generalization errorε > 0 for a fixed agent i , no algorithm can achieve a reduction in the number of oracle calls by afactor larger than the total number of agents ε -close –in a suitable sense– to agent i .(ii)
We next study a naive strategy based on weighted gradient averaging algorithms, coined all-forone , that matches this lower bound, at the cost of high communication and storage requirements.(iii)
We then propose in Section 5 a parallel extension of the simple weighted gradient averagingalgorithm that yields an efficient algorithm for collaborative generalization error minimizationproblems.
In this algorithm, agents compute stochastic gradients at their local estimate, and broadcastit to other agents who may use these to update their own estimates.
For x k = ( x k 1 , . . .
, x kN ) where x kiis the local estimate of agent i at iteration k , updates of the ALL - FOR - ALL algorithm write as:
[Equation]
where g k =
( g k 1 , ..., g kN ) for an unbiased stochastic gradient g ki of function f i , a step size η , and acarefully chosen symmetric matrix W .
Agent i thus uses stochastic gradients that are doubly biased,as gradients of a “wrong function” f j instead of f i computed at a “wrong location” x kj instead of x ki .
Interestingly, note that the A LL - FOR - ALL algorithm is not a gossip algorithm per se (see e.g. [38]),since the matrix W is not doubly-stochastic: gradients are not aggregated with weights that sum to 1 .
Moreover, W depends on the distance between local agents distributions, and thus requires eitherprior information on the local distributions, or estimating these distances as a pre-processing step.(iv)
We finally study in Section 6 the impact of estimating, based on a very limited number of samples,the matrix W to use in the A LL - FOR - ALL algorithm.
Under a mixture model assumption on theagents, we obtain that for a bounded – up to logarithmic factors – number of samples per agent, anyarbitrary small generalization can be reached, with an optimal collaboration speedup in terms of thenumber of agents in each mixture of the mixture model.
Related works Federated Learning is a paradigm in machine learning where training is done collaboratively among several agents, taking into account privacy constraints [27, 30, 34, 44].
A central taskis the training of a common model for all agents, for which both centralized approaches orchestratedby a server and decentralized approaches with no central coordinator [37] have been considered.
Thealgorithms we propose in this paper are well suited for a decentralized implementation.
As observed in Hanzely et al.
[24], training a common model for all users can lead to poor generalization on certain tasks such as e.g. next-word prediction.
To improve both accuracy and fairness,personalized models thus need to be learnt for each agent [32, 36, 46].
Approaches to personalizationinclude fine-tuning [10, 28], transfer learning techniques [17, 42, 44], using shared-representationmodels [11].
Personalization in FL can also be formulated as the training of local models with aregularization term that enforces collaboration between users [24] or with a meta-learning approach[9, 22, 26].
We refer the interested reader to Kulkarni et al.
[31] for a broader survey of Personalized
Federated Learning.
While the goal of personalization is to minimize local generalization errors, the above cited worksdo not provide theoretical guarantees over the sample complexity to obtain small local errors, butinstead control errors on a regularized problem, in terms of communication rounds or full gradientsused, and not in terms of samples used.
Deng et al.
[13], Mansour et al.
[33] among others providegeneralization errors under a statistical learning framework that depend on VC-dimensions andon distances between each local data distribution and the mixture of all datasets.
Donahue and
Kleinberg [14, 15] study the bias-variance trade-off between collaboration and personalization formean estimation in a game-theoretic framework.
Beaussart et al.
[3], Chayti et al.
[8], Grimberget al.
[23] also concurrently frame personalization as a stochastic optimization problem with biased gradients and are the works closest to ours.
They consider the training of a single agent with biasedgradients from another group of agents dedicated to this agent, and obtain performance guarantees interms of distance between individual function
f i and the average N − 1 (cid:80) j f j .
In contrast, we obtainmore general performance bounds based on distance bounds between all pairs of functions f i , f j (orequivalently, pairs of local distributions), in the case where all agents desire to minimize their localobjective; our “bias assumption” is also milder.
In addition, we prove matching lower bounds, andstudy under a mixture model the statistical efficiency of our approach.
Problem Statement and Assumptions
We now detail our objectives and the necessary technical assumptions.
We consider general stochasticgradient methods and formulate our problem, assumptions and algorithms accordingly.
2.1 Problem setting
Let D i for 1 ⩽
i ⩽ N be a probability distribution on a set Ξ (agent i ’s local distribution, not itsempirical distribution), ℓ : R d × Ξ → R a loss function.
We assume that the function f i that agent iaims at minimizing is the generalization error on agent
i ’s local distribution:
[Equation]
We coin this problem as collaborative generalization error minimization (GEM) .
At every iterationk = 1 , 2 , . . .
, agent i may access unbiased i.i.d.
estimates g ki ( x ) of ∇ f i ( x ) :
[Equation]
Counting the number of stochastic gradients used in the whole set of agents to reach a precisionε for f i thus reduces to computing the number of samples required from all agents to obtain localgeneralization error ε for agent i .
To specify the information shared between agents via access tostochastic gradients, we define the following oracle, that lets at every iteration all agents sample astochastic gradient.
After K oracle queries, each agent will have sampled K stochastic gradientsfor a total of NK in the whole set of agents.
Let (cid:8) ( ξ k 1 , . . . , ξ kN ) , k ⩾ 0 (cid:9) a sequence of i.i.d.
randomvariables of law D 1 × · · · × D N .
Given the initial shared knowledge S 0 , at iterations k = 1 , 2 , . . .
For all
⩽
j ⩽
N , agent j samples
ξ
kj
chooses some y kj ∈ R d as a S k − 1 -measurable function.
The shared memory is extended:
S k = S k − 1 ∪ (cid:8) g kj ( y kj ) , ξ kj , 1 ⩽ j ⩽
N
(cid:9) .
Agent j outputs x kj as a S k -measurable function.
For fixed target precision ε > 0 , the objective is to find, using T ε samples from all agents in totalcorresponding to K ε =
T ε /N oracle calls -, models with local generalization error ε .
Throughout thepaper, we assume that each function f
i is minimized over R d , and we denote by x ⋆i such a minimizer.
We further consider the following two standard assumptions.
Assumption 1 (Noise) .
There exists σ 2 > 0 such that for all 1 ⩽ i ⩽ N and x ∈ R d ,
[Equation]
Assumption 2 (Regularity) .
Functions
f i are µ -strongly convex and L -smooth [7].
2.2 Distribution-based distances
We first introduce extensions of classical Integral Probability Metrics (IPMs, [40]) to multivariatefunctions, i.e. pseudo-distances on the set of probability measures parameterized by a set H offunctions, fixed in the sequel.
Definition 1.
For H a set of functions from Ξ to R d and D , D ′ two probability distributions on Ξ , wedefine:
[Equation]
where ξ ∼ D and ξ ′ ∼ D ′ .
This family of pseudo-distances contains a large number of standard distances between distributions,including total variation (with the set of 1 -locally bounded functions, functions that send any ballof radius 1 in a ball of radius 1 ), the Wassertein distance (with the set of 1-Lipschitz functions),maximum mean discrepancies (with the unit ball of a RKHS), or even a simple distance betweenmeans of the distributions (with the set of 1 -Lipschitz affine functions), developed further in Section 6.
Assumption 3 (Distribution-based dissimilarities) .
For some non-negative weights ( b ij ) 1 ⩽ i,j ⩽ N , wehave d H ( D i , D j ) ⩽ b ij for all 1 ⩽ i, j ⩽ N .
We further assume that either of the following holds.
(Weak dissimilarities).
For all 1 ⩽ i ⩽ N , (cid:0) ξ ∈ Ξ (cid:55)→ ∇ x ℓ ( x ⋆i , ξ ) (cid:1) ∈ H .
(Strong dissimilarities).
For all x ∈ R d , (cid:0) ξ ∈ Ξ (cid:55)→ ∇ x ℓ ( x, ξ ) (cid:1) ∈ H .
The “weak dissimilarities” assumption is of course easier to satisfy than the “strong” version, and ourresults will ultimately depend only on the weak assumption.
Under Assumption 3 (weak version) and
Assumption 2, we have f i ( x ⋆j ) − f i ( x ⋆i ) ⩽ b 2 ij / (2 µ ) , which motivates our use of distribution-baseddissimilarity assumptions.
Notation: in the rest of the paper, variables t or T denote the number of stochastic gradients g kisampled (or data item sampled from personal distribution) from all agents, while variables k or Kdenote the iterates of the algorithms (or equivalently the number of oracle calls made).
Information-theoretic lower bound on the sample complexity
In this section, we prove lower bounds on the total number of stochastic gradients required from allagents, to reach ε -generalization for a given agent.
Our lower bounds apply to collaborative GEM, i.e.functions ( f i ) 1 ⩽
i
⩽
N of the form (1), for shared loss function ℓ and user distributions D 1 , . .
, D N .
An oracle ϕ : R N × d → I is a random function that answers some ϕ ( x ) ∈
I where I is an informationset, for every query x ∈ R N × d .
We adapt the definitions of Agarwal et al.
[1] of sample complexityfor SGD to our personalization problem.
Formally, the first-order oracle we defined in Section 2and that we write as ϕ (cid:0) ( D i ) i
=1 ,...,N , ℓ (cid:1) for shared loss function ℓ and user distributions D 1 , . .
, D N ,returns for x ∈ R N × d :
[Equation]
where ξ i ∼ D i .
Given distributions and a loss function (cid:0) ( D i ) i , ℓ (cid:1) , we denote by M the set of allmethods M =
( M K ) K ⩾ 0 : for any K ⩾ 0 , M K makes K oracle calls from oracle ϕ (cid:0) ( D i ) i , ℓ (cid:1)(while using T =
NK stochastic gradient samples from all agents), and returns x Ki ∈ R d for agenti as a measurable function of the K oracle calls.
For a set D of couples of distributions and lossfunction (cid:0) ( D j ) j , ℓ (cid:1) defining functions ( f i ) 1 ⩽
i
⩽
N , we are interested in lower-bounding:
[Equation]
where K εi (cid:16) M , (cid:0) ( D j ) j , ℓ (cid:1)(cid:17) is the number of oracle calls required to reach generalization error ε > 0for agent i , and writes as:
[Equation]
We now define the set D we consider for our lower bounds.
Let b =
( b ij ) 1 ⩽ i,j ⩽ N be non-negativeweights that verify the triangle inequality – namely, b ij ⩽
b ik + b kj for all i, j, k –, and let r, µ, L, σ >0 .
D Lµ ( r, b, σ ) is the set of all (cid:0) ( D i )
1 ⩽ i ⩽ N , ℓ (cid:1) , such that the functions f i parameterized by thesetuples of distributions and shared loss function verify Assumptions 1, 2 and 3 for σ 2 , µ, L > 0 andb ,
such that ∥ x ⋆i ∥
⩽ r for all 1 ⩽
i ⩽ N , and such that f i ( x ⋆j ) − f i ( x ⋆i )
⩽ b 2 ij / (2 µ ) .
We use thenotation a ( · ) =
Ω( b ( · )) for ∃ C > 0 such that a ( · ) ⩾ Cb ( · ) .
Theorem 1 (IT lower bound) .
Let ε ∈
(0 , 1 / 16) , ( b ij ) verifying the triangle inequality, r, σ > 0 .
Assume that the function set H contains the all 1 -Lipschitz affine functions and that d H ⩽ d TV .
For some constant C > 0 independent of the problem and any i ∈ { 1 , . . . , N } :inf M∈ M sup (( D j ) j ,ℓ )
∈ D L =1 /r 2 µ =1 /r 2 ( r,b,σ )
K εi (cid:16) M , (cid:0) ( D j ) j , ℓ (cid:1)(cid:17) =
Ω (cid:32) r
σε
N
εi
b 4 ) (cid:33) , where N εi ( b 2 ) = (cid:80)
j 1 { b 2 ij ⩽ ε } is the number of agents j verifying b 2 ij ⩽ ε .
The proof of this lower bound (Appendix B) builds on lower bounds based on Fano’s inequality [19]for stochastic gradient descent [1] or for information limited statistical estimation [18, 47], adaptedto personalization.
Theorem 1 states that, given the knowledge of ( b ij ) , σ 2 , µ = 1 and L = 1 , thereexist difficult instances of the problem that satisfy all three Assumptions 1, 2 and 3, such that thenumber of oracle calls needed to obtain a generalization error of ε for an agent i is lower-bounded bythe right hand side of the equation in Theorem 1.
The factor Cσ 2 r 2 ε − 1 is reminiscent of stochastic gradient descent, and is present in Agarwal et al.[1]: without cooperation, this is the sample complexity of SGD for a fixed agent.
Cooperation appearsin the factor 1 / N εi ( b/ 4) : the sample complexity is inversely proportional to the number of agentsj that have distributions similar to that of i .
One cannot hope for better than a linear collaborationspeedup proportional to agents ε/ 4 -close to i in terms of the distance d H .
Theorem 1 is a worst-caselower bound, so that a collaboration speedup could be leveraged even for small ε , but this wouldrequire making stronger additional assumptions.
The A LL - FOR - ONE algorithm: parallel weighted gradient averagings
After providing lower complexity bounds in Theorem 1, we present in this section a naive algorithmicapproach based on weighted gradient averagings (WGA), that proves to be sample-optimal.
Eachagent i keeps N shared local models x k 1 , . . .
, x kN , where x kj estimates x ⋆j at iteration k (the knowledgeof x kj needs to be shared by all agents).
At each iteration k , when a sample ξ kj is obtained at agent j ,it is used by that agent to compute unbiased estimates of ∇ f j ( x ki ) for all i ∈ [ N ] .
The iterates of the
WGA algorithm write as, where λ ij ⩾ 0 are such that (cid:80)
j λ ij = 1 for all 1 ⩽
i ⩽ N :
[Equation]
for some step size η > 0 .
We call this algorithm that consists in performing N parallel WGAalgorithms A LL - FOR - ONE (AFO), since every iteration of each gradient averaging for a given nodei requires all the other nodes to compute one stochastic gradient for i .
WGA is thus equivalent totraining models on the mixture of distributions ( D j ) j with weights ( λ ij ) j for all i .
Theorem 2.
Let ( x ki ) 1 ⩽
i ⩽ N,k ⩾ 0 be generated with (2) , and assume that Assumptions 1, 2 and 3(strong version) hold.
For any K ⩾ 0 and 1 ⩽ i ⩽ N , and for η as in Equation (6) ,
[Equation]
Let ε > 0 .
For a specific choice of λ ij =
[Equation]
for a number of oracle calls of:
[Equation]
where N ε i is previously defined in Theorem 1.
Since the oracle complexity of the WGA algorithm matches that of our lower bound, this provesthat our lower bound is optimal.
However, this algorithm may be difficult to use in practice: (i) the choice of λ ij is an explicit function of distribution distances ( b ij ) (defined in Assumption 3) thatcan be (statistically speaking) as hard to compute as solving our optimization problem; and (ii) thememory requirements and computation/communication costs of WGA can be prohibitive for large Nand large ε (they scale with N εi for agent i ).
Note that the strong version of Assumption 3 used in
Theorem 2 can be replaced by a more classical uniform bound of the form ∥∇ f i − ∇ f j ∥ ⩽ b ij .
We first begin by solving this latter issue – an algorithmic one – in the next section, by introducingand studying the A LL - FOR - ALL algorithm.
We discuss (i) in Section 6, where we provide scenariiover which statistical theoretical guarantees can be derived on the error made by estimating thesedistribution distances using only a few samples.
The A LL - FOR - ALL algorithm
In this section, we present the A LL - FOR - ALL algorithm (AFA), an adaptation of the weighted gradientaveraging algorithm.
For 1 ⩽ i ⩽ N , initialize x 0 i = x 0 ∈ R d .
At iteration k , let x
ki ∈ R
d
beagent
i ’s current estimate of x ⋆i , and denote x k
= ( x ki ) 1 ⩽ i ⩽ N ∈ R N × d .
For a step size η > 0 anda matrix W ∈ R N × N with non-negative entries (remarkably and as discussed later, W will notnecessarily verify (cid:80) j W ij = 1 ), iterates of the all-for-all algorithm are generated with Algorithm 1.
In Theorem 3, we control the averaged local generalization error amongst all agents:
[Equation]
Theorem 3 (A LL - FOR - ALL algorithm) .
Let K > 0 , η > 0 , and W a matrix of the form W = ΛΛ ⊤for some stochastic matrix Λ = ( λ ij ) 1 ⩽
i,j ⩽ N .
Assume that Assumptions 1, 2 and 3 (weak version)hold.
The iterates ( x ki ) k ⩾ 0 , 1 ⩽
i ⩽ N generated with Algorithm 1 verify, for η as in Equation (7) :
[Equation]
As for the AFO algorithm, we can deduce from this result the number of oracle calls required by the
A LL - FOR - ALL algorithm to reach an averaged ε -generalization, under the idealistic setting where thedistribution-based distances b ij are accessible.
Corollary 1.
Let ε > 0 .
Under the same assumptions as in Theorem 3, for a choice of matrix
W = ΛΛ ⊤ where λ ij = 1 { b 2 ij/µ<ε }
N εi ( b 2 /µ ) ,
the A LL - FOR - ALL algorithm (Algorithm 1) returns ( x Ki )
1 ⩽ i ⩽ Nsatisfying 1 N (cid:80)
Ni =
1 f i ( x K ε i ) − f i ( x ⋆i ) ⩽ ε , for a number of oracle calls satisfying:
[Equation]
Denoting K
ε
i ) the oracle complexity of the WGA algorithm - that matches the lower bound -, weobserve that the A LL - FOR - ALL algorithm reaches an averaged ε -generalization with an number oforacle calls K ε , of K ε ⩽ 1 N (cid:80) i K ε ( i ) .
The speedup in comparison with a no-collaboration strategy(all agents locally performing SGD) is 1 N (cid:80) i 1 N εi ( b 2 /µ ) :
the mean of all local speedups.
Remark 1.
In Theorem 3, as its proof shows, the quantities b 2 ij / (2 µ ) in the last term can in fact bereplaced by the quantities
f i ( x ⋆j ) − f i ( x ⋆i ) , that control how well the optimal model for j generalizesfor i , and the bias induced by A LL - FOR - ALL iterations is a weighted average of these quantities.
Note that in our lower bound (Theorem 1), we enforce that the functions considered are required tosatisfy f i ( x ⋆j ) − f i ( x ⋆i ) ⩽ b 2 ij / (2 µ ) .
We believe this notion of function proximity that we leverage tobe the weakest achievable in our setting; no prior work uses such a mild proximity assumption.
Perhaps surprisingly, matrix W is in general not a gossip matrix ( i.e. such that W 1 = 1 ): agent idoes not aggregate a convex combination of stochastic gradients, but a combination with scalarsthat do not necessarily sum to 1.
We thus cannot say that the ALL - FOR - ALL algorithm acts as if, inparallel, each agent i trains a model on the mixture of distributions D j with weights W ij .
In fact, asthe analysis shows below, agent i trains a model on the mixture of distributions, with weights λ ij , ifΛ is a stochastic square root of matrix W ( ΛΛ ⊤ = W ), as in the AFO algorithm.
In order to accountfor inter-dependencies between agents that do not directly share information, the all-for-all gradientfiltering uses weights W ij to aggregate information, instead of λ ij .
Propagating information using amatrix W , that induces a similarity graph G W on { 1 , . .
. , N } , such that ( ij ) ∈
E W
if W ij > 0 , isquite natural [4, 43]; yet, ours is the first analysis to give such precise generalization error bounds,through the use of a stochastic optimization framework.
In comparison to Theorem 3, the classical personalized FL approaches that consider personalizedlocal models of the form x i =
¯ x − δ i , where ¯ x is some global quantity shared by all agents, perturbed(and personalized) by some local quantity δ i ( e.g. averaging between local and a global models), canbe seen as the special instances where, for all i , we have λ ii = 1 − α
i
and λ ij =
α i N − 1 if i ̸
=
j forsome α i , and leads to bias terms of the form 1 N (cid:80)
i α
i N − 1 (cid:80) j ̸
= i b ij .
Full and naive collaboration (asingle model trained for all users) corresponds to λ ij = 1 /N for all i, j , and leads to a bias termof 1 N 2 (cid:80) i,j b ij .
The degrees of freedom offered by our matrix W (and by coefficients λ ij ) enablepairwise agent adaptation, and tighter generalization guarantees and bias/variance tradeoffs.
Proof sketch of Theorem 3.
Since brutally analyzing convergence of the iterates ( x k ) generated withx k +
= x
k − WG k seems impossible due to both gradient biases and model biases between agents,we study these iterates through the introduction of a different but related problem.
This approach isin fact similar to some decentralized optimization ones, where a dual problem or a related energyfunction is often introduced [21, 38], upon which well-studied algorithms are applied.
The relatedproblem we formulate is different from and more flexible than all the different personalized FLproblems in the literature [24, 41], that consider regularization terms that enforce consensus.
Forλ =
( λ ij ) 1 ⩽
i,j ⩽ N a stochastic matrix (such that for all 1 ⩽
i ⩽ N , we have (cid:80) Nj =1 λ ij = 1 ), let f Λbe defined as:
[Equation]
where ¯ f = 1 N (cid:80) i f i .
Gradient descent on f Λ writes as y k
+
=
y k − η Λ ⊤ ∇ ¯ f
(Λ y k ) , where∇ ¯ f ( x ) = 1 N (cid:0) ∇ f i ( x i ) (cid:1) 1 ⩽
i ⩽ N for any x ∈ R N × d .
Importantly, notice that denoting x k = Λ y k andsince W = ΛΛ ⊤ , we have the recursion x k +
=
x
k − ηW ∇ ¯ f ( x k ) , making an analysis of the iterates( x k ) possible.
In our case, we however use stochastic gradients given by our oracle.
The full gradient∇ ¯ f ( x ) is thus replaced by ( g ki ( x i ))
i .
Defining ( y k ) k with the recursion: y k +
=
y k − η ((Λ G k
( y k ) i ) i ,initialized at y
= x 01 = . . .
=
y 0 N = x 0 N we have x k = Λ y k , for all k ⩾ 0 , where ( x k ) is generatedusing Algorithm 1.
As a consequence, controlling in Theorem 3 the function values 1 N (cid:80) i f i ( x ki ) isequivalent to controlling f Λ ( y k )
(these two quantities are equal).
The bias-variance trade-off thuswrites as, where y Λ minimizes f Λ and x ⋆ = ( x ⋆i ) 1 ⩽
i ⩽ N :
[Equation]
The rest of the proof, deferred to Appendix D, consists in relating these two terms to Λ and b .
After providing the optimization tools and results to answer for the shortcomings of weighted gradientaveragings, we now turn to quantifying the impact of the use of estimated values ˆ b ij of b ij , in orderto close the loop.
Estimation of d H ( D i , D j ) as a pre-processing step
The sample complexity of estimating the distances d
H
( D i , D j ) depends on the complexity of thefunction space H .
While the estimation of Wassertein or total variation distances are usually hard(in O (1 /S 1 /d ) where d is the ambient dimension and S the number of samples available for theestimation, see e.g. [45]), maximum mean discrepancy (MMD) distances often exhibit lower samplecomplexities in O (1 / √ S )
Moreover, explicit assumptions on the loss function can also providelow sample complexities, as shown below for quadratic loss functions.
In order to formulate statistical results for the estimation of the pairwise distribution-based distancesd H ( D i , D j ) , we need to make additional structural assumptions, on both H and the distributions.
Inspired by Collins et al.
[11], we focus on analyzing an instance of our general GEM setting forquadratic losses and linear models, under which the generalization error of a given agent i writes as:
[Equation]
where z i = ( a i , b i ) is a random variable on R d × R .
The stochastic gradients thus write as∇ x ℓ ( x, ξ i )
=
( a ⊤ i x − b i ) a i for z i =
( a i , b i ) , and are thus linear functions of ξ i
=
z i z ⊤ i .
Hence,
Assumption 3 (weak version) is satisfied for H the set of D ⋆ -Lipschitz and affine functions, where
D ⋆ bounds all ∥ x ⋆ i ∥ for 1 ⩽
i ⩽ N , leading to:
[Equation]
We make the following assumption on the law D
i of the random variables ξ i : they are non-isotropicsubgaussian random variables, that thus benefit from concentration inequalities that are dimensionindependent [20, 29].
Assumption 4.
For some non-negative symetric matrix Σ and all 1 ⩽ i ⩽ N , ξ i are centered andΣ -subgaussian:
[Equation]
and we denote as ν 2 the largest eigenvalue of Σ , and d eff =
∥ Σ ∥ 2 ν 2 its effective dimension.
Importantly, note that d eff can be arbitrarily smaller than the ambient dimension - for the MNISTdataset, d eff is less than 3, while the ambient dimension is 712 [20].
Depending on a smaller dimensionis also an assumption that Collins et al.
[11] use in their work by exploiting shared representations.
We now formulate a structural assumption on the set of agents: there are M clusters C 1 , . .
, C M of Cagents each (to ease notations, with a total number of agents N = MC ).
Within each cluster, agentsdistributions share the same objective, and clusters are “well-separated”.
These models are popularfor modelling population heterogeneity and provide a formal framework for clustering problems; werefer the interested reader to Melnykov and Maitra [35] for a detailed survey on the subject.
Assumption 5 (Well-separated clusters of agents) .
For M, C ⩾ 1 , N writes as N = MC and thereexists {C 1 , . .
, C M } a partition of { 1 , . . .
, N } , µ 1 , . .
. , µ m such that for all 1 ⩽ m ⩽ M , |C m | =
Cand for all i, j ∈ C m , we have E [ ξ i ] =
E [ ξ j ] = µ m .
We denote ∆ 2 = min m ̸
= m ′ ∥
µ m − µ m ′ ∥ 2 andassume that ∆ 2 > 0 .
When distribution-based distances were given (as in Corollary 1), Algorithm 1 achieved the optimalcollaboration speedup, linear in 1 /C under Assumption 5 and for small enough target precision ε .
The cluster model is thus the natural baseline for our problem.
In the case where agents estimate withwhom to collaborate as we do in the sequel, reaching this collaboration speedup of 1 /C will henceprove the effectiveness of the approach.
We assume that agents possess a limited number of samples.
More precisely, for 1 ⩽
i ⩽ N and
S, K ⩾ 1 , agent i possesses K + S i.i.d.
samples of drawn from D i , S of which are dedicated toestimating who to collaborate with, the K remaining dedicated to the optimization process i.e. torunning A LL - FOR - ALL iterations for a number K of oracle calls.
For 1 ⩽ i ⩽ N , let ˆ µ i be an estimation of E [ ξ i ] made with S i.i.d.
samples ξ i, 1 , . . . , ξ i,S , and for1 ⩽ i, j ⩽ N let ˆ b ij be the following estimation of d H ( D i , D j ) :
[Equation]
Computing these distances can be done using only ˜ O ( N ) communications (rather than the N 2communications of a naive approach) by performing randomized gossip communications [6] on thecomplete graph.
Theorem 4 (A LL - FOR - ALL with estimated biases) .
Assume that Assumptions 1, 2, 3 (for someunknown biases b ij ), 4 and 5 hold.
Under the setting described, let ˆΛ be the stochastic matrix withentries
[Equation]
for some u > 0 that verifies u ⩾
[Equation]
d eff S .
The A LL - FOR - ALL algorithm with W =
ˆΛˆΛ ⊤ outputs( x Ki )
1 ⩽ i ⩽ N verifying, where b max = max i,j b ij :
[Equation]
where the mean is taken over both biases estimates ( ˆ b ij ) and gradient estimates ( g k i ) .
Corollary 2.
Under the same assumptions as Theorem 4 and for ε > 0 , the A LL - FOR - ALL algorithmwith estimated biases as described above reaches an averaged generalization error of ε as long as:
[Equation]
Forgetting about the logarithmic factors, only a bounded number of local samples for each user ( Sand K ) are required to reach an averaged arbitrarily small generalization error ε > 0 , in the limitwith an arbitrary large number of agents ( N and C ).
Indeed, due to our regularity assumptions, K– the number of samples kept for the optimization problem – is required only to be of order κ , thecondition number of the problem.
The number of samples S used for estimating the biases is requiredto be of order ν 21 d eff / ∆ 2 , the “signal-to-noise” ratio of our mixture model [35, 39], a natural quantityto depend on.
Corollary 2 hence shows that the optimal collaboration speedup is achieved, up tologarithmic factors: in order to reach an arbitrary small generalization error ε > 0 , are only requiredconstant orders for S and K (the number of samples locally available) if the number of agents islarge enough i.e. if N = ˜Ω( M/ε ) , where M is the number of clusters i.e. we have a linear speedupin the clusters population.
We numerically illustrate our theory in Appendix A on synthetic datasets,with clustered agents (as in this section) in more general settings.
Conclusion
In this paper, we quantified in terms of function and distribution biases, stochastic gradient noise,target precision ε > 0 and functions regularity parameters, the benefit of collaboration betweenagents for shared minimization using stochastic gradient algorithms.
Our lower bound (Theorem 1)states that, under prior knowledge on the distances between local distributions, the collaborativespeedup can be linear only in the first phase of the optimization when the generalization error islarge compared to the distances between distributions.
More specifically, for a given agent
i , thecollaboration speedup is linear in the number of agents that are ε -close to i .
Moreover, we showthat the A LL - FOR - ONE algorithm allows such a speedup and is thus sample optimal.
However, thisalgorithm requires high computation and communication capacities, a drawback that can be mitigatedby the use of a novel algorithm called A LL - FOR - ALL , that benefits from the same collaborationspeedup while being cheaper to deploy.
Finally, we studied the impact of estimating distancesbetween distributions as a pre-processing step to the optimization phase; under a mixture modelassumptions on the agents, we obtain an optimal collaboration speedup.
Extending our results –lower and upper complexity bounds – to other regularity assumptions and Section 6 to more generalsettings, as well as incorporating local steps in the A LL - FOR -
ALL algorithm (even though to lightencommunications, unbiased compressors could here be used, as our analysis encompasses these) areinteresting questions left for future work.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[No] This work
is mainly theoretical, and we do not believe it to have any negative societal impact.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] , except when proving
similar results, we do not repeat the full proof but only sketch the main differences.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[No]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[No] The experiments we ran were consistent through the runs,
and are illustrations of our theory.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[No]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]