On Sample Optimality in Personalized Collaborative and Federated Learning
Abstract
In personalized federated learning, each member of a potentially large set of agents aims to train a model minimizing its loss function averaged over its local data distribution.
We study this problem under the lens of stochastic optimization, focusing on a scenario with a large number of agents, that each possess very few data samples from their local data distribution.
Specifically, we prove novel matching lower and upper bounds on the number of samples required from all agents to approximately minimize the generalization error of a fixed agent.
We provide strategies matching these lower bounds, based on a gradient filtering approach: given prior knowledge on some notion of distance between local data distributions, agents filter and aggregate stochastic gradients received from other agents, in order to achieve an optimal bias-variance trade-off.
Finally, we quantify the impact of using rough estimations of the distances between local distributions of agents, based on a very small number of local samples.
Introduction
A central task in federated learning [31, 40] is the training of a common model from local data sets held by individual agents.
A typical application is when users ( e g. mobile phones, hospitals) want to make predictions ( e.g. next-word prediction, treatment prescriptions), but each has access to very few data samples, hence the need for collaboration.
As highlighted by many recent works ( e.g. Hanzely et al.
[28], Mansour et al.
[39]), while training a global model yields better statistical efficiency on the combined 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 to this generalization issue is the training of personalized models, a midway between a shared model between agents and models trained locally without any coordination.
An ideal approach would take the best of both worlds: increased statistical efficiency by using more samples, while keeping local generalization errors low.
This raises the fundamental question: what is the optimal bias/variance tradeoff between personalization and coordination, and how can it be achieved?
We formulate the personalized federated learning problem as follows, studying it under the lens of stochastic optimization [5].
Consider N ∈ N ∗ agents denoted by integers
⩽
i ⩽ N , each desiring to minimize its own local function f
i : R d → R , while sharing their stochastic gradients.
Since only a limited number of samples are locally available, we focus on stochastic gradient descent -like algorithms, 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 required to reach small generalization error, agents thus need to use stochastic gradients from other agents, that are 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 using some weights λ j into (cid:80)
j λ j g kj , in order to achieve some bias/variance trade-off.
1.1 Contributions and outline of the paper
In this paper, we consider an oracle model where at each step k = 1 , 2 , . .
, all agents may draw a sample according to their local distribution.
We aim at computing the number of stochastic gradients sampled from all agents, required to reach a small generalization error, in terms of biases (distances between functions or distributions), regularity, and noise 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 a factor 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 averaging algorithm that yields an efficient algorithm for collaborative generalization error minimization problems.
In this algorithm, agents compute stochastic gradients at their local estimate, and broadcast it to other agents who may use these to update their own estimates.
For x k = ( x k 1 , . . .
, x kN ) where x ki is 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 a carefully 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. [45]), 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 either prior 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 the agents, we obtain that for a bounded – up to logarithmic factors – number of samples per agent, any arbitrary small generalization can be reached, with an optimal collaboration speedup in terms of the number of agents in each mixture of the mixture model.
1.2 Related works
Federated Learning is a paradigm in machine learning where training is done collaboratively among several agents, taking into account privacy constraints [31, 36, 40, 51].
A central task is the training of a common model for all agents, for which both centralized approaches orchestrated by a server and decentralized approaches with no central coordinator [44] have been considered.
The algorithms we propose in this paper are well suited for a decentralized implementation.
As observed in Hanzely et al.
[28], 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 [38, 43, 54].
Approaches to personalization include fine-tuning [10, 34], transfer learning techniques [19, 49, 51], using shared-representation models [11].
Personalization in FL can also be formulated as the training of local models with a regularization term that enforces collaboration between users [28] or with a meta-learning approach [9, 25, 30].
We refer the interested reader to Kulkarni et al.
[37] for a broader survey of Personalized
Federated Learning.
While the goal of personalization is to minimize local generalization errors, the above cited works do not provide theoretical guarantees over the sample complexity to obtain small local errors, but instead control errors on a regularized problem, in terms of communication rounds or full gradients used, and not in terms of samples used.
Deng et al.
[14], Mansour et al.
[39] among others provide generalization errors under a statistical learning framework that depend on VC-dimensions and on distances between each local data distribution and the mixture of all datasets.
Donahue and
Kleinberg [16, 17] study the bias-variance trade-off between collaboration and personalization for mean estimation in a game-theoretic framework.
Beaussart et al.
[3], Chayti et al.
[8], Grimberg et al.
[27] 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 biased gradients from another group of agents dedicated to this agent, and obtain performance guarantees in terms of distance between individual function
f i and the average N − 1 (cid:80) j f j .
In contrast, we obtain more general performance bounds based on distance bounds between all pairs of functions
f i , f j (or equivalently, pairs of local distributions), in the case where all agents desire to minimize their local objective; our “bias assumption” is also milder.
In addition, we prove matching lower bounds, and study under a mixture model the statistical efficiency of our approach.
Concurrently, [15] use similar dissimilarity notions (Integral Probability Metrics, albeit with respect to different function spaces) to show upper-bounds that closely resemble ours of Theorem 3, when training over mixtures of distributions (without algorithmic solutions as the all-for-all algorithm), together with insights on who to collaborate with for agents.
These results, obtained in a different framework than ours (hypothesis and VC-dimension bounds, rather than our stochastic optimization framework), fall into the frame of our lower bound (and match it up to constant factors), form an orthogonal line of work.
Finally, data-heterogeneity has long been a challenge in Federated optimization, as for instance noted in the analyses and performances of the Local SGD algorithm [33, 53].
Many algorithmic solutions have been proposed to counterweight this effect [32, 42] (non-exhaustive list).
Yet this line of work studies the effect of data-heterogeneity on the convergence guarantees of FL algorithms that train one global model, irrespectively of the local generalization property of this trained global model.
Our work is orthogonal, and focuses on data-heterogeneity as a challenge for statistical meaning (local generalization) of the model(s) trained, as opposed to related works that study data-heterogeneity as a challenge in distributed or federated learning to design fast and scalable algorithms.
Putting into perspective these two views on the challenge data-heterogeneity in FL seems however necessary, and stresses its importance.
Problem Statement and Assumptions
We now detail our objectives and the necessary technical assumptions.
We consider general stochastic gradient 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 its empirical distribution), ℓ : R d × Ξ → R a loss function.
We assume that the function f i that agent i aims 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 iteration k = 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 local generalization error ε for agent i .
To specify the information shared between agents via access to stochastic gradients, we define the following oracle, that lets at every iteration all agents sample a stochastic gradient.
After K oracle queries, each agent will have sampled K stochastic gradients for 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.
random variables of law D 1 × · · · × D N .
Given the initial shared knowledge S 0 , at iterations k = 1 , 2 , . . .
For all 1 ⩽ 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 the paper, 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, [47]) to multivariate functions, i.e. pseudo-distances on the set of probability measures parameterized by a set H of functions, fixed in the sequel.
Definition 1.
For H a set of functions from Ξ to R d and D , D ′ two probability distributions on Ξ , we define:
[Equation]
where ξ ∼ D and ξ ′ ∼ D ′ .
d H is a pseudo-distance on the set of probability measures on Ξ .
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 ball of 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 between means 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 , we have 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 . 2.
(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 our results 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-based dissimilarity assumptions.
Notation: in the rest of the paper, variables t or T denote the number of stochastic gradients g ki sampled (or data item sampled from personal distribution) from all agents, while variables k or K denote 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 all agents, 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 information set, for every query x ∈ R N × d .
We adapt the definitions of Agarwal et al.
[1] of sample complexity for SGD to our personalization problem.
Formally, the first-order oracle we defined in Sectionand 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 : ϕ (cid:0) ( D i ) i , ℓ
[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 all methods 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 agent i as a measurable function of the K oracle calls.
For a set D of couples of distributions and loss function (cid:0) ( D j ) j , ℓ (cid:1) defining functions ( f i
[Equation]
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 ε >for 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-negative weights 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 these tuples of distributions and shared loss function verify Assumptions 1, 2 and 3 for σ 2 , µ, L > 0 and b , 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 the notation 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 any i ∈ { 1 , . . . , N } :
[Equation]
where N εi ( b 2 4 µ ) = (cid:80)
j 1 { b 2 ij ⩽ 4 µε } 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 [21] for stochastic gradient descent [1] or for information limited statistical estimation [20, 55], adapted to personalization.
Theorem 1 states that, given the knowledge of ( b ij ) , σ 2 , µ = 1 and L = 1 , there exist difficult instances of the problem that satisfy all three Assumptions 1, 2 and 3, such that the number of oracle calls needed to obtain a generalization error of ε for an agent i is lower-bounded by the 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 appears in the factor 1 / N εi ( b/ 4) : the sample complexity is inversely proportional to the number of agents j that have distributions similar to that of i .
One cannot hope for better than a linear collaboration speedup proportional to agents 4 µε -close to i in terms of the distance d H .
Theorem 1 is a worst-case lower bound, so that a collaboration speedup could be leveraged even for small ε , but this would require 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 algorithmic approach based on weighted gradient averagings (WGA), that proves to be sample-optimal.
Each agent i keeps N shared local models x k 1 , . .
, x kN , where x kj estimates x ⋆j at iteration k (the knowledge of 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 WGA algorithms
A LL - FOR - ONE (AFO), since every iteration of each gradient averaging for a given node i requires all the other nodes to compute one stochastic gradient for i .
WGA is thus equivalent to training 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(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]
satisfies
[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 proves that 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) that can be (statistically speaking) as hard to compute as solving our optimization problem; and (ii) the memory requirements and computation/communication costs of WGA can be prohibitive for large N and 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 introducing and studying the A LL - FOR - ALL algorithm.
We discuss (i) in Section 6, where we provide scenarii over which statistical theoretical guarantees can be derived on the error made by estimating these distribution distances using only a few samples.
The A LL - FOR - ALL algorithm
Algorithm
In this section, we present the A LL - FOR - ALL algorithm (AFA), an adaptation of the weighted gradient averaging algorithm.
For 1 ⩽ i ⩽ N , initialize x 0 i = x 0 ∈ R d .
At iteration k , let x ki ∈ R d be agent i ’s current estimate of x ⋆i , and denote x k = ( x ki )
⩽
i ⩽ N ∈ R N × d .
For a step size η > 0 and a matrix W ∈ R N × N with non-negative entries (remarkably and as discussed later, W will not necessarily 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 the distribution-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 =
[Equation] [Equation]
N εi
[Equation]
b 2 /µ ) ,
the A LL - FOR - ALL algorithm (Algorithm 1) returns ( x Ki )
1 ⩽ i ⩽ N satisfying 1 N (cid:80) Ni
=
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 -, we observe that the A LL - FOR - ALL algorithm reaches an averaged ε -generalization with an number of oracle 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 [Equation]
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 be replaced by the quantities
f i ( x ⋆ j ) − f i ( x ⋆ i ) , that control how well the optimal model for j generalizes for 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 to satisfy f i ( x ⋆j ) − f i ( x ⋆i ) ⩽ b 2 ij / (2 µ ) .
We believe this notion of function proximity that we leverage to be 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 i does not aggregate a convex combination of stochastic gradients, but a combination with scalars that do not necessarily sum to 1.
We thus cannot say that the ALL - FOR - ALL algorithm acts as if, in parallel, each agent i trains a model on the mixture of distributions D j with weights W ij .
In fact, as the 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 account for inter-dependencies between agents that do not directly share information, the all-for-all gradient filtering uses weights W ij to aggregate information, instead of λ ij .
Propagating information using a matrix W , that induces a similarity graph G W on { 1 , . .
. , N } , such that ( ij ) ∈
E W
if W ij > 0 , is quite natural [4, 50]; 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 personalized local 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), can be seen as the special instances where, for all i , we have λ ii = 1 − α i and λ ij =
α i N − 1 if i ̸
= j for some α i , and leads to bias terms of the form 1 N (cid:80) i
α
i N − 1 (cid:80)
j ̸
= i b ij [14, 25, 39].
Full and naive collaboration (a single model trained for all users) corresponds to λ ij = 1 /N for all i, j , and leads to a bias term of 1 N 2 (cid:80) i,j b ij .
The degrees of freedom offered by our matrix W (and by coefficients λ ij ) enable pairwise 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 with x 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 is in fact similar to some decentralized optimization ones, where a dual problem or a related energy function is often introduced [23, 45], upon which well-studied algorithms are applied.
The related problem we formulate is different from and more flexible than all the different personalized FL problems in the literature [28, 48], 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 and since W = ΛΛ ⊤ , we have the recursion x k +1 =
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 generated using Algorithm 1.
As a consequence, controlling in Theorem 3 the function values 1 N (cid:80) i f i ( x ki ) is equivalent to controlling f Λ ( y k ) (these two quantities are equal).
The bias-variance trade-off thus writes 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 gradient averagings, we now turn to quantifying the impact of the use of estimated values ˆ b ij of b ij , in order to 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 the function 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 the estimation, see e.g. [52]), maximum mean discrepancy (MMD) distances often exhibit lower sample complexities in O (1 / √ S )
Moreover, explicit assumptions on the loss function can also provide low sample complexities, as shown below for quadratic loss functions.
Yet, the results presented in this section can be generalized beyond linear models with squared losses, as long as concentration inequalities for controlling how far empirical distributions are from the true distribution in terms of distance d H .
In order to formulate statistical results for the estimation of the pairwise distribution-based distances d 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 for quadratic 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-isotropic subgaussian random variables, that thus benefit from concentration inequalities that are dimensionindependent [22, 35].
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 MNIST dataset, d eff is less than 3, while the ambient dimension is 712 [22].
Depending on a smaller dimension is 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 C agents each (to ease notations, with a total number of agents N = MC ).
Within each cluster, agents distributions share the same objective, and clusters are “well-separated”.
These models are popular for modelling population heterogeneity and provide a formal framework for clustering problems; we refer the interested reader to Melnykov and Maitra [41] for a detailed survey on the subject.
Assumption 5 (Well-separated clusters of agents) .
For M, C ⩾ 1 , N writes as N = MC and there exists {C 1 , . .
, C M } a partition of { 1 , . . .
, N } , µ 1 , . .
. , µ m such that for all 1 ⩽ m ⩽ M , |C m | =
C and for all i, j ∈ C m , we have E [ ξ i ] =
E [ ξ j ] = µ m .
We denote ∆ 2 = min m ̸
= m ′ ∥
µ m − µ m ′ ∥ 2 and assume that ∆ 2 > 0 .
When distribution-based distances were given (as in Corollary 1), Algorithm 1 achieved the optimal collaboration 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 with whom to collaborate as we do in the sequel, reaching this collaboration speedup of 1 /C will hence prove 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 to estimating who to collaborate with, the K remaining dedicated to the optimization process i.e. to running 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 for 1 ⩽ 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 Ncommunications of a naive approach) by performing randomized gossip communications [6] on the complete graph. Theorem 4.
Under Assumptions 1,2,3 and 5, and for ε > 0 , the A LL - FOR - ALL algorithm with estimated biases as just described reaches an averaged generalization error of ε as long as the number of pre-used samples S and of oracle calls K satisfy:
[Equation]
Theorem 5 (A LL - FOR - ALL with estimated biases) .
Assume that Assumptions 1, 2, 3 (for some unknown biases b ij ), 4 and 5 hold.
Under the setting described, let ˆΛ be the stochastic matrix with entries
[Equation]
for some u > 0 that verifies u ⩾ 4 ν 2 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 : E
[Equation]
where the mean is taken over both biases estimates (ˆ b ij ) and gradient estimates ( g ki ) .
Corollary 2.
Under the same assumptions as Theorem 5 and for ε > 0 , the A LL - FOR - ALL algorithm with 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 ( S and K ) are required to reach an averaged arbitrarily small generalization error ε > 0 , in the limit with 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 κ , the condition number of the problem.
The number of samples S used for estimating the biases is required to be of order ν 21 d eff / ∆ 2 , the “signal-to-noise” ratio of our mixture model [41, 46], a natural quantity to depend on.
Corollary 2 hence shows that the optimal collaboration speedup is achieved, up to logarithmic factors: in order to reach an arbitrary small generalization error ε > 0 , are only required constant orders for S and K (the number of samples locally available) if the number of agents is large enough i.e. if N = ˜Ω( M/ε ) , where M is the number of clusters i.e. we have a linear speedup in the clusters population.
We numerically illustrate our theory in Appendix A on synthetic datasets, with clustered agents (as in this section), as well as in a setting where agents are distributed according to a more general “distribution of agent”.
A closely related work [26] also studies a model where agents verify a cluster structure as described in Assumption 5 for quadratic losses and linear models.
Yet, we highlight several differences between their approach and ours.
First, Ghosh et al.
[26] perform an online clustering of the agents, as opposed to our pre-training hierarchical approach.
While the results we obtain in Theorem 5 and Corollaryand those of Ghosh et al.
[26] have the same linear speedup in the number of agents, ours require no initialization condition.
Finally, our algorithm is decentralized, thus leading to improved scalability (especially in terms of the number of clusters) and privacy [13], if of interest.
Finally, not being restricted to clusters in the analysis of the A LL - FOR - ALL algorithm leads to a better collaboration speedup and fairness (in the sense that performance does not impact a few agents) in a non-clustered scenario, where an approach based on clusters would be highly non-optimal for agents that are at the border of the inferred clusters.
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 between agents for shared minimization using stochastic gradient algorithms.
Our lower bound (Theorem 1) states that, under prior knowledge on the distances between local distributions, the collaborative speedup can be linear only in the first phase of the optimization when the generalization error is large compared to the distances between distributions.
More specifically, for a given agent i , the collaboration speedup is linear in the number of agents that are ε -close to i .
Moreover, we show that the A LL - FOR - ONE algorithm allows such a speedup and is thus sample optimal.
However, this algorithm requires high computation and communication capacities, a drawback that can be mitigated by the use of a novel algorithm called A LL - FOR - ALL , that benefits from the same collaboration speedup while being cheaper to deploy.
Finally, we studied the impact of estimating distances between distributions as a pre-processing step to the optimization phase; under a mixture model assumptions 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 general settings, as well as incorporating local steps in the A LL - FOR -
ALL algorithm (even though to lighten communications, unbiased compressors could here be used, as our analysis encompasses these) are interesting questions left for future work.
See [24] for extensions to convex-smooth functions (not necessarily strongly convex), convex-Lipschitz functions and to asynchronous gradient oracles.
Aknowledgments This work was supported by ANR-19-P3IA-0001(PRAIRIE 3IA Institute).
We also acknowledge support from the MSR-INRIA joint centre
[-We first notice that x ⋆j = 1 d (cid:0) 12 + δ j α l (cid:1) 
-]
[-and ¯ f is L/N -smooth and µ/N -strongly convex, f Λ is L/N -relatively smooth and µ/N -relatively strongly convex [2] with respect to 12 ∥ y ∥ 2 W = 
-]
[-E (cid:20) e C ψ 2 β 2 S 2 (cid:80) Ss =
-]