Introduction

The success of learning algorithms trained with stochastic gradient descent (SGD) variants, dating back to the seminal AlexNet architecture (Krizhevsky et al., 2012)-arguably initiating the 'deep learning revolution '-and empirically verified comprehensively in Sutskever et al. (2013), emphasizes the importance of incorporating simple momentum for achieving rapid convergence in neural network learning. Although more complex momentum (or acceleration) algorithms have been proposed (Nesterov, 1983;Bubeck et al., 2015;Liu & Wright, 2015;Van Scoy et al., 2017;Allen-Zhu, 2018;Cyrus et al., 2018;Jain et al., 2018, Jin et al., 2018) demonstrating faster convergence rates than plain SGD on standard classes of loss functions, Polyak's original simple momentum update (Polyak, 1964) defies theory by remaining highly effective in practice and remains a popular choice for many applications. Despite several studies analyzing the performance of stochastic momentum methods (Flammarion & Bach, 2015;Loizou & Richtárik, 2017;Gadat et al., 2018;Kidambi et al., 2018;Yan et al., 2018;Can et al., 2019;Gitman et al., 2019;Liu et al., 2020;Loizou & Richtárik, 2020;Sebbouh et al., 2021) a gap persists between existing theoretical guarantees and their superior practical performance. We aim to bridge this gap by analyzing the properties of simple stochastic momentum methods in the context of quadratic optimization.

Given an n × d matrix A and a length n vector b, the linear least squares problem min

is one of the most fundamental problems in optimization. One approach to solving (1.1) is Polyak's heavy ball momentum (HBM) (Polyak, 1964), also called 'standard' or 'classical' momentum. (HBM) updates the parameter estimate x k for the solution x * as

where α k and β k are the step-size and momentum parameters, respectively. This is equivalent to the update

The gradient of the objective (1.1) is easily computed to be ∇f (x) = A T (Ax-b), and when A T A has finite condition number κ = λ max /λ min , where λ max and λ min are the largest and smallest eigenvalues of A T A, (HBM) with properly chosen constant step-size and momentum parameters α k = α and β k = β provably attains the optimal linear rate

When β k = 0, (HBM) reduces to the standard gradient descent algorithm which, with the optimal choice of step-sizes α k , converges at a sub-optimal linear rate

On large-scale problems, computing the gradient ∇f (x) can be prohibitively expensive. For problems such as (1.1), it is common to replace applications of the gradient with a minibatch stochastic gradient

where S k contains B indices drawn independently with replacement from {1, 2, . . . , n} where, at each draw, p j is the probability that an index j is chosen. Note that this sampling strategy ensures E[∇f S k (x)] = ∇f (x).

We denote by minibatch-heavy ball momentum (Minibatch-HBM) the following algorithm: starting from initial conditions x 1 = x 0 , iterate until convergence

( M i n i b a t c h -H B M )

In the case of (1.1), the minibatch stochastic gradient can be written

where a T j is the jth row of A and b j is the jth entry of b. A..u/p01o, 1.1 Throughout, we will assume that, for some η ≥ 1, the sampling probabilities are such that ( 1 . 3 ) where • F represents the matrix Frobenious norm.

Re/a(2 1.2 If rows are sampled proportionally to their squared norm p j = a j 2 / A 2 F , we have η = 1. If rows are sampled i.i.d. uniformly, p j = 1/n for j = 1, . . . , n, (1.3) is satisfied with η = n max j a j 2 / A 2 F . Applied to the problem (1.1), plain SGD (B = 1, β = 0) with an appropriately chosen step-size (p j = a j 2 / A 2 F ) is equivalent to the randomized Kaczmarz (RK) algorithm if importance weighted sampling (η = 1) is used (Needell et al., 2014). The standard version of RK extends the original cyclic Kaczmarz algorithm (Kaczmarz, 1937) and, as proved in Strohmer & Vershynin (2008), converges in a number of iterations scaling with d κ, where κ = λ ave /λ min is the smoothed condition number of A T A, and λ ave and λ min are the average and smallest eigenvalues of A T A, respectively. Note the important relationship between the smoothed and the standard condition numbers:

This relationship implies that, when n ≥ d √ κ, RK at least matches (up to constants) the performance of (HBM). If κ κ or n + d √ κ then RK can significantly outperform (HBM), at least in terms of total number of row products.

While the number of row products of RK is reduced compared to (HBM), the number of iterations required to converge is increased. In practice, running times are not necessarily directly associated with the number of row products, and instead depend on other factors such as communication and memory access patterns. These costs often scale with the number of iterations, so it is desirable to understand whether the iteration complexity of (Minibatch-HBM) can be reduced to that of (HBM).

Contributions

We aim to make precise the observation that stochastic momentum affords acceleration in the minibatch setting. An illustration of this phenomenon is provided in Fig. 1. Our main theoretical result is a proof Gradient descent with heavy ball momentum (HBM) allows for accelerated convergence over gradient descent (GD). SGD allows for lower per iteration costs, and the use of batching (Minibatch-SGD) reduces the variance of the iterates. While batching and momentum are often used simultaneously (Minibatch-HBM), convergence guarantees have remained elusive, even for quadratic objectives (1.1). In this paper we prove that, on such objectives, (Minibatch-HBM) converges linearly at the same rate as (HBM), provided the batch size is sufficiently large in a precise sense.

that the linear rate of convergence of (HBM) can be matched by (Minibatch-HBM), provided the batch size is larger than a critical size. Informally, this result can be summarized as follows:

T#eo(e/ 1.3 Consider (Minibatch-HBM) applied to a strongly convex quadratic objective (1.1) with stochastic gradients (1.2) whose sampling probabilities satisfy (1.3) with parameter η ≥ 1. Suppose that the minimizer x * satisfies Ax * = b. Then, with the same fixed step-size and momentum parameters α k = α, β k = β as (HBM), there exists a constant C > 0 such that, if κ is sufficiently large and B ≥ Cηd log(d) κ√ κ, the (Minibatch-HBM) iterates converge in expected norm at least at a linear rate 1 -1/ √ κ.

For a more precise statement of Theorem 1.3, see Corollary 3.4.

Re/a(2 1. 4 The bound for B does not depend directly on n, and in many situations B n. In these cases, (Minibatch-HBM) offers a provable speedup over (HBM). Theorem 3.2 provides a more finegrained relationship between the momentum and step-size parameters than Theorem 1.3. In particular, it shows that it is always possible to take B < n and have (Minibatch-HBM) converge similar to (HBM), albeit at the cost of slowed convergence compared to the rate of convergence of (HBM) with the optimal α, β.

Owing to the equivalence between SGD on convex quadratics and the RK algorithm (Needell et al., 2014) our convergence guarantees give a provable iteration complexity √ κ for RK-type algorithms. Our analysis method is quite general, and can be used to certify fast linear convergence for various forms of momentum beyond heavy ball momentum; in Appendix A we illustrate the generality of the approach by proving an analogous convergence result for a minibatch Nesterov's accelerated gradient method in the setting of linear regression.

Literature review

In the remainder of this section, we provide an overview of state-of-art existing results for row-sampling methods for solving (consistent) linear squares problems. A summary is given in Table 1.

Randomized Kaczmarz. A number of improvements to the standard RK algorithm have been proposed. Liu and Wright (Liu & Wright, 2015) introduce an accelerated randomized Kaczmarz (ARK) method that, through the use of Nesterov's acceleration, can achieve a faster rate of convergence

Ta3&e 1 Runtime comparisons for row-sampling iterative methods for solving a consistent linear least squares problem (1.1) to constant accuracy when A ∈ R n×d for large condition number κ. Constants and a logarithmic dependence on the accuracy parameter are suppressed. Here κ and κ are the regular and smoothed condition numbers of A T A. Due to practical considerations such as parallelization, data movement, caching, energy efficiency, etc., the real-world cost of an iteration does not necessarily scale linearly with the number of row products Algorithm Iterations # row prods/iter. References (HBM) √ κ n Hestenes et al. (1952), Polyak (1964), Liesen & Strako4 (2013) SGD/RK d κ 1 Strohmer & Vershynin (2008), Needell et al. (2014) ARK d √ κ 1 Liu & Wright (2015) (Minibatch-HBM) √ κ B d log(d) κ√ κ (this paper)

compared to RK. However, their rate is still sub-optimal compared to the rate attained by (HBM). Moreover, ARK is less able to take advantage of potential sparsity in the data matrix A than the standard RK algorithm and (Minibatch-HBM). This issue is partially addressed by a special implementation of ARK for sparse matrices, but is still of concern for particularly sparse matrices.

Minibatching in the setting of the RK has been studied extensively in 'small' batch regimes (Needell & Tropp, 2014;Needell & Ward, 2016;Moorman et al., 2020). These works view minibatching as a way to reduce the variance of iterates and improve on the standard RK algorithm. In general, the bounds for the convergence rates for such algorithms are complicated, but can improve on the convergence rate of RK by up to a factor of B in the best case. This 'best case' improvement, however, can only be attained for small B; indeed, RK reduces to standard gradient descent in the deterministic gradient limit. In contrast to these works, we study minibatching in the RK method as a necessary algorithmic structure for unlocking the fast convergence rate of HBM. Minibatch-HBM. Several recent works provide theoretical convergence guarantees for (Minibatch-HBM). Loizou & Richtárik (2017, 2020) show that (Minibatch-HBM) can achieve a linear rate of convergence for solving convex linear regression problems. However, the linear rate they show is slower than the rate of (deterministic) (HBM) in the same setting. Gitman et al. (2019) establish local convergence guarantees for the (Minibatch-HBM) method for general strongly convex functions and appropriate choice of the parameters α k and β k . Liu et al. (2020) show that (Minibatch-HBM) converges as fast as SGD for smooth strongly convex and nonconvex functions. Under the assumption that the stochastic gradients have uniformly bounded variance, Can et al. (2019) provide a number of convergence guarantees for stochastic HBM. In particular, it is shown that the same fast rate of convergence can be attained for full-batch quadratic objectives with bounded additive noise. The results of Can et al. (2019) however do not apply to the setting of RK, where the variance of the stochastic gradients necessarily grows proportionally to the squared norm of the full gradient.

Jain et al. (2018) demonstrate that (Minibatch-HBM) with a batch size of B = 1 provably fails to achieve faster convergence than SGD. They acknowledged that the favorable empirical results of (Minibatch-HBM), such as those found in Sutskever et al. (2013), should be seen as an 'artifact' of large minibatching, where the variance of stochastic gradients is sufficiently reduced that the deterministic convergence behavior of (HBM) dominates. In this paper, we aim to precisely quantify this observation by providing a characterization of the minimal batch size required for (Minibatch-HBM) to achieve fast linear convergence comparable to that of (HBM).

In concurrent work, Lee et al. (2022) analyze the dynamics of (Minibatch-HBM) applied to quadratic objectives corresponding to a general class of random data matrices. Their results show that when the batch size is sufficiently large, (Minibatch-HBM) converges like its deterministic counterpart, but convergence is necessarily slower for smaller batch sizes. The batch size requirement of Lee et al. (2022) is a factor of κ better than what we obtain (see Theorem 1.3). However, while our analysis makes no assumptions on A, Lee et al. (2022) requires certain invariance assumptions on the singular vectors of A. It would be interesting to understand whether the extra factor of κ in our bound can be improved or whether is a necessary artifact of the lack of assumptions on A. Stochastic Nesterov's Accelerated Gradient (SNAG). Several recent works have analyzed the theoretical convergence properties of stochastic Nesterov's accelerated gradient (SNAG) methods and their variants in both strongly convex and nonconvex settings (Ghadimi & Lan, 2016;Can et al., 2019;Vaswani et al., 2019;Aybat et al., 2020). Can et al. (2019) andAybat et al. (2020) demonstrate accelerated convergence guarantees of SNAG method variants to a neighborhood of the solution for problems with uniformly bounded noise. Additionally, Ghadimi & Lan (2016) provide convergence guarantees for SNAG variants in nonconvex settings. Vaswani et al. (2019) show that SNAG methods achieve accelerated convergence rates to the solution for over-parameterized machine learning models under the assumption of the strong gradient growth condition, where the ℓ 2 norm of the stochastic gradients is assumed to be bounded by the norm of the gradient. Our contributions imply that the strong gradient growth conditions hold in the consistent under-parameterized least squares setting for the stochastic minibatch gradient, and demonstrating that this condition implies acceleration for (Minibatch-HBM), in addition to Nesterov momentum. Acceleration techniques have also been integrated with the variance reduction techniques to achieve optimal convergence rate guarantees for finite-sum problems (Frostig et al., 2015;Defazio, 2016;Allen-Zhu, 2018;Lin et al., 2018;Zhou et al., 2019). Ma et al. (2018) established critical batch size for SGD to retain the convergence rate of deterministic gradient descent method where as in this work we establish the critical batch size for (Minibatch-HBM) to retain the convergence properties of (HBM).

Notation

We denote vectors using lower case roman letters and matrices using upper case roman letters. We use x ∈ R d to denote the variables of the optimization problem and x * to denote the minimizer of f (x). We use • to represent the Euclidean norm for vectors and operator norm for matrices and • F to represent the matrix Frobenious norm. Throughout, A will be an n × d matrix and b a length n vector. The eigenvalues of A T A are denoted λ 1 , λ 2 , . . . , λ d , and we write λ max , λ min , and λ ave for the largest, smallest, and average eigenvalue of A T A. The regular and smoothed condition numbers κ and κ of A T A are respectively defined as κ = λ max /λ min and κ = λ ave /λ min . All logarithms are natural, and we denote complex numbers by i and Euler's number 2.718 . . . by e.

Preliminaries

In this section we provide an overview of (HBM) analysis and important statements from random matrix theory that are used in proving the convergence of (Minibatch-HBM).

Standard analysis of heavy ball momentum for quadratics

We review the standard analysis for heavy ball momentum (HBM) in the setting of strongly convex quadratic optimization problems (1.1) (see Recht (2010)). Here and henceforth, we take α k = α and β k = β as constants.

First, we re-write the (HBM) updates as

so, by definition, the (HBM) updates satisfy

This can be written more concisely as

where T is the transition matrix taking us from the error vectors at steps k and k -1 to the error vectors at steps k + 1 and k. Repeatedly applying this relation, we find

from which we obtain the error bound

The assumption x 1 = x 0 allows us to write

The difficulty in analyzing (HBM) compared to plain gradient descent (when β = 0) lies in the fact T k need not provide a useful upper bound for T k . Indeed, while T k = T k for symmetric matrices, this is not necessarily the case for nonsymmetric matrices. To get around this issue, it is common to bound the spectral radius ρ(T) = max j {|λ j |} and use Gelfand's formula ρ(T) = lim k→∞ T k 1/k to derive the rate of convergence (Recht, 2010).

To bound the spectral radius of T, note that T is orthogonally similar to a block diagonal matrix consisting of 2 × 2 components {T j }. Specifically,

, where

for each j = 1, 2, . . . , n, U is a certain orthogonal matrix (see Recht ( 2010)), and {λ j } are the eigenvalues of A T A. For each j = 1, 2, . . . , n, the eigenvalues z ± j of T j are easily computed to be

and are therefore nonreal if and only if

In this case, the magnitude of both the eigenvalues of T j is

Here we have used that

2 whenever (2.5) holds. Therefore, provided

we have that

( 2 . 7 )

We would like to choose √ β = ρ(T) as small as possible subject to the condition that (2.6) holds.foot_2 Note that (2.6) is equivalent to the condition

which we can rewrite as

( 2 . 9 )

As noted at the start of this section, this gives an asymptotic rate of convergence √ β.

A closer look at the quadratic case

To derive a bound for stochastic (HBM) at finite k, it is desired to understand the eigendecomposition of the matrix T from (2.4) more carefully. Thus, we might aim to diagonalize T as

where U is the previously described orthogonal matrix rotating T into block diagonal form (2.4) and D and C are block-diagonal matrices with 2 × 2 blocks {D j } and {C j } where, for each j = 1, . . . , d, T j is diagonalized as

Given such a factorization, we would have T k = (UC)D k (UC) -1 . Then, since that U is unitary, we would obtain the bound

is the condition number of the eigenvector matrix UC.

However, if α and β are chosen as in (2.9), T is defective (that is, does not have a complete basis of eigenvectors) and no such diagonalization exists. 2 To avoid this issue, we perturb the choices of α and β, and define, for some ( ∈ (0, λ min ), L = λ max + ( and ℓ = λ min -( .

( 2 . 1 1 )

ensures that (2.6) holds and that T is diagonalizable. Indeed, since we can write z ± j = a j ± ib j for a j , b j ∈ R with b j = 0, it is easily verified that the (up to a scaling of the eigenvectors) eigendecomposition

We clearly have D = max j |z ± j | = √ β, so the (HBM) iterates satisfy the convergence guarantee

( 2 . 1 4 )

Note that M(α, β), L and ℓ each depend on λ max , λ min and ( . The dependency of M(α, β) on these values is through C, which, up to a unitary scaling by U, is the eigenvector matrix of the transition matrix T. We can bound M(α, β) by the following lemma.

Le//a 2.1 For any ( ∈ (0, λ min ) set ℓ = λ min -( and L = λ max + ( and choose α = 4/(

, where UC is the eigenvector matrix for T. Then,

α ( (( + λ maxλ min ) .

Proof. In order to bound M(α, β), we note that the block diagonal structure of C implies that C = max{ C j : j = 1, . . . , d} and C -1 = max{ C -1 j : j = 1, . . . , d}. By construction, the specified values of α and β ensure that (1+β -αλ i ) 2 < 4β for all i = 1, 2 . . . , d, i.e. condition (2.5). This implies |z ± j | = √ β, so we easily compute

By direct computation, we find

To bound C -1 j we first note that the condition (2.5) is also equivalent to

Since L ≥ λ max we therefore have the bound

.

The result follows by combining the above expressions.

ON THE FAST CONVERGENCE OF MINIBATCH HBM 1407 2.3 Lemmas from nonasymptotic random matrix theory Before we prove our main result, we need to introduce tools from nonasymptotic random matrix theory which are crucial components of the proof. P(opo.101o, 2.2 Consider a finite sequence {W k } of independent random matrices with common dimension d 1 × d 2 . Assume that E[W i ] = 0 and W i ≤ W for each index i and introduce the random matrix

Let v(Z) be the matrix variance statistic of the sum:

Then,

The bound on E[ Z ] is Theorem 6.1.1 in Tropp (2015), and the bound on E[ Z 2 ] follows from equation 6.1.6 in Tropp (2015) and the fact E[max i W i 2 ] ≤ W. Equation 6.1.6 in Tropp (2015) comes from applying Theorem A.1 in Chen et al. (2012) to the Hermitian dilation of Z. Under the stated conditions, the logarithmic dependence on the dimension is necessary (Tropp, 2015).

We will also use a theorem on products of random matrices from Huang et al.

Main results

We are now prepared to analyze (Minibatch-HBM) applied to strongly convex least squares problems of the form (1.1). We begin by considering the case of consistent linear systems, i.e. systems for which b is in the column span of A. In Section 3.1 we then provide an analogous result for inconsistent least squares problems. We begin with a useful technical lemma that bounds the batch size required to ensure that a certain random matrix is near it's expectation.

Le//a 3.1 Define W j = B -1 (-p -1 j a j a T j + A T A) and let

where S is a list of B indices each chosen independently according to (1.3). Then,

Proof. Since the sampling probabilities satisfy (1.3), a j a T j = a j 2 ≤ ηp j A 2 F . Then, since η ≥ 1,

Next, observe that

Using that E[(p j ) -1 a j a T j ] = A T A and a j 2 ≤ ηp j A 2 F , we find that

Here we write M 1 M 2 if M 2 -M 1 is positive semi-definite. Note that M 1 M 2 implies the largest eigenvalue of M 2 is greater than the largest eigenvalue of M 1 . Therefore, using that 0 E W T j W j followed by the fact 0 A T A A 2 F I,

Thus, since W j is symmetric and the samples in S are i.i.d., we obtain a bound for the variance statistic

Together with (3.1) and (3.2), Theorem 2.2 implies

The first term is bounded by δ/2 when

whereas the second term is bounded by δ/2 when

The result follows by taking the max of these quantities.

Our main result is the following theorem.

T#eo(e/ 3.2 Consider (Minibatch-HBM) applied to a strongly convex quadratic objective (1.1) with stochastic gradients (1.2) whose sampling probabilities satisfy (1.3). Fix parameters α and β satisfying

Then, for all k > 0, assuming that the minimizer x * satisfies Ax * = b, the (Minibatch-HBM) iterates satisfy

where M(α, β) is the condition number of eigenvector matrix UC for T defined in (2.10).

Re/a(2 3.3 For k ≤ k * , (k * ) k/k * ≤ max{e, k}. Thus, (Minibatch-HBM) converges at at rate nearly √ β for k ≤ max{d, k * }. For k > max{d, k * }, the convergence is still linear for k * sufficiently large, but at a rate slower than √ β. Specifically, (Minibatch-HBM) converges at a rate

Proof. Due to assumption of consistency, we have that Ax * = b. Therefore, we can write the minibatch gradient (1.2) as . 4 ) Define the random matrix

Then, analogously to (2.2), the (Minibatch-HBM) iterates satisfy the recurrence

where Y S k is the stochastic transition matrix at iteration k. After k iterations, the error satisfies

so our goal is to bound the norm of the random matrix

This is a product of random matrices, so we may hope to apply Proposition 2.3. However, while E[Y S i ] = T, where T is the deterministic transition matrix (2.2), E[Y S i ] is not necessarily bounded by √ β. Thus, to apply Proposition 2.3 we will instead consider

where

and U and C are the matrices from (2.10). Then, as desired,

Thus, if we can guarantee that the variances { E[ X S i -E[X S i ] 2 ]} are not too large, we can apply Proposition 2.3 to obtain a rate similar to (HBM). Towards this end, note that

where W j is as in Lemma 3.1. This and the fact that (3.6) where W = j∈S j W j . Using Lemma 3.1 and (3.6), we have E X S i -EX S i 2 ≤ δ provided that the batch size B satisfies

Applying Proposition 2.3 to the product (3.5) with the parameters

gives the bound

Set δ 2 = β log(k * )/(2k * ) so that 2v = k log(k * )/k * . This gives the desired expression for B. Moreover, we then have

Thus, we find that

giving the desired bound for the iterates.

The expressions for the required batch size in Theorem 3.2 is somewhat complicated, but can be simplified in the large condition number limit.

Co(o&&a(y 3.4 Fix c ∈ (0, 2). There exist parameters α and β and a constant C > 0 (depending on c) such that, for all κ sufficiently large, the (Minibatch-HBM) iterates converge in expected norm at least at a linear rate 1c/ √ κ provided that B ≥ Cηd log(d) κ√ κ.

Proof. Suppose ( = c 1 λ min for c 1 ∈ (0, 1). Then, using that the definitions of L and ℓ from (2.11) imply that L/ℓ = (λ max + c 1 λ min )/(λ minc 1 λ min ) = κ/(1c 1 ) + c 1 /(1c 1 ), we have

we have

Therefore, if we take c 1 = 1 -((c + 6)/8) 2 and c 2 = (2c)/4 we have that, for κ sufficiently large,

Using Lemma 2.1 we have that

This, with the fact that β = O(1), implies

and

Thus, the bound on B becomes B ≥ O(ηd log(d) κ√ κ).

Inconsistent least squares problems

Our results can be extended to inconsistent systems. On such systems, the stochastic gradients at the optimal point x * need not equal zero, even though ∇f (x * ) = 0. As a result, stochastic gradient methods will only converge to within a convergence horizon of the minimizer x * .

Re/a(2 3. T#eo(e/ 3.6 In the setting of Theorem 3.2, define r = Ax * -b and ) = max i |r i |/ a i . Let the batch size B satisfy the conditions in Theorem 3.2. Then, provided k * is chosen so that δ = 2 log(k * )/(k * log(1/β)) < 1, the (Minibatch-HBM) iterates satisfy

Re/a(2 3.7 The term in R containing r scales with 1/ √ B whereas the term containing ) scales with 1/B. Thus, when B is large, the term in R involving ) becomes small relative to the term involving r .

Re/a(2 3.8 In the large κ limit considered in Corollary 3.4

Proof of Theorem 3.6. Since b = Ax * -r, our stochastic gradients for inconsistent systems satisfy

where M S k is as in the proof of Theorem 3.2, and define

We therefore find the iterates satisfy the update formula

Thus, after k iterations, the error satisfies

The first term is identical to the case r = 0, so we have

Here we have used the triangle inequality and definition of operator norm followed by the linearity of expectation and independence of the minibatch draws across iterations.

As in the proof of Theorem 3.2,

Therefore, assuming all minibatch draws {S j } are identically distributed, we have

It remains to bound E[ r S 1 ], and to do so we again turn to the matrix Bernstein inequality. Define the sum

and note that, with ) = max i |r i |/ a i , and the assumption (1.3) on the sampling probabilities,

By direct computation we also observe that

Therefore, applying Theorem 2.2 we obtain the bound

Combining everything gives the desired result.

Numerical results

We conduct numerical experiments on quadratic objectives (1.1) to illustrate the performance of (Minibatch-HBM). Throughout this section, we use the value

as a heuristic for the batch size needed to observe linear convergence at a rate similar to that of (HBM). This value is obtained from Theorem 3.2 by making several simplifying assumptions. Specifically, we drop the dependence on M(α, β) which results from a change of basis followed by a return to the original basis (which we believe is likely an artifact of our analysis approach) and replace k * / log(k * ) by 1/ log(1/β) = O( √ κ). In all experiments we use n = 10 6 , d = 10 2 and set ( = λ min /10 3 . Each experiment is repeated 100 times and the median and 5th to 95th percentile range for each algorithm/parameter choice are plotted.

Row-norm and uniform sampling

In this example, we study the dependence of (Minibatch-HBM) on the sampling probabilities. Our bounds are sharpest when p j ∝ a j 2 , but in practice it is common to use uniform sampling p j = 1/n to avoid the need for computing row-norms which requires accessing the entire data matrix. We take A = DG, where G is an n × d matrix whose entries are independently 1 with probability 1/10 or 0 with probability 9/10 and D is an n × n diagonal matrix whose diagonal entries are 1 with probability 9/10 and 10 with probability 1/10. Thus, uniform sampling probabilities p j = 1/n satisfy (1.3) provided η ≥ n max j a j 2 / A 2 F ≈ 23. We use a planted solution x * with i.i.d. standard normal entries. In Fig. 2 we show the convergence of (Minibatch-HBM) with row norm sampling and uniform sampling at several values of B. As expected, row norm sampling works better than uniform sampling for a fixed value of B. However, since the norms of rows are not varying too significantly, the convergence rates are still comparable. See Needell et al. (2014) for a further discussion on sampling probabilities in the context of RK and SGD.

Sensitivity to batch size

The fact (Minibatch-HBM) exhibits accelerated convergence is an artifact of batching, and we expect different behavior at different batch sizes. In fact, we have already observed this phenomenon on the previous example. In Theorem 3.2, we provide an upper bound on the required batch size depending on spectral properties of A T A such as κ = λ ave /λ min and κ = λ max /λ min . To study whether the stated dependence on such quantities is reasonable, we construct a series of synthetic problems with prescribed spectra.

F1g. 2. Median and 5th to 95th percentile error norm x k -x * of (Minibatch-HBM) for row norm sampling and uniform sampling at varying values of batch size B. For reference, we also show the convergence of (HBM) and the (HBM) bound (2.14).

We construct problems A = UΣV T by choosing the singular vectors U and V uniformly at random, and selecting singular values {) i } with exponential or algebraic decay. For exponential decay we use the squared singular valuesfoot_5

and for algebraic decay we use the squared singular values

In both cases, the condition number of the A T A is κ and ρ determines how fast the singular values of A decay. We again use a planted solution x * with i.i.d. standard normal entries.

In Figs 3 and 4 we report the results of our experiments. Here we consider consistent equations with condition numbers κ = 10, 30, 100. For each value of κ, we generate two problems according to (4.1) with ρ = 0.1 and ρ = 0.8 and two problems according to (4.2) with ρ = 2 and ρ = 1. In Fig. 3 we run (Minibatch-HBM) (row norm sampling) with B = cB * for c = 10 -3 , 10 -2 , 10 -1 , 10 0 and show the convergence as a function of the iterations k. In Fig. 4 we run (Minibatch-HBM) for a fixed number of iterations, and show the convergence as a function of the batch size B. This empirically illustrates that when the batch size is near B * , the rate of convergence is nearly that of (HBM). computation. 5 We also note that, while the error in the iterates stagnates at different points, the value of the objective function is quite similar in all cases, nearly matching the residual norm of the true least squares solution.

Computational tomography

One of the most prevalent applications of Kaczmarz-like methods is in tomographic image reconstruction, notably in medical imaging. In X-ray tomography (e.g. CT scans), X-rays are passed through an object and the intensity of resulting X-ray beam is measured. This process is repeated for numerous known orientations of the X-ray beams relative to the object of interest. With a sufficient number of F1g. 5. Median and 5th to 95th percentile error norm and residual norm of (Minibatch-HBM) for varying values of batch size B on an inconsistent problem. For reference, we also show the convergence of (HBM) as well as optimal residual norm. measurements, it is possible to reconstruct a 'slice' of the interior of the object of interest. In theory, this reconstruction involves solving a large, sparse consistent linear system.

In this example we consider the performance of (Minibatch-HBM) on a tomography problem corresponding to a parallel beam geometry scanner with 128 sensor pixels. Measurements are taken for 360 degrees of rotation at half-degree increments, totaling 720 measurements. The goal is to reconstruct a 64 × 64 pixel image. This results in a measurement matrix of dimensions (720 • 128) × (64 • 64) = 92160 × 4096, which we construct using the ASTRA toolbox (van Aarle et al., 2015).

We employ a planted solution of a walnut, which we aim to recover from the resulting measurement data. Uniform sampling is employed due to the roughly similar norms of all rows. The step-size α and momentum parameter β are chosen so that (HBM) converges at a reasonable rate and so that B * = 407 is not too large relative to n. In Fig. 6 we report the convergence of (HBM) and (Minibatch-HBM), along with the resulting images recovered by the algorithms after k = 500 iterations.

Conclusion

We provided a first analysis of accelerated convergence of minibatch heavy ball momentum method for quadratics using standard choices of the momentum step-sizes. Our proof method involves a refined quantitative analysis of the convergence proof for (deterministic) heavy ball momentum, combined with matrix concentration results for sums and products of independent matrices. Our proof technique is general, and also can be used to verify the accelerated convergence of a minibatch version of Nesterov's acceleration for quadratics, using the constant step-size parameters suggested in Nesterov's original paper. An interesting direction for future work is to combine the simple minibatch momentum algorithm with an adaptive gradient update such as AdaGrad (McMahan & Streeter, 2010;Duchi et al., 2011;Ward et al., 2019;Défossez et al., 2020), to potentially learn the momentum parameters α k and β k adaptively. Such an analysis would also shed light on the convergence of ADAM (Kingma & Adam, 2014), an extension of momentum which combines adaptive gradients and momentum in a careful way to achieve state-of-the-art performance across various large-scale optimization problems. Again, T is unitarily similar to a block diagonal matrix whose blocks are

and it is easy to see the eigenvalues of T j are

Rather than aiming to optimize the parameters α and β as we did for (HBM), we will simply use the standard choices of parameters suggested in Nesterov ( 2013):

By direct computation, we find 4β

and therefore that

Thus, the (NAG) iterates satisfy the convergence guarantee

where M(α, β) is the eigenvector condition number of the transition matrix T (note that this value is different from the value for (HBM) bounded in Lemma 2.1). We now provide a bound, analogous to Theorem 3.2, for minibatch-NAG, the algorithm obtained by using the stochastic gradients (1.2) in (NAG).

T#eo(e/ A.1 Set ℓ = λ min , L = λ max and define α = 1/L and β = ( √ L/ℓ -1)/( √ L/ℓ + 1). For any

Then, for all k > 0, assuming that the minimizer x * satisfies Ax * = b, the (Minibatch-HBM) iterates satisfy

The proof of Theorem A.1 is almost the same as the proof identical to the proof of Theorem 3.2, so we skip repeated parts.

Proof. For NAG we have that Thus, using the submultiplicativity of the operator norm, analogous to (3.6), we have that

Again, using Lemma 3.1 and (3.6), we see that E X S i -EX S i 2 ≤ δ provided that the batch size B satisfies B ≥ 8eη log(2d) max 5 A 2 F A 2 α 2 M(α, β) 2 δ -2 , (20|A 4 F α 2 M(α, β) 2 δ -2 ) 1/2 . Using the same choice of δ 2 = β log(k * )/(2k * ) we get the desired bound.