Introduction

The success of learning algorithms trained with stochastic gradient descent (SGD) variants, dating back to the seminal AlexNet architecture [22]-arguably initiating the "deep learning revolution"-and empirically verified comprehensively in [45], 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 [1,3,6,17,18,26,37,48], demonstrating faster convergence rates than plain SGD on standard classes of loss functions, Polyak's original simple momentum update [39] 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 [4,10,12,14,20,27,28,29,41,51], 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 a 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) [39], 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

).

(HBM)

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

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

(Minibatch-HBM)

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.

Assumption 1 Throughout, we will assume that, for some η ≥ 1, the sampling probabilities are such that

where ∥ • ∥ F represents the matrix Frobenious norm. Gradient descent with heavy ball momentum (HBM) allows for accelerated convergence over gradient descent (GD). Stochastic gradient descent (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.

Remark 1 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 stochastic gradient descent (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 [34]. The standard version of RK extends the original cyclic Kaczmarz algorithm [19] and, as proved in [44], 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 Figure 1. Our main theoretical result is a proof 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:

Theorem 1 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 For a more precise statement of Theorem 1, see Corollary 1.

Remark 2

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 4 provides a more fine-grained relationship between the momentum and step-size parameters than Theorem 1. 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 randomized Kaczmarz algorithm [34] our convergence guarantees give a provable iteration complexity √ κ for randomized Kaczmarz 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 [26] introduce an accelerated randomized Kaczmarz (ARK) method which, through the use of Nesterov's acceleration, can achieve a faster rate of convergence 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 randomized Kaczmarz has been studied extensively in "small" batch regimes [32,35,36]. 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 randomized Kaczmarz 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 and Richtárik [28,29] 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. [14] 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. [27] 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. [4] 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. [4] however do not apply to the setting of randomized Kaczmarz, where the variance of the stochastic gradients necessarily grows proportionally to the squared norm of the full gradient.

Jain et al. [17] 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 [45], 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. [23] 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 [23] is a factor of κ better than what we obtain (see Theorem 1). However, while our analysis makes no assumptions on A, [23] 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 [2,4,13,49]. Aybat et al. [2] and Can et al. [4] demonstrate accelerated convergence guarantees of SNAG method variants to a neighborhood of the solution for problems with uniformly bounded noise. Additionally, Ghadimi and Lan [13] provide convergence guarantees for SNAG variants in nonconvex settings. Vaswani et al. [49] 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 [1,7,11,25,52].

Ma et al. [30] 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 a 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 [40]). 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 non-symmetric 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 [40].

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 [40]), 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 non-real if and only if

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

Here we have used that

we have that

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

λ j for all j = 1, 2, . . . , n, 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 T j = C j D j C -1 j . 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

where M(α, β ) = ∥C∥∥C -1 ∥ = ∥(UC)∥∥(UC) -1 ∥ 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 ),

(2.11)

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

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. .

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. □

Lemmas from non-asymptotic random matrix theory

Before we prove our main result, we need to introduce tools from non-asymptotic random matrix theory which are crucial components of the proof.

Proposition 2 Consider a finite sequence {W k } of independent random matrices with common dimension d 1 × d 2 . Assume that

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 [46], and the bound on E[∥Z∥ 2 ] follows from equation 6.1.6 in [46] and the fact E[max i ∥W i ∥ 2 ] ≤ W . Equation 6.1.6 in [46] comes from applying Theorem A.1 in [5] to the Hermitian dilation of Z. Under the stated conditions, the logarithmic dependence on the dimension is necessary [46].

We will also use a theorem on products of random matrices from [16]:

Proposition 3 (Corollary 5.4 in [16]) Consider an independent sequence of d × d random matrices X 1 , . . . , X k , and form the product

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 which bounds the batch size required to ensure that a certain random matrix is near it's expectation.

Lemma 2 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 iid, we obtain a bound for the variance statistic

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

The first term is bounded by δ /2 when B ≥ 8eη∥A∥ 2 F ∥A∥ 2 log(2d)δ -2 whereas the second term is bounded by δ /2 when B ≥ 16eη∥A∥ 2 F log(2d)δ -1 . The result follows by taking the max of these quantities. □ Our main result is the following theorem.

Theorem 4 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).

Remark 3 For k ≤ k * , (k * ) k/k * ≤ max{e, k}. Thus, Minibatch-HBM converges at at rate nearly β for k ≤ k * . For k > k * , the convergence is still linear for k * sufficiently large, but at a rate slower than β . Specifically, Minibatch-HBM converges at a rate ( β ) 1-δ where δ = 2 log(k * )/(k * log(1/β )).

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

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 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 3 we will instead consider

where X S i = (UC) -1 Y S i (UC) 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 3 to obtain a rate similar to HBM. Towards this end, note that

where W j is as in Lemma 2. This and the fact that

where W = ∑ j∈S j W j . Using Lemma 2 and (3.6), we have

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

gives the bound

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 4 is somewhat complicated, but can be simplified in the large condition number limit.

Corollary 1 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 1 -c/ √ κ 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 )/(λ min -c 1 λ min ) = κ/(1 -c 1 ) + c 1 /(1 -c 1 ), we have

we have

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

Using Lemma 1 we have that

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

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 * .

Remark 4 As shown in [33, Theorem 2.1], the RK iterates converge to within an expected radius √ d κσ of the least squares solution, where σ = max i |r i |/∥a i ∥. The minibatch-RK algorithm from [32] improves the convergence horizon by roughly a factor of √ B.

Theorem 5 In the setting of Theorem 4, define r = Ax * -b and σ = max i |r i |/∥a i ∥. Let the batch size B satisfy the conditions in Theorem 4. Then, provided k * is chosen so that δ = 2 log(k * )/(k * log(1/β )) < 1, the Minibatch-HBM iterates satisfy

Remark 5 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∥.

Remark 6 In the large

Proof Theorem 5 Since = Ax * -r, our stochastic gradients for inconsistent systems satisfy

where M S k is as in the proof of Theorem 4, 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 4,

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

where, in the third inequality, we have used the fact that (k * ) j/k * ≥ k * whenever j ≥ k * . Now, note that, by assumptionfoot_2 ,

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 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 4 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

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. 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 a 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 iid standard normal entries. In Figure 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 [34] for a further discussion on sampling probabilities in the context of RK and SGD.

Row-norm and uniform sampling

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 4, 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.

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_3

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 iid standard normal entries.

In Figures 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 Figure 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 Figure 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).

Inconsistent systems

For inconsistent systems, stochastic gradients at the optimum will not be zero and convergence is only possible to within the so-called convergence horizon. In this experiment we sample A as in the previous experiment using (4.1) with κ = 50 and ρ = 0.5. We take b = Ax + ε, where x is has iid standard normal entries and ε is drawn uniformly from the hypersphere of radius 10 -5 . Thus, the minimum residual norm ∥b -Ax * ∥ is around 10 -5 . We use several different values of B.

The results of the experiment are shown in Figure 5. For larger batch sizes, the convergence horizon of Minibatch-HBM gets smaller. In particular, when the batch size is increased by a factor of 100, the convergence horizon is decreased by about a factor of 10. This aligns with the intuition that the convergence horizon should depend on √ B. For reference, we also show the convergence for standard RK. Note RK requires more iterations to converge, although each iteration involves significantly less 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 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 [47].

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 Figure 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 [8,9,31,50], to potentially learn the momentum parameters α k and β k adaptively. Such an analysis would also shed light on convergence of ADAM [21], 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.

A. Analysis of Nesterov's acceleration for quadratics

Another common approach to accelerating the convergence of gradient descent is Nesterov's accelerated gradient descent (NAG). NAG uses iterates

or equivalently,

Therefore, a computation analogous to the above computation for HBM shows that the NAG iterates satisfy the transition relation 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 [38]:

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 1). We now provide a bound, analogous to Theorem 4, for minibatch-NAG, the algorithm obtained by using the stochastic gradients (1.2) in (NAG). Then, for all k > 0, assuming that the minimizer x * satisfies Ax * = b, the Minibatch-HBM iterates satisfy

The proof of Theorem 6 is almost the same as the proof identical to the proof of Theorem 4, 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 2 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. □