50
SEGA: Variance Reduction via Gradient Sketching * Filip Hanzely Konstantin Mishchenko Peter Richt´ arik § October 19, 2018 Abstract We propose a randomized first order optimization method—SEGA (SkEtched GrAdient)— which progressively throughout its iterations builds a variance-reduced estimate of the gradient from random linear measurements (sketches) of the gradient obtained from an oracle. In each iteration, SEGA updates the current estimate of the gradient through a sketch-and-project op- eration using the information provided by the latest sketch, and this is subsequently used to compute an unbiased estimate of the true gradient through a random relaxation procedure. This unbiased estimate is then used to perform a gradient step. Unlike standard subspace de- scent methods, such as coordinate descent, SEGA can be used for optimization problems with a non-separable proximal term. We provide a general convergence analysis and prove linear convergence for strongly convex objectives. In the special case of coordinate sketches, SEGA can be enhanced with various techniques such as importance sampling, minibatching and accelera- tion, and its rate is up to a small constant factor identical to the best-known rate of coordinate descent. * Accepted to NIPS 2018. King Abdullah University of Science and Technology, Kingdom of Saudi Arabia King Abdullah University of Science and Technology, Kingdom of Saudi Arabia § King Abdullah University of Science and Technology, Kingdom of Saudi Arabia — School of Mathematics, University of Edinburgh, United Kingdom — Moscow Institute of Physics and Technology, Russia 1 arXiv:1809.03054v2 [math.OC] 18 Oct 2018

Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

  • Upload
    others

  • View
    0

  • Download
    0

Embed Size (px)

Citation preview

Page 1: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

SEGA: Variance Reduction via Gradient Sketching∗

Filip Hanzely† Konstantin Mishchenko‡ Peter Richtarik§

October 19, 2018

Abstract

We propose a randomized first order optimization method—SEGA (SkEtched GrAdient)—

which progressively throughout its iterations builds a variance-reduced estimate of the gradient

from random linear measurements (sketches) of the gradient obtained from an oracle. In each

iteration, SEGA updates the current estimate of the gradient through a sketch-and-project op-

eration using the information provided by the latest sketch, and this is subsequently used to

compute an unbiased estimate of the true gradient through a random relaxation procedure.

This unbiased estimate is then used to perform a gradient step. Unlike standard subspace de-

scent methods, such as coordinate descent, SEGA can be used for optimization problems with

a non-separable proximal term. We provide a general convergence analysis and prove linear

convergence for strongly convex objectives. In the special case of coordinate sketches, SEGA can

be enhanced with various techniques such as importance sampling, minibatching and accelera-

tion, and its rate is up to a small constant factor identical to the best-known rate of coordinate

descent.

∗Accepted to NIPS 2018.†King Abdullah University of Science and Technology, Kingdom of Saudi Arabia‡King Abdullah University of Science and Technology, Kingdom of Saudi Arabia§King Abdullah University of Science and Technology, Kingdom of Saudi Arabia — School of Mathematics,

University of Edinburgh, United Kingdom — Moscow Institute of Physics and Technology, Russia

1

arX

iv:1

809.

0305

4v2

[m

ath.

OC

] 1

8 O

ct 2

018

Page 2: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Contents

1 Introduction 4

1.1 Gradient sketching . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4

1.2 Related work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5

1.3 Outline . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6

1.4 Notation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6

2 The SEGA Algorithm 6

2.1 SEGA as a variance-reduced method . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7

2.2 SEGA versus coordinate descent . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8

3 Convergence of SEGA for General Sketches 9

3.1 Smoothness assumptions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 9

3.2 Main result . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 9

4 Convergence of SEGA for Coordinate Sketches 11

4.1 Defining D: samplings . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11

4.2 Non-accelerated method . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12

4.3 Accelerated method . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13

5 Experiments 14

5.1 Comparison to projected gradient . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 14

5.2 Comparison to zeroth-order optimization methods . . . . . . . . . . . . . . . . . . . 14

5.3 Subspace SEGA: a more aggressive approach . . . . . . . . . . . . . . . . . . . . . . . 15

6 Conclusions and Extensions 16

6.1 Conclusions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16

6.2 Extensions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16

A Proofs for Section 3 21

A.1 Proof of Theorem 3.3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22

A.2 Proof of Lemma A.3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23

A.3 Proof of Lemma A.4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24

B Proofs for Section 4 24

B.1 Technical Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24

B.2 Proof of Theorem 4.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25

B.3 Proof of Corollary 4.3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26

B.4 Accelerated SEGA with arbitrary sampling . . . . . . . . . . . . . . . . . . . . . . . . 26

2

Page 3: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

B.4.1 Proof of Corollary 4.5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 30

B.5 Proof of Lemma B.3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31

B.6 Proof of Lemma B.4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31

C Subspace SEGA: a More Aggressive Approach 32

C.1 The algorithm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32

C.2 Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34

C.3 Main result . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 36

C.4 Optimal choice of B and Sk . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 37

C.5 The conclusion of subspace SEGA . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 38

D Simplified Analysis of SEGA 1 38

D.1 Technical Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 39

D.2 Proof of Theorem D.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 40

E Simplified Analysis of SEGA II 41

E.1 Two lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 42

E.2 Proof of Theorem D.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 43

F Extra Experiments 44

F.1 Evolution of Iterates: Extra Plots . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 44

F.2 Experiments from Section 5 with empirically optimal stepsize . . . . . . . . . . . . . 46

F.3 Experiment: comparison with randomized coordinate descent . . . . . . . . . . . . . 47

F.4 Experiment: large-scale logistic regression . . . . . . . . . . . . . . . . . . . . . . . . 49

G Frequently Used Notation 50

3

Page 4: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

1 Introduction

Consider the optimization problem

minx∈Rn

F (x)def= f(x) +R(x), (1)

where f : Rn → R is smooth and µ–strongly convex, and R : Rn → R ∪ {+∞} is a closed convex

regularizer. In some applications, R is either the indicator function of a convex set or a sparsity-

inducing non-smooth penalty such as group `1-norm. We assume that, as in these two examples,

the proximal operator of R, defined as

proxαR(x)def= argmin

y∈Rn

{R(y) +

1

2α‖y − x‖2B

},

is easily computable (e.g., in closed form). Above we use the weighted Euclidean norm ‖x‖Bdef=

〈x, x〉1/2B , where 〈x, y〉Bdef= 〈Bx, y〉 is a weighted inner product associated with a positive definite

weight matrix B. Strong convexity of f is defined with respect to the geometry induced by this

inner product and norm1.

1.1 Gradient sketching

In this paper we design proximal gradient-type methods for solving (1) without assuming that

the true gradient of f is available. Instead, we assume that an oracle provides a random linear

transformation (i.e., a sketch) of the gradient, which is the information available to drive the

iterative process. In particular, given a fixed distribution D over matrices S ∈ Rn×b (b ≥ 1 can but

does not need to be fixed), and a query point x ∈ Rn, our oracle provides us the random linear

transformation of the gradient given by

ζ(S, x)def= S>∇f(x) ∈ Rb, S ∼ D. (2)

Information of this type is available/used in a variety of scenarios. For instance, randomized

coordinate descent (CD) methods use oracle (2) with D corresponding to a distribution over standard

basis vectors. Minibatch/parallel variants of CD methods utilize oracle (2) with D corresponding

to a distribution over random column submatrices of the identity matrix. If one is prepared to use

difference of function values to approximate directional derivatives, then one can apply our oracle

model to zeroth-order optimization [8]. Indeed, the directional derivative of f in a random direction

S = s ∈ Rn×1 can be approximated by ζ(s, x) ≈ 1ε (f(x + εs) − f(x)), where ε > 0 is sufficiently

small.

We now illustrate this concept using two examples.

1f is µ–strongly convex if f(x) ≥ f(y) + 〈∇f(y), x− y〉B + µ2‖x− y‖2B for all x, y ∈ Rn.

4

Page 5: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Example 1.1 (Sketches). (i) Coordinate sketch. Let D be the uniform distribution over standard

unit basis vectors e1, e2, . . . , en of Rn. Then ζ(ei, x) = e>i ∇f(x), i.e., the ith partial derivative of f

at x. (ii) Gaussian sketch. Let D be the standard Gaussian distribution in Rn. Then for s ∼ D we

have ζ(s, x) = s>∇f(x), i.e., the directional derivative of f at x in direction s.

1.2 Related work

In the last decade, stochastic gradient-type methods for solving problem (1) have received unprece-

dented attention by theoreticians and practitioners alike. Specific examples of such methods are

stochastic gradient descent (SGD) [45], variance-reduced variants of SGD such as SAG [46], SAGA [10],

SVRG [22], and their accelerated counterparts [26, 1]. While these methods are specifically designed

for objectives formulated as an expectation or a finite sum, we do not assume such a structure.

Moreover, these methods utilize a fundamentally different stochastic gradient information: they

have access to an unbiased estimator of the gradient. In contrast, we do not assume that (2) is an

unbiased estimator of ∇f(x). In fact, ζ(S, x) ∈ Rb and ∇f(x) ∈ Rn do not even necessarily belong

to the same space. Therefore, our algorithms and results should be seen as complementary to the

above line of research.

While the gradient sketch ζ(S, x) does not immediatey lead to an unbiased estimator of the

gradient, SEGA uses the information provided in the sketch to construct an unbiased estimator of

the gradient via a sketch-and-project process. Sketch-and-project iterations were introduced in [16]

in the contex of linear feasibility problems. A dual view uncovering a direct relationship with

stochastic subspace ascent methods was developed in [17]. The latest and most in-depth treatment

of sketch-and-project for linear feasibility is based on the idea of stochastic reformulations [44].

Sketch-and-project can be combined with Polyak [31, 30] and Nesterov momentum [15], extended

to convex feasibility problems [32], matrix inversion [19, 18, 15], and empirical risk minimization

[14, 13]. Connections to gossip algorithms for average consensus were made in [29, 28].

The line of work most closely related to our setup is that on randomized coordinate/subspace

descent methods [36, 17]. Indeed, the information available to these methods is compatible with

our oracle for specific distributions D. However, the main disadvantage of these methods is that

they are not able to handle non-separable regularizers R. In contrast, the algorithm we propose—

SEGA—works for any regularizer R. In particular, SEGA can handle non-separable constraints even

with coordinate sketches, which is out of range of current coordinate descent methods. Hence, our

work could be understood as extending the reach of coordinate and subspace descent methods from

separable to arbitrary regularizers, which allows for a plethora of new applications. Our method is

able to work with an arbitrary regularizer due to its ability to build an unbiased variance-reduced

estimate of the gradient of f throughout the iterative process from the random linear measurements

thereof provided by the oracle. Moreover, and unlike coordinate descent, SEGA allows for general

sketches from essentially any distribution D.

5

Page 6: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Another stream of work on designing gradient-type methods without assuming perfect access

to the gradient is represented by the inexact gradient descent methods [9, 11, 47]. However, these

methods deal with deterministic estimates of the gradient and are not based on linear transforma-

tions of the gradient. Therefore, this second line of research is also significantly different from what

we do here.

1.3 Outline

We describe SEGA in Section 2. Convergence results for general sketches are described in Section 3.

Refined results for coordinate sketches are presented in Section 4, where we also describe and analyze

an accelerated variant of SEGA. Experimental results can be found in Section 5. We also include

here experiments with a subspace variant of SEGA, which is described and analyzed in Appendix C.

Conclusions are drawn and potential extensions outlined in Section 6. A simplified analysis of SEGA

in the case of coordinate sketches and for R ≡ 0 is developed in Appendix D (under standard

assumptions as in the main paper) and E (under alternative assumptions). Extra experiments for

additional insights are included in Appendix F.

1.4 Notation

We introduce notation when and where needed. For convenience, we provide a table of frequently

used notation in Appendix G.

2 The SEGA Algorithm

In this section we introduce a learning process for estimating the gradient from the sketched infor-

mation provided by (2); this will be used as a subroutine of SEGA.

Let xk be the current iterate, and let hk be the current estimate of the gradient of f . We

then query the oracle, and receive new gradient information in the form of the sketched gradient

(2). At this point, we would like to update hk based on this new information. We do this using

a sketch-and-project process [16, 17, 44]: we set hk+1 to be the closest vector to hk (in a certain

Euclidean norm) satisfying (2):

hk+1 = arg minh∈Rn

‖h− hk‖2B

subject to S>k h = S>k∇f(xk). (3)

The closed-form solution of (3) is

hk+1 = hk −B−1Zk(hk −∇f(xk)) = (I−B−1Zk)h

k + B−1Zk∇f(xk), (4)

6

Page 7: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Algorithm 1: SEGA: SkEtched GrAdient Method

1 Initialize: x0, h0 ∈ Rn; B � 0; distribution D;

stepsize α > 0

2 for k = 1, 2, . . . do

3 Sample Sk ∼ D4 gk = hk + θkB

−1Zk(∇f(xk)− hk)5 xk+1 = proxαR(xk − αgk)6 hk+1 = hk + B−1Zk(∇f(xk)− hk)

Figure 1: Iterates of SEGA and CD

where Zkdef= Sk

(S>k B−1Sk

)†S>k . Notice that hk+1 is a biased estimator of ∇f(xk). In order to

obtain an unbiased gradient estimator, we introduce a random variable2 θk = θ(Sk) for which

ED [θkZk] = B. (5)

If θk satisfies (5), it is straightforward to see that the random vector

gkdef= (1− θk)hk + θkh

k+1 (4)= hk + θkB

−1Zk(∇f(xk)− hk) (6)

is an unbiased estimator of the gradient:

ED[gk]

(5)+(6)= ∇f(xk). (7)

Finally, we use gk instead of the true gradient, and perform a proximal step with respect to R.

This leads to a new randomized optimization method, which we call SkEtched GrAdient Method

(SEGA). The method is formally described in Algorithm 1. We stress again that the method does

not need the access to the full gradient.

2.1 SEGA as a variance-reduced method

As we shall show, both hk and gk are becoming better at approximating ∇f(xk) as the iterates xk

approach the optimum. Hence, the variance of gk as an estimator of the gradient tends to zero,

which means that SEGA is a variance-reduced algorithm. The structure of SEGA is inspired by the

JackSketch algorithm introduced in [13]. However, as JackSketch is aimed at solving a finite-sum

optimization problem with many components, it does not make much sense to apply it to (1).

Indeed, when applied to (1) (with R = 0, since JackSketch was analyzed for smooth optimization

only), JackSketch reduces to gradient descent. While JackSketch performs Jacobian sketching

(i.e., multiplying the Jacobian by a random matrix from the right, effectively sampling a subset

2Such a random variable may not exist. Some sufficient conditions are provided later.

7

Page 8: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

of the gradients forming the finite sum), SEGA multiplies the Jacobian by a random matrix from

the left. In doing so, SEGA becomes oblivious to the finite-sum structure and transforms into the

gradient sketching mechanism described in (2).

2.2 SEGA versus coordinate descent

We now illustrate the above general setup on the simple example when D corresponds to a distri-

bution over standard unit basis vectors in Rn.

Example 2.1. Let B = Diag(b1, . . . , bn) � 0 and let D be defined as follows. We choose Sk = ei

with probability pi > 0, where e1, e2, . . . , en are the unit basis vectors in Rn. Then

hk+1 (4)= hk + e>i (∇f(xk)− hk)ei, (8)

which can equivalently be written as hk+1i = e>i ∇f(xk) and hk+1

j = hkj for j 6= i. Note that hk+1

does not depend on B. If we choose θk = θ(Sk) = 1/pi, then

ED [θkZk] =n∑i=1

pi1

piei(e

>i B−1ei)

−1e>i =n∑i=1

eie>i

1/bi= B

which means that θk is a bias-correcting random variable. We then get

gk(6)= hk +

1

pie>i (∇f(xk)− hk)ei. (9)

In the setup of Example 2.1, both SEGA and CD obtain new gradient information in the form of

a random partial derivative of f . However, the two methods process this information differently,

and perform a different update:

(i) While SEGA allows for arbitrary proximal term, CD allows for separable proximal term only [48,

27, 12].

(ii) While SEGA updates all coordinates in every iteration, CD updates a single coordinate only.

(iii) If we force hk = 0 in SEGA and use coordinate sketches, the method transforms into CD.

Based on the above observations, we conclude that SEGA can be applied in more general settings

for the price of potentially more expensive iterations3. For intuition-building illustration of how

SEGA works, Figure 1 shows the evolution of iterates of both SEGA and CD applied to minimizing a

simple quadratic function in 2 dimensions. For more figures of this type, including the composite

case where CD does not work, see Appendix F.1.

In Section 4 we show that SEGA enjoys the same theoretical iteration complexity rates as CD,

up to a small constant factor. This remains true when comparing state-of-the-art variants of

CD utilizing importance-sampling, parallelism/mini-batching and acceleration with the appropriate

corresponding variants of SEGA.

3Forming vector g and computing the prox.

8

Page 9: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Remark 2.2. Nontrivial sketches S and metric B might, in some applications, bring a substantial

speedup against the baseline choices mentioned in Example 2.1. Appendix C provides one setting

where this can happen: there are problems where the gradient of f always lies in a particular d-

dimensional subspace of Rn. In such a case, suitable choice of S and B leads to O(nd

)–times faster

convergence compared to the setup of Example 2.1. In Section 5.3 we numerically demonstrate this

claim.

3 Convergence of SEGA for General Sketches

In this section we state a linear convergence result for SEGA (Algorithm 1) for general sketch

distributions D under smoothness and strong convexity assumptions.

3.1 Smoothness assumptions

We will use the following general version of smoothness.

Assumption 3.1 (Q-smoothness). Function f is Q-smooth with respect to B, where Q � 0 and

B � 0. That is, for all x, y, the following inequality is satisfied:

f(x)− f(y)− 〈∇f(y), x− y〉B ≥1

2‖∇f(x)−∇f(y)‖2Q, (10)

Assumption 3.1 is not standard in the literature. However, as Lemma A.1 states, for B = I

and Q = M−1, Assumption 3.1 is equivalent to M-smoothness (see Assumption 3.2), which is a

common assumption in modern analysis of CD methods. Hence, our assumption is more general

than the commonly used assumption.

Assumption 3.2 (M-smoothness). Function f is M-smooth for some matrix M � 0. That is, for

all x, y, the following inequality is satisfied:

f(x) ≤ f(y) + 〈∇f(y), x− y〉+1

2‖x− y‖2M. (11)

Assumption 3.2 is fairly standard in the CD literature. It appears naturally in various application

such as empirical risk minimization with linear predictors and is a baseline in the development of

minibatch CD methods [43, 40, 38, 41]. We will adopt this notion in Section 4, when comparing

SEGA to coordinate descent. Until then, let us consider the more general Assumption 3.1.

3.2 Main result

We are now ready to present one of the key theorems of the paper, which states that the iterates

of SEGA converge linearly to the optimal solution.

9

Page 10: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Theorem 3.3. Assume that f is Q–smooth with respect to B, and µ–strongly convex. Choose

stepsize α > 0 and Lyapunov parameter σ > 0 so that

α (2(C−B) + σµB) ≤ σED [Z] , αC ≤ 1

2(Q− σED [Z]) , (12)

where Cdef= ED

[θ2kZk

]. Fix x0, h0 ∈ dom(F ) and let xk, hk be the random iterates produced by

SEGA. Then

E[Φk]≤ (1− αµ)kΦ0,

where Φk def= ‖xk−x∗‖2B +σα‖hk−∇f(x∗)‖2B is a Lyapunov function and x∗ is the solution of (1).

Note that the convergence of the Lyapunov function Φk implies both xk → x∗ and hk → ∇f(x∗).

The latter means that SEGA is variance reduced, in contrast to CD in the proximal setup with non-

separable R, which does not converge to the solution.

To clarify on the assumptions, let us mention that if σ is small enough so that Q−σED [Z] � 0,

one can always choose stepsize α satisfying

α ≤ min

{λmin(ED [Z])

λmax(2σ−1(C−B) + µB),λmin(Q− σED [Z])

2λmax(C)

}(13)

and inequalities (12) will hold. Therefore, we get the next corollary.

Corollary 3.4. If σ < λmin(Q)λmax(ED[Z]) , α satisfies (13) and k ≥ 1

αµ log Φ0

ε , then E[‖xk − x∗‖2B

]≤ ε.

As Theorem 3.3 is rather general, we also provide a simplified version thereof, complete with a

simplified analysis (Theorem D.1 in Appendix D). In the simplified version we remove the proximal

setting (i.e., we set R = 0), assume L–smoothness4, and only consider coordinate sketches with

uniform probabilities. The result is provided as Corollary 3.5.

Corollary 3.5. Let B = I and choose D to be the uniform distribution over unit basis vectors in

Rn. If the stepsize satisfies

0 < α ≤ min

1− Lσn

2Ln,

1

n(µ+ 2(n−1)

σ

) ,

then ED[Φk+1

]≤ (1− αµ)Φk, therefore the iteration complexity is O(nL/µ).

Remark 3.6. In the fully general setting, one might choose α to be bigger than bound (13), which

depends on eigen properties of matrices ED [Z] ,C,Q,B, leading to a better overall complexity ac-

cording to Corollary 3.4. However, in the simple case with B = I, Q = I and Sk = eik with uniform

probabilities, bound (13) is tight.

4The standard L–smoothness assumption is a special case of M–smoothness for M = LI, and hence is less general

than both M–smoothness and Q–smoothness with respect to B.

10

Page 11: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

CD SEGA

Nonaccelerated method

importance sampling, b = 1

Trace(M)µ log 1

ε [36] 8.55 · Trace(M)µ log 1

ε

Nonaccelerated method

arbitrary sampling

(maxi

vipiµ

)log 1

ε [43] 8.55 ·(

maxivipiµ

)log 1

ε

Accelerated method

importance sampling, b = 11.62 ·

∑i

√Mii√µ log 1

ε [3] 9.8 ·∑i

√Mii√µ log 1

ε

Accelerated method

arbitrary sampling1.62 ·

√maxi

vip2iµ

log 1ε [20] 9.8 ·

√maxi

vip2iµ

log 1ε

Table 1: Complexity results for coordinate descent (CD) and our sketched gradient method (SEGA),

specialized to coordinate sketching, for M–smooth and µ–strongly convex functions.

4 Convergence of SEGA for Coordinate Sketches

In this section we compare SEGA with coordinate descent. We demonstrate that, specialized to

a particular choice of the distribution D (where S is a random column submatrix of the identity

matrix), which makes SEGA use the same random gradient information as that used in modern

state-of-the-art randomized CD methods, SEGA attains, up to a small constant factor, the same

convergence rate as CD methods.

Firstly, in Section 4.2 we develop SEGA with arbitrary “coordinate sketches” (Theorem 4.2).

Then, in Section 4.3 we develop an accelerated variant of SEGA in a very general setup known as

arbitrary sampling (see Theorem B.5) [43, 42, 39, 40, 6]. Lastly, Corollary 4.3 and Corollary 4.5 pro-

vide us with importance sampling for both nonaccelerated and accelerated method, which matches

up to a constant factor cutting-edge coordinate descent rates [43, 3] under the same oracle and

assumptions5. Table 1 summarizes the results of this section. We provide a dedicated analysis for

the methods from this section in Appendix B.

We now describe the setup and technical assumptions for this section. In order to facilitate a

direct comparison with CD (which does not work with non-separable regularizer R), for simplicity

we consider problem (1) in the simplified setting with R ≡ 0. Further, function f is assumed to be

M–smooth (Assumption 3.2) and µ–strongly convex.

4.1 Defining D: samplings

In order to draw a direct comparison with general variants of CD methods (i.e., with those analyzed

in the arbitrary sampling paradigm), we consider sketches in (3) that are column submatrices of

5There was recently introduced a notion of importance minibatch sampling for coordinate descent [20]. We state,

without a proof, that SEGA with block coordinate sketches allows for the same importance sampling as developed in

the mentioned paper.

11

Page 12: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

the identity matrix: S = IS , where S is a random subset (aka sampling) of [n]def= {1, 2, . . . , n}.

Note that the columns of IS are the standard basis vectors ei for i ∈ S and hence

Range (S) = Range (ei : i ∈ S) .

So, distribution D from which we draw matrices is uniquely determined by the distribution of

sampling S. Given a sampling S, define p = (p1, . . . , pn) ∈ Rn to be the vector satisfying pi =

P (ei ∈ Range (S)) = P (i ∈ S), and P to be the matrix for which

Pij = P ({i, j} ⊆ S) .

Note that p and P are the probability vector and probability matrix of sampling S, respec-

tively [40]. We assume throughout the paper that S is proper, i.e., we assume that pi > 0 for all

i. State-of-the-art minibatch CD methods (including the ones we compare against [43, 20]) utilize

large stepsizes related to the so-called ESO Expected Separable Overapproximation (ESO) [40] pa-

rameters v = (v1, . . . , vn). ESO parameters play a key role in SEGA as well, and are defined next.

Assumption 4.1 (ESO). There exists a vector v satisfying the following inequality

P ◦M � Diag(p)Diag(v), (14)

where ◦ denotes the Hadamard (i.e., element-wise) product of matrices.

In case of single coordinate sketches, parameters v are equal to coordinate-wise smoothness

constants of f . An extensive study on how to choose them in general was performed in [40]. For

notational brevity, let us set Pdef= Diag(p) and V

def= Diag(v) throughout this section.

4.2 Non-accelerated method

We now state the convergence rate of (non-accelerated) SEGA for coordinate sketches with arbitrary

sampling of subsets of coordinates. The corresponding CD method was developed in [43].

Theorem 4.2. Assume that f is M–smooth and µ–strongly convex. Denote Ψk def= f(xk)−f(x∗)+

σ‖hk‖2P−1. Choose α, σ > 0 such that

σI− α2(VP−1 −M) � γµσP−1, (15)

where γdef= α− α2 maxi{ vipi } − σ. Then the iterates of SEGA satisfy E

[Ψk]≤ (1− γµ)kΨ0.

We now give an importance sampling result for a coordinate version of SEGA. We recover, up

to a constant factor, the same convergence rate as standard CD [36]. The probabilities we chose are

optimal in our analysis and are proportional to the diagonal elements of matrix M.

12

Page 13: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Corollary 4.3. Assume that f is M–smooth and µ–strongly convex. Suppose that D is such that

at each iteration standard unit basis vector ei is sampled with probability pi ∝ Mii. If we choose

α = 0.232Trace(M) , σ = 0.061

Trace(M) , then E[Ψk]≤(

1− 0.117µTrace(M)

)kΨ0.

The iteration complexities provided in Theorem 4.2 and Corollary 4.3 are summarized in Table 1.

We also state that σ, α can be chosen so that (15) holds, and the rate from Theorem 4.2 coincides

with the rate from Table 1.

Remark 4.4. Theorem 4.2 and Corollary 4.3 hold even under a non-convex relaxation of strong

convexity – Polyak- Lojasiewicz inequality: µ(f(x) − f(x∗)) ≤ 12‖∇f(x)‖22. Therefore, SEGA also

converges for a certain class of non-convex problems. For an overview on different relaxations of

strong convexity, see [23].

4.3 Accelerated method

In this section, we propose an accelerated (in the sense of Nesterov’s method [33, 34]) version of

SEGA, which we call ASEGA. The analogous accelerated CD method, in which a single coordinate

is sampled in every iteration, was developed and analyzed in [3]. The general variant utilizing

arbitrary sampling was developed and analyzed in [20].

Algorithm 2: ASEGA: Accelerated SEGA

1 Initialize: x0 = y0 = z0 ∈ Rn; h0 ∈ Rn; S; parameters α, β, τ, µ > 0

2 for k = 1, 2, . . . do

3 xk = (1− τ)yk−1 + τzk−1

4 Sample Sk = ISk , where Sk ∼ S, and compute gk, hk+1 according to (4), (6)

5 yk = xk − αP−1gk

6 zk = 11+βµ(zk + βµxk − βgk)

The method and analysis is inspired by [2]. Due to space limitations and technicality of the

content, we state the main theorem of this section in Appendix B.4. Here, we provide Corollary 4.5,

which shows that Algorithm 2 with single coordinate sampling enjoys, up to a constant factor, the

same convergence rate as state-of-the-art accelerated coordinate descent method NUACDM of Allen-

Zhu et al. [3].

Corollary 4.5. Let the sampling be defined as follows: S = {i} with probability pi ∝√

Mii, for

i ∈ [n]. Then there exist acceleration parameters and a Lyapunov function Υk such that f(yk) −f(x∗) ≤ Υk and

E[Υk]≤ (1− τ)kΥ0 =

(1−O

( √µ∑

i

√Mii

))kΥ0.

The iteration complexity guarantees provided by Theorem B.5 and Corollary 4.5 are summarized

in Table 1.

13

Page 14: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Figure 2: Convergence of SEGA and PGD on synthetic problems with n = 500. The indicator “Xn” in the label

indicates the setting where the cost of solving linear system is Xn times higher comparing to the oracle call. Recall

that a linear system is solved after each n oracle calls. Stepsizes 1/λmax(M) and 1/(nλmax(M)) were used for PGD

and SEGA, respectively.

5 Experiments

In this section we perform numerical experiments to illustrate the potential of SEGA. Firstly, in

Section 5.1, we compare it to projected gradient descent (PGD) algorithm. Then in Section 5.2, we

study the performance of zeroth-order SEGA (when sketched gradients are being estimated through

function value evaluations) and compare it to the analogous zeroth-order method. Lastly, in Sec-

tion 5.3 we verify the claim from Remark 3.6 that in some applications, particular sketches and

metric might lead to a significantly faster convergence. In the experiments where theory-supported

stepsizes were used, we obtained them by precomputing strong convexity and smoothness measures.

5.1 Comparison to projected gradient

In this experiment, we illustrate the potential superiority of our method to PGD. We consider the

`2 ball constrained problem (R is the indicator function of the unit ball) with the oracle providing

the sketched gradient in the random Gaussian direction. As we mentioned in the introduction, a

method moving in the gradient direction (analogue of CD), will not converge due to the proximal

nature of the problem. Therefore, we can only compare against the projected gradient. However, in

order to obtain the full gradient, one needs to gather n sketched gradients and solve a linear system

to recover the gradient. To illustrate this, we choose 4 different quadratic problems, according to

Table 2 in the appendix. We stress that these are synthetic problems generated for the purpose of

illustrating the potential of our method against a natural baseline. Figure 2 compares SEGA and

PGD under various relative cost scenarios of solving the linear system compared to the cost of the

oracle calls. The results show that SEGA significantly outperforms PGD as soon as solving the linear

system is expensive, and is as fast as PGD even if solving the linear system comes for free.

5.2 Comparison to zeroth-order optimization methods

In this section, we compare SEGA to the random direct search (RDS) method [5] under a zeroth-order

oracle for unconstrained optimization. For SEGA, we estimate the sketched gradient using finite

14

Page 15: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Figure 3: Comparison of SEGA and randomized direct search for various problems. Theory supported stepsizes were

chosen for both methods. 500 dimensional problem.

Figure 4: Comparison of SEGA with sketches from a correct subspace versus coordinate sketches naiveSEGA. Stepsize

chosen according to theory. 1000 dimensional problem.

differences. Note that RDS is a randomized version of the classical direct search method [21, 24, 25].

At iteration k, RDS moves to argmin(f(xk + αksk), f(xk − αksk), f(xk)

)for a random direction

sk ∼ D and a suitable stepszie αk. For illustration, we choose f to be a quadratic problem based

on Table 2 and compare both Gaussian and coordinate directions. Figure 3 shows that SEGA

outperforms RDS.

5.3 Subspace SEGA: a more aggressive approach

As mentioned in Remark 3.6, well designed sketches are capable of exploiting structure of f and lead

to a better rate. We address this in detail Appendix C where we develop and analyze a subspace

variant of SEGA.

To illustrate this phenomenon in a simple setting, we perform experiments for problem (1) with

f(x) = ‖Ax−b‖2, where b ∈ Rd and A ∈ Rd×n has orthogonal rows, and with R being the indicator

function of the unit ball in Rn. That is, we solve the problem

min‖x‖2≤1

‖Ax− b‖2.

We assume that n � d. We compare two methods: naiveSEGA, which uses coordinate sketches,

and subspaceSEGA, where sketches are chosen as rows of A. Figure 4 indicates that subspaceSEGA

outperforms naiveSEGA roughly by the factor nd , as claimed in Appendix C.

15

Page 16: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

6 Conclusions and Extensions

6.1 Conclusions

We proposed SEGA, a method for solving composite optimization problems under a novel stochastic

linear first order oracle. SEGA is variance-reduced, and this is achieved via sketch-and-project

updates of gradient estimates. We provided an analysis for smooth and strongly convex functions

and general sketches, and a refined analysis for coordinate sketches. For coordinate sketches we

also proposed an accelerated variant of SEGA, and our theory matches that of state-of-the-art

CD methods. However, in contrast to CD, SEGA can be used for optimization problems with a

non-separable proximal term. We develop a more aggressive subspace variant of the method—

subspaceSEGA—which leads to improvements in the n� d regime. In the Appendix we give several

further results, including simplified and alternative analyses of SEGA in the coordinate setup from

Example 2.1. Our experiments are encouraging and substantiate our theoretical predictions.

6.2 Extensions

We now point to several potential extensions of our work.

Speeding up the general method. We believe that it should be possible to extend ASEGA to

the general setup from Theorem 3.3. In such a case, it might be possible to design metric B and

distribution of sketches D so as to outperform accelerated proximal gradient methods [35, 4].

Biased gradient estimator. Recall that SEGA uses unbiased gradient estimator gk for updat-

ing the iterates xk in a similar way JacSketch [13] or SAGA [10] do this for the stochastic finite

sum optimization. Recently, a stochastic method for finite sum optimization using biased gradi-

ent estimators was proven to be more efficient [37]. Therefore, it might be possible to establish

better properties for a biased variant of SEGA. To demonstrate the potential of this approach, in

Appendix F.1 we plot the evolution of iterates for the very simple biased method which uses hk as

an update for line 3 in Algorithm 1.

Applications. We believe that SEGA might work well in applications where a zeroth-order ap-

proach is inevitable, such as reinforcement learning. We therefore believe that SEGA might be

an efficient proximal method in some reinforcement learning applications. We also believe that

communication-efficient variants of SEGA can be used for distributed training of machine learning

models. This is because SEGA can be adapted to communicate sparse model updates only.

16

Page 17: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

References

[1] Zeyuan Allen-Zhu. Katyusha: The first direct acceleration of stochastic gradient methods. In

Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, pages

1200–1205. ACM, 2017.

[2] Zeyuan Allen-Zhu and Lorenzo Orecchia. Linear coupling: An ultimate unification of gradient

and mirror descent. In Innovations in Theoretical Computer Science, 2017.

[3] Zeyuan Allen-Zhu, Zheng Qu, Peter Richtarik, and Yang Yuan. Even faster accelerated coordi-

nate descent using non-uniform sampling. In Proceedings of The 33rd International Conference

on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1110–

1119, 2016.

[4] Amir Beck and Marc Teboulle. A fast iterative shrinkage-thresholding algorithm for linear

inverse problems. SIAM Journal on Imaging Sciences, 2(1):183–202, 2009.

[5] El Houcine Bergou, Peter Richtarik, and Eduard Gorbunov. Random direct search method

for minimizing nonconvex, convex and strongly convex functions. Manuscript, 2018.

[6] Antonin Chambolle, Matthias J. Ehrhardt, Peter Richtarik, and Carola-Bibiane Schoenlieb.

Stochastic primal-dual hybrid gradient algorithm with arbitrary sampling and imaging appli-

cations. SIAM Journal on Optimization, 28(4):27832808, 2018.

[7] Chih-Chung Chang and Chih-Jen Lin. LibSVM: A library for support vector machines. ACM

transactions on intelligent systems and technology (TIST), 2(3):27, 2011.

[8] Andrew R Conn, Katya Scheinberg, and Luis N Vicente. Introduction to derivative-free opti-

mization, volume 8. Siam, 2009.

[9] Alexandre d’Aspremont. Smooth optimization with approximate gradient. SIAM Journal on

Optimization, 19(3):1171–1183, 2008.

[10] Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. SAGA: A fast incremental gradient

method with support for non-strongly convex composite objectives. In Advances in Neural

Information Processing Systems, pages 1646–1654, 2014.

[11] Olivier Devolder, Francois Glineur, and Yurii Nesterov. First-order methods of smooth convex

optimization with inexact oracle. Mathematical Programming, 146(1-2):37–75, 2014.

[12] Olivier Fercoq and Peter Richtarik. Accelerated, parallel and proximal coordinate descent.

SIAM Journal on Optimization, (25):1997–2023, 2015.

17

Page 18: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

[13] Robert M Gower, Peter Richtarik, and Francis Bach. Stochastic quasi-gradient methods:

Variance reduction via Jacobian sketching. arXiv preprint arXiv:1805.02632, 2018.

[14] Robert Mansel Gower, Donald Goldfarb, and Peter Richtarik. Stochastic block BFGS: squeez-

ing more curvature out of data. In 33rd International Conference on Machine Learning, pages

1869–1878, 2016.

[15] Robert Mansel Gower, Filip Hanzely, Peter Richtarik, and Sebastian Stich. Accelerated

stochastic matrix inversion: general theory and speeding up BFGS rules for faster second-

order optimization. arXiv:1802.04079, 2018.

[16] Robert Mansel Gower and Peter Richtarik. Randomized iterative methods for linear systems.

SIAM Journal on Matrix Analysis and Applications, 36(4):1660–1690, 2015.

[17] Robert Mansel Gower and Peter Richtarik. Stochastic dual ascent for solving linear systems.

arXiv preprint arXiv:1512.06890, 2015.

[18] Robert Mansel Gower and Peter Richtarik. Linearly convergent randomized iterative methods

for computing the pseudoinverse. arXiv:1612.06255, 2016.

[19] Robert Mansel Gower and Peter Richtarik. Randomized quasi-Newton updates are linearly

convergent matrix inversion algorithms. SIAM Journal on Matrix Analysis and Applications,

38(4):1380–1409, 2017.

[20] Filip Hanzely and Peter Richtarik. Accelerated coordinate descent with arbitrary sampling

and best rates for minibatches. arXiv preprint arXiv:1809.09354, 2018.

[21] Robert Hooke and Terry A Jeeves. “Direct search” solution of numerical and statistical prob-

lems. Journal of the ACM (JACM), 8(2):212–229, 1961.

[22] Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive vari-

ance reduction. In Advances in Neural Information Processing Systems, pages 315–323, 2013.

[23] Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal-

gradient methods under the Polyak-Lojasiewicz condition. In Joint European Conference on

Machine Learning and Knowledge Discovery in Databases, pages 795–811. Springer, 2016.

[24] Tamara G Kolda, Robert Michael Lewis, and Virginia Torczon. Optimization by direct search:

New perspectives on some classical and modern methods. SIAM Review, 45(3):385–482, 2003.

[25] Jakub Konecny and Peter Richtarik. Simple complexity analysis of simplified direct search.

arXiv preprint arXiv:1410.0390, 2014.

[26] Hongzhou Lin, Julien Mairal, and Zaid Harchaoui. A universal catalyst for first-order opti-

mization. In Advances in Neural Information Processing Systems, pages 3384–3392, 2015.

18

Page 19: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

[27] Qihang Lin, Zhaosong Lu, and Lin Xiao. An accelerated proximal coordinate gradient method.

In Advances in Neural Information Processing Systems, pages 3059–3067, 2014.

[28] Nicoas Loizou and Peter Richtarik. Accelerated gossip via stochastic heavy ball method. In

56th Annual Allerton Conference on Communication, Control, and Computing, 2018.

[29] Nicolas Loizou and Peter Richtarik. A new perspective on randomized gossip algorithms. In

IEEE Global Conference on Signal and Information Processing (GlobalSIP), pages 440–444,

2016.

[30] Nicolas Loizou and Peter Richtarik. Linearly convergent stochastic heavy ball method for

minimizing generalization error. In NIPS Workshop on Optimization for Machine Learning,

2017.

[31] Nicolas Loizou and Peter Richtarik. Momentum and stochastic momentum for stochastic

gradient, Newton, proximal point and subspace descent methods. arXiv:1712.09677, 2017.

[32] Ion Necoara, Peter Richtarik, and Andrei Patrascu. Randomized projection methods for convex

feasibility problems: conditioning and convergence rates. arXiv:1801.04873, 2018.

[33] Yurii Nesterov. A method of solving a convex programming problem with convergence rate

O(1/k2). Soviet Mathematics Doklady, 27(2):372–376, 1983.

[34] Yurii Nesterov. Introductory lectures on convex optimization: A basic course. Kluwer Academic

Publishers, 2004.

[35] Yurii Nesterov. Smooth minimization of nonsmooth functions. Mathematical Programming,

103:127–152, 2005.

[36] Yurii Nesterov. Efficiency of coordinate descent methods on huge-scale optimization problems.

SIAM Journal on Optimization, 22(2):341–362, 2012.

[37] Lam M. Nguyen, Jie Liu, Katya Scheinberg, and Martin Takac. SARAH: A novel method

for machine learning problems using stochastic recursive gradient. In Proceedings of the 34th

International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning

Research, pages 2613–2621. PMLR, 2017.

[38] Zheng Qu and Peter Richtarik. Coordinate descent with arbitrary sampling I: Algorithms and

complexity. Optimization Methods and Software, 31(5):829–857, 2016.

[39] Zheng Qu and Peter Richtarik. Coordinate descent with arbitrary sampling I: Algorithms and

complexity. Optimization Methods and Software, 31(5):829–857, 2016.

[40] Zheng Qu and Peter Richtarik. Coordinate descent with arbitrary sampling II: Expected

separable overapproximation. Optimization Methods and Software, 31(5):858–884, 2016.

19

Page 20: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

[41] Zheng Qu, Peter Richtarik, Martin Takac, and Olivier Fercoq. SDNA: Stochastic dual Newton

ascent for empirical risk minimization. In Proceedings of The 33rd International Conference on

Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1823–1832.

PMLR, 2016.

[42] Zheng Qu, Peter Richtarik, and Tong Zhang. Quartz: Randomized dual coordinate ascent with

arbitrary sampling. In Advances in Neural Information Processing Systems, pages 865–873,

2015.

[43] Peter Richtarik and Martin Takac. On optimal probabilities in stochastic coordinate descent

methods. Optimization Letters, 10(6):1233–1243, 2016.

[44] Peter Richtarik and Martin Takac. Stochastic reformulations of linear systems: algorithms

and convergence theory. arXiv:1706.01108, 2017.

[45] H. Robbins and S. Monro. A stochastic approximation method. Annals of Mathematical

Statistics, 22:400–407, 1951.

[46] Nicolas Le Roux, Mark Schmidt, and Francis Bach. A stochastic gradient method with an

exponential convergence rate for finite training sets. In Advances in Neural Information Pro-

cessing Systems, pages 2663–2671, 2012.

[47] Mark Schmidt, Nicolas L Roux, and Francis R Bach. Convergence rates of inexact proximal-

gradient methods for convex optimization. In Advances in Neural Information Processing

Systems, pages 1458–1466, 2011.

[48] Shai Shalev-Shwartz and Tong Zhang. Proximal stochastic dual coordinate ascent. arXiv

preprint arXiv:1211.2717, 2012.

20

Page 21: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Appendix

A Proofs for Section 3

Lemma A.1. Suppose that B = I and f is twice differentiable. Assumption 3.1 is equivalent to

Assumption 3.2 for Q = M−1.

Proof: We first establish that Assumption 3.1 implies Assumption 3.2. Summing up (10) for

(x, y) and (y, x) yields

〈∇f(x)−∇f(y), x− y〉 ≥ ‖∇f(x)−∇f(y)‖2Q.

Using Cauchy Schwartz inequality we obtain

‖x− y‖Q−1 ≥ ‖∇f(x)−∇f(y)‖Q.

By the mean value theorem, there is z ∈ [x, y] such that ∇f(x)−∇f(y) = ∇2f(z)(x− y). Thus

‖x− y‖Q−1 ≥ ‖x− y‖∇2f(z)Q∇2f(z).

The above is equivalent to(∇2f(z)

)− 12 Q−1

(∇2f(z)

)− 12 �

(∇2f(z)

) 12 Q

(∇2f(z)

) 12

Note that for any M′ � 0 we have M′ �M−1 if and only if M � I. Thus(∇2f(z)

)− 12 Q−1

(∇2f(z)

)− 12 � I,

which is equivalent to Q−1 � ∇2f(z). To establish the other direction, denote φ(y) = f(y) −〈∇f(x), y〉. Clearly, x is minimizer of φ and therefore we have

φ(x) ≤ φ(x−M−1∇f(y)) ≤ φ(y)− 1

2‖∇f(y)‖2M−1 ,

which is exactly (10) for Q = M−1.

Lemma A.2. For B � 0 and Zkdef= Sk(S

>k B−1Sk)

†S>k , then

Z>k B−1Zk = Zk. (16)

Proof: It is a property of pseudo-inverse that for any matrices A,B it holds ((AB)†)> =

(B>A>)†, so Z>k = Zk. Moreover, we also know for any A that A†AA† = A† and, thus,

Z>k B−1Zk = Sk(S>k B−1Sk)

†S>k B−1Sk(S>k B−1Sk)

†S>k = Sk(S>k B−1Sk)

†S>k = Zk.

21

Page 22: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

A.1 Proof of Theorem 3.3

We first state two lemmas which will be crucial for the analysis. They characterize key properties

of the gradient learning process (4), (6) and will be used later to bound expected distances of both

hk+1 and gk from ∇f(x∗). The proofs are provided in Appendix A.2 and A.3 respectively

Lemma A.3. For all v ∈ Rn we have

ED[‖hk+1 − v‖2B

]= ‖hk − v‖2B−ED[Z] + ‖∇f(xk)− v‖2ED[Z]. (17)

Lemma A.4. Let Cdef= ED

[θ2Z

]. Then for all v ∈ Rn we have

ED[‖gk − v‖2B

]≤ 2‖∇f(xk)− v‖2C + 2‖hk − v‖2C−B.

For notational simplicity, it will be convenient to define Bregman divergence between x and y:

Df (x, y)def= f(x)− f(y)− 〈∇f(y)), x− y〉B

We can now proceed with the proof of Theorem 3.3. Let us start with bounding the first term in

the expression for Φk+1. From Lemma A.4 and strong convexity it follows that

ED[‖xk+1 − x∗‖2B

]= ED

[‖ proxαR(xk − αgk)− proxαR(x∗ − α∇f(x∗))‖2B

]≤ ED

[‖xk − αgk − (x∗ − α∇f(x∗))‖2B

]= ‖xk − x∗‖2B − 2αED

[(gk −∇f(x∗))>B(xk − x∗)

]+α2ED

[‖gk −∇f(x∗)‖2B

]≤ ‖xk − x∗‖2B − 2α(∇f(xk)−∇f(x∗))>B(xk − x∗)

+2α2‖∇f(xk)−∇f(x∗)‖2C + 2α2‖hk −∇f(x∗)‖2C−B≤ ‖xk − x∗‖2B − αµ‖xk − x∗‖2B − 2αDf (xk, x∗)

+2α2‖∇f(xk)−∇f(x∗)‖2C + 2α2‖hk −∇f(x∗)‖2C−B.

Using Assumption 3.1 we get

−2αDf (xk, x∗) ≤ −α‖∇f(xk)−∇f(x∗)‖2Q.

As for the second term in Φk+1, we have by Lemma A.3

ασED[‖hk+1 −∇f(x∗)‖2B

]= ασ‖hk −∇f(x∗)‖2B−ED[Z] + ασ‖∇f(xk)−∇f(x∗)‖2ED[Z]

Combining it into Lyapunov function Φk,

Φk+1 ≤ (1− αµ)‖xk − x∗‖2B + ασ‖hk −∇f(x∗)‖2B−ED[Z] + 2α2‖hk −∇f(x∗)‖2C−B+ασ‖∇f(xk)−∇f(x∗)‖2ED[Z] + 2α2‖∇f(xk)−∇f(x∗)‖2C − α‖∇f(xk)−∇f(x∗)‖2Q.

22

Page 23: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

To see that this gives us the theorem’s statement, consider first

ασED [Z] + 2α2C− αQ = 2α(αC− 12(Q− σED [Z])) ≤ 0,

so we can drop norms related to ∇f(xk)−∇f(x∗). Next, we have

ασ(B− ED [Z]) + 2α2(C−B) = α (α(2(C−B) + σµB)− ED [Z]) + σα(1− αµ)B

≤ σα(1− αµ)B,

which follows from our assumption on α.

A.2 Proof of Lemma A.3

Proof: Keeping in mind that Z>k = Zk and (B−1)> = B−1, we first write

ED[‖hk+1 − v‖2B

](8)= ED

[∥∥∥hk + B−1Zk(∇f(xk)− hk)− v∥∥∥2

B

]= ED

[∥∥∥(I−B−1Zk)

(hk − v) + B−1Zk(∇f(xk)− v)∥∥∥2

B

]= ED

[∥∥∥(I−B−1Zk)

(hk − v)∥∥∥2

B

]+ ED

[∥∥∥B−1Zk(∇f(xk)− v)∥∥∥2

B

]+2(hk − v)>ED

[(I−B−1Zk

)>BB−1Zk

](∇f(xk)− v)

= (hk − v)>ED[(

I−B−1Zk)>

B(I−B−1Zk

)](hk − v)

+(∇f(xk)− v)>ED[ZkB

−1BB−1Zk]

(∇f(xk)− v)

+2(hk − v)>ED[Zk − ZkB

−1Zk]

(∇f(xk)− v).

By Lemma A.2 we have ZkB−1Zk = Zk, so the last term in the expression above is equal to 0. As

for the other two, expanding the matrix factor in the first term leads to

ED[(

I−B−1Zk)>

B(I−B−1Zk

)]= ED

[(I− ZkB

−1)B(I−B−1Zk

)]= ED

[B− ZkB

−1B−BB−1Zk + ZkB−1BB−1Zk

]= B− ED [Zk] .

We, thereby, have derived

ED[‖hk+1 − v‖2B

]= (hk − v)> (B− ED [Zk]) (hk − v)

+(∇f(xk)− v)>ED[ZkB

−1Zk]

(∇f(xk)− v)

= ‖hk − v‖2B−ED[Z] + ‖∇f(xk)− v‖2ED[Z].

23

Page 24: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

A.3 Proof of Lemma A.4

Proof: Throughout this proof, we will use without any mention that Z>k = Zk.

Writing gk − v = a+ b, where adef= (I− θkB−1Zk)(h

k − v) and bdef= θkB

−1Zk(∇f(xk)− v), we

get ‖gk‖2B ≤ 2(‖a‖2B + ‖b‖2B). Using Lemma A.2 and the definition of θk yields

ED[‖a‖2B

]= ED

[‖(I− θkB−1Zk

)(hk − v)‖2B

]= (hk − v)>ED

[(I− θkZkB−1

)B(I− θkB−1Zk

)](hk − v)

= (hk − v)>ED[(

B− θkZkB−1B−BθkB−1Zk + θ2

kZkB−1BB−1Zk

)](hk − v)

= (hk − v)>ED[(

B− 2B + θ2kZk

)](hk − v)

= ‖hk − v‖2ED[θ2Z]−B.

Similarly, the second term in the upper bound on gk can be rewritten as

ED[‖b‖2B

]= ED

[‖θkB−1Zk(∇f(xk)− v)‖2B

]= (∇f(xk)− v)>ED

[θ2kZkB

−1BB−1Zk]

(∇f(xk)− v)

= ‖∇f(xk)− v‖2C.

Combining the pieces, we get the claim.

B Proofs for Section 4

B.1 Technical Lemmas

We first start with an analogue of Lemma A.4 allowing for a norm different from ‖ · ‖B. We

remark that matrix Q′ in the lemma is not to be confused with the smoothness matrix Q from

Assumption 3.1.

Lemma B.1. Let Q′ � 0. The variance of gk as an estimator of ∇f(xk) can be bounded as follows:

1

2ED[‖gk‖2Q′

]≤ ‖hk‖2

P−1(P◦Q′)P−1−Q′ + ‖∇f(xk)‖2P−1(P◦Q′)P−1 . (18)

Proof: Denote Sk to be a matrix with columns ei for i ∈ Range (Sk). We first write

gk = hk − P−1SkS>k h

k︸ ︷︷ ︸a

+ P−1SkS>k∇f(xk)︸ ︷︷ ︸b

.

24

Page 25: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Let us bound the expectation of each term individually. The first term is equal to

ED[‖a‖2Q′

]= ED

[∥∥∥(I− P−1SkS>k

)hk∥∥∥2

Q′

]= (hk)>ED

[(I− P−1SkS

>k

)>Q′(I− P−1SkS

>k

)]hk

= (hk)>ED[(

Q′ − P−1SkS>k Q′ −Q′SkS

>k P−1

)]hk

+(hk)>ED[(

P−1SkS>k Q′SkS

>k P−1

)]hk

= (hk)>(P−1(P ◦Q′)P−1 −Q′

)hk.

The second term can be bounded as

ED[‖b‖2Q′

]= ED

[∥∥∥P−1S>k∇f(xk)Sk

∥∥∥2

Q′

]= ED

[‖∇f(xk)‖2

P−1SkS>k Q′SkS

>k P−1

]= ‖∇f(xk)‖2

P−1(P◦Q′)P−1

It remains to combine the two bounds.

We also state the analogue of Lemma A.3, which allows for a different norm as well.

Lemma B.2. For all diagonal D � 0 we have

ED[‖hk+1‖2D

]= ‖hk‖2

D−PD+ ‖∇f(xk)‖2

PD. (19)

Proof: Denote Sk to be a matrix with columns ei for i ∈ Sk. We first write

hk+1 = hk − SkS>k h

k + SkS>k∇f(xk).

Therefore

ED[‖hk+1‖2D

]= ED

[∥∥∥(I− SkS>k )hk + SkS

>k∇f(xk)

∥∥∥2

D

]= ED

[∥∥∥(I− SkS>k )hk

∥∥∥2

D

]+ ED

[∥∥∥SkS>k∇f(xk)∥∥∥2

D

]+2ED

[hk>

(I− SkS>k )DSkS

>k∇f(xk)

]= ‖hk‖2

D−PD+ ‖∇f(xk)‖2

PD.

B.2 Proof of Theorem 4.2

Proof: Throughout the proof, we will use the following Lyapunov function:

Ψk def= f(xk)− f(x∗) + σ‖hk‖2P−1 .

25

Page 26: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Following similar steps to what we did before, we obtain

E[Ψk+1

] (11)

≤ f(xk)− f(x∗) + αE[〈∇f(xk), gk〉

]+α2

2E[‖gk‖2M

]+ σE

[‖hk+1‖2

P−1

]= f(xk)− f(x∗)− α‖∇f(xk)‖22 +

α2

2E[‖gk‖2M

]+ σE

[‖hk+1‖2

P−1

](18)

≤ f(xk)− f(x∗)− α‖∇f(xk)‖22 + α2‖∇f(xk)‖2P−1(P◦M)P−1 + α2‖hk‖2

P−1(P◦M)P−1−M

+σE[‖hk+1‖2

P−1

].

This is the place where the ESO assumption comes into play. By applying it to the right-hand side

of the bound above, we obtain

E[Ψk+1

] (14)

≤ f(xk)− f(x∗)− α‖∇f(xk)‖22 + α2‖∇f(xk)‖2VP−1 + α2‖hk‖2

VP−1−M

+σE[‖hk+1‖2

P−1

](19)= f(xk)− f(x∗)− α‖∇f(xk)‖22 + α2‖∇f(xk)‖2

VP−1 + α2‖hk‖2VP−1−M

+σ‖∇f(xk)‖22 + σ‖hk‖2P−1−I

= f(xk)− f(x∗)−(α− α2 max

i

vipi− σ

)‖∇f(xk)‖22

+‖hk‖2α2(VP−1−M)+σ(P−1−I).

Due to Polyak- Lojasiewicz inequality, we can further upper bound the last expression by(1−

(α− α2 max

i

vipi− σ

)(f(xk)− f(x∗)) + ‖hk‖2

α2(VP−1−M)+σ(P−1−I).

To finish the proof, it remains to use (15).

B.3 Proof of Corollary 4.3

The claim was obtained by choosing carefully α and σ using numerical grid search. Note that by

strong convexity we have I � µDiag(M)−1, so we can satisfy assumption (15). Then, the claim

follows immediately noticing that we can also set V = Diag(M) while maintaining(α− α2 max

i

Mii

pi− σ

)≥ 0.117

Trace(M).

B.4 Accelerated SEGA with arbitrary sampling

Before establishing the main theorem, we first state two technical lemmas which will be crucial

for the analysis. First one, Lemma B.3 provides a key inequality following from (6). The second

one, Lemma B.4, analyzes update (5) and was technically established throughout the proof of

Theorem 4.2. We include a proof of lemmas in Appendix B.5 and B.6 respectively.

26

Page 27: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Lemma B.3. For every u ∈ Rn we have

β〈∇f(xk+1), zk − u〉 − βµ

2‖xk+1 − u‖22

≤ β2 1

2E[‖gk‖22

]+

1

2‖zk − u‖22 −

1 + βµ

2E[‖zk+1 − u‖22

](20)

Lemma B.4. Letting η(v, p)def= maxi

√vipi

, we have

f(xk+1)− E[f(yk+1)

]+ ‖hk‖2

α2(VP−3−P−1MP−1)≥(α− α2η(v, p)2

)‖∇f(xk)‖2

P−1 . (21)

Now we state the main theorem of Section 4.3, providing a convergence rate of ASEGA (Algo-

rithm 2) for arbitrary minibatch sampling. As we mentioned, the convergence rate is, up to a

constant factor, same as state-of-the-art minibatch accelerated coordinate descent [20].

Theorem B.5. Assume M–smoothness and µ–strong convexity and that v satisfies (14). Denote

Υk def=

2

75

η(v, p)−2

τ2

(E[f(yk)

]− f(x∗)

)+

1 + βµ

2E[‖zk − x∗‖22

]+ σE

[‖hk‖2

P−2

]and choose

c1 = max

(1, η(v, p)−1

õ

mini pi

)(22)

α =1

5η(v, p)2(23)

β =2

75τη(v, p)2(24)

σ = 5β2 (25)

τ =

√4

9·54 η(v, p)−4µ2 + 875η(v, p)−2µ− 2

75η(v, p)−2µ

2(26)

Then, we have

E[Υk]≤(1− c−1

1 τ)k

Υ0.

Proof: The proof technique is inspired by [2]. First of all, let us see what strong convexity of f

gives us:

β(f(xk+1)− f(x∗)

)≤ β〈∇f(xk+1), xk+1 − x∗〉 − βµ

2‖x∗ − xk+1‖22.

Thus, we are interested in finding an upper bound for the scalar product that appeared above. We

have

β〈∇f(xk+1), zk − u〉 − βµ

2‖xk+1 − u‖22 + σE

[‖hk+1‖2

P−2

](20)

≤ β2 1

2E[‖gk‖22

]+

1

2‖zk − u‖22 −

1 + βµ

2E[‖zk+1 − u‖22

]+ σE

[‖hk+1‖2

P−2

].

27

Page 28: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Using the Lemmas introduced above, we can upper bound the norms of gk and hk+1 by using norms

of hk and ∇f(xk) to get the following:

β2 1

2E[‖gk‖22

]+ σE

[‖hk+1‖2

P−2

](19)

≤ β2 1

2E[‖gk‖22

]+ σ‖hk‖2

P−2−P−1 + σ‖∇f(xk)‖2P−1

(18)

≤ β2‖hk‖2P−1−I + β2‖∇f(xk)‖2

P−1 + σ‖hk‖2P−2−P−1 + σ‖∇f(xk)‖2

P−1 .

Now, let us get rid of ∇f(xk) by using the gradients property from Lemma B.4:

β2 1

2E[‖gk‖22

]+ σE

[‖hk+1‖2

P−2

](21)

≤ β2‖hk‖2P−1−I +

(β2 + σ

) f(xk+1)− f(yk+1) + ‖hk‖2α2(VP−3−P−1MP−1)

α− α2η(v, p)2+ σ‖hk‖2

P−2−P−1

= ‖hk‖2β2(P−1−I)+ (β2+σ)α2

α−α2η(v,p)2(VP−3−P−1MP−1)+σ(P−2−P−1)

+β2 + σ

α− α2η(v, p)2(f(xk+1)− E

[f(yk+1)

])

≤ ‖hk‖2β2P−1+

(β2+σ)α2

α−α2η(v,p)2VP−3+σ(P−2−P−1)

+β2 + σ

α− α2η(v, p)2(f(xk+1)− E

[f(yk+1)

]).

Plugging this into the bound with which we started the proof, we deduce

β〈∇f(xk+1), zk − u〉 − βµ

2‖xk+1 − u‖22 + σE

[‖hk+1‖2

P−2

]≤ ‖hk‖2

β2P−1+(β2+σ)α2

α−α2η(v,p)2VP−3+σ(P−2−P−1)

+β2 + σ

α− α2η(v, p)2(f(xk+1)− E

[f(yk+1)

]) +

1

2‖zk − u‖22 −

1 + βµ

2E[‖zk+1 − u‖22

].

Recalling our first step, we get with a few rearrangements

β(f(xk+1)− f(x∗)

)≤ β〈∇f(xk+1), xk+1 − x∗〉 − βµ

2‖x∗ − xk+1‖22

= β〈∇f(xk+1), xk+1 − zk〉+ β〈∇f(xk+1), zk − x∗〉 − βµ

2‖x∗ − xk+1‖22

=(1− τ)β

τ〈∇f(xk+1), yk − xk+1〉+ β〈∇f(xk+1), zk − x∗〉 − βµ

2‖x∗ − xk+1‖22

≤ (1− τ)β

τ

(f(yk)− f(xk+1)

)+ ‖hk‖2

β2P−1+(β2+σ)α2

α−α2η(v,p)2VP−3+σ(P−2−P−1)

+β2 + σ

α− α2η(v, p)2(f(xk+1)− E

[f(yk+1)

]) +

1

2‖zk − x∗‖22

−1 + βµ

2E[‖zk+1 − x∗‖22

]− σE

[‖hk+1‖2

P−2

].

28

Page 29: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Let us choose σ, β such that for some constant c2 (which we choose at the end) we have

c2σ = β2, β =α− α2η(v, p)2

(1 + c−12 )τ

.

Consequently, we have

α− α2η(v, p)2

(1 + c−12 )τ2

(E[f(yk+1)

]− f(x∗)

)+

1 + βµ

2E[‖zk+1 − x∗‖22

]+ σE

[‖hk+1‖2

P−2

]≤ (1− τ)

α− α2η(v, p)2

(1 + c−12 )τ2

(f(yk)− f(x∗)

)+

1

2‖zk − x∗‖22

+‖hk‖2(P−1−(1−c2)I+

(1+c2)α2

α−α2η(v,p)2VP−2

)σP−1

Let us make a particular choice of α, so that for some constant c3 (which we choose at the end) we

can obtain the equations below:

α =1

c3η(v, p)2⇒ α− α2η(v, p)2 =

c3 − 1

c23

η(v, p)−2,α2

α− α2η(v, p)2=

1

(c3 − 1)η(v, p)2.

Thus

c3−1c23

η(v, p)−2

(1 + c−12 )τ2

(E[f(yk+1)

]− f(x∗)

)+

1 + βµ

2E[‖zk+1 − x∗‖22

]+ σE

[‖hk+1‖2

P−2

]≤ (1− τ)

c3−1c23

η(v, p)−2

(1 + c−12 )τ2

(f(yk)− f(x∗)

)+

1

2‖zk − x∗‖22

+‖hk‖2(P−1−(1−c2)I+

(1+c2)

(c3−1)η(v,p)2VP−2

)σP−1

.

Using the definition of η(v, p), one can see that the above gives

c3−1c23

η(v, p)−2

(1 + c−12 )τ2

(E[f(yk+1)

]− f(x∗)

)+

1 + βµ

2E[‖zk+1 − x∗‖22

]+ σE

[‖hk+1‖2

P−2

]≤ (1− τ)

c3−1c23

η(v, p)−2

(1 + c−12 )τ2

(f(yk)− f(x∗)

)+

1

2‖zk − x∗‖22 + ‖hk‖2(

P−1−(1−c2)I+1+c2c3−1

I)σP−1

.

To get the convergence rate, we shall establish(1− c2 −

1 + c2

c3 − 1

)c1I � τP−1 (27)

and

1 + βµ ≥ 1

1− τ. (28)

To this end, let us recall that

β =c3 − 1

c22

η(v, p)−2τ−1 1

1 + c−12

.

29

Page 30: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Now we would like to set equality in (28), which yields

0 = τ2 +c3 − 1

c22

η(v, p)−2 1

1 + c−12

µτ − c3 − 1

c22

η(v, p)−2 1

1 + c−12

µ = 0.

This, in turn, implies

τ =

√(c3−1c22

)2η(v, p)−4 1

(1+c−12 )

2µ2 + 4 c3−1c22

η(v, p)−2 11+c−1

2

µ− c3−1c22

η(v, p)−2 11+c−1

2

µ

2

= O

√c3 − 1

c22

1√1 + c−1

2

η(v, p)−1√µ

.

Notice that for any c ≤ 1 we have√c2+4c−c

2 ≤√c and therefore

τ ≤

√c3 − 1

c22

η(v, p)−1 1√1 + c−1

2

õ. (29)

Using this inequality and a particular choice of constants, we can upper bound P−1 by a matrix

proportional to identity as shown below:

τP−1(29)

√c3 − 1

c22

η(v, p)−1 1√1 + c−1

2

√µP−1

√c3 − 1

c22

η(v, p)−1 1√1 + c−1

2

õ

mini piI

(22)

√c3 − 1

c22

1√1 + c−1

2

c1I

(∗)�

(1− c2 −

1 + c2

c3 − 1

)c1I,

which is exactly (27). Above, (∗) holds for choice c3 = 5 and c2 = 15 . It remains to verify that (23),

(24), (25) and (26) indeed correspond to our derivations.

We also mention, without a proof, that acceleration parameters can be chosen in general such

that c1 can be lower bounded by constant and therefore the rate from Theorem B.5 coincides with

the rate from Table 1. Corollary 4.5 is in fact a weaker result of that type.

B.4.1 Proof of Corollary 4.5

It suffices to verify that one can choose v = Diag(M) in (14) and that due to pi ∝√

Mii we have

c1 = 1.

30

Page 31: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

B.5 Proof of Lemma B.3

Proof: Firstly (6), is equivalent to

zk+1 = argminz

ψk(z)def=

1

2‖z − zk‖22 + β〈gk, z〉+

βµ

2‖z − xk+1‖22.

Therefore, we have for every u

0 = 〈∇ψk(zk+1), zk+1 − u〉

= 〈zk+1 − zk, zk+1 − u〉+ β〈gk, zk+1 − u〉+ βµ〈zk+1 − xk+1, zk+1 − u〉. (30)

Next, by generalized Pythagorean theorem we have

〈zk+1 − zk, zk+1 − u〉 =1

2‖zk − zk+1‖22 −

1

2‖zk − u‖22 +

1

2‖u− zk+1‖22 (31)

and

〈zk+1 − xk+1, zk+1 − u〉 =1

2‖xk+1 − zk+1‖22 −

1

2‖xk+1 − u‖22 +

1

2‖u− zk+1‖22. (32)

Plugging (31) and (32) into (30) we obtain

β〈gk, zk − u〉 − βµ

2‖xk+1 − u‖22

≤ β〈gk, zk − zk+1〉 − 1

2‖zk − zk+1‖22 +

1

2‖zk − u‖22 −

1 + βµ

2‖zk+1 − u‖22

(∗)≤ β2

2‖gk‖22 +

1

2‖zk − u‖22 −

1 + βµ

2‖zk+1 − u‖22.

The step marked by (∗) holds due to Cauchy-Schwartz inequality. It remains to take the expectation

conditioned on xk+1 and use (7).

B.6 Proof of Lemma B.4

Proof: The shortest, although not the most intuitive, way to write the proof is to put matrix

factor into norms. Apart from this trick, the proof is quite simple consists of applying smoothness

31

Page 32: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

followed by ESO:

E[f(yk+1)

]− f(xk+1)

(11)

≤ −αE[〈∇f(xk), P−1gk〉

]+α2

2E[‖P−1gk‖2M

]= −α‖∇f(xk)‖2

P−1 +α2

2E[‖gk‖P−1MP−1

](18)

≤ −α‖∇f(xk)‖2P−1 + α2‖∇f(xk)‖2

P−1(P◦P−1MP−1)P−1

+α2‖hk‖2P−1(P◦P−1MP−1)P−1−P−1MP−1

= −α‖∇f(xk)‖2P−1 + α2‖∇f(xk)‖2

P−2(P◦M)P−2

+α2‖hk‖2P−2(P◦M)P−2−P−1MP−1

(14)

≤ −α‖∇f(xk)‖2P−1 + α2‖∇f(xk)‖2

VP−3

+α2‖hk‖2VP−3−P−1MP−1

≤ −(α− α2 max

i

vip2i

)‖f(xk)‖2

P−1 + α2‖hk‖2VP−3−P−1MP−1 .

C Subspace SEGA: a More Aggressive Approach

In this section we describe a more aggressive variant of SEGA, one that exploits the fact that the

gradients of f lie in a lower dimensional subspace if this is indeed the case.

In particular, assume that F (x) = f(x) +R(x) and

f(x) = φ(Ax),

where A ∈ Rm×n6. Note that ∇f(x) lies in Range(A>). There are situations where the dimension

of Range(A>)

is much smaller than n. For instance, this happens when m� n. However, standard

coordinate descent methods still move around in directions ei ∈ Rn for all i. We can modify the

gradient sketch method to force our gradient estimate to lie in Range(A>), hoping that this will

lead to faster convergence.

C.1 The algorithm

Let xk be the current iterate, and let hk be the current estimate of the gradient of f . Assume

that the sketch S>k∇f(xk) is available. We can now define hk+1 through the following modified

6Strong convexity is not compatible with the assumption that A does not have full rank, so a different type of

analysis using Polyak- Lojasiewicz inequality is required to give a formal justification. However, we proceed with the

analysis anyway to build the intuition why this approach leads to better rates.

32

Page 33: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

sketch-and-project process:

hk+1 = arg minh∈Rn

‖h− hk‖2B

subject to S>k h = S>k∇f(xk), (33)

h ∈ Range(A>).

Before proceeding further, we note that there are such sketches and metric (as discussed in Sec-

tion C.4) which keep h ∈ Range(A>)

implicitly, and therefore one might omit the extra constraint

in such case. In fact, the mentioned sketches also lead to a faster convergence, which is the main

takeaway from this section.

Standard arguments reveal that the closed-form solution of (33) is

hk+1 = H(hk −B−1Sk(S

>k HB−1Sk)

†S>k (Hhk −∇f(xk))), (34)

where

Hdef= A>(ABA>)†AB (35)

is the projector onto Range(A>). A quick sanity check reveals that this gives the same formula as

(4) in the case where Range(A>)

= Rn. We can also write

hk+1 = Hhk −HB−1Zk(Hhk −∇f(xk)) =(I−HB−1Zk

)Hhk + HB−1Zk∇f(xk), (36)

where

Zkdef= Sk(S

>k HB−1Sk)

†S>k . (37)

Assume that θk is chosen in such a way that

ED [θkZk] = B.

Then, the following estimate of ∇f(xk)

gkdef= Hhk + θkHB−1Zk(∇f(xk)−Hhk) (38)

is unbiased, i.e. ED[gk]

= ∇f(xk). After evaluating gk, we perform the same step as in SEGA:

xk+1 = proxαR(xk − αgk).

By inspecting (33), (35) and (38), we get the following simple observation.

Lemma C.1. If h0 ∈ Range(A>), then hk, gk ∈ Range

(A>)

for all k.

33

Page 34: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Consequently, if h0 ∈ Range(A>), (34) simplifies to

hk+1 = hk −HB−1Sk(S>k HB−1Sk)

†S>k (hk −∇f(xk)) (39)

and (38) simplifies to

gkdef= hk + θkHB−1Zk(∇f(xk)− hk). (40)

Example C.2 (Coordinate sketch). Consider B = I and the choice of D given by S = ei with

probability pi > 0. Then we can choose the bias-correcting random variable as θ = θ(s) = wipi

, where

widef= ‖Hei‖22 = e>i Hei. Indeed, with this choice, (5) is satisfied. For simplicity, further choose

pi = 1/n for all i. We then have

hk+1 = hk − e>i hk − e>i ∇f(xk)

wiHei =

(I− Heie

>i

wi

)hk +

Heie>i

wi∇f(xk) (41)

and (40) simplifies to

gkdef= (1− θk)hk + θkh

k+1 = hk + nHeie>i

(∇f(xk)− hk

). (42)

C.2 Lemmas

All theory provided in this subsection is, in fact, a straightforward generalization of our non-

subspace results. The reader can recognize similarities in both statements and proofs with that of

previous sections.

Lemma C.3. Define Zk and H as in equations (37) and (35). Then Zk is symmetric, ZkHB−1Zk =

Zk, H2 = H and HB−1 = B−1H>.

Proof: The symmetry of Zk follows from its definition. The second statement is a corollary of the

equations ((A1A2)†)> = (A>2 A>1 )† and A†1A1A†1 = A†1, which are true for any matrices A1,A2.

Finally, the last two rules follow directly from the definition of H and the property A†1A1A†1 = A†1.

Lemma C.4. Assume hk ∈ Range(A>). Then

ED[‖hk+1 − v‖2B

]= ‖hk − v‖2B−ED[Z] + ‖∇f(xk)− v‖2ED[Z]

for any vector v ∈ Range(A>).

34

Page 35: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Proof: By Lemma C.3 we can rewrite HB−1 as B−1H>, so

ED[‖hk+1 − v‖2B

](36)= ED

[∥∥∥hk −HB−1Zk(hk −∇f(xk))− v

∥∥∥2

B

]= ED

[∥∥∥(I−HB−1Zk)

(hk − v) + HB−1Zk(∇f(xk)− v)∥∥∥2

B

]= ED

[∥∥∥(I−B−1H>Zk

)(hk − v) + HB−1Zk(∇f(xk)− v)

∥∥∥2

B

]= ED

[∥∥∥(I−B−1H>Zk

)(hk − v)

∥∥∥2

B

]+ ED

[∥∥∥HB−1Zk(∇f(xk)− v)∥∥∥2

B

]+2(hk − v)>ED

[(I−B−1H>Zk

)>BHB−1Zk

](∇f(xk)− v)

= (hk − v)>ED[(

I−B−1H>Zk

)>B(I−HB−1Zk

)](hk − v)

+(∇f(xk)− v)>ED[ZkB

−1H>BHB−1Zk

](∇f(xk)− v)

+2(hk − v)>ED[BHB−1Zk − ZkHHB−1Zk

](∇f(xk)− v). (43)

By Lemma C.3 we have

ZkHHB−1Zk = ZkHB−1Zk = Zk,

so the last term in (43) is equal to 0. As for the other two, expanding the matrix factor in the first

term leads to(I−B−1H>Zk

)>B(I−HB−1Zk

)=

(I− ZkHB−1

)B(I−HB−1Zk

)= B− ZkHB−1B−BB−1H>Zk + ZkHB−1BHB−1Zk

= B− ZkH−H>Zk + Zk.

Let us mention that H(hk − v) = hk − v and (hk − v)>H> = (hk − v)> as both vectors hk and v

belong to Range(A>). Therefore,

(hk − v)>ED[B− ZkH−H>Zk + Zk

](hk − v) = (hk − v)> (B− ED [Zk]) (hk − v).

It remains to consider

ED[ZkB

−1H>BHB−1Zk

]= ED

[ZkHB−1BHB−1Zk

]= ED [Zk] .

We, thereby, have derived

ED[‖hk+1 − v‖2B

]= (hk − v)> (B− ED [Zk]) (hk − v)

+(∇f(xk)− v)>ED[ZkB

−1Zk]

(∇f(xk)− v)

= ‖hk − v‖2B−ED[Zk] + ‖∇f(xk)− v‖2ED[Z].

35

Page 36: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Lemma C.5. Suppose hk ∈ Range(A>)

and gk is defined by (38). Then

ED[‖gk − v‖2B

]≤ ‖hk − v‖2C−B + ‖∇f(xk)− v‖2C (44)

for any v ∈ Range(A>), where

Cdef= ED

[θ2Z

]. (45)

Proof: Writing gk−v = a+b, where adef= (I−θkHB−1Zk)(h

k−v) and bdef= θkHB−1Zk(∇f(xk)−v),

we get ‖gk‖2B ≤ 2(‖a‖2B + ‖b‖2B). By definition of θk,

ED[‖a‖2B

]= ED

[‖(I− θkHB−1Zk

)(hk − v)‖2B

]= (hk − v)>ED

[(I− θkZkB−1H

)B(I− θkHB−1Zk

)](hk − v)

= (hk − v)>ED[(

B− θkZkB−1HB−BθkHB−1Zk + θ2kZkB

−1HBHB−1Zk)]

(hk − v).

According to Lemma C.3, HB−1 = B−1H and ZkHB−1Zk = Zk, so

ED[‖a‖2B

]= (hk − v)>ED

[(B− θkZkH− θkH>Zk + θ2

kZk

)](hk − v)

= ‖hk − v‖2ED[θ2Z]−B,

where in the last step we used the assumption that hk and v are from Range(A>)

and H is the

projector operator onto Range(A>).

Similarly, the second term in the upper bound on gk can be rewritten as

ED[‖b‖2B

]= ED

[‖θkHB−1Zk(∇f(xk)− v)‖2B

]= (∇f(xk)− v)>ED

[θ2kZkB

−1H>BHB−1Zk

](∇f(xk)− v)

= ‖∇f(xk)− v‖2ED[θ2kZk].

Combining the pieces, we get the claim.

C.3 Main result

The main result of this section is:

Theorem C.6. Assume that f is Q–smooth, µ–strongly convex, and that α > 0 is such that

α (2(C−B) + σµB) ≤ σED [Z] , αC ≤ 1

2(Q− σED [Z]) . (46)

If we define Φk def= ‖xk − x∗‖2B + σα‖hk −∇f(xk)‖2B, then E

[Φk]≤ (1− αµ)kΦ0.

Proof: Having established Lemmas C.3, C.4 and C.5, the proof follows the same steps as the

proof of Theorem 3.3.

36

Page 37: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

C.4 Optimal choice of B and Sk

Let us now slightly change the value of θk that we use in the algorithm. Instead of seeking for

θk giving ED [θkZk] = B, we will use the one that gives ED [θkZk] = BH. This will steal lead to

ED[gk]

= ∇f(xk) and, if f is strongly-convex, we can still show the convergence rate of Theo-

rem C.6. Although the strong convexity assumption is simplistic, the new idea results in a surprising

finding.

Let a1, . . . , am be the columns of A> and U ∈ Rd×n be a matrix that transforms these columns

into an orthogonal basis of ddef= Rank(A) vectors. Set B = U>U. Then, 〈ai, aj〉B = 0 for any

i 6= j. Assume for simplicity, that ‖ai‖B 6= 0 for i ≤ d and ‖ai‖B = 0 for i > d. This is always true

up to permutation of a1, . . . , am. Choose also Sk ∈ Rn equal to ξidef= Bai‖ai‖B with i sampled with

probability pi > 0, and θk = p−1i . Clearly, one has

ED [θkZk] =d∑i=1

pip−1i ξi(ξ

>i HB−1ξi)

†ξ>i =d∑i=1

ξi‖ai‖2B(a>i BHB−1Bai)†ξ>i .

Since ai lies in Range(A>), we have Hai = ai, which gives

ED [θkZk] =d∑i=1

ξi‖ai‖2B(a>i Bai)†ξ>i =

d∑i=1

ξiξ>i . (47)

By definition of B,

(ABA>)† = (diag(‖ai‖2B))† =

d∑i=1

‖ai‖−2B eie

>i .

Thus,

BH = BA>(ABA>)AB =

d∑i=1

(Bai)>Bai

‖ai‖2B= ED [θkZk] ,

so we have achieved our goal. Note that if h0 ∈ Range(A>), we have hk ∈ Range

(A>)

even

without implicitly enforcing it in (33). Therefore, the method can be seen as SEGA with a smart

choice of both sketches and metric in which we project.

To show how the choice of B and of the sketches provided above improves the rate, let us take

a closer look at the conditions of Theorem C.6. We have

C(45)= ED

[θ2Z

] (47)=

d∑i=1

pip−2i ξiξ

>i =

d∑i=1

p−1i ξiξ

>i .

If we assume that σ ≤ 2/µ, then the first bound on α simplifies to

α(2(C−B) + σµB) ≤ 2αC ≤ σED [Z] = σd∑i=1

piξiξ>i ,

37

Page 38: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

where the second part needs to be verified by choosing α to be small enough. For this it is sufficient

to take α ≤ σmax p−2i as every summand ξiξ

>i in the expression for C is positive definite. As

for the second condition, it is enough to choose σ ≤ λmax(Q)2λmin(ED[Z]) and α ≤ λmax(Q)

4λmin(C) . Note that

ξiξ>i ≤ ‖ξi‖22I, so for uniform sampling with pi = 1

d and uniform Q–smoothness with Q = 1LI we

get the following condition on α:

α ≤ min

d2,

1

4Ldmaxi ‖ξi‖22

}.

In particular, choosing σ = min{

2µ ,

λmax(Q)2λmin(ED[Z])

}= min

{2µ ,

d2Lmaxi ‖ξi‖22

}, we get the requirement

α ≤ min

{2

µd2,

1

4Ldmaxi ‖ξi‖22

}.

Typically, d � 1µ , so the leading term in the maximum above is the second one and we get O

(1d

)requirement instead of previous O

(1n

).

C.5 The conclusion of subspace SEGA

Let us recall that gk = hk + θkB−1Zk(∇f(xk) − hk). A careful examination shows that when we

reduce θk from O(n) to O(d), we put more trust in the value of hk with the benefit of reducing

the variance of gk. This insight points out that a practical implementation of the algorithm may

exploit the fact that hk learns the gradient of f by using smaller θk.

It is also worth noting that SEGA is a stationary point algorithm regardless of the value of θk.

Indeed, if one has xk = x∗ and hk = ∇f(x∗), then gk = ∇f(x∗) for any θk. Therefore, once we get

a reasonable hk, it is well grounded to choose gk to be closer to hk. This argument is also supported

by our experiments.

Finally, the ability to take bigger stepsizes is also of high interest. One can think of extending

other methods in this direction, especially if interested in applications with a small rank of matrix A.

D Simplified Analysis of SEGA 1

In this section we consider the setup from Example 2.1 with B = I uniform probabilities: pi = 1/n

for all i. We now state the main complexity result.

Theorem D.1. Let B = I and choose D to be the uniform distribution over unit basis vectors in

Rn. Choose σ > 0 and define

Φk def= ‖xk − x∗‖22 + σα‖hk‖22,

where {xk, hk}k≥0 are the iterates of the gradient sketch method. If the stepsize satisfies

0 < α ≤ min

1− Lσn

2Ln,

1

n(µ+ 2(n−1)

σ

) , (48)

38

Page 39: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

then ED[Φk+1

]≤ (1− αµ)Φk. This means that

k ≥ 1

αµlog

1

ε⇒ E

[Φk]≤ εΦ0.

In particular, if we let σ = n2L , then α = 1

(4L+µ)n satisfies (48), and we have the iteration

complexity

n

(4 +

1

κ

)κ log

1

ε= O(nκ),

where κdef= L

µ is the condition number.

This is the same complexity as NSync [43] under the same assumptions on f . NSync also needs

just access to partial derivatives. However, NSync uses variable stepsizes, while SEGA can do the

same with fixed stepsizes. This is because SEGA learns the direction gk using past information.

D.1 Technical Lemmas

Since f is L–smooth, we have

‖∇f(xk)‖22 ≤ 2L(f(xk)− f(x∗)). (49)

On the other hand, by µ–strong convexity of f we have

f(x∗) ≥ f(xk) + 〈∇f(xk), x∗ − xk〉+µ

2‖x∗ − xk‖22. (50)

Lemma D.2. The variance of gk as an estimator of ∇f(xk) can be bounded as follows:

ED[‖gk‖22

]≤ 4Ln(f(xk)− f(x∗)) + 2(n− 1)‖hk‖22. (51)

Proof: In view of (9), we first write

gk = hk − 1

pie>i h

kei︸ ︷︷ ︸a

+1

pie>i ∇f(xk)ei︸ ︷︷ ︸

b

,

and note that pi = 1/n for all i. Let us bound the expectation of each term individually. The first

term is equal to

ED[‖a‖22

]= ED

[∥∥∥hk − ne>i hkei∥∥∥2

2

]= ED

[∥∥∥(I− neie>i)hk∥∥∥2

2

]= (hk)>ED

[(I− neie>i

)> (I− neie>i

)]hk

= (n− 1)‖hk‖22.

39

Page 40: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

The second term can be bounded as

ED[‖b‖22

]= ED

[∥∥∥ne>i ∇f(xk)ei

∥∥∥2

2

]= n2

n∑i=1

1

n(e>i ∇f(xk))2

= n‖∇f(xk)‖22= n‖∇f(xk)−∇f(x∗)‖22

(49)

≤ 2Ln(f(xk)− f(x∗)),

where in the last step we used L–smoothness of f . It remains to combine the two bounds.

Lemma D.3. For all v ∈ Rn we have

ED[‖hk+1‖22

]=

(1− 1

n

)‖hk‖22 +

1

n‖∇f(xk)− v‖22. (52)

Proof: We have

ED[‖hk+1‖22

](8)= ED

[∥∥∥hk + e>ik(∇f(xk)− hk)eik∥∥∥2

2

]= ED

[∥∥∥(I− eike>ik

)hk + eike

>ik∇f(xk)

∥∥∥2

2

]= ED

[∥∥∥(I− eike>ik

)hk∥∥∥2

2

]+ ED

[∥∥∥eike>ik∇f(xk)∥∥∥2

2

]= (hk)>ED

[(I− eike

>ik

)> (I− eike

>ik

)]hk(∇f(xk))>ED

[(eike

>ik

)>eike>ik

]∇f(xk)

= (hk)>ED[I− eike

>ik

]hk + (∇f(xk))>ED

[eike

>ik

]∇f(xk)

=

(1− 1

n

)‖hk‖22 +

1

n‖∇f(xk)‖22.

D.2 Proof of Theorem D.1

We can now write

ED[‖xk+1 − x∗‖22

]= ED

[‖xk − αgk − x∗‖22

]= ‖xk − x∗‖22 + α2ED

[‖gk‖22

]− 2α〈ED

[gk], xk − x∗〉

(7)= ‖xk − x∗‖22 + α2ED

[‖gk‖22

]− 2α〈∇f(xk), xk − x∗〉

(50)

≤ (1− αµ)‖xk − x∗‖22 + α2ED[‖gk‖22

]− 2α(f(xk)− f(x∗)).

Using Lemma D.2, we can further estimate

ED[‖xk+1 − x∗‖22

]≤ (1− αµ)‖xk − x∗‖22

+2α(2Lnα− 1)(f(xk)− f(x∗)) + 2(n− 1)α2‖hk‖22.

40

Page 41: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Let us now add σαED[‖hk+1‖22

]to both sides of the last inequality. Recalling the definition of the

Lyapunov function, and applying Lemma A.3, we get

ED[Φk+1

]≤ (1− αµ)‖xk − x∗‖22 + 2α(2Lnα− 1)(f(xk)− f(x∗)) + 2(n− 1)α2‖hk‖22

+σα

(1− 1

n

)‖hk‖22 +

σα

n‖∇f(xk)‖22

(49)

≤ (1− αµ)‖xk − x∗‖22 + 2α

(2Lnα+

n− 1

)︸ ︷︷ ︸

I

(f(xk)− f(x∗))

+

(1− 1

n+

2(n− 1)α

σ

)︸ ︷︷ ︸

II

σα‖hk‖22.

Let us choose α so that I ≤ 0 and II ≤ 1 − αµ. This leads to the bound (48). For any α > 0

satisfying this bound we therefore have ED[Φk+1

]≤ (1 − αµ)Φk, as desired. Lastly, as we have

freedom to choose σ, let us pick it so as to maximize the upper bound on the stepsize.

E Simplified Analysis of SEGA II

In this section we consider the setup from Example 2.1 with arbitrary non-uniform probabilities:

pi > 0 for all i. We provide a simplified analysis of SEGA in this scenario. However, we will do this

under slightly different assumptions. In particular, we shall assume that smoothness and strong

convexity of f are measured with respect to the same norm.

In this setup, as we shall see, uniform probabilities are optimal. That is, uniform probabilities

are identical to the importance sampling probabilities. We note that this would be the case even

for standard coordinate descent under these assumptions, as follows from the results in [43].

Let G = Diag(g1, . . . , gn) � 0 and assume that

‖∇f(x)−∇f(y)‖G−1 ≤ L‖x− y‖G

and7

f(x) ≥ f(y) + 〈∇f(y), x− y〉+µ

2‖x− y‖2G

for all x, y ∈ Rn. These two assumptions combined lead to the following inequalities:

f(y) + 〈∇f(y), x− y〉+µ

2‖x− y‖2G ≤ f(x) ≤ f(y) + 〈∇f(y), x− y〉+

L

2‖x− y‖2G.

We define gk as before, but change the method to:

xk+1 = xk − αG−1gk (53)

We now state the main complexity result.

7Note that in the strong convexity inequality below the scalar product is without any additional metric unlike in

other sections.

41

Page 42: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Theorem E.1. Choose σ > 0 and define Φk def= ‖xk − x∗‖2G + σα‖hk‖2

Diag(

1gipi

), where {xk, hk}k≥0

are the iterates of the gradient sketch method. If the stepsize satisfies

0 < α ≤ mini

{pi

(1

µ+ L− σ

2

),

pi2σ (1− pi) + 2Lµ

µ+L

}, (54)

then ED[Φk+1

]≤(

1− αµ 2Lµ+L

)Φk. This means that

k ≥ L+ µ

2αLµlog

1

ε⇒ E

[Φk]≤ εΦ0.

In particular, if we choose gi = 1 and pi = 1n for all i, then if we set σ = 1

2L , we can choose stepsize

α = 3L−µ4Ln(L+µ) , and obtain the rate 2L+2µ

3L−µ n(Lµ + 1

)log 1

ε ≤ 2n(Lµ + 1

)log 1

ε .

E.1 Two lemmas

Lemma E.2. Let d1, . . . , dn > 0. The variance of gk as an estimator of ∇f(xk) can be bounded as

follows:

ED[‖gk‖2Diag(di)

]≤ 2‖hk‖2

Diag(di

1−pipi

) + 2‖∇f(xk)‖2Diag

(dipi

). (55)

Proof: In view of (9), we first write

gk = hk − 1

pie>i h

kei︸ ︷︷ ︸a

+1

pie>i ∇f(xk)ei︸ ︷︷ ︸

b

.

Let us bound the expectation of each term individually. The first term is equal to

ED[‖a‖2G−1

]= ED

[∥∥∥∥hk − 1

pie>i h

kei

∥∥∥∥2

Diag(di)

]

= ED

[∥∥∥∥(I− 1

pieie>i

)hk∥∥∥∥2

Diag(di)

]

= (hk)>ED

[(I− 1

pieie>i

)>Diag(di)

(I− 1

pieie>i

)]hk

= (hk)>ED[(

Diag(di)−2dipieie>i +

dip2i

eie>i

)]hk

=

n∑i=1

di

(1

pi− 1

)(hki )

2.

The second term can be bounded as

ED[‖b‖2Diag(di)

]= ED

[∥∥∥∥ 1

pie>i ∇f(xk)ei

∥∥∥∥2

Diag(di)

]=

n∑i=1

dipi

(e>i ∇f(xk))2.

It remains to combine the two bounds.

42

Page 43: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Lemma E.3. For all v ∈ Rn and d1, . . . , dn > 0 we have

ED[‖hk+1 − v‖2Diag(di)

]= ‖hk − v‖2Diag(di(1−pi)) + ‖∇f(xk)− v‖2Diag(dipi)

. (56)

Proof: We have

ED[‖hk+1 − v‖2Diag(di)

](8)= ED

[∥∥∥hk + e>i (∇f(xk)− hk)ei − v∥∥∥2

Diag(di)

]= ED

[∥∥∥(I− eie>i)

(hk − v) + eie>i (∇f(xk)− v)

∥∥∥2

Diag(di)

]= ED

[∥∥∥(I− eie>i)

(hk − v)∥∥∥2

Diag(di)

]+ ED

[∥∥∥eie>i (∇f(xk)− v)∥∥∥2

Diag(di)

]= (hk − v)>ED

[(I− eie>i

)>Diag(di)

(I− eie>i

)](hk − v)

+(∇f(xk)− v)>ED[(eie

>i )>Diag(di)eie

>i

](∇f(xk)− v)

= (hk − v)>ED[Diag(di)− dieie>i

](hk − v)

+(∇f(xk)− v)>ED[dieie

>i

](∇f(xk)− v)

= ‖hk − v‖2Diag(di(1−pi)) + ‖∇f(xk)− v‖2Diag(dipi).

E.2 Proof of Theorem D.1

Proof: Since f is L–smooth and µ–strongly convex, we have the inequality

〈∇f(x)−∇f(y), x− y〉 ≥ µL

µ+ L‖x− y‖2G +

1

µ+ L‖∇f(x)−∇f(y)‖2G−1 .

In particular, we will use it for x = xk and y = x∗:

〈∇f(xk), x∗ − xk〉 ≤ − µL

µ+ L‖x− x∗‖2G −

1

µ+ L‖∇f(xk)‖2G−1 . (57)

We can now write

ED[‖xk+1 − x∗‖2G

](53)= ED

[‖xk − αG−1gk − x∗‖2G

]= ‖xk − x∗‖2G + α2ED

[‖G−1gk‖2G

]− 2α〈ED

[gk], xk − x∗〉

(7)= ‖xk − x∗‖2G + α2ED

[‖gk‖2G−1

]+ 2α〈∇f(xk), x∗ − xk〉

(57)

≤(

1− αµ 2Lµ+L

)‖xk − x∗‖2G + α2ED

[‖gk‖2G−1

]− 2α

µ+L‖∇f(xk)‖2G−1 .

Using Lemma E.2 to bound ED[‖gk‖2G−1

], we can further estimate

ED[‖xk+1 − x∗‖2G

]≤

(1− αµ 2L

µ+L

)‖xk − x∗‖2G + 2α2‖∇f(xk)‖2

Diag(

1pigi

)− 2αµ+L‖∇f(xk)‖2G−1 + 2α2‖hk‖2

Diag

(1−pipigi

).

43

Page 44: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Let us now add σαED

[‖hk+1‖2

Diag(

1gipi

)]

to both sides of the last inequality. Recalling the defini-

tion of the Lyapunov function, and applying Lemma E.3 with v = 0 and di = 1gipi

, we get

ED[Φk+1

]≤

(1− αµ 2L

µ+L

)‖xk − x∗‖2G + 2α2‖∇f(xk)‖2

Diag(

1pigi

) +(ασ − 2α

µ+L

)‖∇f(xk)‖2G−1

+(2α2 + ασ)‖hk‖2Diag

(1−pipigi

)≤

(1− αµ 2L

µ+L

)‖xk − x∗‖2G + σα‖hk‖2

Diag(( 2ασ

+1) 1−pipigi

)+‖∇f(xk)‖2

Diag(

2α2

pigi+σαgi− 2α

(µ+L)gi

).If we now choose α > 0 such that

pi+ σ − 2

µ+ L≤ 0,

(2α

σ+ 1

)(1− pi) ≤ 1− αµ 2L

µ+ L,

then we get the recursion

ED[Φk+1

]≤(

1− αµ 2Lµ+L

)Φk ≤ (1− αµ)Φk.

F Extra Experiments

F.1 Evolution of Iterates: Extra Plots

Here we show some additional plots similar to Figure 1, which we believe help to build intuition

about how the iterates of SEGA behave. We also include plots for biasSEGA, which uses biased

estimators of the gradient instead. We found that the iterates of biasSEGA often behave in a more

stable way, as could be expected given the fact that they enjoy lower variance. However, we do not

have any theory supporting the convergence of biasSEGA; this is left for future research.

44

Page 45: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Figure 5: Evolution of iterates of

SEGA, CD and biasSEGA (updates

made via hk+1 instead of gk).

Figure 6: Iterates of SEGA, CD and

biasSEGA (updates made via hk+1

instead of gk). Different starting

point.

Figure 7: Iterates of projected

SEGA, projected CD (which do not

converge) and projected biasSEGA

(updates made via hk+1 instead of

gk). The constraint set is repre-

sented by the shaded region.

45

Page 46: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

F.2 Experiments from Section 5 with empirically optimal stepsize

In the experiments in Section 5, we worked with quadratic functions of the form

f(x)def=

1

2x>Mx− b>x,

where b is a random vector with independent entries from N (0, 1) and Mdef= UΣU> according to

Table 2 for U obtained from QR decomposition of random matrix with independent entries from

N (0, 1). For each problem, the starting point was chosen to be a vector with independent entries

from N (0, 1).

Type Σ

1 Diagonal matrix with first n/2 components equal to 1 and the rest equal to n

2 Diagonal matrix with first n− 1 components equal to 1 and the remaining one equal to n

3 Diagonal matrix with i–th component equal to i

4 Diagonal matrix with components coming from uniform distribution over [0, 1]

Table 2: Spectrum of M.

The results are provided in Figures 8-10. They include zeroth-order experiments and the sub-

space version of SEGA.

Figure 8: Counterpart to Figure 2 – convergence illustration of SEGA and PGD. The indicator “Xn” in the label

stands for the setting when the cost of solving linear system is Xn times higher comparing to the oracle call. Recall

that a linear system is solved after each n oracle calls. Empirically best stepsizes were used both PGD and SEGA.

Figure 9: Counterpart to Figure 3 – comparison of SEGA and randomized direct search for a various problems.

Empirically best stepsizes were used for both methods.

46

Page 47: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Figure 10: Counterpart to Figure 4 – comparison of SEGA with sketches from a correct subspace versus naive SEGA.

Optimal (empirically) stepsize chosen.

F.3 Experiment: comparison with randomized coordinate descent

In this section we numerically compare the results from Section 4 to analogous results for coordinate

descent (as indicated in Table 1). We consider the ridge regression problem on LibSVM [7] data, for

both primal and dual formulation. For all methods, we have chosen parameters as suggested from

theory Figure 11 shows the results. We can see that in all cases, SEGA is slower to the corresponding

coordinate descent method, but still is competitive. We however observe only constant times

difference in terms of the speed, as suggested by Table 1.

47

Page 48: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Figure 11: Comparison of SEGA and ASEGA with corresponding coordinate descent methods for

R = 0.

48

Page 49: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

Figure 12: Comparison of SEGA with CD on logistic regression problem with similar stepsizes.

F.4 Experiment: large-scale logistic regression

In this experiment, we set B to be identity matrix and compare CD to SEGA with coordinate sketches,

both with uniform sampling and with similar stepsizes. The problem considered is logistic regression

with `2 penalty:

minx∈Rn

1

m

m∑i=1

log(

1 + exp(−bia>i x))

2‖x‖22,

where ai and bi are data-dependent. Clearly, this regularizer is separable, so we can easily apply

both methods. The value of µ was chosen to be of order 1m in both experiments. Here we use

real-world large scale datasets from the LIBSVM [7] library, a summary can be found in Table 3.

To make it clear whether CD and SEGA converge with the same speed if given similar stepsizes, we

use stepsize 1L for CD and 1

dL for SEGA. The results can be found in Figure 12.

Dataset m n L µ

Epsilon 400000 2000 0.25 2.5 · 10−5

Covtype 581012 54 21930585. 25 10−1

Table 3: Description of the datasets used in our logistic regression experiments. Constants m,

n, L and µ denote respectively the size of the training set, the number of features, the Lipschitz

constant, and the value of `2 penalty.

49

Page 50: Filip Hanzelyy Konstantin Mishchenkoz Peter Richt arikx ... · above line of research. While the gradient sketch (S;x) does not immediatey lead to an unbiased estimator of the gradient,

G Frequently Used Notation

Basic

E [·], P (·) Expectation / Probability

〈·, ·〉B, ‖ · ‖B Weighted inner product and norm: 〈x, y〉B = x>By; ‖x‖B =√〈x, x〉B

ei i-th vector from the standard basis

I Identity matrix

λmax(·), λmin(·) Maximal eigenvalue / minimal eigenvalue

f Objective to be minimized over set Rn (1)

R Regularizer (1)

x∗ Global optimum

L Lipschitz constant for ∇fQ Smoothness matrix (10)

M Smoothness matrix, equal to Q−1 for B = I (11)

µ Strong convexity constant

SEGA

D Distribution over sketch matrices S

S Sketch matrix (3)

ED [·] Expectation over the choice of S

b Random variable such that S ∈ Rn×b

ζ(S, x) Sketched gradient at x (2)

Z S(S>B−1S

)†S>

θ Random variable for which ED [θZ] = B (5)

C ED[θ2Z

]Thm 3.3

h, g Biased and unbiased gradient estimators (4), (6)

α Stepsize

Φ Lyapunov function Thm 3.3,

σ Parameter for Lyapunov function Thm 3.3, 4.2

Extra Notation for Section 4

p, P Probability vector and matrix

v vector of ESO parameters (14)

P, V Diag(p),Diag(v)

γ α− α2 maxi{ vipi } − σ Thm 4.2

y, z Extra sequences of iterates for ASEGA

τ, β Parameters for ASEGA

Ψ,Υ Lyapunov functions Thm 4.2, B.5

η(v, p) maxi√vipi

Table 4: Summary of frequently used notation.

50