arXiv CSopen access

Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

machine learning, deep learning, neural networks

arXiv:2605.18609v1 [cs.LG] 18 May 2026

Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration Sachin Garg University of Michigan [email protected]

Michał Dereziński University of Michigan [email protected]

Abstract Accelerating stochastic gradient methods with classical momentum schemes, such as Polyak’s heavy ball, has proven highly successful in training large-scale machine learning models, particularly when combined with the hardware acceleration of large mini-batch computations. Yet, the effect of classical momentum on stochastic minibatch optimization has been poorly understood theoretically, with prior works requiring strong noise assumptions and extremely large mini-batches. In this work, we develop a general theory of stochastic momentum acceleration for optimizing over quadratics in the interpolation regime, a popular abstraction for studying deep learning dynamics which also includes classical methods such as randomized Kaczmarz and coordinate descent. Our framework encompasses both heavy ball and Nesterov-style momentum, allows for arbitrary mini-batch sizes, and makes minimal assumptions on the stochastic noise. In particular, we show that acceleration from classical momentum is directly proportional to the gradient mini-batch size (up to a natural saturation point), thereby enabling perfect parallelization of mini-batch computations. Our theory also provides a simple choice for the momentum parameter, which is shown to be effective empirically.

Introduction

Consider the following optimization task which naturally arises in training machine learning models: min F (w) := Eξ∼D f (w; ξ) , (1) w∈Rd

where w represents the model parameters and ξ ∼ D represents the stochastic noise. Problem (1) provides a powerful abstraction for designing and analyzing algorithms for model training, particularly in the context of deep learning and large language models [31, 24, 14, 54]. When given oracle access to the full gradient g(w) = ∇F (w) of the function F , problem (1) is very well understood. Here, the simplest strategy is to iteratively follow the direction

of steepest descent, wt+1 = wt − α g(wt ).

(2)

When F is smooth and strongly convex, this Gradient Descent algorithm (GD) attains linear convergence, i.e., it reaches an ϵ-approximate solution after O(κF log(1/ϵ)) iterations, where κF is an appropriate condition number of the deterministic function F . In 1964, Polyak [44] showed that the convergence rate of GD can be improved by introducing momentum into its descent trajectory: wt+1 = wt − α g(wt ) + β (wt − wt−1 ).

(3)

With an appropriate choice of step size α and momentum parameter β, this Heavy Ball √ Momentum algorithm (HBM) reaches an ϵ-approximate solution after O( κF log(1/ϵ)) iterations. In 1983, Nesterov [40] proposed a refinement of this scheme, known as Nesterov Accelerated Gradient (NAG), which essentially incorporates one additional gradient step to both wt and wt−1 in the momentum. For large-scale problems, instead of the full gradient we often only have access to stochasPm 1 tic gradients ∇f (w, ξ). By aggregating m such gradients into the vector g(w) = m i=1 ∇f (w; ξi ), the GD update (2) turns into the popular mini-batch Stochastic Gradient Descent algorithm (SGD, [47]), which has become a workhorse of modern machine learning. The success of SGD in deep learning has been partly attributed to the fact that it is particularly effective in the interpolation regime [4, 38, 3], i.e., when all of the functions f (·, ξ) can be simultaneously minimized (the model fits perfectly to the training data). In this setting, for smooth and strongly convex objectives, SGD attains linear convergence with O(κf log(1/ϵ)) iteration complexity, where κf is a corresponding condition number of the stochastic function f . The use of large mini-batch size m allows us to leverage parallel computing architectures while reducing the variance of the stochastic gradient [34], however its effect on the iteration complexity of vanilla SGD is limited [7] and not fully understood in practice [18]. A natural strategy for improving the performance of mini-batch SGD is to incorporate the classical momentum update (3) of Polyak, or its Nesterov counterpart. This approach has proven highly successful empirically, particularly when paired with large mini-batches [50, 24, 48, 16], and it inspired state-of-the-art training algorithms such as Adam [30] as well as many other more elaborate acceleration schemes for SGD [22, 55, 28]. However, this has led to a fundamental theory-practice gap: While there is significant theoretical progress on accelerating SGD through carefully designed schemes [11, 23, 1, 2, 27, 26, 52, 35, 6], the SGD algorithms used in practice still largely rely on the simple and classical momentum based on HBM, whose remarkable effectiveness remains poorly understood. Initial efforts in that direction were done by Loizou and Richtarik [37], who showed linear convergence of SGD with HBM, but without attaining any acceleration. Bollapragada, Chen and Ward [5] showed that HBM-accelerated mini-batch SGD can exhibit faster convergence on quadratic tasks, with related and follow-up works by [32, 41, 51, 58]. However, in order to attain this acceleration, all of these prior works require either strong additional assumptions on the stochastic noise or prohibitively large mini-batch sizes that depend heavily on problem parameters. 2

1.1

Our contributions

In this work, we address the above theory-practice gap by developing a general framework for analyzing the convergence of momentum-based stochastic processes, which encompasses mini-batch SGD with both HBM and NAG acceleration for stochastic quadratic minimization in the interpolation regime. Our analysis leads to a strikingly simple conclusion about the role of classical momentum in large-scale stochastic optimization: Classical momentum enables perfect parallelization for mini-batch SGD in the interpolation regime. Main result. Concretely, our main technical result (Theorem 1) shows that properly tuned SGD using mini-batch size m and either HBM or NAG momentum attains O (κf /m + √ κf ) log(κf /ϵ) iteration complexity in the interpolation regime. Note that for m = 1 (no mini-batching) this result essentially recovers the previously mentioned linear convergence for SGD, but as we increase m, the iteration complexity improves proportionally to 1/m √ until it saturates at an accelerated κf rate. Thus, up to the point of saturation, the overall work (number of stochastic gradients computed) remains constant, allowing us to fully exploit the embarrassingly parallel mini-batch computations. Such perfect parallelization does not in general hold for mini-batch SGD without acceleration, even over quadratics in the interpolation regime. Thus, our analysis provides an explanation for the success of momentum acceleration in allowing deep learning optimizers such as Adam to benefit from large mini-batches and highly parallelizable GPU architectures. Applications. Our analysis framework can also be directly applied to accelerate several popular SGD algorithms for solving large-scale regression tasks, such as randomized Kaczmarz [49] and coordinate descent [20]. In particular, we use our theory to develop a simple momentum-accelerated block coordinate descent algorithm, and we evaluate it on several kernel ridge regression tasks. Our numerical results confirm the perfect parallelization phenomenon, and they also show that a simple choice of β = 1 − 1/m, as motivated by our theory, already leads to meaningful acceleration. We also compare classical momentum with an optimally tuned state-of-the-art acceleration scheme for block coordinate descent [11], showing that, surprisingly, in certain large mini-batch regimes classical momentum yields substantially better performance.

1.2

Overview of our techniques

Limitations of prior approaches. The convergence behavior of a momentum-accelerated stochastic method follows a dynamical system that is characterized by a product of nonsymmetric random matrices. Existing matrix concentration results can be used to bound the spectral norm of a product of random matrices in terms of the product of the spectral norms of the expected values of the matrices [25]. Unfortunately, this strategy proves vacuous for momentum acceleration, as the product of norms diverges to infinity even when the norm of the product converges to zero, because of the lack of symmetry. To address this, [5] diagonalize the matrices so that they can use matrix product concentration effectively. 3

However, passing from the original matrices to the diagonalized form incurs a penalty factor that appears both in the convergence bound and in the required mini-batch size, and can blow up independently of the condition number of the problem. In fact, the penalty can even be infinite, as the underlying matrices may not be diagonalizable. To get around this, [41] diagonalize the matrices to a Jordan form, however this still does not avoid a potentially large penalty factor in both the convergence bound and the mini-batch size. Our approach. In order to get around these issues, we abandon both the diagonalization step and matrix product concentration. Instead, we use the Schur decomposition together with a matrix recursion to bound the original matrix product directly. Our analysis consists of two key components, following a bias-variance decomposition of the expected error. 1. Bias: Convergence of expected iterates via Schur decomposition. For the expected iterates, it suffices to analyze the product of certain deterministic non-symmetric matrices. Here, we start by rearranging the matrices into a block-diagonal form, and carefully analyze both their real and complex eigenvalues in order to bound the spectral radius. Next, we turn this into a spectral norm bound for the matrix product. In this step, instead of using diagonalization as done in prior work, we rely on the Schur decomposition [43], thereby avoiding condition number factors in our bound. 2. Variance: Convergence of expected error via matrix recursion. In order to control the random iterates, we must bound the stochastic error incurred at each iteration of the process. Instead of relying on black-box matrix concentration results, we do this by deriving a recursion for the expectation of the product of the random matrices, which relates its spectral norm to all monomials of the individual block matrices from the block-diagonal representations derived in the bias analysis. Unfolding this recursion proves quite technical as it requires separately handling real and complex eigenvalues of the diagonal blocks. Our analysis is not specific to one algorithm, but rather provides a general framework for effectively characterizing the convergence behavior of stochastic momentum methods. We illustrate this by stating our results for a class of stochastic processes which includes both HBM and NAG as special cases, and can be instantiated to provide convergence guarantees for a variety of SGD-type algorithms.

Further related work

In addition to the classical HBM and NAG momentum, a variety of more sophisticated momentum-based methods have been proposed in the literature [11, 23, 1, 2, 27, 26, 52, 35, 6]. Concurrently, substantial theoretical progress has been made toward understanding the convergence properties of stochastic gradient descent (SGD) equipped with various forms of momentum. Alongside these developments, a consistent empirical finding is the accelerated convergence achieved by incorporating classical momentum schemes such as HBM or NAG into

mini-batch SGD [50, 24, 48, 16]. Notwithstanding this strong empirical evidence, the theoretical understanding of the mechanisms underlying such acceleration remains incomplete. In particular, even in the setting of quadratic objectives, a comprehensive explanation of when and why these classical momentum methods yield acceleration is still lacking. A substantial body of work has analyzed the behavior of mini-batch SGD with HBM and NAG [21, 58, 41, 51, 5, 37, 8, 57, 17, 29, 15]. Many of these studies establish that, in general, SGD with O(1) size mini-batches and with such momentum does not admit provable improvements over its non-momentum counterpart [29, 36, 26, 56], attributing observed empirical gains primarily to variance reduction induced by large mini-batches. Despite these negative results, more recent works have revisited this question under restrictive settings, including large mini-batches, carefully tuned step sizes and momentum parameters, and specific assumptions on the data and noise [5, 51, 41, 58, 32]. For instance, [58] demonstrate accelerated convergence of randomized linear solvers with HBM via adaptive parameter tuning; however, their analysis relies on distributional assumptions on the sketching matrices that may not hold for commonly used subsampling schemes, and their guarantees do not explicitly quantify the improvement in convergence rates. Similarly, [32] establish acceleration for mini-batch SGD only in an asymptotic regime and under additional assumptions such as random, orthogonally invariant data matrices. The works most closely related to ours are [5, 41, 51]. In particular, [5] establish accelerated convergence of mini-batch SGD with HBM under the requirement that the batch √ size satisfies m = Ω̃(κf κF ), along with additional stable diagonalization assumptions. In a different setting, [51] prove acceleration for inconsistent linear systems, but require a significantly larger batch size, specifically m = Ω̃(κ3f ). More recently, [41] show that mini-batch √ d SGD with HBM has iteration complexity Õ κ + mϵ to reach ϵ-accuracy, showing sublinear convergence that improves with mini-batch size. They do not consider the interpolation regime, but rather, rely on other assumptions such as anisotropic gradient noise and carefully designed step-size schedules. Since we are focused on the interpolation regime and attaining linear convergence, their results are largely incomparable to ours.

Stochastic momentum acceleration framework

Notation. We represent vectors by lowercase boldface letters and matrices by uppercase boldface letters. For symmetric matrices, we use Loewner ordering and say A ⪯ B if B − A is a positive semidefinite (psd) matrix. ∥A∥ denotes √ spectral norm, ∥A∥F denotes Frobenius norm of matrix A. For a psd matrix A, ∥x∥A = x⊤ Ax denotes the Mahalanobis norm of x. We let A† denote the Moore-Penrose pseudoinverse of A. We use Õ, Ω̃ notation to hide logarithmic factors. Stochastic contraction process. We start by defining a class of stochastic dynamical systems called stochastic contraction processes [9]. This framework can be instantiated to represent the dynamics of many stochastic algorithms, as described below, and thus serves as a useful abstraction for our theory. 5

Definition 1 For a random psd matrix Π such that 0 ⪯ Π ⪯ I and scalar α ∈ [0, 1], we let A(∆; α) := (I − α Π)∆ define a stochastic contraction with average rate Π̄ = E[Π]. We call Π the stochastic rate matrix of A. Given a sequence of independent At with the same Π̄ and α, we call ({At }t≥0 , α) a stochastic contraction process. To see how this definition corresponds to an SGD algorithm, consider the minimization problem (1) over F (w) = Eξ∼D [f (w; ξ)], where for every ξ, function f (w; ξ) is a convex quadratic with respect to w that is minimized at some w∗ ∈ Rd . As long as function f (·; ξ) is L-smooth, i.e., ∥∇f (w, ξ) − ∇f (v, ξ)∥ ≤ L∥w − v∥, then the evolution of the error vector ∆t = wt − w∗ resulting from the SGD update (2) with g(wt ) = L1 ∇f (wt ; ξt ) follows a stochastic contraction process: α 2 ∆t+1 = At (∆t ; α), where At (∆; α) = I − ∇ f (wt ; ξt ) ∆. (4) L The stochastic rate matrix defining the stochastic contraction process At is Πt = L1 ∇2 f (wt ; ξt ), and its average rate is Π̄ = L1 ∇2 F (w∗ ). Note that the Hessian is the same everywhere since both f (·; ξt ) and F are assumed to be quadratic. Stochastic momentum process. We next show how to introduce momentum into this process by incorporating information from two preceding steps. We define these dynamics in sufficient generality to capture both HBM and NAG momentum as special cases. Definition 2 Given a stochastic contraction process ({At }t≥0 , α), initial vectors ∆−1 , ∆0 , and parameters β, ω ≥ 0, define the corresponding stochastic momentum process recursively as: h i ∆t+1 = At (∆t ; α) + β At (∆t ; ω) − At−1 (∆t−1 ; ω) . (5) Note that if we set the momentum parameter β to zero, then we recover the usual SGD dynamics of a stochastic contraction process. On the other hand, if we set β > 0 but ω = 0, then the momentum term simplifies to β[∆t −∆t−1 ], which precisely corresponds to Polyak’s heavy-ball momentum (HBM). On the other hand, by choosing ω = α we introduce one more SGD step to the momentum term, thereby recovering Nesterov’s momentum (NAG). For example, we can use this framework to model HBM acceleration for the SGD process {At }t≥0 from (4) by letting ∆t = wt − w∗ ,

and

wt+1 = wt −

α ∇f (wt ; ξt ) + β(wt − wt−1 ). L

It is easy to verify that such a sequence {∆t }t≥0 is a stochastic momentum process associated with ({At }t≥0 , α), momentum parameter β, and ω = 0. A similar derivation follows immediately for Nesterov’s momentum by setting ω = α.

Example 1: Randomized Kaczmarz. As a concrete example of an algorithm that falls directly under our framework, we consider the randomized Kaczmarz (RK) algorithm [49], along with its block versions [39]. Let A ∈ Rn×d and y ∈ Rn define a linear regression problem with the loss function F (w) = ∥Aw − y∥2 . Also, let D be a distribution over subsets S ⊆ {1, ..., n}, and consider the following update: wt+1 =

argmin ∥w − wt ∥2 = wt − A†St (ASt wt − ySt ),

St ∼ D,

(6)

w: ASt w=ySt

where ASt denotes the submatrix consisting of rows indexed by St . In particular, when S ∼ D generates single-sample sets such that Pr[S = {i}] ∝ ∥Ai,: ∥2 , then this recovers classical randomized Kaczmarz. When there exists w∗ such that Aw∗ = y (i.e., in the interpolation regime), the update (6) can be cast as a stochastic contraction process ({At }t≥0 , 1) by letting Πt = A†St ASt and ∆t = wt − w∗ . In particular, for randomized Kaczmarz, a simple calculation shows that this process has an average rate Π̄ = E[Πt ] = A⊤ A/∥A∥2F . Similar or better rates have been derived for block versions of Kaczmarz [13, 11], and other extensions such as sketch-and-project [12, 46, 45]. By applying momentum via Definition 2 to ({At }t≥0 , 1), we can precisely recover convergence dynamics of the error vector of the accelerated Kaczmarz algorithm considered by [5]. Example 2: Coordinate Descent. Another application of our framework comes in the context of solving positive semidefinite linear systems, such as in kernel ridge regression. For a positive definite matrix K ∈ Rn×n and y ∈ Rn , consider the (block) coordinate descent (CD, [33]) update: wt+1 = wt − I⊤St (KSt ,St )† (Kwt − y)St ,

St ∼ D,

(7)

which only updates the coordinates of wt indexed by St . Letting w∗ denote the solution of the linear system Kw = y, we can cast (7) as a stochastic contraction process ({At }t≥0 , 1) by 1/2 1/2 choosing ∆t = K1/2 (wt − w∗ ) and Πt = KSt (KSt ,St )† KSt . Here, the canonical distribution D for single-sample coordinate descent is Pr[S = {i}] ∝ Ki,i , which yields average rate Π̄ = K/tr(K). Convergence rate of a stochastic contraction process. The above examples are of base algorithms, i.e., without momentum acceleration. Before we analyze their accelerated variants through the stochastic momentum framework, we establish their base convergence rate via our general formalism. Proposition 1 Given a stochastic contraction process ({At }t≥0 , α) with average rate Π̄, and an initial vector ∆0 ∈ range(Π̄), consider the sequence ∆t+1 = At (∆t ; α). Then, we have: 1 t ∥∆0 ∥2 , where κ = α/λ+ E ∥∆t ∥2 ≤ 1 − min (Π̄). κ We will refer to κ as the stochastic condition number associated with the process ({At }t≥0 , α). For instance, by plugging in our concrete examples of (single-sample) randomized Kaczmarz and coordinate descent [49], we recover the classical convergence guarantees for these methods, with stochastic condition numbers κRK = ∥A∥2F ∥A† ∥2 and κCD = tr(K)∥K† ∥. 7

Effect of mini-batching. The stochastic contraction framework is naturally amenable to mini-batching of the updates via simple averaging of the contractions. It is thus natural to ask how mini-batching affects the convergence dynamics. Definition 3 Given a stochastic contraction A with average rate Π̄, let A[m] denote the average of m independent copies of A, i.e., m

[m]

1 X (i) A (∆; α), (∆; α) := m i=1

where

A(i) ∼ A.

(8)

For example, if ({At }t≥0 , α) is a stochastic contraction process for the SGD algorithm Pm (i) [m] 1 given by (4), then ({At }t≥0 , α) corresponds to mini-batch SGD where g(wt ) = mL i=1 ∇f (wt , ξt ). Note that mini-batching is different from choosing a larger block set St in randomized Kaczmarz or coordinate descent, since those schemes do more than just averaging the samples (such as inverting the matrices ASt and KSt ,St ). In fact, combining blocking and minibatching in RK and CD is a very natural strategy for making these methods scalable to parallel architectures, as discussed later. Proposition 2 If A is a stochastic contraction with average rate Π̄, then A[m] is also a stochastic contraction with the same average rate Π̄. Moreover, if Π[m] is the stochastic rate matrix associated with A[m] , then it follows that: h 2 i 1 [m] E Π −Π̄ ⪯ · Π̄(I − Π̄). (9) m Proposition 2, together with Proposition 1, illustrates why mini-batching of SGD does not in general lead to faster convergence rates, since the average rate matrix Π̄ (and thus also the associated stochastic condition number) remains the same regardless of m. However, minibatching does reduce the variance of the stochastic noise matrix, as shown in (9), which will be crucial for the stochastic momentum dynamics. We defer the proofs of Propositions 1 and 2 to Appendix D.

Main result

We are now ready to state our main result, which characterizes the convergence dynamics of a stochastic momentum process (Definition 2) corresponding to stochastic contractions with mini-batching (Definitions 1 and 3). Crucially, our result makes no assumptions on the mini-batch size m or the stochastic noise, and applies to both heavy-ball and Nesterov momentum. Theorem 1 Consider a stochastic contraction process ({At }t≥0 , α) with α ∈ [0, 1], average rate Π̄, and condition number κ = α/λ+ min (Π̄). For m ≥ 1, ∆0 ∈ range(Π̄), let ∆−1 = 0,

[m]

{∆t }t≥0 be a stochastic momentum process for ({At }t≥0 , α) with ω ∈ [0, α] and momentum parameter   if m ≤ c1 , 0 √ c1 β = 1 − 2m if m ∈ [c1 , c2 κ],  √  1 − 2c2c1√κ if m ≥ c2 κ, where c1 , c2 > 0 are absolute constants. Then, there are absolute constants C, c > 0 such that: n m 1 ot 2 3 ,√ ∥∆0 ∥2 for any t ≥ 1. E ∥∆t ∥ ≤ Ct 1 − c min κ κ √ Remark 1 In particular, we get E ∥∆t ∥2 ≤ ϵ∥∆0 ∥2 after t = O((κ/m + κ) log(κ/ϵ)) √ iterations, compared to t = O(κ log(1/ϵ)) without momentum. Thus, for m ∈ [1, Θ( κ)], momentum enables perfect parallelization up to a logarithmic factor (which is caused by the t3 term in the bound). √ Remark 2 In the regime m ∈ [1, Θ( κ)], our result yields a simple problem-independent momentum parameter β = 1 − Θ(1/m). Numerical experiments confirm that this choice works well in practice.

4.1

Proof sketch

First, without loss of generality we can assume that α = 1 (otherwise, we can simply rescale Π̄ appropriately). Also, note that the regime of 1 ≤ m ≤ c1 follows immediately from [m] [m] Proposition 1. Next, let {Πt }t≥0 be the stochastic noise matrices associated with {At }t≥0 . From Proposition 2, these matrices satisfy the variance bound (9). In fact, this is the only property we need out of the mini-batching, so to simplify the notation, we will drop the superscript and consider a sequence {Πt }t≥0 of matrices such that E (Πt −Π̄)2 ⪯ m1 Π̄(I− Π̄). With these conventions, after simple calculations the stochastic momentum process (5) can be expressed as the following recursion: ∆t+1 = (1 + β)I − (1 + βω) Πt ∆t − β I − ω Πt−1 ∆t−1 . (10) For notational inconvenience, we restrict ourselves to NAG acceleration in this proof sketch (i.e. ω = α = 1). The detailed proof in the appendix covers HBM and NAG momentum simultaneously. Let Π̄ = VΛV⊤ be the eigendecomposition of Π̄, where Λ = diag(λ1 , λ2 , · · · , λd ) and VV⊤ = V⊤ V = I. Let rank(Π̄) = r ≤ d, so that λi = 0 for i > r. Following prior works [26, 5, 41], we study the dynamics of stochastic process ∆t in (10) via the following two-step matrix transition rule: ⊤ ⊤ V ∆t−1 V ∆t (1 + β)(I − V⊤ Πt V) −β(I − V⊤ Πt−1 V) · (11) = V⊤ ∆t−2 V⊤ ∆t−1 I 0 {z } | :=Yt

We break down the proof of Theorem 1 into two main components, starting with the convergence of the expected iterates (bias), and then analyzing the effect of the stochastic noise (variance). 4.1.1

Bias: Convergence of ∥E∆t ∥

The convergence behavior of the expected iterates is characterized by a simpler transition rule: ⊤ ⊤ V ∆0 (1 + β)(I − Λ) −β(I − Λ) V ∆t t , where T := E[Yt ] = . (12) E =T V⊤ ∆−1 I 0 V⊤ ∆t−1 The first thing that comes to mind is to use Tt ∥ ≤ ∥T∥t , but, unfortunately, for nonzero β the spectral norm of T is greater than 1, making this naive idea ineffective. However, the spectral radius of T is less than 1 and this paves the way ahead, suggesting a two step approach: Step 1: Bound the spectral radius of Tt , denoted as ρ(Tt ); Step 2: Using structural properties of T, bound Tt ∥ in terms of ρ(Tt ). Step 1. By carefully permuting rows and columns of T, we can transform it into a blockdiagonal matrix with 2 by 2 blocks on the diagonal. In particular, the following lemma holds: Lemma 1 (Block-diagonalization) There exists a permutation matrix P such that, (1 + β)(1 − λi ) −β(1 − λi ) ⊤ PTP = diag(T1 , T2 , · · · , Td ), where Ti := . 1 0 Due to Lemma 1, taking maximum over the spectral radius of each Ti will upper bound ρ(T). Interestingly, the eigenvalues of Ti depend on interaction between λi and β, and can (1−β)2 be either real or complex depending upon whether λi ≤ (1+β) 2 . Let γi denote the eigenvalue of Ti with the larger magnitude, so that |γi | equals the spectral radius of Ti . We prove the following result: Lemma 2 (Spectral radius of Ti ) Let β = 1 − ϕ1 with ϕ ≥ 2. Then, ( 2

|γi | ≤

2 1 − ϕ1 (1 − λi ), if λi ≥ (1−β) (i.e., when γi is complex), (1+β)2 2 1 − ϕλ2 i (1 − λi )2 , otherwise (i.e., when γi is real).

It is easy to see that the maximum is attained at i = r, thus yielding an upper bound on ρ(T).

Step 2. To go from ρ(Tt ) to Tt ∥, previous work [5] assumes that T is diagonalizable, i.e., there exists an invertible matrix C such that CTC−1 is a diagonal matrix. This assumption implies, Tt ∥ ≤ C∥ · C−1 ∥ · ρ(Tt ). 2

, then T is not Unfortunately, T need not be diagonalizable. In fact, if for any i, λi = (1−β) (1+β)2 diagonalizable. More problematically, it can be shown that n ϕ (1 − β)2 o , where δ = min λi − ∥C∥ · ∥C−1 ∥ = Ω . i δ (1 + β)2 To avoid this blow up, we abandon the diagonalization approach and rely on a fundamental matrix decomposition result in linear algebra, known as Schur’s decomposition. Lemma 3 (Schur’s decomposition [43]) For every i, there exists a unitary matrix Ui such that γi1 xi H Ui Ti Ui = , 0 γi2 where γi1 and γi2 are the two eigenvalues of Ti , 1 ≤ |xi | ≤ 3 and UH i denotes the hermitian of Ui . Using Schur’s decomposition it readily follows that for any i, t t−1 ∥Tti ∥ = ∥UH . i Ti Ui ∥ ≤ O(t) · |γi |

Using the above in conjunction with Lemma 2 we show that: ϕ (1 − κ1 ), the transition rule (11) Theorem 2 (Convergence of ∥E∆t ∥) For ρ = 1 − 2κ satisfies, ∥E∆t ∥2 ≤ O(t2 ) · ρt−2 ∥∆0 ∥2 . 4.1.2

Variance: Convergence of E∥∆t ∥2

Note that (11) implies, Y ⊤ t−1 V⊤ ∆t V ∆0 ⊤ P = ·P PYt−j P . V⊤ ∆t−1 V⊤ ∆−1 j=0 | {z }

Σt

To establish convergence of E∥∆t ∥ , it suffices to bound ∥E[Σ⊤t Σt ]∥. Since Σt is a product of non-symmetric random matrices, upper bounding E[Σ⊤t Σt ] is challenging. Prior work [5] uses matrix product concentration [25] and require careful tuning of the momentum parameter to keep T diagonalizable while controlling ∥C∥ · ∥C−1 ∥. This leads to mini-batch requirements scaling as ∥C∥2 ·∥C−1 ∥2 , which can become arbitrarily large near non-diagonalizability, and introduces additional condition number dependence in m. Consequently, their guarantees degrade rapidly unless the mini-batch size is sufficiently large. These limitations stem from the reliance on product concentration results and diagonalization. In contrast, we avoid both and directly bound ∥E[Σ⊤t Σt ]∥. 11

Recursive approach to bound ∥E[Σ⊤t Σt ]∥. Let Σ̄t := E[Σt ] denote the matrix mean and Mt := E[Σ⊤t Σt ] − Σ̄⊤t Σ̄t denote the matrix variance of Σt . Moreover, let Ȳ = E[PYt P⊤ ] and Rt := P(Yt − E[Yt ])P⊤ . Then, for any pt−1 ≥ E[Σ⊤t−1 Σt−1 ] , we have the following matrix recursion: Mt ⪯ pt−1 · E[R⊤t Rt ] + Ȳ⊤ Mt−1 Ȳ.

(13)

Note that, Ȳ and Σ̄t are block-diagonal due to Lemma 1. If one can show that E[R⊤t Rt ] is upper bounded by block-diagonal matrices, then via the matrix recursion, so is Mt . We show this in the following lemma. Lemma 4 Let P be the permutation matrix from Lemma 1. Then we have, λ (1 − λ ) λ (1 − λ ) i i i i ⊤ where Wi = diag E[Rt Rt ] ⪯ 4 · diag(W1 , W2 , · · · , Wn ), , . m m Using Lemma 4 and the repeated use of (13), it follows that both Mt and E[Σ⊤t Σt ] are also upper bounded by block-diagonal matrices. By taking a maximum over those blocks, we prove the following recursive expression for E[Σ⊤t Σt ] . Lemma 5 For any t ≥ 1, if there exist p0 , p1 , p2 , · · · , pt−1 such that E[Σ⊤k Σk ] ≤ pk for 0 ≤ k ≤ t − 1, then there exists pt such that for qi = 4λi (1 − λi )/m, ⊤

E[Σt Σt ] ≤ pt ≤ max

1≤i≤r

(Ti ) Tti ⊤

+ qi ·

t−1 X

pj · (T⊤i )t−1−j Tt−1−j i

j=0

Unfolding the recursion. The next step is to unfold the recursive expression for pt . By coupling with a simpler recursion, we show that, defining ρt := max1≤i≤r |γi | + qi · ℓi (t) , pt ≤ ρtt +

t−1 X ρjt · max (T⊤i )t−j Tt−j , i j=0

1≤i≤r

(14)

for a carefully defined function ℓi (t), see (25). We can now use the bias analysis (Lemmas 2 and 3) to obtain tight bounds for the terms (T⊤i )t−j Tt−j . The final challenging task is to i bound ρt . Upper bounding ρt . We provide a fine-grained analysis of ρt by carefully analyzing |γi | + qi · ℓi (t) for every single i ≤ r. Recall that, depending upon the interaction between λi and β, the eigenvalue γi can be real or complex (Lemma 2). 2

(1−β) Case 1: λi − (1+β) This is the easy case where γi is complex, 2 is sufficiently positive. and we get

|γi | + qi · ℓi (t) < 1 −

1 , 32ϕ

for ϕ ≤

m . C

(15)

Case 2: λi − (1−β) is close to zero. In this case, we enter the regime of near non(1+β)2 diagonalizability of T. A crude analysis here would lead to a blow up factor in m. However, 2 we are able to avoid this by exploiting the specific form of ℓi (t) and combine it with λi ≈ (1−β) (1+β)2 to show that (15) still holds. 2

Case 3: λi − (1−β) is sufficiently negative. Previous works such as [5, 41, 32] select (1+β)2 large values for the momentum parameter β and for the mini-batch size m, because they 2 . Our crucial insight is that, for smaller values of β, enforce the condition that λi ≥ (1−β) (1+β)2 2

(1−β) the quantity λi − (1+β) 2 can be negative and sufficiently bounded away from 0 for some i. In fact, this is the case that determines the convergence rate for E∥∆t ∥2 and allows us to handle SGD with classical momentum even under small mini-batches (see details in Appendix B). In particular, in this case we show that for ϕ ≤ m , C

|γi | + qi · ℓi (t) ≤ 1 −

ϕλi . 4

(16)

Combining (15), (16) and maximizing over i, bounds ρt . Substituting in (14) and boundbounds pt . Noting the relation between m and ϕ finishes the proof of ing (T⊤i )t−j Tt−j i Theorem 1.

Experiments

For our experiments, we consider kernel ridge regression on four benchmark datasets sourced from Scikit-learn [42] and OpenML [53] (more details and results in Appendix C). For each dataset, we solve the resulting psd linear system (K + λI)w = y, where K is the kernel matrix and λ ≥ 0 is the ridge parameter, using randomized block coordinate descent (7) with uniform sampling as the base update. We then incorporate mini-batching by averaging m independently generated update vectors in each iteration. Note that this mini-batching is separate from block sampling, which is less readily parallelizable due to the block inverse step. Each of the m update vectors in the mini-batch is produced using an independently sampled random block. Finally, we consider four different configurations of the algorithm, depending on the type of momentum acceleration: 1. CD: Vanilla block coordinate descent (7) using mini-batches of size m and varying block sizes. 2. CDpp: For comparison, we include a tuned state-of-the-art acceleration scheme for CD (i.e., not classical momentum) based on a Lyapunov-based approach with two auxiliary sequences [19, 10, 11]. 3. CD+NAG (β = 1 − 1/m): CD with classical NAG momentum (i.e., ω = 1 in our framework), using a simple choice of momentum parameter motivated by our theory. 4. CD+NAG (β = adaptive): We implement a heuristic strategy that learns the near-optimal β for any m by executing an adaptive binary search procedure at runtime (see Appendix C). 13

Block size=5

Block size=10

105 104 103

1 2 4 10

100 500 1 2 4 10

Mini-batch size (m)

100 500

No. of row sampled processed

(a) Covtype: iterations vs mini-batch size 108

Block size=1

Block size=5

Block size=10

106 105 1 2 4 10

100 500 1 2 4 10

Mini-batch size (m)

100 500

CDpp

105

Block size=10

Block size=25

Block size=50

104

103 1 2 4 10 50 250 1000 1 2 4 10 50 250 1000 1 2 4 10 50 250 1000

Mini-batch size (m)

(b) California Housing: iterations vs mini-batch size

107

104

No. of iterations to convergence

Block size=1

No. of row sampled processed

No. of iterations to convergence

106

109

Block size=10

Block size=25

Block size=50

108 107 106 105 104

1 2 4 10 50 250 1000 1 2 4 10 50 250 1000 1 2 4 10 50 250 1000

Mini-batch size (m)

(d) California Housing: row samples vs minibatch size

CD+NAG (β = 1 − 1/m)

CD+NAG (β = adaptive)

Figure 1: Convergence of block coordinate descent with NAG momentum and baselines. Top row: iterations to reach error 10−7 . Bottom row: total work, measured by the number of sampled rows. We also separately compare several variants of classical momentum by interpolating between CD+HBM (ω = 0) and CD+NAG (ω = 1), see Fig. 2. Although our theory (Theorem 1) guarantees identical convergence rates for all ω ∈ [0, 1], empirical performance can depend on the choice of ω, with larger values often yielding improved results. This observation is also noted in the remark following Lemma 7. Furthermore, as shown in Fig. 2, all variants of classical momentum perform very similarly; nevertheless, ω = 1 consistently achieves the best results, which is why we chose NAG momentum in our main experiments. In Fig. 1(a)-(b) we can see that, as expected, mini-batching by itself does not improve the convergence of CD. (Due to this phenomenon, we ran CD only up to mini-batch size m = 32, and extrapolated for larger values of m.) However, incorporating momentum into the update unlocks significant speed-up from mini-batching. Remarkably, while the specialized acceleration scheme of CDpp tends to attain better acceleration than CD+NAG for m = 1, the classical momentum does a better job of exploiting large mini-batches. This can be seen even more clearly in Fig. 1(c)-(d), which shows total work performed by the algorithm. For CD+NAG, the total work stays constant as we increase m, until reaching a saturation point, which allows for perfect parallelization of mini-batches. Notably, this perfect parallelization

Block size=10

Block size=20

104

103 1 2 4 10

250 1 2 4 10

Mini-batch size (m)

250 1 2 4 10

250

No. of row sampled processed

No. of iterations to convergence

Block size=1

105

(a) Abalone: iterations vs mini-batch size CD

CDM : ω = 0 (HBM)

CDM : ω = 0.25

108

Block size=1

Block size=10

Block size=20

107 106 105 104

1 2 4 10

250 1 2 4 10

Mini-batch size (m)

250 1 2 4 10

250

(b) Abalone: row samples vs mini-batch size CDM : ω = 0.50

CDM : ω = 0.75

CDM : ω = 1 (NAG)

Figure 2: Convergence behavior for CDM with different ω across block sizes. Left column: iterations to reach error 10−10 . Right column: total work, measured by the number of sampled rows. phenomenon is not observed for CDpp. Next, comparing CD+NAG with β = 1 − 1/m to its adaptive variant, we observe that in the perfect parallelization regime the simple choice of momentum parameter is essentially √ optimal. For larger m, as suggested by our theory, β should be capped at 1 − Θ(1/ κ), which is why the adaptive selection procedure becomes more effective.

Conclusions

We develop a theoretical framework for quantifying the effect of classical momentum schemes like heavy ball on the convergence of mini-batch SGD. Our results demonstrate that momentum enables perfect parallelization of mini-batch computations, which is confirmed by numerical experiments.

Acknowledgments This work was supported in part by NSF CAREER Grant CCF-233865 and a Google ML and Systems Junior Faculty Award. The work was done in part while MD was visiting the Simons Institute for the Theory of Computing.

References [1] Seth J Alderman, Roan W Luikart, and Nicholas F Marshall. Randomized kaczmarz with geometrically smoothed momentum. SIAM Journal on Matrix Analysis and Applications, 45(4):2287–2313, 2024.

[2] Zeyuan Allen-Zhu. Katyusha: The first direct acceleration of stochastic gradient methods. Journal of Machine Learning Research, 18(221):1–51, 2018. [3] Amit Attia, Matan Schliserman, Uri Sherman, and Tomer Koren. Fast last-iterate convergence of sgd in the smooth interpolation regime. arXiv preprint arXiv:2507.11274, 2025. [4] Mikhail Belkin, Daniel J Hsu, and Partha Mitra. Overfitting or perfect fitting? risk bounds for classification and regression rules that interpolate. Advances in neural information processing systems, 31, 2018. [5] Raghu Bollapragada, Tyler Chen, and Rachel Ward. On the fast convergence of minibatch heavy ball momentum. IMA Journal of Numerical Analysis, 45(3):1397–1424, 2025. [6] Sébastien Bubeck, Yin Tat Lee, and Mohit Singh. A geometric alternative to nesterov’s accelerated gradient descent. arXiv preprint arXiv:1506.08187, 2015. [7] Richard H Byrd, Gillian M Chin, Jorge Nocedal, and Yuchen Wu. Sample size selection in optimization methods for machine learning. Mathematical programming, 134(1):127– 155, 2012. [8] Bugra Can, Mert Gurbuzbalaban, and Lingjiong Zhu. Accelerated linear convergence of stochastic momentum methods in wasserstein distances. In International Conference on Machine Learning, pages 891–901. PMLR, 2019. [9] Michał Dereziński and Xiaoyu Dong. Last-iterate convergence of randomized kaczmarz and sgd with greedy step size. arXiv preprint arXiv:2604.09909, 2026. [10] Michal Dereziński, Daniel LeJeune, Deanna Needell, and Elizaveta Rebrova. Finegrained analysis and faster algorithms for iteratively solving linear systems. Journal of Machine Learning Research, 26(144):1–49, 2025. [11] Michał Dereziński, Deanna Needell, Elizaveta Rebrova, and Jiaming Yang. Randomized kaczmarz methods with beyond-krylov convergence. SIAM Journal on Matrix Analysis and Applications, 46(4):2558–2588, 2025. [12] Michał Dereziński and Elizaveta Rebrova. Sharp analysis of sketch-and-project methods via a connection to randomized singular value decomposition. SIAM Journal on Mathematics of Data Science, 6(1):127–153, 2024. [13] Michał Dereziński and Jiaming Yang. Solving dense linear systems faster than via preconditioning. In Proceedings of the 56th Annual ACM Symposium on Theory of Computing, pages 1118–1129, 2024.

[14] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pretraining of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers), pages 4171–4186, 2019. [15] Nicolas Flammarion and Francis Bach. From averaging to acceleration, there is only a step-size. In Conference on learning theory, pages 658–695. PMLR, 2015. [16] Jingwen Fu, Bohan Wang, Huishuai Zhang, Zhizheng Zhang, Wei Chen, and Nanning Zheng. When and why momentum accelerates sgd: An empirical study. arXiv preprint arXiv:2306.09000, 2023. [17] Sébastien Gadat, Fabien Panloup, and Sofiane Saadane. Stochastic heavy ball. 2018. [18] Noah Golmant, Nikita Vemuri, Zhewei Yao, Vladimir Feinberg, Amir Gholami, Kai Rothauge, Michael W Mahoney, and Joseph Gonzalez. On the computational inefficiency of large batch sizes for stochastic gradient descent. arXiv preprint arXiv:1811.12941, 2018. [19] Robert Gower, Filip Hanzely, Peter Richtárik, and Sebastian U Stich. Accelerated stochastic matrix inversion: general theory and speeding up bfgs rules for faster secondorder optimization. Advances in Neural Information Processing Systems, 31, 2018. [20] Robert M Gower and Peter Richtárik. Randomized iterative methods for linear systems. SIAM Journal on Matrix Analysis and Applications, 36(4):1660–1690, 2015. [21] Kanan Gupta, Jonathan W Siegel, and Stephan Wojtowytsch. Nesterov acceleration despite very noisy gradients. Advances in Neural Information Processing Systems, 37:20694–20744, 2024. [22] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842– 1850. PMLR, 2018. [23] Deren Han and Jiaxin Xie. On pseudoinverse-free randomized methods for linear systems: Unified framework and acceleration. Optimization Methods and Software, pages 1–36, 2025. [24] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016. [25] De Huang, Jonathan Niles-Weed, Joel A Tropp, and Rachel Ward. Matrix concentration for products. Foundations of Computational Mathematics, 22(6):1767–1799, 2022.

[26] Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerating stochastic gradient descent. arXiv preprint arXiv:1704.08227, 3:8, 2017. [27] Chi Jin, Praneeth Netrapalli, and Michael I Jordan. Accelerated gradient descent escapes saddle points faster than gradient descent. In Conference on learning theory, pages 1042–1085. PMLR, 2018. [28] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. [29] Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham Kakade. On the insufficiency of existing momentum schemes for stochastic optimization. In 2018 Information Theory and Applications Workshop (ITA), pages 1–9. IEEE, 2018. [30] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. [31] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012. [32] Kiwon Lee, Andrew Cheng, Elliot Paquette, and Courtney Paquette. Trajectory of minibatch momentum: batch size saturation and convergence in high dimensions. Advances in Neural Information Processing Systems, 35:36944–36957, 2022. [33] Dennis Leventhal and Adrian S Lewis. Randomized methods for linear constraints: convergence rates and conditioning. Mathematics of Operations Research, 35(3):641– 654, 2010. [34] Mu Li, Tong Zhang, Yuqiang Chen, and Alexander J Smola. Efficient mini-batch training for stochastic optimization. In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 661–670, 2014. [35] Ji Liu and Stephen Wright. An accelerated randomized kaczmarz algorithm. Mathematics of Computation, 85(297):153–178, 2016. [36] Nicolas Loizou and Peter Richtárik. Linearly convergent stochastic heavy ball method for minimizing generalization error. arXiv preprint arXiv:1710.10737, 2017. [37] Nicolas Loizou and Peter Richtárik. Momentum and stochastic momentum for stochastic gradient, newton, proximal point and subspace descent methods. Computational Optimization and Applications, 77(3):653–710, 2020. [38] Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of sgd in modern over-parametrized learning. In International Conference on Machine Learning, pages 3325–3334. PMLR, 2018. 18

[39] Deanna Needell and Joel A Tropp. Paved with good intentions: analysis of a randomized block kaczmarz method. Linear Algebra and its Applications, 441:199–221, 2014. [40] Yurii Nesterov. A method for solving the convex programming problem with convergence rate o (1/k2). In Dokl akad nauk Sssr, volume 269, page 543, 1983. [41] Rui Pan, Yuxing Liu, Xiaoyu Wang, and Tong Zhang. Accelerated convergence of stochastic heavy ball method under anisotropic gradient noise. arXiv preprint arXiv:2312.14567, 2023. [42] Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit-learn: Machine learning in python. the Journal of machine Learning research, 12:2825–2830, 2011. [43] Robert J Plemmons. Matrix analysis (roger a. horn and charles r. johnson), 1988. [44] Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr computational mathematics and mathematical physics, 4(5):1–17, 1964. [45] Pratik Rathore, Zachary Frangella, Sachin Garg, Shaghayegh Fazliani, Michal Derezinski, and Madeleine Udell. Turbocharging gaussian process inference with approximate sketch-and-project. Advances in Neural Information Processing Systems, 38:141691– 141728, 2026. [46] Pratik Rathore, Zachary Frangella, Jiaming Yang, Michał Dereziński, and Madeleine Udell. Have askotch: A neat solution for large-scale kernel ridge regression. arXiv preprint arXiv:2407.10070, 2024. [47] Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400–407, 1951. [48] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4510–4520, 2018. [49] Thomas Strohmer and Roman Vershynin. A randomized kaczmarz algorithm with exponential convergence. Journal of Fourier Analysis and Applications, 15(2):262–278, 2009. [50] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning, pages 1139–1147. pmlr, 2013. [51] Kejie Tang, Weidong Liu, Yichen Zhang, and Xi Chen. Acceleration of stochastic gradient descent with momentum by averaging: finite-sample rates and asymptotic normality. arXiv preprint arXiv:2305.17665, 2023. 19

[52] Bryan Van Scoy, Randy A Freeman, and Kevin M Lynch. The fastest known globally convergent first-order method for minimizing strongly convex functions. IEEE Control Systems Letters, 2(1):49–54, 2017. [53] Joaquin Vanschoren, Jan N Van Rijn, Bernd Bischl, and Luis Torgo. Openml: networked science in machine learning. ACM SIGKDD Explorations Newsletter, 15(2):49–60, 2014. [54] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. [55] Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409.11321, 2024. [56] Runzhe Wang, Sadhika Malladi, Tianhao Wang, Kaifeng Lyu, and Zhiyuan Li. The marginal value of momentum for small learning rate sgd. arXiv preprint arXiv:2307.15196, 2023. [57] Yan Yan, Tianbao Yang, Zhe Li, Qihang Lin, and Yi Yang. A unified analysis of stochastic momentum methods for deep learning. arXiv preprint arXiv:1808.10396, 2018. [58] Yun Zeng, Deren Han, Yansheng Su, and Jiaxin Xie. On adaptive stochastic heavy ball momentum for solving linear systems. SIAM Journal on Matrix Analysis and Applications, 45(3):1259–1286, 2024.

Convergence analysis for expected iterates

We start by recalling the matrix transition rule (11), ⊤ ⊤ V ∆t (1 + β)I − (1 + βω)V⊤ Πt V −β · (I − ωV⊤ Πt−1 V) V ∆t+1 · = V⊤ ∆t−1 V⊤ ∆t I 0 {z } | =Yt

where ω ∈ [0, 1] and ∆−1 = 0. As range(Πt ) ⊂ range(Π̄), we note that if ∆0 , ∆1 ∈ range(Π̄), then for all t, ∆t ∈ range(Π̄). In this section, we analyze the update rule (11) in expectation. Let T denote EYt . We have, (1 + β)I − (1 + βω)V⊤ Π̄V −β · (I − ωV⊤ Π̄V) T= . I 0 It is useful to transform T into a block-diagonal form. The following lemma guarantees that a permutation matrix P exists such that T can be transformed into a block-diagonal matrix.

Lemma 6 (Block-diagonalization via permutation) There exists a permutation matrix P such that, PTP⊤ = diag(T1 , T2 , · · · , Tn ), (1 + β) − (1 + βω)λi −β(1 − ωλi ) where Ti = , and 1 0  ⊤  v1 ∆t+1 ⊤  v1⊤ ∆t    V ∆t+1 ···  P = ⊤  . V ∆t vd⊤ ∆t+1  vd⊤ ∆t The resulting block-diagonal structure helps us throughout the analysis. As (PTP⊤ )t = PTt P⊤ and ∥PTt P⊤ ∥ = ∥Tt ∥, we need to upper bound ∥(PTP⊤ )t ∥. Due to the blockdiagonal form of PTP⊤ , we have ∥(PTP⊤ )t ∥ = maxi ∥Tti ∥. For any fixed i, let γi1 , γi2 ∈ C denote the two eigenvalues of Ti . We always take γi1 to be the eigenvalue with larger magnitude. In the next lemma, we provide a tight upper bound on the magnitude of eigenvalues of Ti . Lemma 7 (Bounding spectral radius of transition matrix) Let ω ∈ [0, 1] and β = 1 − ϕ1 for some ϕ ≥ 2. Then, the following holds: If Ti has complex eigenvalues then, 1 |γi2 |2 = |γi1 |2 = 1 − (1 − ωλi ). ϕ If Ti has real and equal eigenvalues then, 1 (1 − ωλi ). |γi2 |2 = |γi1 |2 = 1 − ϕ If Ti has real and distinct eigenvalues then, ϕλi 2 |γi2 |2 < |γi1 |2 < 1 − (1 − ωλi )2 . 2 Proof We have, (1 + β) − (1 + βω)λi −β(1 − ωλi ) Ti = 1 0 It is easy to verify that the eigenvalues of Ti , i.e., γi1 and γi2 are determined by the following expression. q 2 1 (γi1 , γi2 ) = (1 − λi ) + β(1 − ωλi ) ± (1 − λi ) + β(1 − ωλi ) − 4β(1 − ωλi ) . 2 The expression inside the square root can be negative, zero, or positive. So we consider three cases: 21

2 1. (1 − λi ) + β(1 − ωλi ) − 4β(1 − ωλi ) < 0: In this case, both γi1 and γi2 are complex and have the same magnitude. We get, 1 (1 − ωλi ). (17) |γi2 |2 = |γi1 |2 = β(1 − ωλi ) = 1 − ϕ 2 2 2 = 14 (1 − λi ) + β(1 − = γi2 2. (1 − λi ) + β(1 − ωλi ) − 4β(1 − ωλi ) = 0. In this case, γi1 2 ωλi ) = β(1 − ωλi ) . Therefore, 2 2 = = γi1 γi2

1 1− (1 − ωλi ). ϕ

(18)

2 3. (1 − λi ) + β(1 − ωλi ) − 4β(1 − ωλi ) > 0. The larger eigenvalue is given as, 1 γi1 = (1 − λi ) + β(1 − ωλi ) ± 2

2 (1 − λi ) + β(1 − ωλi ) − 4β(1 − ωλi ) .

(19)

1−λi , Let ui = 1−ωλ i

r 2 1 4β γi1 = ui + β + ui + β − (1 − ωλi ) 2 1 − ωλi s 1 2β(1 − λi ) 4β 2 2 = − ui + β + ui + β + (1 − ωλi ) 2 1 − ωλi 1 − ωλi s 1 2β(1 + λi ) 2 2 = ui + β + ui + β − (1 − ωλi ) 2 1 − ωλi s 2β(1 − λi ) 4βλi 1 2 2 ui + β + ui + β − − (1 − ωλi ) = 2 1 − ωλi 1 − ωλi r 4βλi 1 = ui + β + (ui − β)2 − (1 − ωλi ) 2 1 − ωλi q 4βλi Let x = ui − β, y = 1−ωλ , then we have x = yr , where r = 2x > 2. Now, if x, y ∈ R, 2 y i p x > y > 0 and x ≤ yr for some r > 2, then x2 − y 2 ≤ x − yr . Therefore, 2 y 1 ui + β + ui − β − (1 − ωλi ) 2 r y2 = ui − (1 − ωλi ) 4x y2 = 1− (1 − ωλi ). 4(1 − β)

γi1 <

The last inequality holds as ui ≤ 1 and x ≤ 1 − β. Therefore, we get ϕβλi (1 − ωλi ) γi1 < 1 − 1 − ωλi Now for β ≥ 0.5, we get ϕλi γi1 < 1 − (1 − ωλi ) 2

(20)

Combining the relations (17), (18) and (20) we finish the proof.

Remark 3 Lemma 7 implies that NAG momentum can perform slightly better than HBM as ω = 1 for NAG and ω = 0 for HBM. We demonstrate this phenomenon empirically in Fig 2, where we compare iteration complexity of block coordinate descent with momentum for various values of ω. We are now ready for the analysis of expected iterates. As ∆t ∈ range(Π̄) for all t, if for some i, vi ∈ Null(Π̄), then vi⊤ ∆t = 0. Let r denote the rank of Π̄ and note that vi⊤ ∆t = 0 , √1 , where C for all t and i > r. Let ϕ be any fixed constant satisfying 2 ≤ ϕ ≤ min m C 3 λr is some absolute constant (to be determined later) and m is as it appears in Proposition 2. Let the momentum parameter β = 1 − ϕ1 . We prove the following result: Theorem 3 Let ω ∈ [0, 1]. Then, the matrix transition rule (11) satisfies, ∥E∆t+1 ∥2 ≤ 17 · (3t + 2)2 · ρt−1 .∥∆0 ∥2 , where ρ = 1 − ϕλ2 r (1 − ωλr ). We will need the following auxiliary lemmas to finish the proof of Theorem 3. Lemma 8 (Two-step matrix transition rule) We have ∥E∆t+1 ∥2 ≤ 17 · max ∥Tti ∥ · ∥∆0 ∥2 . 1≤i≤r

Proof Let t ≥ 1. The update rule (11) satisfies, ⊤ ⊤ V ∆t+1 V ∆t ⊤ EP = PTP · EP . V⊤ ∆t V⊤ ∆t−1 Due to the block diagonal structure of PTP⊤ , we get ⊤ ⊤ vi ∆t+1 vi ∆t E = Ti · E ⊤ vi⊤ ∆t vi ∆t−1 23

Recursively applying the above relation we get ⊤ ⊤ vi ∆t vi ∆t+1 t = Ti · E ⊤ E vi ∆t−1 vi⊤ ∆t Therefore, ∥Evi⊤ ∆t+1 ∥2 ≤ ∥Evi⊤ ∆t+1 ∥2 + ∥Evi⊤ ∆t ∥2 ≤ ∥Tti ∥2 · ∥Evi⊤ ∆1 ∥2 + ∥vi⊤ ∆0 ∥2 .

(21)

As by notation ∆−1 = 0, we have Evi⊤ ∆1 = 1 + β − (1 + βω)λi vi⊤ ∆0 ∥Evi⊤ ∆1 ∥2 ≤ 16∥vi⊤ ∆0 ∥2 . Substituting the bound for ∥Evi⊤ ∆1 ∥2 in (21) we get ∥Evi⊤ ∆t+1 ∥2 ≤ 17∥Tti ∥2 · ∥vi⊤ ∆0 ∥2 . Now summing over i and noting that vi⊤ ∆t = 0 for all i > r, we get ∥EV⊤ ∆t+1 ∥2 ≤ 17 · max ∥Tti ∥ · ∥V⊤ ∆0 ∥2 . 1≤i≤r

Nothing that V has orthonormal columns finishes the proof. To complete the proof of Theorem 3, we need to bound the spectral norm of Tti . Note that Lemma 7 bounds the spectral radius of Tti , not the spectral norm. For bounding the spectral norm, we use the following fundamental matrix decomposition result. For any complex valued matrix U, let UH denote the conjugate transpose of U. Lemma 9 (Schur’s decomposition, Theorem 2.3.1 in [43]) For every i, there exists a unitary matrix Ui such that γi1 xi H , Ui Ti Ui = 0 γi2 where γi1 and γi2 are the two eigenvalues of Ti and xi is a constant such that 1 ≤ |xi | ≤ 3. The unitary matrix, Ui is explicitly given as 1 γi1 −1 Ui = p · , 1 γ̄i1 1 + |γi1 |2 where γ̄i1 denotes complex conjugate of γi1 . Furthermore, for any k ≥ 1, k P j k−1−j γi1 xi k−1 γi2 γi1 H k j=0 Ui Ti Ui = k 0 γi2

Corollary 1 If γi1 ̸= γi2 , then xi |xi

k−1 X

j k−1−j k−1 1−(γi2 /γi1 )k = xi γi1 · 1−(γi2 /γi1 ) , implying j=0 γi2 γi1

Pk−1

j k−1−j γi2 γi1 | < |xi | · |γi1 |k ·

j=0

Corollary 2 If γi1 = γi2 , then xi |xi

3 . |γi1 − γi2 |

Pk−1

k−1 X

j k−1−j k−1 = xi γi1 · k, implying, j=0 γi2 γi1

j k−1−j γi2 γi1 | < |xi | · |γi1 |k ·

j=0

k . |γi1 |

Corollary 3 For any i, k k−1 ∥Tki ∥ = ∥UH . i Ti Ui ∥ ≤ (3k + 2) · |γi1 |

Finishing proof of Theorem 3. Using Corollary 3 we get ∥Tti ∥2 ≤ (3t + 2)2 · |γi1 |2(t−1) . 2 Furthermore, if λ is such that (1−λ )+β(1−ωλ ) −4β(1−ωλr ) < 0, then we have |γi1 |2 ≤ r r r 2 1 1 − 2ϕ (1 − ωλr ) for i ≤ r. On the other hand if, (1 − λi ) + β(1 − ωλi ) − 4β(1 − ωλi ) < 0, then |γi1 |2 < 1− ϕλ2 r (1−ωλr ) for i ≤ r. Now the statement follows by noting that ϕ < √1λr .

Convergence in ℓ2-norm with mini-batch averaging

In this section, we aim to recover a result similar to Theorem 3 but in ℓ2 norm. Recall the transformed update rule, ⊤ ⊤ V ∆t+1 V ∆t ⊤ P = PYt P · P , V⊤ ∆t V⊤ ∆t−1 As vi⊤ ∆t = 0 for i > r, we can restrict the dynamics to the top-r dimensional space. Let Q ∈ R2r×2d consist of first 2r rows of 2d × 2d identity matrix. We get, ⊤ ⊤ V ∆t+1 V ∆t ⊤ ⊤ ⊤ QP = QPYt P Q · QP V⊤ ∆t V⊤ ∆t−1 Recursively we get, ⊤ j=t−1 ⊤ Y V ∆t+1 V ∆1 ⊤ ⊤ QP = QPYt−j P Q · QP . ⊤ V ∆t V⊤ ∆0 j=0

Therefore, t t−1 h Y Y i E∥∆t+1 ∥2 ≤ 17 · E QPYj⊤ P⊤ Q⊤ QPYt−j P⊤ Q⊤ · ∥∆0 ∥2 . j=1

j=0

h Q Q i t t−1 ⊤ ⊤ ⊤ ⊤ ⊤ Our task is to bound E . We prove the j=1 QPYj P Q j=0 QPYt−j P Q following theorem, which is only a restatement of Theorem 1. 25

Theorem 4 (General result) Let ϕ satisfy 2 ≤ ϕ ≤ min update rule (11) with β = 1 − ϕ1 satisfies,

m C

, 3√1λr

for C > 76800. Then,

E∥∆t ∥2 ≤ 170t3 · ρt−1 · ∥∆0 ∥2 , where ρ = 1 − ϕλ4 r . Note that Theorem 1 follows from Theorem 4 by taking c1 = 2 · 76800, c2 = c1 /3.

B.1

Proof of Theorem 4

We prove a series of results, building towards the proof of Theorem 4. Let Zt = V⊤ Πt V, and Z̄ = EZt . Note that Z̄ = Λ, where Λ = diag(λ1 , λ2 , · · · λd ). We have, E(Zt − Z̄)2 = E(Z2t ) − Z̄2 = E(Z2t ) − Λ2 = V⊤ E(Π2t )V − Λ2 As by Definition 1, Πt ⪯ I, we always have E(Π2t ) ⪯ Π̄, implying E(Zt − Z̄)2 ⪯ Λ − Λ2 . However, upper bounding E(Π2t ) by E Πt can be quite suboptimal, for instance, when Πt is itself an average of multiple independent contractions. For proving a tighter bound on the second moment of Zt , we use Proposition 2. Due to Proposition 2, we have, 2 2 1 1 · V⊤ Π̄(I − Π̄)V = · Λ(I − Λ) E Zt − Z̄ = V⊤ E Πt −Π̄ V ⪯ m m

(22)

Lemma 10 Let P be the permutation matrix as in Lemma 6. Then, we have, P · E (Yt − EYt )⊤ (Yt − EYt ) · P⊤ ⪯ 4 · diag(W1 , W2 , · · · , Wn ), 2 i ) ω λi (1−λi ) . where Wi = diag λi (1−λ , m m Proof We have Zt = V⊤ Πt V, Zt−1 = V⊤ Πt−1 V and Z̄ = EZt = EZt−1 . We have, (1 + β)I − (1 + βω)Zt −β(I − ωZt−1 ) (1 + β)I − (1 + βω)Z̄ −β(I − ω Z̄) Yt = , EYt = , I 0 I 0 and,

(1 + βω) Z̄ − Zt ) 0 (1 + βω) Z̄ − Zt ) −ωβ(Z̄ − Zt−1 ) E[(Yt − EYt ) (Yt − EYt )] = E · −ωβ(Z̄ − Zt−1 ) 0 0 0 (1 + βω)2 (Zt − Z̄)2 −ωβ(1 + βω)(Zt − Z̄)(Zt−1 − Z̄) =E −ωβ(1 + βω)(Zt−1 − Z̄)(Zt − Z̄) ω 2 β 2 (Zt−1 − Z̄)2 2 (1 + βω)2 · E Zt − Z̄ 0 = . 0 ω 2 β 2 · E(Zt−1 − Z̄)2 ⊤

Using relation (22) we get, 4 Λ(I − Λ) 0 · E[(Yt − EYt ) (Yt − EYt )] ⪯ 0 ω 2 Λ(I − Λ) m ⊤

The last step is to observe that permuting the rows and columns using permutation matrix P yields the result. In the next result, we construct a recursive framework that will help us in upper bounding the spectral norm of the product of random matrices Yt . As already exploited in Lemma 7 and Lemma 8, it is convenient to do the analysis on the transformed matrices (by using permutation P). Consider the following notations: Σt =

t−1 Y

⊤

QPYt−j P Q ,

Σ̄t =

j=0

t−1 Y

QPEYt P⊤ Q⊤ , Rt = QP(Yt − EYt )P⊤ Q⊤ , Ȳ = QPEYt P⊤ Q⊤ ,

j=0

where P is the permutation matrix such that PEYt P⊤ = diag(T1 , T2 , · · · , Tn ). We prove the following lemma: Lemma 11 For any t ≥ 1, there exist p0 , p1 , ..., pt such that p0 = 1 and, E[Σ⊤t Σt ] ≤ pt ≤ max

1≤i≤r

(T⊤i )t Tti +

t−1 4λi (1 − λi ) X · pj · (T⊤i )t−1−j Tt−1−j . i m j=0

Proof We provide a proof by doing induction over t. Note that t = 0 forms the base case, and as Σ0 = I, the base case hold for p0 = 1. Now for k < t, let E[Σ⊤k Σk ] ≤ pk , for pk satisfying the the inequality in statement of Lemma 11, and let Mk = E[Σ⊤k Σk ] − Σ̄⊤k Σ̄k . For Σt we have, E[Σ⊤t Σt ] = E QPYt⊤ P⊤ Q⊤ E[Σ⊤t−1 Σt−1 ]QPYt P⊤ Q⊤ = E QPEYt⊤ P⊤ Q⊤ E[Σ⊤t−1 Σt−1 ]QPEYt P⊤ Q⊤ + E R⊤t E[Σ⊤t−1 Σt−1 ]Rt = E Ȳ⊤ E[Σ⊤t−1 Σt−1 ]Ȳ + E R⊤t E[Σ⊤t−1 Σt−1 ]Rt = Σ̄⊤t Σ̄t + Ȳ⊤ E[Σ⊤t−1 Σt−1 ] − Σ̄⊤t−1 Σ̄t−1 Ȳ + E R⊤t E[Σ⊤t−1 Σt−1 ]Rt Rearranging the terms, we get, E[Σ⊤t Σt ] − Σ̄⊤t Σ̄t = E R⊤t E[Σ⊤t−1 Σ⊤t−1 ]Rt + Ȳ⊤ E[Σ⊤t−1 Σt−1 ] − Σ̄⊤t−1 Σ̄t−1 Ȳ. {z } | Mt

By induction hypothesis E[Σ⊤t−1 Σt−1 ] ≤ pt−1 , we get, Mt ⪯ pt−1 · E[R⊤t Rt ] + Ȳ⊤ Mt−1 Ȳ, E[Σ⊤t Σt ] ⪯ Σ̄⊤t Σ̄t + pt−1 · E[R⊤t Rt ] + Ȳ⊤ Mt−1 Ȳ, 27

(23) (24)

Note that E[R⊤t Rt ] = M1 . Using Lemma 10, it follows that M1 is upper bounded by a block-diagonal matrix. Let us define D1 := 4 · diag(W1 , W2 , · · · Wr ) where Wi′ s are as given in Lemma 10, and for k > 1, Dk := pk−1 · D1 + Ȳ⊤ Dk−1 Ȳ. Note that D1 is blockdiagonal and M1 ⪯ D1 . Let’s proceed inductively and assume that Dk−1 is block-diagonal and Mk−1 ⪯ Dk−1 . By definition, Dk is block-diagonal as Ȳ is block-diagonal and by induction hypothesis Dk−1 is block-diagonal. Furthermore, by (23), Mk ⪯ Dk as M1 ⪯ D1 and Mk−1 ⪯ Dk−1 by induction hypothesis. Therefore, by induction, ∀k, Mk ⪯ Dk , and D′k s are block-diagonal satisfying Dk = pk−1 · D1 + Ȳ⊤ Dk−1 Ȳ. We now write the recursive expression for diagonal blocks in Dt . Let the ith diagonal block of Ȳ be Ti , and ith diagonal block of Dt be Dt,i . This yields, Dt,i = pt−1 · D1,i + T⊤i Dt−1,i Ti . Unfolding the above recursion, we get, Dt,i = pt−1 · D1,i + T⊤i (pt−2 · D1,i + T⊤i Dt−2,i Ti )Ti = pt−1 · D1,i + pt−2 · T⊤i D1,i Ti + (T⊤i )2 Dt−2,i T2i . Continuing the recursion, we get t−1 X

Dt,i =

pj · (T⊤i )t−1−j D1,i Tt−1−j , or i

j=0 t−2 X

Dt−1,i =

pj · (T⊤i )t−2−j D1,i Tt−2−j . i

j=0

Substituting in (24), we bound EΣ⊤t Σt as ⊤

EΣt Σt ⪯ max

1≤i≤r

(Ti ) Tti + pt−1 · D1,i + ⊤

t−2 X

pj · (T⊤i )t−1−j D1,i Tt−1−j i

j=0

⪯ max (T⊤i )t Tti + 1≤i≤r

t−1 X

pj · (T⊤i )t−1−j D1,i Tt−1−j . i

j=0

Using triangle inequality and norm submultiplicativity, we get, ⊤

EΣt Σt ≤ pt = max

1≤i≤r

(Ti ) Tti ⊤

+ ∥D1,i ∥ ·

t−1 X

pj · (T⊤i )t−1−j Tt−1−j i

j=0

The proof follows by substituting for ∥D1,i ∥ using Lemma 10. Now the task is to upper bound the recursive expression for pt in Lemma 11. In the next lemma, we explicitly bound the spectral norm of (T⊤i )k Tki for any i and any k ≥ 1. 28

Lemma 12 For any k ≥ 0, (T⊤i )k Tki ≤ 11 · |γi1 |2k · h2i (k). where hi (k) = min

, 2 |γi1 | |γi1 −γi2 |

if γi1 ̸= γi2 , and hi (k) = |γki1 | if γi1 = γi2 .

Proof We use Lemma 9 to upper bound ∥Tki ∥2F . Note that, (T⊤i )k Tki ≤ ∥Tki ∥2F = |γi1 |2k + |xi |2 |gi (k)|2 + |γi2 |2k , P k j k−1−j i2 /γi1 ) | where gi (k) = k−1 . We have, |gi (k)| = |γi1 |k · min |γki1 | , |1−(γ j=0 γi1 γi2 |γi1 −γi2 | hi (k). Note that hi (k) is always larger than 1. Moreover, as |xi | ≤ 3, we get,

≤ |γi1 |k ·

(T⊤i )k Tki ≤ 11 · |γi1 |2k · h2i (k).

Now suppose t > 0 and our aim is to upper bound pt+1 . In the following analysis, we construct an upper bound for pt+1 , by defining ft+1 , which itself has a recursive definition, but a simpler one compared to pt+1 . We will use Lemma 11 and Lemma 12 to build this recursion. For convenience, consider the following notations: αi (t) := (T⊤i )t Tti , qi :=

n o 4λi (1 − λi ) , ℓi (t) := max 11 · |γi1 |k · h2i (k) , 0≤k≤t m

(25)

where hi (k) is defined in Lemma 12. We have p0 = 1, and pt+1 = max

1≤i≤r

⊤

(Ti )

t+1

Tt+1 i

+ qi ·

t X

⊤

pj · (Ti )

t−j

Tt−j i

j=0

t n o X < max αi (t + 1) + qi · pj · |γi1 |t−j · ℓi (t − j) 1≤i≤r

j=0

t n o X ≤ max αi (t + 1) + qi · ℓi (t) · pj · |γi1 |t−j . 1≤i≤r

j=0

Informed by the above relations, we construct upper bounds for pk for 0 ≤ k ≤ t+1, denoted by fk and defined as follows: f0 = 1,

fk+1 = max αi (k + 1) + qi · ℓi (t) · 1≤i≤r

k X j=0

k−j

fj · |γi1 |

for 0 ≤ k ≤ t.

Note that pk+1 ≤ fk+1 for all 0 ≤ k ≤ t. For any fixed k such that 0 ≤ k ≤ t, k X fk+1 = max αi (k + 1) + qi · ℓi (t) · fj · |γi1 |k−j 1≤i≤r

j=0

= max αi (k + 1) + qi · ℓi (t) · fk + |γi1 | · qi · ℓi (t) · 1≤i≤r

k−1 X

fj · |γi1 |k−1−j

j=0

{z ∗

}

Now note, irrespective of the index i, the term marked as ∗ is always upper bounded by fk . Therefore, fk+1 < max αi (k + 1) + |γi1 | + qi · ℓi (t) · fk 1≤i≤r ≤ max |γi1 | + qi · ℓi (t) · fk + max αi (k + 1) 1≤i≤r

1≤i≤r

≤ ρ · fk + max αi (k + 1). 1≤i≤r

where we let ρ := max1≤i≤r |γi1 | + qi · ℓi (t) . Now, using the above recursive relation we can upper bound ft+1 as, ft+1 ≤ ρ · ft + max αi (t + 1) 1≤i≤r

≤ ρ · ft−1 + ρ · max αi (t) + max αi (t + 1) 1≤i≤r

≤ ρt+1 +

t X j=0

1≤i≤r

ρj · max αi (t + 1 − j). 1≤i≤r

(26)

We need to bound ρ. In the following lemma, we provide a bound on the term |γi1 | + qi · ℓi (t) for any i. For simplicity, we provide the proof for ω ∈ {0, 1} (which covers both HBM and NAG). The proof for ω ∈ (0, 1) is completely analogous and follows by using a more general version of Lemma 14. , √1 for C > 76800 and ω ∈ [0, 1]. Then for any Lemma 13 Let 2 ≤ ϕ ≤ min m C 3 λr 1 ≤ i ≤ r, |γi1 | + r · ℓi (t) ≤ 1 −

ϕλr . 4

Proof As we are proving the result for a fixed i, we drop the subscript i for notational convenience. The proof proceeds in cases: 1. Case 1: γ1 is complex. We know from Lemma 7 that |γ1 |2 = 1 − ϕ1 (1 − ωλ). We get 1 4λ(1 − λ) |γ1 | + q · ℓ(t) < 1 − + · ℓ(t). 2ϕ m 30

(27)

o n o n 2 . Let k ∗ be the Recall ℓ(t) = max0≤k≤t 11 · |γ1 |k · h2 (k) and h(k) = min |γk1 | , |γ1 −γ 2| ∗

index such that ℓ(t) = 11 · |γ1 |k · h2 (k ∗ ). We get, 1 44λ(1 − λ) 2 ∗ ∗ |γ1 | + r · ℓ(t) < 1 − + · h (k ) · |γ1 |k . 2ϕ m Now look at h2 (k ∗ ), h2 (k ∗ ) = min

n (k ∗ )2

o 4 . , |γ1 |2 |γ1 − γ2 |2

Due to Lemma 7, we know 2 1 4β(1 − ωλ) − (1 − λ) + β(1 − ωλ) 4 (1 − λ) = · (4βδ − (1 + βδ)2 (1 − λ) , 4

|γ1 − γ2 |2 =

2 2 . Also, as β = 1 − ϕ1 ≥ 0.5, we have |γ11 |2 ≤ 1−ωλ ≤ 1−λ . Therefore, where δ = 1−ωλ 1−λ

n o 1 44λ 16 ∗ |γ1 | + q · ℓ(t) < 1 − + · min 2(k ∗ )2 , · |γ1 |k (28) 2 2ϕ m 4βδ − (1 + βδ) (1 − λ) Let ϵ = 4βδ − (1 + βδ)2 (1 − λ). We get, 4βδ ϵ − 2 (1 + βδ) (1 + βδ)2 ϵ 4βδ + λ=1− (1 + βδ)2 (1 + βδ)2 (1 − βδ)2 ϵ = + 2 (1 + βδ) (1 + βδ)2

1−λ=

(29)

We consider two subcases: 2

(a) Subcase-1: ϵ > (1−βδ) . Using (29) we get, (1+βδ)2 λ 1 <1+ < 2. 2 4βδ − (1 + βδ) (1 − λ) (1 + βδ)2 (b) Subcase-2:

(1−βδ)2 ϵ ≤ (1+βδ) 2.

Using (29) we have λ ≤

1 1 + (1+βδ) 2

(30)

(1−βδ)2 . (1+βδ)2

Note that

1 δ ∈ {1, 1−λ } and b ≥ 0.5. Using item 1 from Lemma 14 with x = δ we get,

λ≤ 1+

(1 − β)2 1 3 < , 2 2 (1 + β) (1 + β) 2(2ϕ − 1)2 31

(as β = 1 −

1 ). ϕ

We get, ∗

λ · (k ∗ )2 · |γ1 |k ≤

3 ∗ · (k ∗ )2 · |γ1 |k . 2 2(2ϕ − 1)

∗

The quantity (k ∗ )2 · |γ1 |k is maximized at k ∗ = log−2|γ1 | , implying ∗

λ · (k ∗ )2 · |γ1 |k ≤

3 4 1 · · , 2(2ϕ − 1)2 (1 − |γ1 |)2 e2 ∗

2 and |γ1 |k = e−2 , where in the last inequality we used the relation log−2|γ1 | < 1−|γ 1| 1 , we have (1−|γ1 1 |)2 < 4ϕ2 . Therefore, with e being Euler’s number. As |γ1 | < 1 − 2ϕ ∗

λ · (k ∗ )2 · |γ1 |k ≤

3 16ϕ2 3 16ϕ2 · ≤ · 2 < 2. 2(2ϕ − 1)2 e2 4ϕ2 e

(31)

Depending on the above two subcases, we substitute (30) and (31) in (28) to prove 1 44 · 32 + . |γ1 | + q · ℓ(t) < 1 − 2ϕ m m If ϕ ≤ 5632 , we get,

|γ1 | + q · ℓ(t) < 1 −

1 . 4ϕ

(32)

2. Case 2: γ1 is real and γ1 = γ2 . In this case we have, λ=

(1 − βδ)2 (1 − β)2 1 ≤ = , 2 2 (1 + βδ) (1 + β) (2ϕ − 1)2

where the last inequality follows from item 2 from Lemma 14. The proof is exactly 2 same to subcase-2 of the first case (complex γ1 ). Using Lemma 7 we have, γ1 = 1 − ϕ1 (1 − ωλ). Therefore, 1 4λ(1 − λ) |γ1 | + q · ℓ(t) < 1 − + · ℓ(t). 2ϕ m n o Recall ℓ(t) = max0≤k≤t 11 · |γ1 |k · h2 (k) and h(k) = |γk1 | . Let k ∗ be the index such

∗

that ℓ(t) = 11 · |γ1 |k · h2 (k ∗ ). We get, 1 44λ(1 − λ) (k ∗ )2 ∗ + · · |γ1 |k |γ1 | + q · ℓ(t) < 1 − 2 2ϕ m |γ1 | r 1 176λ 1 ∗ 2 k∗ 2 < 1− + · (k ) · |γ1 | as γ1 > (1 − λ) 1 − 2ϕ m ϕ 2 1 176 1 16ϕ < 1− + · · 2 2 2ϕ m (2ϕ − 1) e 1 200 < 1− + . 2ϕ m 32

m Therefore, for ϕ < 400 , we get

|γ1 | + q · ℓ(t) < 1 −

1 . 4ϕ

(33)

3. Case 3: γ1 is real with γ1 > γ2 . Using Lemma 7, we have γ1 < 1 − ϕλ (1 − ωλ). 2 Therefore, ϕλ 4λ(1 − λ) γ1 + q · ℓ(t) < 1 − + · ℓ(t). 2 m o n o n 2 . Let k ∗ be Here again, ℓ(t) = max0≤k≤t 11 · |γ1 |k · h2 (k) and h(k) = min γk1 , γ1 −γ 2

∗

the index such that ℓ(t) = 11 · |γ1 |k · h2 (k ∗ ). We get, ϕλ 44λ(1 − λ) 2 ∗ ∗ γ1 + r · ℓ(t) < 1 − + · h (k ) · γ1k , 2 m

(34)

where, 2

∗

h (k ) = min

n (k ∗ )2 γ12

o 4 , . (γ1 − γ2 )2

We have 2 (1 − λ) + β(1 − ωλ) − 4β(1 − ωλ) = (1 − λ) (1 + βδ)2 (1 − λ) − 4βδ ,

(γ1 − γ2 )2 =

where δ = 1−ωλ . Let ϵ = (1 + βδ)2 (1 − λ) − 4βδ. Note that this yields 1−λ λ=

(1 − βδ)2 ϵ − . (1 + βδ)2 (1 + βδ)2

(1−βδ) ′ Moreover, let a(λ) = λ − (1+βδ) 2 and let λ denote the zero of a(λ) in (0, 1). It turns out that a(λ) is an increasing function for all λ satisfying a(λ) < 0. As λ in our case satisfies a(λ) < 0, we get

(1 − β)2 1 1 λ<λ < = ≤ . 2 2 (1 + β) (2ϕ − 1) 9 ′

(35)

where we used item 2 from Lemma 14 again for the second last inequality. The last inequality used ϕ ≥ 2. Consider the following two subcases:

(1−βδ) (a) Subcase 1: 0 < ϵ ≤ (1+βδ) 2 . We get,

λ≥ 1−

(1 − βδ)2 1 . (1 + βδ)2 (1 + βδ)2

Using item 3 from Lemma 14 we get λ>

(1 − β)2 1 1 = > . 2(1 + β)2 2(2ϕ − 1)2 8ϕ2

(36)

We know that (see proof of Lemma 7), 1 3 2 1 (1 − λ) + β(1 − ωλ) > (1 + β)(1 − λ) > (1 − λ) > (by (35)) γ1 > 2 2 4 3 Therefore, ϕλ 44λ(1 − λ) (k ∗ )2 k∗ + · 2 · γi γ1 + q · ℓ(t) < 1 − 2 m γ1 ϕλ 100λ 2 ∗ < 1− + · (k ∗ )2 · γ1k (as γ1 > ). 2 m 3 ∗

As (k ∗ )2 · γ1k is maximized at k ∗ = log−2γ1 and γ1 < 1 − ϕλ , we get 2 ϕλ 100λ 3 γ1 + q · ℓ(t) < 1 − + · 2 2. 2 m ϕλ Finally using (36) we get, 1 2400 γ1 + q · ℓ(t) < 1 − + 16ϕ m 1 <1− , 32ϕ

(37)

m where last inequality follows for ϕ ≤ 76800 . 2

(1−βδ) 2 (b) Subcase 2: ϵ > (1+βδ) 2 . In this case we can simply lower bound (γ1 − γ2 ) . Note that, (1 − βδ)2 1 λ< 1− · . (1 + βδ)2 (1 + βδ)2 2

Using item 4 from Lemma 14 with x = δ we get (1−βδ) ≥ 12 · (1−β) . Therefore, (1+βδ)2 (1+β)2 (γ1 − γ2 )2 = (1 − λ) · ϵ (1 − βδ)2 > (1 − λ) · (1 + βδ)2 (1 − β)2 1 ≥ (1 − λ) · = (1 − λ) · . 2 (1 + β) (2ϕ − 1)2 34

Substituting the above in (34) we get, ϕλ 176λ + · (2ϕ − 1)2 2 m ϕλ 704λϕ2 + <1− 2 m

γ1 + q · ℓ(t) < 1 −

m For ϕ < 2816 we get,

γ1 + q · ℓ(t) < 1 −

ϕλ . 4

(38)

Combining the inequalities (31), (32), (37) and (38) we get for ϕ ≤ min{ m , √1 } and C 3 λr C > 76800, |γ1 | + r · ℓ(t) ≤ 1 −

ϕλ . 4

This finishes the proof.

Finishing the proof of Theorem 4. We are now ready to complete the proof of Theorem 4. Recall the bound on pt+1 in (26). We had, t+1

pt+1 ≤ ft+1 < ρ

t X

ρj · max αi (t + 1 − j) i

j=0

Due to Lemma 13 we know, ρ<1−

ϕλr . 4

Furthermore, αi (t + 1 − j) = (T⊤i )t+1−j Tt+1−j , and we know from Lemma 9 that i 2 (T⊤i )t+1−j Tt+1−j ≤ 3(t + 1 − j) + 2 · |γi1 |t−j i < 3(t + 1)2 + 2 · |γi1 |t−j < 3(t + 1)2 + 2 · ρt−j . Therefore, pt+1 < ρt+1 + (t + 1) · 3(t + 1)2 + 2) · ρt < 2(t + 1) · 3(t + 1)2 + 2 · ρt < 10(t + 1)3 · ρt . h Q Q i t t−1 ⊤ ⊤ ⊤ ⊤ ⊤ Note that pt upper bounds E QPY P Q QPY P Q and this fint−j j j=1 j=0 ishes the proof of Theorem 4. The following helper lemma is used in the proof of Lemma 13 for ω = 0. A more general variant can be used to prove Lemma 13 for any ω ∈ (0, 1). The proof of Lemma 14 follows by a routine calculus exercise. 35

Lemma 14 Let x ≥ 1, 12 ≤ β < 1, and define (1 − βx)2 (1 − βx)2 1 1 , g(x) = 1 − f (x) = 1 + (1 + βx)2 (1 + βx)2 (1 + βx)2 (1 + βx)2 1. Over the set of all x ≥ 1 satisfying 1 − x1 ≤ f (x), the function f (x) is maximized at 2 x = 1. In particular, f (x) < 32 · (1−β) . (1+β)2 2

(1−βx) (1−βx) (1−β) 2. Over the set of all x ≥ 1 satisfying 1 − x1 = (1+βx) 2 , we have (1+βx)2 ≤ (1+β)2 . 2

(1−βx) 4 (1−β) 3. Over the set of all x ≥ 1 satisfying 1 − x1 < (1+βx) 2 , inf x g(x) > 5 · (1+β)2 . 2

4. Over all the set of all x ≥ 1 satisfying 1 − x1 < g(x), inf x (1−βx) > 12 · (1−β) . (1+βx)2 (1+β)2

Additional experiments

In this section, we provide additional details about our experimental setup, and include numerical results on two datasets to complement the results in Section 5. Finally, we also explore the dependence of classical momentum on the ω parameter, which interpolates between HBM and NAG.

C.1

Experimental setup

We transform the data matrices through StandardScaler from sklearn.preprocessing. In particular, after sampling n = 2048 rows from the dataset, we perform the following transformation, scaler = StandardScaler() X_scaled = scaler.fit_transform(X) K = rbf_kernel(X_scaled, gamma=0.1) We consider the following values for λ and ϵ tolerance in our experiment: 1. Covtype, λ = 0.5, ϵ = 10−7 . 2. California Housing, λ = 0.16, ϵ = 10−7 . 3. Abalone, λ = 20, ϵ = 10−10 . 4. Phoneme, λ = 10−10 , ϵ = 10−6 . All experiments were conducted on a MacBook Air (M3) with 16 GB of memory.

Block size=10

Block size=20

104 103 102

1 2 4 10

250 1 2 4 10

Mini-batch size (m)

250 1 2 4 10

250

No. of iterations to convergence

Block size=1

105

Block size=1

108

Block size=10

107 106 105 104

1 2 4 10

250 1 2 4 10

Mini-batch size (m)

250 1 2 4 10

250

CDpp

Block size=10

Block size=25

Block size=50

105 104 103 1 2 4 10 50 250 1000 1 2 4 10 50 250 1000 1 2 4 10 50 250 1000

Mini-batch size (m)

(b) Phoneme: iterations vs mini-batch size

Block size=20 No. of row sampled processed

No. of row sampled processed

(a) Abalone: iterations vs mini-batch size

106

1010

Block size=10

Block size=25

Block size=50

109 108 107 106 105 104

1 2 4 10 50 2501000 1 2 4 10 50 2501000 1 2 4 10 50 2501000

Mini-batch size (m)

(d) Phoneme: row samples vs mini-batch size

CD + NAG (β = 1 − 1/m)

CD + NAG (β = adaptive)

Figure 3: Convergence of block coordinate descent with NAG momentum and baselines. Top row: iterations to reach error 10−10 for Abalone and 10−6 for Phoneme. Bottom row: total work, measured by the number of sampled rows.

C.2

Plots on Abalone and Phoneme datasets

In Fig 3, we compare, CD, CDpp, CD+NAG (β = 1 − 1/m), and CD+NAG (β = adaptive) on two additional datasets. The results on Abalone dataset follow the same pattern as observed in Fig 1 for California Housing and Covtype datasets. However, the results on Phoneme dataset look different, as CDpp performs better than CD+NAG. We provide insights into this phenomenon in the following paragraph. Comparison between CDpp and CD+NAG. Under certain incoherenceassumptions on the data matrix K, the condition number of Π̄ can be bounded as O nκ̄k k , where k is the block-size and κ̄k denotes a certain stochastic condition number quantity (see [11] for details). In this complexity of CDpp with mini-batch size m can be shown qcase, the iteration p κ̄k to scale as Õ · nk + κ̄k nk , whereas the iteration complexity of CD+NAG scales as m p Õ κ̄mk · nk + κ̄k nk . Therefore, when m < κ̄k , CDpp can converge faster than CD+NAG, but when m ≥ κ̄k , the perfect parallelization phenomenon takes over and CD+NAG becomes faster. While this discussion requires some additional assumptions on the data, it provides some intuition for why one type of acceleration may be preferable over the other, depending

on the conditioning properties of the dataset.

C.3

Adaptive momentum

Next, we describe the implementation of CD+NAG (β = adaptive), where we select the momentum parameter β adaptively during the course of the optimization. Consider solving the positive definite linear system Kw = y using block coordinate descent with NAG momentum, and with mini-batch size m (here, for simplicity we absorb the regularizer λ into K). Our aim is to adaptively find ϕ∗ ∈ [2, m] such that β ∗ = 1 − ϕ1 is a near-optimal momentum parameter for mini-batch size m. We provide the full heuristic logic for adaptive momentum tuning in Algorithm 1. The matrix K̃ and vector b̃ are simply a few randomly sampled rows from K and b, used to approximate the true residuals. In all our experiments, we sampled 100 rows randomly from the data matrices for approximating the true residual.

Algorithm 1 Adaptive momentum tuning for CD+NAG Require: K, K̃, b, b̃, m, maximum iterations T Require: Warm-up length Twarm , check interval Tcheck 1: Initialize w−1 = 0, x−1 = 0 1 2: Set initial momentum β ← 1 − m ▷ Warm-up phase 3: for t = 0, . . . , Twarm − 1 do 4: xt+1 ← mini-batch CD step(wt , m) 5: wt+1 ← xt+1 + β(xt+1 − xt ) 6: end for

▷ Adaptive phase 1 7: β+ = 1 − m 8: β− = 1 − 12 9: (w+,t , x+,t ) ← (wt , xt ) 10: (w−,t , x−,t ) ← (wt , xt ) 1−β

11: while t < T and 1−β− > 2 do +

12: for i ∈ {+, −} do 13: xi,t+1 ← mini-batch CD step(wi,t , m) 14: wi,t+1 ← xi,t+1 + βi (xi,t+1 − xi,t ) 15: if t mod Tcheck = 0 then 16: ri ← ∥K̃xi,t+1 − b̃∥/∥b̃∥ 17: R ← r+ /r− 18: if R < 1 then p ← p + 1 19: else p ← 0 20: end if 21: if R ≥ 3 then ▷ Large momentum is too aggressive 1 1 , β ← 1 − 22: ϕ+ ← 12 1−β + ϕ+ + 23: w+,t+1 ← w−,t+1 , x+,t+1 ← x−,t+1 24: p←0 25: else if R ≤ 1/2 or p > 20 then ▷ Small momentum is too conservative 1 2 β− ← 1 − ϕ− 26: ϕ− ← 1−β− , 27: w−,t+1 ← w+,t+1 , x−,t+1 ← x+,t+1 28: p←0 29: end if 30: end if 31: end for 32: t←t+1 33: end while 34: β ← β+ 35: (wt , xt ) ← (w+,t , x+,t ) 36: while t < T do ▷ CD+NAG with optimal β 37: xt+1 ← mini-batch CD step(wt , m) 38: wt+1 ← xt+1 + β(xt+1 − xt ) 39: t←t+1 39 40: end while 41: return wt

Proofs omitted from Section 3

Proof of Proposition 1 As ∆t+1 = At (∆t ; α) = (I − α Πt )∆t , we have E∥∆t+1 ∥2 = ∆⊤t E[(I − α Πt )2 ]∆t ≤ ∆⊤t E[(I − α Πt )]∆t = ∆⊤t (I − αΠ̄)∆t .

(since 0 ⪯ Πt ⪯ I and 0 < α ≤ 1)

As ∆0 ∈ range(Π̄), all ∆t ∈ range(Π̄) with probability 1. Therefore, we get 1 E∥∆t+1 ∥2 ≤ 1 − ∥∆0 ∥2 , κ where κ = α/λ+ min (Π̄). Proof of Proposition 2 It is easy to see that m

[m]

1 X (i) = Π , m i=1

where Π(i) is the stochastic rate matrix associated with A(i) . Therefore, h i 2 E (Π[m] −Π̄)2 = (Π[m] )2 + Π̄ − E[Π[m] ] · Π̄ − Π̄ · E[Π[m] ] 2 2 = Π[m] − Π̄ m m X 1 1 X 2 (i) 2 E Π(i) Π(j) − Π̄ = 2· E Π + 2· m i=1 m i,j=1,i̸=j 2 1 1 2 · E Π(1) − · Π̄ m m 1 ⪯ · Π̄(I − Π̄). m

Record · ID 200500 · SHA-256 08522086b2b4884e
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.