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 [30, 39] 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.
[27], Mansour et al.
[38]), 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. [44]), 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 [30, 35, 39, 50].
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 [43] have been considered.
The algorithms we propose in this paper are well suited for a decentralized implementation.
As observed in Hanzely et al.
[27], 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 [37, 42, 53].
Approaches to personalization include fine-tuning [10, 33], transfer learning techniques [18, 48, 50], 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 [27] or with a meta-learning approach [9, 24, 29].
We refer the interested reader to Kulkarni et al.
[36] 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.
[38] 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 [15, 16] 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.
[26] 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.
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 [32, 52].
Many algorithmic solutions have been proposed to counterweight this effect [31, 41] (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)
[Equation]
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, [46]) 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 :
[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 ) 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 some constant C > 0 independent of the problem and 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 [20] for stochastic gradient descent [1] or for information limited statistical estimation [19, 54], 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) ,
E
[Equation]
Let ε > 0 .
For a specific choice of λ ij = 1 { b 2 ij<ε/ 2 }
N
[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 algorithm 1: Step size η > 0 , matrix W ∈ R N × N , initialization x 01 = . . .
= x 0 N ∈ R d ( x 0 i at agent i ).
2: for k = 0 , 1 , 2 , . . .
K − 1 do 3:
Agents 1 ⩽ j ⩽ N compute g kj ( x kj ) and broadcast it to all agents i such that W ij > 0 .
For i = 1 , . . .
, N , update: x k +
[Equation]
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] [Equation]
N εi ( 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 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 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, 49]; 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, 24, 38].
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 [22, 44], 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 [27, 47], 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. [51]), 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 [21, 34].
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 [21].
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 [40] 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 (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 :
[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 4 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 [40, 45], 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 [25] 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.
[25] perform an online clustering of the agents, as opposed to our pre-training hierarchical approach.
While the results we obtain in Theorem 4 and Corollaryand those of Ghosh et al.
[25] 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 [23] 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
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] 2.
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 experimental 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 they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments 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] 4.
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]
[Figure]
To test the robustness of our theory, we build toy problems from synthetic datasets, placing ourselves in the scenario we considered throughout the paper: a large number of agents with heterogeneous data, that each have too few samples available from their local data distribution in order to reach a small generalization error on their own.
In Figure 1, we consider N = 500 agents, and a quadratic loss function ℓ ( x, ξ = ( a, b ))
=
12 ( a ⊤ x − b ) 2 , for x, a ∈ R d ( d = 100 ) and µ ∈ R d .
For i = 1 , . . . , 500 , the distribution D i of ξ i
=
( a i a ⊤ i , a i b i ) as a centered Gaussian random variable of covariance matrix Σ
i for a i , and b i is the sign of a ⊤ i u for some fixed u ∈ R d , flipped with probability 0 .
In both figures, each agents have 10 samples available for the optimization phase ( K = 10 oracle calls), corresponding to a total number of samples used of N × K = 5000 .
We computed and showed the 500 steps of all ten oracle calls, each step corresponding to the use of the stochastic gradient of a single agent.
The dotted lines represent our baselines.
The blue one is the A LL - FOR - ALL algorithm with matrix W exactly as in Corollary 1 with b ij = ∥
Σ
i − Σ j ∥ .
The orange dotted line consists in the no-collaboration baseline: each agents performs SGD on its own without sharing information (corresponds to W = I N ).
The green dotted line corresponds to the “single-model” approach without personalization: one model is trained for all agents, using SGD and all samples from all agents (corresponding to W = 1 n 11 ⊤ ).
The choice of the algorithm for the single model approach without collaboration is in fact unimportant, since all algorithms would reach the same asymptotic bias here.
The full lines (red, violet and brown) correspond to estimating the pairwise distances from empirical distributions (as in Section 6), using respectively S = 1 , S = 3 and S = 5 samples.
In Figure 1(a), we consider M = 10 (unkonwn) clusters C 1 , . . .
, C M .
All
i ∈ C m have the same covariance matrix Σ m , equal to I d / √ d
+ e m e ⊤ m , where e m is the m -th element of the canonical basis of R d .
In Figure 1(b), Σ i =
Diag( u ( i ) 1 , . . .
, u ( d ) 1 ) / √ d where the ( u ( i ) ℓ ) are i.i.d.
uniformly distributed in [0 , 1] .
Performing rough estimations of the pairwise distance between agents’ local distributions thus appears to be quite robust in both our settings.
In the “cluster” setting, this was predicted by our theory, and the numerical results are compelling.
In the “ i.i.d. ” setting, using very few samples for the estimation also appears to be very efficient.
B Proof of our lower-bound (Theorem 1) B.1 General framework to prove lower bounds [1]
The idea is that, when optimizing a function f ( x ) =
E [ ℓ ( x, ξ )] and finding a good approximation of a minimizer x ⋆ , we learn some information on the distribution D over which samples are drawn.
In order to prove lower bounds, we construct a loss function ℓ , and distributions D α 1 , . . .
, D αN , where α is a random parameter.
We argue that minimizing the objective function up to a certain precision gives a good estimator (quantified) of the random seeds α .
Then, using Fano inequality, we bound the efficiency of such an estimator in terms of number of oracle calls, obtaining a lower bound on the sample complexity.
This approach is inspired by Agarwal et al.
[1], who prove IT-lower bounds for stochastic gradient descent.
We adapt their proof technique to the personalized and multi-agent setting.
Constructing difficult loss functions For any two functions f, g : R d → R , we define the discrepancy measure ρ ( f, g ) as:
[Equation]
which is a pseudo metrics.
Now, for a finite set V of parameters, let G ( δ ) =
(cid:8) g δα , α ∈ V (cid:9) be a set of functions indexed by V , that depend on δ (fixed in the set).
The dependency in δ of each g α ∈ G ( δ ) is left implicit in the following subsections.
We define:
[Equation]
Minimizing is Bernoulli parameters identification
The two following lemmas justify that optimizing a function g α ∈ G ( δ ) to a precision of order ψ ( δ ) is more difficult than estimating the parameter α .
Lemma 1 (Agarwal et al.
For any x ∈ R d , there can be at most one function g α in G ( δ ) such that:
[Equation]
Lemma 2 (Agarwal et al.
Assume that for some fixed but unknown α ∈ V there exists a method
M K based on the data ϕ
= { X 1 , ..., X K } that returns x K (function of ϕ ) satisfying an error of:
[Equation]
where the mean is taken over the randomness of both the oracle Φ , the method M K and α ∈ V if random.
Then, there exists a hypothesis test ˆ α : ϕ → V such that:
[Equation]
Suppose now that the parameter α in the previous Lemma is chosen uniformly at random in V .
Let ˆ α : ϕ
→ V be a hypothesis test estimating α .
By Fano inequality [12], we have:
[Equation]
where I (cid:0) ϕ, α (cid:1) is the mutual information between ϕ and α , that we need to upper-bound.
Combining
Fano inequality with Lemmas 1 and 2, fixing a target error ε = ψ ( δ ) , we obtain a lower bound on the
K ε the number of oracle calls required to reach an ε generalization error:
[Equation]
where ϕ K ε is the information contained in K ε oracle calls.
If we have an equality of the form
I (cid:0) ϕ K ε , α (cid:1) =
K ε I (cid:0) ϕ 1 , α (cid:1) , this gives:
[Equation]
Playing with the different parameters δ, α, V gives lower bounds.
We refer the interested reader to
Chapter 2 in Cover and Thomas [12] for Fano inequality and mutual information.
B.2 Applying this to prove Theorem
For simplicity, assume that r 2 = d and σ 2 = 1 .
Let δ > 0 a free parameter.
Let V =
{ α 1 , . . .
, α L } ⊂ {− 1 , 1 } d be a subset of the hypercube such that for all k ̸ = l ,
[Equation]
i.e. V is a d/ 4 -packing of the hypercube.
We know that we can set |V| ⩾ (2 / √ e ) d/ 2 .
Without loss of generality, we prove a lower bound in the case where the agent that desires to minimize its local function is indexed by 1 .
Let:
[Equation]
for x, ξ ∈ R d and, for fixed δ > 0 and any α ∈ V :
[Equation]
We keep the same notations as last subsection ( ψ ( δ ) , ρ ).
We have:
[Equation]
leading to ψ ( δ ) ⩾ δ 2 / 4 since V is a d/ 4 -packing of the hypercube.
For any i = 1 , . . .
, N , let D i be the probability distribution on { 0 , 1 } d of the following random variable:
[Equation]
where δ i = ( δ − b i 1 ) + , where s + = max(0 , s ) for s ∈ R , k is taken uniformly at random in { 1 , . . .
, d } , ( ϵ k ) is the canonical basis of R d , and Ber ( p ) is a Bernoulli random variable, independent of k .
The mutual information is thus, in our case:
[Equation]
where we use the fact that I (cid:0) Ber (cid:0) 12 + 1 b 1 i ⩽ δ
α k δ i (cid:1) , α k (cid:1) ⩽ C 1 δ 2 for some constant C 1 > 0 , for δ i ⩽ 1 / 4 .
Setting the target precision as ε = δ 2 / 4 , we obtain:
[Equation]
The loss function and distributions built verify our regularity assumptions for µ = 1 /d , L = 1 /d , noise σ 2 ⩽ 1 .
We first verify that for all 1 ⩽ j, k ⩽ N , we have f j ( x ⋆k )
− f j ( x ⋆j ) ⩽ b 2 kj .
We first notice that x ⋆j = 1 d (cid:0) 12 + δ j α l (cid:1)
⩽ l ⩽ d , so that:
[Equation]
since the weights b verify the triangle inequality.
Under the assumptions of Theorem 1 on H , we have, in terms of distribution-based distances:
[Equation]
The minimum of each g α is attained at x α = 12 + δα , we thus need to assume that r is of order √ d , and a rescaling leads to the dependency in r .
The dependency in σ 2 for σ 2 > 1 is obtained by taking
D ′ i
= Ber(1 /σ 2 )
σ 2 D i .
In this case, we have a noise amplitude of order σ 2 instead of order 1 , and a factor 1 /σ 2 appears in the mutual information.
C Proof of Theorem
Proof of Theorem 2.
We begin by proving the following descent lemma.
Lemma 3.
Let F : R d → R be L -smooth and µ -strongly convex and G : R d → R be differentiable.
Consider the iterates generated by:
[Equation]
where E (cid:2) g k | y k (cid:3) = ∇ G ( y k ) and E (cid:104)(cid:13)(cid:13) g k − ∇ G ( y k ) (cid:13)(cid:13) 2 |
y k (cid:105) ⩽
σ 2 g .
Then, we have, where y ⋆ minimizes F , as long as η ⩽ 1 /L :
[Equation]
Proof Lemma 3.
We use smoothness of F :
[Equation]
Then, using E (cid:104)(cid:13) (cid:13) g k
(cid:13) (cid:13) 2 (cid:105) ⩽
[Equation]
(cid:13)
[Equation]
G ( y k ) (cid:13) (cid:13) 2 (cid:105) +
[Equation]
2 g , and − η E (cid:2)
[Equation]
g k
[Equation]
∇ F ( y k ) ⟩ (cid:3)
[Equation]
− η
[Equation]
(cid:2)
[Equation]
we obtain that:
[Equation]
Finally, we conclude using − η 2 E (cid:104)(cid:13)(cid:13) ∇ F ( y k ) (cid:13)(cid:13)
[Equation]
⩽ η E
[Equation]
F ( y k
[Equation]
− F ( y ⋆ ) (cid:3) and η < 1 /L .
Let i ∈ [ N ] .
To prove Theorem 2, we now use Lemma 3 to study the sequence y k
= x ki with F = f i ,
G = (cid:80) j λ ij f j := f λ and g k = (cid:80) j λ ij g kj ( x ki ) .
For all x ∈ R d , using Assumption 3:
[Equation]
⩽
[Equation]
(cid:88)
[Equation]
=
[Equation]
λ
[Equation]
b
[Equation]
, since ∇ x ℓ ( x, · ) ∈ H and d H ( D i , D j ) ⩽ b ij .
Then, using the independence (conditionally on x ki ) of the ( g kj ) j and E (cid:2) g kj ( x ki ) | x ki (cid:3) = ∇ f j ( x ki ) :
E   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:88)
j
[Equation]
⩽ (cid:88)
j λ 2 ij σ 2 .
Consequently,
E (cid:2)
f i ( x k +1 i ) −
[Equation]
⋆i
[Equation]
f
[Equation]
η
N (cid:88)
[Equation]
=
[Equation]
λ
[Equation]
b
[Equation]
ij + η 2 σ 2 L
N (cid:88)
j = λ 2 ij .
Writing H
k
= (1 − ηµ ) − k E (cid:2)
f i ( x ki ) − f i ( x ⋆i ) (cid:3) , unrolling the recursion leads to:
[Equation]
Finally, using (cid:80) k<K (
[Equation]
ηµ
[Equation]
k ⩽ (1 −
[Equation]
K
[Equation]
, we have:
[Equation]
Using (1 − ηµ ) K ⩽ e − Kηµ , we optimize of η .
For
[Equation]
we obtain:
[Equation]
leading to the first part of Theorem 2.
For the second part, we simply plug the expression of λ ij in the proven formula.
D Proof of Theorem
We recall that for a stochastic matrix Λ , we defined
[Equation]
Then, y Λ is defined as a minimizer of f Λ , and we write x ⋆ =
( x ⋆ 1 , . . .
, x ⋆N ) where x ⋆i is the minimizer of f i .
We first begin with the following simple lemmas.
Lemma 4.
If Assumption 3 (weak version) holds, then for all i, j = 1 , . .
, N , we have:
[Equation]
Proof.
Using strong-convexity of f i and ∇ f i ( x ⋆i ) :
[Equation]
where the last inequality is deduced using the weak version of Assumption 3.
Lemma 5.
If Λ is a stochastic matrix,
[Equation]
Proof.
Writing the optimality of y Λ gives:
[Equation]
where we used convexity of each f i .
Then, subtracting ¯ f ( x ⋆ ) and using stochasticity of Λ :
[Equation]
We are now armed to prove Theorem 3.
Proof.
We have the following bias-variance decomposition, where the inequality is a consequence of
Lemma 5:
[Equation]
We thus need to upper-bound the optimization term f Λ ( y k ) − f Λ ( y Λ ) .
We recall that y k verifies the recursion:
[Equation]
for
[Equation]
that verifies:
[Equation]
The function f Λ is however not necessarily strongly convex.
However, since ∇ 2 f Λ ( y ) =
Λ ⊤ ∇ 2 ¯ f (Λ y )Λ
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 =
y ⊤ Wy .
Note also that the spectral radius of W is 1 , since Λ is stochastic.
Instead of using stochastic Bregman gradient descent ( e.g. Dragomir et al.
[17]), we use Lemma 6 that we prove at the end of the paper: classical SGD that naturally generalizes to relative smoothness and strong convexity assumptions, when the mirror map is quadratic.
This leads to:
[Equation]
for a choice of stepsizes of:
[Equation]
concluding the proof.
E Proof of Theorem
We first start by recalling that, for a Σ -subgaussian random variable ξ , using Theorem 1 from [28], we have for any t ⩾ 0 :
[Equation]
Consequently, for u ⩾ max( ν 2 d eff , ∥ Σ ∥ 22 /ν 2 ) , we have:
[Equation]
Remarking that ∥ Σ ∥ 22 ν 2 = ν 2 (cid:80) k ν 4 k ν 4 ⩽ ν 2 (cid:80) k ν 2 k ν 2 = ν 2 d eff , where ν 2 k are the eigenvalues of Σ , this condition on u is in fact u ⩾ ν 2 d eff .
Proof of our Theorem.
From Theorem 3, we have, conditionally on the S samples used in estimating biases:
[Equation]
We hence need to bound E (cid:104)(cid:80) i,j ˆ λ 2 ij (cid:105) and E (cid:104)(cid:80) i,j ˆ λ ij b 2 ij (cid:105) , and start by assuming that u ⩾ 4 S ν 2 d eff .
First, denoting b max = max i,j b ij , we have:
[Equation]
Using a triangle inequality, that gives us 2 ∥ ˆ µ i − ˆ µ j − E
[ˆ µ i − ˆ µ j ] ∥⩾
[Equation]
ij − 2ˆ b 2 ij ,
P (cid:16)
[Equation]
ij
[Equation]
u |
[Equation]
2 ij > 4 u (cid:17) ⩽ P (cid:16)
[Equation]
µ i −
[Equation]
µ
[Equation]
− E [ˆ µ i − ˆ µ j ] ∥ 2 ⩾
[Equation]
(cid:17) .
Then, since each ξ i and ξ j are Σ -subgaussian (and independent), ˆ µ i − ˆ µ j − E
[ˆ µ i − ˆ µ j ] is 4Σ /S
subgaussian 2 , so that using our assumption on u , we can use Theorem 1 of Hsu et al.
[Equation]
Then, for 1 ⩽ i ⩽ N , let ˆ N i
=
(cid:80) Nj =1 1 { ˆ b ij ⩽ u } .
Fix 1 ⩽
m ⩽ M .
We have:
[Equation]
, so that E (cid:104)(cid:80) i ∈C M
[Equation]
1 ⩽ j
[Equation]
N
[Equation] [Equation]
=
[Equation]
m 1ˆ N i
[Equation]
⩽ 1 C + 2 Ce − Su 8 ν 2 , concluding the proof.
We then prove the resulting corollary by taking u = ∆ / 4 , and the condition on u translates into
S ⩾
[Equation]
F SGD under strong-convexity and smoothness assumptions
We recall the following well-known result.
2 indeed, using the subgaussian norm ∥·∥ ψ 2 , for real-valued independent random variables X 1 , . . .
, X S and any β > 0 , E (cid:104) e β 1 S (cid:80) Ss =1 X s (cid:105) = (cid:81) s E (cid:104)
e βS X s (cid:105) ⩽ (cid:81) s E (cid:20)
e C ψ 2 β 2 S 2 ∥
X s ∥ ψ 2 (cid:21) =
E (cid:20) e C ψ 2 β 2 S 2 (cid:80) Ss =
∥
X s ∥ ψ 2 (cid:21) , so that (cid:13)(cid:13)(cid:13) 1 S (cid:80) Ss =1 X s (cid:13)(cid:13)(cid:13) ψ
2 ⩽ 1 S 2 (cid:80) Ss =1 ∥
X s ∥ ψ 2 , and we apply this to the random variables ˆ µ i ⊤ y
Lemma 6 (SGD, s.c. and smooth) .
Define ∥ x ∥ 2 A =
x ⊤ Ax for some non-negative and symmetric matrix A .
Let f : X → R µ -relatively strongly convex and L -relatively smooth with respect to 12 ∥ x ∥ 2 A .
Let ( f t , g t ) t ⩾ 0 be first order oracle calls such that for all t ⩾ 0 : ∀ x ∈ X
[Equation]
for some σ > 0 .
Let L A be the largest eigenvalue of A , and assume that L A ⩽ 1 (our result generalizes to any L A ).
Let ( x t ) t ⩾ 0 be generated with:
[Equation]
for a fixed stepsize 12 L ⩾ η > 0 , and assume that all the iterates lie in X .
Assume that f is minimized over X at some interior point x ⋆ .
We have for any T > 0 :
[Equation]
For fixed T > 0 , setting η = min (cid:0) 1 / (2 L ) , 1 µT ln( f 0 µ 2 T Lσ 2 ) (cid:1) gives:
[Equation]
Thus, for fixed target precision ε > 0 , using stepsize η ε = min (cid:0) µε 2 Lσ 2 , 12 L (cid:1) and setting T ε = ⌈ ln (cid:0) ε − 1 ( f
[Equation]
x 0 ) − f ( x ⋆ )) (cid:1) 1 η ε µ ⌉ , we have:
[Equation]
with a number of oracle calls
T
[Equation]
Proof.
For some t ⩾ 0 , denoting f t = E (cid:2)
f ( x t +1 )
− f ( x ⋆ ) (cid:3) , using relative smoothness, unbiasedness of the stochastic gradients and then relative strong convexity:
[Equation]
Using relative strong convexity of f , we have:
[Equation]
yielding, for η < 1 / ( LL A )
[Equation]
Then, for some T > 0 and since L A ⩽ 1 , sum the above inequality multiplied by (1 − ηµ ) − t − 1 :
[Equation]
leading to the desired result.