Local SGD for non-i.i.d. data Konstantin Mishchenko Work done - - PowerPoint PPT Presentation

Local SGD for non-i.i.d. data

Konstantin Mishchenko

Work done together with Ahmed Khaled and Peter Richtárik

Problem

Convex

min

x

1 M

M

X

m=1

fm(x)

Problem

Convex In practice, usually a neural network

min

x

1 M

M

X

m=1

fm(x)

Problem

min

x

1 M

M

X

m=1

fm(x)

fm(x) = Eξfm(x; ξ)

Local SGD

min

x

1 M

M

X

m=1

fm(x)

xm

t+1 =

( ˆ xt+1, if t mod H = 0 xm

t γrfm(xm t ; ξm t ), otherwise

Local SGD

min

x

1 M

M

X

m=1

fm(x)

xm

t+1 =

(

1 M

PM

j=1(xj t γrfj(xj t; ξj t )), if t mod H = 0

xm

t γrfm(xm t ; ξm t ), otherwise

Local SGD

min

x

1 M

M

X

m=1

fm(x)

xm

t+1 =

( ˆ xt+1, if t mod H = 0 xm

t γrfm(xm t ; ξm t ), otherwise

H = 1 − → minibatch SGD

H = T − → one-shot averaging

Local GD

xm

t+1 = xm t γrfm(xm t )

The Variance of Local GD

σ2

f def

= 1 M

M

X

m=1

krfm(x∗)k2

xm

t+1 = xm t γrfm(xm t )

Analysis difficulties in local GD

xm

t+1 = xm t γrfm(xm t )

ˆ xt

def

= 1 M

M

X

m=1

xm

t

Analysis difficulties in local GD

xm

t+1 = xm t γrfm(xm t )

Vt

def

= 1 M

M

X

m=1

kxm

t ˆ

xtk2

ˆ xt

def

= 1 M

M

X

m=1

xm

t

gt

def

= 1 M

M

X

m=1

rfm(xm

t )

Theorem

Choose H such that H 

√ T √ M , then γ = √ M 4L √ T  1 4HL, and hence

f(ˆ xT ) f(x∗)  8Lkx0 x∗k2 p MT + 3Mσ2

fH2

2LT . To get a convergence rate of 1/ p MT we can choose H = O(T 1/4M −3/4), which implies a total number of Ω(T 3/4M 3/4) communication steps. If a rate of 1/ p T is desired instead, we can choose a larger H = O(T 1/4).

Plots

Plots

Local SGD

xm

t+1 = xm t γrfm(xm t ; ξm t )

Eξkrfm(x; ξ) rfm(x)k2  σ2

Eξkrfm(x; ξ) rfm(x)k2  4LDfm(x, x∗) + 2σ2

Local SGD

xm

t+1 = xm t γrfm(xm t ; ξm t )

Eξkrfm(x; ξ) rfm(x)k2  4LDfm(x, x∗) + 2σ2

σdif

def

= 1 M

M

X

m=1

Eξkrfm(x∗, ξ)k2

Local SGD

xm

t+1 = xm t γrfm(xm t ; ξm t )

Eξkrfm(x; ξ) rfm(x)k2  4LDfm(x, x∗) + 2σ2

Df(x, y) = f(x) f(y) hrf(y), x yi

σdif

def

= 1 M

M

X

m=1

Eξkrfm(x∗, ξ)k2

Theorem

Choose H such that H ≤

p T p M , then γ = p M 8L p T ≤ 1 8HL and

Ef(ˆ xT ) − f(x⇤) ≤ 32Lkˆ

x0x∗k2 p MT

+

5σ2

dif

2L p MT + σ2

difM(H1)2

4LT

.

Theorem

Choose H such that H ≤

p T p M , then γ = p M 8L p T ≤ 1 8HL and

Ef(ˆ xT ) − f(x⇤) ≤ 32Lkˆ

x0x∗k2 p MT

+

5σ2

dif

2L p MT + σ2

difM(H1)2

4LT

.

Optimal H is H = 1 + bT 1/4M −3/2c

Theorem

Choose H such that H ≤

p T p M , then γ = p M 8L p T ≤ 1 8HL and

Ef(ˆ xT ) − f(x⇤) ≤ 32Lkˆ

x0x∗k2 p MT

+

5σ2

dif

2L p MT + σ2

difM(H1)2

4LT

.

Optimal H is H = 1 + bT 1/4M −3/2c

Improves to H = 1 + bT 1/2M −3/2c

if Ekrfm(x; ξ) rfm(x)k2  σ2

Plot

Open questions

Meta-Learning

We can learn an ”improvable” model

Open questions

Meta-Learning

We can learn an ”improvable” model

min

x

1 m

M

X

m=1

fm(x γrfm(x))

Reference

Better Communication Complexity for Local SGD arXiv:1909.04746 First Analysis of Local GD on Heterogeneous Data arXiv:1909.04715 NeurIPS workshop on Federated Learning http://federated-learning.org/fl-neurips-2019/