Original: Accelerating SGD for Highly Ill-Conditioned Huge-Scale Online Matrix Completion
Revised: Accelerating SGD for Highly Ill-Conditioned Huge-Scale Online Matrix Completion

Original: Abstract
Revised: Abstract

Original: The matrix completion problem seeks to recover a d ⇥ d ground truth matrix of low rank r ⌧ d from observations of its individual elements.
Revised: The matrix completion problem seeks to recover a d ⇥ d ground truth matrix of low rank r ⌧ d from observations of its individual elements.

Original: Real-world matrix completion is often a huge-scale optimization problem, with d so large that even the simplest full-dimension vector operations with O ( d ) time complexity become prohibitively expensive.
Revised: Real-world matrix completion is often a huge-scale optimization problem, with d so large that even the simplest full-dimension vector operations with O ( d ) time complexity become prohibitively expensive.

Original: Stochastic gradient descent (SGD) is one of the few algorithms capable of solving matrix completion on a huge scale, and can also naturally handle streaming data over an evolving ground truth.
Revised: Stochastic gradient descent (SGD) is one of the few algorithms capable of solving matrix completion on a huge scale, and can also naturally handle streaming data over an evolving ground truth.

Original: Unfortunately,
Revised: Unfortunately,

Original: SGD experiences a dramatic slow-down when the underlying ground truth is illconditioned; it requires at least O (  log(1 / ✏ ))
Revised: SGD experiences a dramatic slow-down when the underlying ground truth is illconditioned; it requires at least O (  log(1 / ✏ ))

Original: iterations to get ✏ -close to ground truth matrix with condition number  .
Revised: iterations to get ✏ -close to ground truth matrix with condition number  .

Original: In this paper, we propose a preconditioned version of SGD that preserves all the favorable practical qualities of SGD for huge-scale online optimization while also making it agnostic to  .
Revised: In this paper, we propose a preconditioned version of SGD that preserves all the favorable practical qualities of SGD for huge-scale online optimization while also making it agnostic to  .

Original: For a symmetric ground truth and the Root Mean Square Error (RMSE) loss, we prove that the preconditioned SGD converges to ✏ -accuracy in O (log(1 / ✏ )) iterations, with a rapid linear convergence rate as if the ground truth were perfectly conditioned with  = 1 .
Revised: For a symmetric ground truth and the Root Mean Square Error (RMSE) loss, we prove that the preconditioned SGD converges to ✏ -accuracy in O (log(1 / ✏ )) iterations, with a rapid linear convergence rate as if the ground truth were perfectly conditioned with  = 1 .

Original: In ournumerical experiments, we observe a similar acceleration for ill-conditioned matrix completion under the root mean square error (RMSE) loss, Euclidean distance matrix (EDM) completion under pairwise square loss, and collaborative ﬁltering under the Bayesian Personalized Ranking (BPR) loss.
Revised: In our experiments, we observe a similar acceleration for item-item collaborative ﬁltering on the MovieLens25M dataset via a pair-wise ranking loss, with 100 million training pairs and 10 million testing pairs. [See supporting code at https://github.com/Hong-Ming/ScaledSGD .]

Original:  Introduction
Revised:  Introduction

Original: The matrix completion problem seeks to recover an underlying d ⇥ d ground truth matrix M of low rank r ⌧ d from observations of its individual matrix elements M i,j .
Revised: The matrix completion problem seeks to recover an underlying d ⇥ d ground truth matrix M of low rank r ⌧ d from observations of its individual matrix elements M i,j .

Original: The problem appears most prominently in the context of collaborative ﬁltering and recommendation system, but also numerous other applications.
Revised: The problem appears most prominently in the context of collaborative ﬁltering and recommendation system but also numerous other applications.

Original: In this paper, we focus on the symmetric and positive semideﬁnite variant of the problem, in which the underlying matrix M can be factored as M = ZZ T where the factor matrix Z is d ⇥ r , though that our methods have natural extensions to the nonsymmetric case.
Revised: In this paper, we focus on the symmetric and positive semideﬁnite variant of the problem, in which the underlying matrix M can be factored as M = ZZ T where the factor matrix Z is d ⇥ r , though that our methods have natural extensions to the nonsymmetric case.

Original: We note that the symmetric positive semideﬁnite variant is actually far more common in collaborative ﬁltering, due to the prevalence of item-item models, which enjoy better data (most platforms contain several orders of magnitude more users than items) and more stable recommendations (the similarity between items tends to change slowly over time) than user-user and user-item models.
Revised: We note that the symmetric positive semideﬁnite variant is actually far more common in collaborative ﬁltering, due to the prevalence of item-item models, which enjoy better data (most platforms contain several orders of magnitude more users than items) and more stable recommendations (the similarity between items tends to change slowly over time) than user-user and user-item models.

Original: For the full-scale, online instances of matrix completion that arise in real-world collaborative ﬁltering, stochastic gradient descent or SGD is the only viable algorithm for learning the underlying matrix M .
Revised: For the full-scale, online instances of matrix completion that arise in real-world collaborative ﬁltering, stochastic gradient descent or SGD is the only viable algorithm for learning the underlying matrix M .

Original: The basic idea is to formulate a candidate matrix of the form XX T with respect to a  learned factor matrix X 2 R d ⇥ r , and to minimize a cost function of the form  ( XX T  M ) .
Revised: The basic idea is to formulate a candidate matrix of the form XX T with respect to a 36th Conference on Neural Information Processing Systems (NeurIPS 2022). learned factor matrix X 2 R d ⇥ r , and to minimize a cost function of the form  ( XX T  M ) .

Original: Earlier work used the root mean square error (RMSE) loss k XX T  M k 2 F =
Revised: Earlier work used the root mean square error (RMSE) loss k XX T  M k 2 F =

Original: P i,j ( XX T  M ) 2 i,j , though later work have focused on pairwise losses like the BPR [1] that optimize for ordering and therefore give better recommendations.
Revised: P i,j ( XX T  M ) 2 i,j , though later work have focused on pairwise losses like the BPR [1] that optimize for ordering and therefore give better recommendations.

Original: For the RMSE loss, the corresponding SGD iterations with (rescaled) learning rate ↵ > 0 reads
Revised: For the RMSE loss, the corresponding SGD iterations with (rescaled) learning rate ↵ > 0 reads

Original: [Equation]
Revised: [Equation]

Original: where M ij is the sampled ( i, j ) -th element of the ground truth matrix M , and x i , x j and x i, + , x j, + denote the i -th and j -th rows of the current iterate X t and new iterate X t +1 .
Revised: where M ij is the sampled ( i, j ) -th element of the ground truth matrix M , and x i , x j and x i, + , x j, + denote the i -th and j -th rows of the current iterate X t and new iterate X t +1 .

Original: Pairwise losses like the BPR can be shown to have a similar update equation over three rows of X [1].
Revised: Pairwise losses like the BPR can be shown to have a similar update equation over three rows of X [1].

Original: Given that only two or three rows of X are accessed and updated at any time, SGD is readily accessible to massive parallelization and distributed computing.
Revised: Given that only two or three rows of X are accessed and updated at any time, SGD is readily accessible to massive parallelization and distributed computing.

Original: For very large values of d , the update equation (1) can be run by multiple workers in parallel without locks, with vanishing probability of collision [2].
Revised: For very large values of d , the update equation (1) can be run by multiple workers in parallel without locks, with vanishing probability of collision [2].

Original: The blocks of X that are more frequently accessed together can be stored on the same node in a distributed memory system.
Revised: The blocks of X that are more frequently accessed together can be stored on the same node in a distributed memory system.

Original: Unfortunately, the convergence rate of SGD can sometimes be extremely slow.
Revised: Unfortunately, the convergence rate of SGD can sometimes be extremely slow.

Original: One possible explanation, as many recent authors have pointed out [3–6], is that matrix factorization models are very sensitive to ill-conditioning of the ground truth matrix M .
Revised: One possible explanation, as many recent authors have pointed out [3–6], is that matrix factorization models are very sensitive to ill-conditioning of the ground truth matrix M .

Original: The number of SGD iterations grows at least linearly the condition number  , which here is deﬁned as the ratio between the largest and the r -th largest singular values of M .
Revised: The number of SGD iterations grows at least linearly the condition number  , which here is deﬁned as the ratio between the largest and the r -th largest singular values of M .

Original: Ill-conditioning causes particular concern because most realworld data are ill-conditioned.
Revised: Ill-conditioning causes particular concern because most realworld data are ill-conditioned.

Original: In one widely cited study [7], it was found that the dominant singular value accounts for only ⇡ 80% prediction accuracy, with diversity of individual preferences making up the remainder ill-conditioned singular values.
Revised: In one widely cited study [7], it was found that the dominant singular value accounts for only ⇡ 80% prediction accuracy, with diversity of individual preferences making up the remainder ill-conditioned singular values.

Original: Cloninger et al.
Revised: Cloninger et al.

Original: [8] notes that there are certain applications of matrix completion that have condition numbers as high as  = 10 15 .
Revised: [8] notes that there are certain applications of matrix completion that have condition numbers as high as  = 10 15 .

Original: This paper is inspired by a recent full-batch gradient method called ScaledGD [4, 9] and a closely related algorithm PrecGD [5] in which gradient descent is made immune to ill-conditioning in the ground truth by right-rescaling the full-batch gradient by the matrix ( X T X )  1 .
Revised: This paper is inspired by a recent full-batch gradient method called ScaledGD [4, 9] and a closely related algorithm PrecGD [5] in which gradient descent is made immune to ill-conditioning in the ground truth by right-rescaling the full-batch gradient by the matrix ( X T X )  1 .

Original: Applying this same strategy to the SGD update equation (1) yields the row-wise updates
Revised: Applying this same strategy to the SGD update equation (1) yields the row-wise updates

Original: [Equation]
Revised: [Equation]

Original: in which we precompute and cache the preconditioner P
Revised: in which we precompute and cache the preconditioner P

Original: =
Revised: =

Original: ( X T X )  1 ahead of time 1 , and update it after the iteration as
Revised: ( X T X )  1 ahead of time 1 , and update it after the iteration as

Original: [Equation]
Revised: [Equation]

Original: by making four calls to the Sherman–Morrison rank-1 update formula
Revised: by making four calls to the Sherman–Morrison rank-1 update formula

Original: [Equation]
Revised: [Equation]

Original: This way, the rescaled update equations use just O ( r 2 ) arithmetic operations, which for modest values of r is only marginally more than the O ( r ) cost of the unscaled update equations (1).
Revised: This way, the rescaled update equations use just O ( r 2 ) arithmetic operations, which for modest values of r is only marginally more than the O ( r ) cost of the unscaled update equations (1).

Original: Indeed, the nearest-neighbor algorithms inside most collaborative ﬁlters have exponential complexity with respect to the latent dimensionality r , and so are often implemented with r small enough for (1) and (2) to have essentially the same runtime.
Revised: Indeed, the nearest-neighbor algorithms inside most collaborative ﬁlters have exponential complexity with respect to the latent dimensionality r , and so are often implemented with r small enough for (1) and (2) to have essentially the same runtime.

Original: Here, we observe that the rescaled update equations (2) preserve essentially all of the practical advantages of SGD for huge-scale, online optimization: it can also be run by multiple workers in parallel without locks, and it can also be easily implemented over distributed memory.
Revised: Here, we observe that the rescaled update equations (2) preserve essentially all of the practical advantages of SGD for huge-scale, online optimization: it can also be run by multiple workers in parallel without locks, and it can also be easily implemented over distributed memory.

Original: The only minor difference is that separate copies of P should be maintained by each worker, and resynchronized once differences grow large.
Revised: The only minor difference is that separate copies of P should be maintained by each worker, and resynchronized once differences grow large.

Original: Contributions In this paper, we provide a rigorous proof that the rescaled update equations (1), which we name ScaledSGD, become immune to the effects of ill-conditioning in the underlying ground truth matrix.
Revised: Contributions In this paper, we provide a rigorous proof that the rescaled update equations (1), which we name ScaledSGD, become immune to the effects of ill-conditioning in the underlying ground truth matrix.

Original: For symmetric matrix completion under the root mean squared error (RMSE) loss function, regular SGD is known to have an iteration count of O (  4 · dr log( d/ ✏ )) within a local neighborhood of the ground truth [10].
Revised: For symmetric matrix completion under the root mean squared error (RMSE) loss function, regular SGD is known to have an iteration count of O (  4 · dr log( d/ ✏ )) within a local neighborhood of the ground truth [10].

Original: This ﬁgure is optimal in the dimension d , the rank r , and the ﬁnal accuracy ✏ , but suboptimal by four exponents with respect to condition number  .
Revised: This ﬁgure is optimal in the dimension d , the rank r , and the ﬁnal accuracy ✏ , but suboptimal by four exponents with respect to condition number  .

Original: In contrast, we prove for the same setting that ScaledSGD attains an optimal convergence rate, converging to ✏ -accuracy in O ( dr log( d/ ✏ )) iterations for all values of the condition number  .
Revised: In contrast, we prove for the same setting that ScaledSGD attains an optimal convergence rate, converging to ✏ -accuracy in O ( dr log( d/ ✏ )) iterations for all values of the condition number  .

Original: In fact, our theoretical result predicts that ScaledSGD converges as if the ground truth matrix is perfectly conditioned, with a condition number of  = 1 .
Revised: In fact, our theoretical result predicts that ScaledSGD converges as if the ground truth matrix is perfectly conditioned, with a condition number of  = 1 .

Original: At ﬁrst sight, it appears quite natural that applying the ScaledGD preconditioner to SGD should result in accelerated convergence.
Revised: At ﬁrst sight, it appears quite natural that applying the ScaledGD preconditioner to SGD should result in accelerated convergence.

Original: However, the core challenge of stochastic algorithms like SGD is that each iteration can have substantial variance that “drown out” the expected progress made in the iteration.
Revised: However, the core challenge of stochastic algorithms like SGD is that each iteration can have substantial variance that “drown out” the expected progress made in the iteration.

Original: In the case of ScaledSGD, a rough analysis would suggest that the highly ill-conditioned preconditioner should improve convergence in expectation, but at the cost of dramatically worsening the variance.
Revised: In the case of ScaledSGD, a rough analysis would suggest that the highly ill-conditioned preconditioner should improve convergence in expectation, but at the cost of dramatically worsening the variance.

Original: Surprisingly, we ﬁnd in this paper that the speciﬁc scaling ( X T X )  1 used in ScaledSGD not only does not worsen the variance, but in fact improves it.
Revised: Surprisingly, we ﬁnd in this paper that the speciﬁc scaling ( X T X )  1 used in ScaledSGD not only does not worsen the variance, but in fact improves it.

Original: Our key insight and main theoretical contribution is Lemma 4, which shows that the same mechanism that allows ScaledGD to converge faster (compared to regular GD) also allows ScaledSGD to enjoy reduced variance (compared to regular
Revised: Our key insight and main theoretical contribution is Lemma 4, which shows that the same mechanism that allows ScaledGD to converge faster (compared to regular GD) also allows ScaledSGD to enjoy reduced variance (compared to regular

Original: SGD).
Revised: SGD).

Original: In fact, it is this effect of variance reduction that is responsible for most (  3 out of  4 ) of our improvement over the previous state-of-the-art.
Revised: In fact, it is this effect of variance reduction that is responsible for most (  3 out of  4 ) of our improvement over the previous state-of-the-art.

Original: It turns out that a careful choice of preconditioner can be used as a mechanism for variance reduction , while at the same time also fulﬁlling its usual, classical purpose, which is to accelerate convergence in expectation.
Revised: It turns out that a careful choice of preconditioner can be used as a mechanism for variance reduction , while at the same time also fulﬁlling its usual, classical purpose, which is to accelerate convergence in expectation.

Original: Related work Earlier work on matrix completion analyzed a convex relaxation of the original problem, showing that nuclear norm minimization can recover the ground truth from a few incoherent measurements [11–15].
Revised: Related work Earlier work on matrix completion analyzed a convex relaxation of the original problem, showing that nuclear norm minimization can recover the ground truth from a few incoherent measurements [11–15].

Original: This approach enjoys a near optimal sample complexity but incurs an
Revised: This approach enjoys a near optimal sample complexity but incurs an

Original: O ( d 3 ) per-iteration computational cost, which is prohibitive for a even moderately large d .
Revised: O ( d 3 ) per-iteration computational cost, which is prohibitive for a even moderately large d .

Original: More recent work has focused more on a nonconvex formulation based on Burer and Monteiro [16], which factors the optimization variable as M = XX T where X 2 R d ⇥ r and applies a local search method such as alternating-minimization [17–20], projected gradient descent [21, 22] and regular gradient descent [23–26].
Revised: More recent work has focused more on a nonconvex formulation based on Burer and Monteiro [16], which factors the optimization variable as M = XX T where X 2 R d ⇥ r and applies a local search method such as alternating-minimization [17–20], projected gradient descent [21, 22] and regular gradient descent [23–26].

Original: A separate line of work [27–35] focused on global properties of nonconvex matrix recovery problems, showing that the problem has no spurious local minima if sampling operator satisﬁes certain regularity conditions such as incoherence or restricted isometry.
Revised: A separate line of work [27–35] focused on global properties of nonconvex matrix recovery problems, showing that the problem has no spurious local minima if sampling operator satisﬁes certain regularity conditions such as incoherence or restricted isometry.

Original: The convergence rate of SGD has been well-studied for general classes of functions [36–39].
Revised: The convergence rate of SGD has been well-studied for general classes of functions [36–39].

Original: For matrix completion in particular, Jin et al.
Revised: For matrix completion in particular, Jin et al.

Original: [10] proved that SGD converges towards an ✏ -accurate solution in O ( d  4 log(1 / ✏ )) iterations where  is the condition number of M .
Revised: [10] proved that SGD converges towards an ✏ -accurate solution in O ( d  4 log(1 / ✏ )) iterations where  is the condition number of M .

Original: Unfortunately, this quartic dependence on  makes SGD extremely slow and impractical for huge-scale applications.
Revised: Unfortunately, this quartic dependence on  makes SGD extremely slow and impractical for huge-scale applications.

Original: This dramatic slow down of gradient descent and its variants caused by ill-conditioning has become well-known in recent years.
Revised: This dramatic slow down of gradient descent and its variants caused by ill-conditioning has become well-known in recent years.

Original: Several recent papers have proposed full-batch algorithms to overcome this issue [9, 40, 41], but these methods cannot be used in the huge-scale optimization setting where d is so large that even full-vector operations with O ( d ) time complexity are too expensive.
Revised: Several recent papers have proposed full-batch algorithms to overcome this issue [9, 40, 41], but these methods cannot be used in the huge-scale optimization setting where d is so large that even full-vector operations with O ( d ) time complexity are too expensive.

Original: As a deterministic full-batch method, ScaledGD [9] requires a projection onto the set of incoherent matrices at every iteration in order to maintain rapid convergence.
Revised: As a deterministic full-batch method, ScaledGD [9] requires a projection onto the set of incoherent matrices at every iteration in order to maintain rapid convergence.

Original: Instead our key ﬁnding here is that the stochasticity of SGD alone is enough to keep the iterates as incoherent as the ground truth, which allows for rapid progress to be made.
Revised: Instead our key ﬁnding here is that the stochasticity of SGD alone is enough to keep the iterates as incoherent as the ground truth, which allows for rapid progress to be made.

Original: The second-order method proposed in [41] costs at least O ( d ) per-iteration and has no straightforward stochastic analog.
Revised: The second-order method proposed in [41] costs at least O ( d ) per-iteration and has no straightforward stochastic analog.

Original: PrecGD [5] only applies to matrices that satisﬁes matrices satisfying the restricted isometry property, which does not hold for matrix completion.
Revised: PrecGD [5] only applies to matrices that satisﬁes matrices satisfying the restricted isometry property, which does not hold for matrix completion.

Original:  Background: Linear convergence of SGD
Revised:  Background: Linear convergence of SGD

Original: In our theoretical analysis, we restrict our attention to symmetric matrix completion under the root mean squared error (RMSE) loss function.
Revised: In our theoretical analysis, we restrict our attention to symmetric matrix completion under the root mean squared error (RMSE) loss function.

Original: Our goal is to solve the following nonconvex optimization
Revised: Our goal is to solve the following nonconvex optimization

Original: [Equation]
Revised: [Equation]

Original: in which we assume that the d ⇥ d ground truth ZZ T ⌫ 0 matrix is exactly rank- r , with a ﬁnite condition number
Revised: in which we assume that the d ⇥ d ground truth ZZ T ⌫ 0 matrix is exactly rank- r , with a ﬁnite condition number

Original: [Equation]
Revised: [Equation]

Original: In order to be able to reconstruct ZZ T from a small number of measurements, we will also need to assume that the ground truth has small coherence [42]
Revised: In order to be able to reconstruct ZZ T from a small number of measurements, we will also need to assume that the ground truth has small coherence [42]

Original: [Equation]
Revised: [Equation]

Original: Recall that µ takes on a value from 1 to d/r , with the smallest achieved by dense, orthonormal choices of Z whose rows all have magnitudes of 1 / p d , and the largest achieved by a ground truth
Revised: Recall that µ takes on a value from 1 to d/r , with the smallest achieved by dense, orthonormal choices of Z whose rows all have magnitudes of 1 / p d , and the largest achieved by a ground truth

Original: ZZ T containing a single nonzero element.
Revised: ZZ T containing a single nonzero element.

Original: Assuming incoherence µ =
Revised: Assuming incoherence µ =

Original: O (1) with respect to d , it is a well-known result that all d 2 matrix elements of ZZ T can be perfectly reconstructed from just
Revised: O (1) with respect to d , it is a well-known result that all d 2 matrix elements of ZZ T can be perfectly reconstructed from just

Original: O ( dr log d ) random samples of its matrix elements [12, 43].
Revised: O ( dr log d ) random samples of its matrix elements [12, 43].

Original: This paper considers solving (3) in the huge-scale, online optimization setting, in which individual matrix elements of the ground truth ( ZZ T ) i,j = z
Revised: This paper considers solving (3) in the huge-scale, online optimization setting, in which individual matrix elements of the ground truth ( ZZ T ) i,j = z

Original: Ti z j are revealed one-at-a-time, uniformly at random with replacement, and that a current iterate X is continuously updated to streaming data.
Revised: Ti z j are revealed one-at-a-time, uniformly at random with replacement, and that a current iterate X is continuously updated to streaming data.

Original: We note that this is a reasonably accurate model for how recommendation engines are tuned to user preferences in practice, although the uniformity of random sampling is admittedly an assumption made to ease theoretical analysis.
Revised: We note that this is a reasonably accurate model for how recommendation engines are tuned to user preferences in practice, although the uniformity of random sampling is admittedly an assumption made to ease theoretical analysis.

Original: Deﬁne the stochastic gradient operator as
Revised: Deﬁne the stochastic gradient operator as

Original: [Equation]
Revised: [Equation]

Original: where ( i, j ) ⇠ Unif([ d ] ⇥ [ d ]) , where x i , x j 2 R r are the i -th and j -th rows of X , and the scaling d 2 is chosen that, over the randomness of the sampled index ( i, j ) , we have exactly E [ SG ( X )]
Revised: where ( i, j ) ⇠ Unif([ d ] ⇥ [ d ]) , where x i , x j 2 R r are the i -th and j -th rows of X , and the scaling d 2 is chosen that, over the randomness of the sampled index ( i, j ) , we have exactly E [ SG ( X )]

Original: =
Revised: =

Original: r f ( X ) .
Revised: r f ( X ) .

Original: Then, the classical online SGD algorithm can be written as
Revised: Then, the classical online SGD algorithm can be written as

Original: [Equation]
Revised: [Equation]

Original: where ↵ > 0 .
Revised: where ↵ > 0 .

Original: [Equation]
Revised: [Equation]

Original: Here, we observe that a single iteration of SGD coincides with full-batch gradient descent in expectation, as in E [ X t +1 | X t ] =
Revised: Here, we observe that a single iteration of SGD coincides with full-batch gradient descent in expectation, as in E [ X t +1 | X t ] =

Original: X t  ↵ r f ( X t ) .
Revised: X t  ↵ r f ( X t ) .

Original: Therefore, assuming that bounded deviations and bounded variances, it follows from standard arguments that the behavior of many iterations of SGD should concentrate about that of full-batch gradient descent X t +1 =
Revised: Therefore, assuming that bounded deviations and bounded variances, it follows from standard arguments that the behavior of many iterations of SGD should concentrate about that of full-batch gradient descent X t +1 =

Original: X t  ↵ r f ( X t ) .
Revised: X t  ↵ r f ( X t ) .

Original: Within a region sufﬁciently close to the ground truth, full-batch gradient descent is well-known to converge at a linear rate to the ground truth [23, 44].
Revised: Within a region sufﬁciently close to the ground truth, full-batch gradient descent is well-known to converge at a linear rate to the ground truth [23, 44].

Original: Within this same region, Jin et al.
Revised: Within this same region, Jin et al.

Original: [10] proved that SGD also converges linearly.
Revised: [10] proved that SGD also converges linearly.

Original: For an incoherent ground truth with µ = O (1) , they proved that SGD with an aggressive choice of step-size is able to recover the ground truth to ✏ - accuracy O (  4 dr log( d/ ✏ )) iterations, with each iteration costing O ( r ) arithmetic operations and selecting 1 random sample.
Revised: For an incoherent ground truth with µ = O (1) , they proved that SGD with an aggressive choice of step-size is able to recover the ground truth to ✏ - accuracy O (  4 dr log( d/ ✏ )) iterations, with each iteration costing O ( r ) arithmetic operations and selecting 1 random sample.

Original: This iteration count is optimal with respect to d, r, and 1 / ✏ , although its dependence on  is a cubic factor (i.e. a factor of  3 ) worse than full-batch gradient descent’s ﬁgure of O (  log(1 / ✏ )) , which is itself already quite bad, given that  in practice can readily take on values of 10 3 to 10 6 .
Revised: This iteration count is optimal with respect to d, r, and 1 / ✏ , although its dependence on  is a cubic factor (i.e. a factor of  3 ) worse than full-batch gradient descent’s ﬁgure of O (  log(1 / ✏ )) , which is itself already quite bad, given that  in practice can readily take on values of 10 3 to 10 6 .

Original: Theorem 1 (Jin, Kakade, and Netrapalli [10]) .
Revised: Theorem 1 (Jin, Kakade, and Netrapalli [10]) .

Original: For Z 2 R d ⇥ r with  max ( Z ) = 1 and f ( X ) =
Revised: For Z 2 R d ⇥ r with  max ( Z ) = 1 and f ( X ) =

Original: k XX T  ZZ T k 2 F and h i ( X ) =
Revised: k XX T  ZZ T k 2 F and h i ( X ) =

Original: k e Ti X k 2 , deﬁne the following
Revised: k e Ti X k 2 , deﬁne the following

Original: [Equation]
Revised: [Equation]

Original: For an initial point X 0 2 R d ⇥ r that satiﬁes f
Revised: For an initial point X 0 2 R d ⇥ r that satiﬁes f

Original: ( X 0 )  12 f max and max i
Revised: ( X 0 )  12 f max and max i

Original: h i ( X 0 )  12 h max , there exists some constant c such that for any learning rate ↵ < c · (  · h max · d 2 log d )  1 , with probability at least 1  T/d 10 , we will have for all t  T iterations of SGD that
Revised: h i ( X 0 )  12 h max , there exists some constant c such that for any learning rate ↵ < c · (  · h max · d 2 log d )  1 , with probability at least 1  T/d 10 , we will have for all t  T iterations of SGD that

Original: [Equation]
Revised: [Equation]

Original: The reason for Theorem 1’s additional  3 dependence beyond full-batch gradient descent is due to its need to maintain incoherence in its iterates.
Revised: The reason for Theorem 1’s additional  3 dependence beyond full-batch gradient descent is due to its need to maintain incoherence in its iterates.

Original: Using standard techniques on martingale concentration, one can readily show that SGD replicates a single iteration of full-batch gradient descent over an epoch of d 2 iterations.
Revised: Using standard techniques on martingale concentration, one can readily show that SGD replicates a single iteration of full-batch gradient descent over an epoch of d 2 iterations.

Original: This results in an iteration count O (  · d 2 log(1 / ✏ )) with an optimal dependence on  , but the entire matrix is already fully observed after collecting d 2 samples.
Revised: This results in an iteration count O (  · d 2 log(1 / ✏ )) with an optimal dependence on  , but the entire matrix is already fully observed after collecting d 2 samples.

Original: Instead,
Revised: Instead,

Original: Jin et al.
Revised: Jin et al.

Original: [10] noted that the variance of SGD iterations is controlled by the step-size ↵ times the maximum coherence µ X = dr · max i,t k e Ti X t k 2 over the iterates X t , X t  1 , . . .
Revised: [10] noted that the variance of SGD iterations is controlled by the step-size ↵ times the maximum coherence µ X = dr · max i,t k e Ti X t k 2 over the iterates X t , X t  1 , . . .

Original: , X 0 .
Revised: , X 0 .

Original: If the iterates can be kept incoherent with µ
Revised: If the iterates can be kept incoherent with µ

Original: X =
Revised: X =

Original: O (1) , then SGD with a more aggressive step-size will reproduce an iteration of full-batch gradient descent after an epoch of just O
Revised: O (1) , then SGD with a more aggressive step-size will reproduce an iteration of full-batch gradient descent after an epoch of just O

Original: ( dr log d ) iterations.
Revised: ( dr log d ) iterations.

Original: The main ﬁnding in Jin et al.
Revised: The main ﬁnding in Jin et al.

Original: [10]’s proof of Theorem 1 is that the stochasticity of SGD is enough to keep the iterates incoherent.
Revised: [10]’s proof of Theorem 1 is that the stochasticity of SGD is enough to keep the iterates incoherent.

Original: This contrasts with full-batch methods at the time, which required an added regularizer [20, 30, 45] or an explicit projection step [9].
Revised: This contrasts with full-batch methods at the time, which required an added regularizer [20, 30, 45] or an explicit projection step [9].

Original: (As pointed out by a reviewer, it was later shown by Ma et al.
Revised: (As pointed out by a reviewer, it was later shown by Ma et al.

Original: [46] that full-batch gradient descent is also able to maintain incoherence without a regularizer nor a projection.)
Revised: [46] that full-batch gradient descent is also able to maintain incoherence without a regularizer nor a projection.)

Original: Unfortunately, maintaining incoherence requires shrinking the step-size by a factor of  , and the actual value of µ X that results is also a factor of  2 worse than the original coherence µ of the ground truth Z .
Revised: Unfortunately, maintaining incoherence requires shrinking the step-size by a factor of  , and the actual value of µ X that results is also a factor of  2 worse than the original coherence µ of the ground truth Z .

Original: The resulting iteration count O (  4 · dr log( d/ ✏ )) is made optimal with respect to d, r, and 1 / ✏ , but only at the cost of worsening its the dependence on the condition number  by another three exponents .
Revised: The resulting iteration count O (  4 · dr log( d/ ✏ )) is made optimal with respect to d, r, and 1 / ✏ , but only at the cost of worsening its the dependence on the condition number  by another three exponents .

Original: Finally, the quality of the initial point X 0 also has a dependence on the condition number  .
Revised: Finally, the quality of the initial point X 0 also has a dependence on the condition number  .

Original: In order to guarantee linear convergence, Theorem 1 requires X 0 to lie in the neighborhood k X 0 X T 

Revised: In order to guarantee linear convergence, Theorem 1 requires X 0 to lie in the neighborhood k X 0 X T

Original:  ZZ T k F <  min ( Z T Z ) =
Revised:  ZZ T k F <  min ( Z T Z ) =

Original: O (   1 ) .
Revised: O (   1 ) .

Original: This dependence on  is optimal, because fullbatch gradient descent must lose its ability to converge linearly in the limit  !
Revised: This dependence on  is optimal, because fullbatch gradient descent must lose its ability to converge linearly in the limit  !

Original: However, the leading constant can be very pessmistic, because the theorem must formally exclude spurious critical points X spur that have r f ( X spur )
Revised: However, the leading constant can be very pessmistic, because the theorem must formally exclude spurious critical points X spur that have r f ( X spur )

Original: = 

Revised: =

Original: but f ( X spur ) > 0 in order to be provably correct.
Revised: but f ( X spur ) > 0 in order to be provably correct.

Original: In practice, it is commonly observed that SGD converges globally, starting from an arbitrary, possibly random initialization [30], at a linear rate that is consistent with local convergence theorems like Theorem 1.
Revised: In practice, it is commonly observed that SGD converges globally, starting from an arbitrary, possibly random initialization [30], at a linear rate that is consistent with local convergence theorems like Theorem 1.

Original: It is now commonly argued that gradient methods can escape saddle points with high probability [47], and so their performance is primarily dictated by local convergence behavior [48, 49].
Revised: It is now commonly argued that gradient methods can escape saddle points with high probability [47], and so their performance is primarily dictated by local convergence behavior [48, 49].

Original:  Proposed algorithm and main result
Revised:  Proposed algorithm and main result

Original: Inspired by a recent full-batch gradient method called ScaledGD [4, 9] and a closely related algorithm PrecGD [5], we proposed the following algorithm
Revised: Inspired by a recent full-batch gradient method called ScaledGD [4, 9] and a closely related algorithm PrecGD [5], we proposed the following algorithm

Original: [Equation]
Revised: [Equation]

Original: where ↵ >
Revised: where ↵ >

Original: [Equation]
Revised: [Equation]

Original: ScaledSGD
Revised: ScaledSGD

Original: [Equation]
Revised: [Equation]

Original: As we mentioned in the introduction, the preconditioner P =
Revised: As we mentioned in the introduction, the preconditioner P =

Original: ( X T X )  1 can be precomputed and cached in a practical implementation, and afterwards efﬁciently updated using the Sherman–
Revised: ( X T X )  1 can be precomputed and cached in a practical implementation, and afterwards efﬁciently updated using the Sherman–

Original: Morrison formula.
Revised: Morrison formula.

Original: The per-iteration cost of ScaledSGD is O ( r 2 ) arithmetic operations and 1 random sample, which for modest values of r is only marginally more than the cost of SGD.
Revised: The per-iteration cost of ScaledSGD is O ( r 2 ) arithmetic operations and 1 random sample, which for modest values of r is only marginally more than the cost of SGD.

Original: Our main result in this paper is that, with a region sufﬁciently close to the ground truth, this simple rescaling allows ScaledSGD to converge linearly to ✏ -accuracy O ( dr log( d/ ✏ )) iterations, with no further dependence on the condition number  .
Revised: Our main result in this paper is that, with a region sufﬁciently close to the ground truth, this simple rescaling allows ScaledSGD to converge linearly to ✏ -accuracy O ( dr log( d/ ✏ )) iterations, with no further dependence on the condition number  .

Original: This iteration count is optimal with respect to d, r, 1 / ✏ , and  , and in fact matches SGD with a perfectly conditioned ground truth  = 1 .
Revised: This iteration count is optimal with respect to d, r, 1 / ✏ , and  , and in fact matches SGD with a perfectly conditioned ground truth  = 1 .

Original: In our numerical experiments, we observe that ScaledSGD converges globally from a random initialization at the same rate as SGD as if  = 1 .
Revised: In our numerical experiments, we observe that ScaledSGD converges globally from a random initialization at the same rate as SGD as if  = 1 .

Original: Theorem 2 (Main) .
Revised: Theorem 2 (Main) .

Original: For Z 2 R d ⇥ r with  max ( Z ) = 1 and f ( X ) =
Revised: For Z 2 R d ⇥ r with  max ( Z ) = 1 and f ( X ) =

Original: k XX T  ZZ T k 2 F and g i ( X ) =
Revised: k XX T  ZZ T k 2 F and g i ( X ) =

Original: e Ti X ( X T X )  1 X T e i , select a radius ⇢ < 1 / 2 and set
Revised: e Ti X ( X T X )  1 X T e i , select a radius ⇢ < 1 / 2 and set

Original: [Equation]
Revised: [Equation]

Original: For an initial point X 0 2 R d ⇥ r that satiﬁes f
Revised: For an initial point X 0 2 R d ⇥ r that satiﬁes f

Original: ( X 0 )  12 f max and max
Revised: ( X 0 )  12 f max and max

Original: i
Revised: i

Original: g
Revised: g

Original: i
Revised: i

Original: ( X 0 )  12 g max , there exists some constant c such that for any learning rate ↵ < c · [( g max + ⇢ ) · d 2 log d ]  1 , with probability at least 1  T/d 10 , we will have for all t  T iterations of ScaledSGD that:
Revised: ( X 0 )  12 g max , there exists some constant c such that for any learning rate ↵ < c · [( g max + ⇢ ) · d 2 log d ]  1 , with probability at least 1  T/d 10 , we will have for all t  T iterations of ScaledSGD that:

Original: [Equation]
Revised: [Equation]

Original: Theorem 2 eliminates all dependencies on the condition number  in Theorem 1 except for the quality of the initial point, which we had already noted earlier as being optimal.
Revised: Theorem 2 eliminates all dependencies on the condition number  in Theorem 1 except for the quality of the initial point, which we had already noted earlier as being optimal.

Original: Our main ﬁnding is that it is possible to maintain incoherence while making aggressive step-sizes towards a highly ill-conditioned ground truth ZZ T .
Revised: Our main ﬁnding is that it is possible to maintain incoherence while making aggressive step-sizes towards a highly ill-conditioned ground truth ZZ T .

Original: In fact, Theorem 2 says that, with high probability, the maximum coherence µ X over of any iterate X t will only be a mild constant factor of ⇡ 16 times worse than the coherence µ of the ground truth ZZ T .
Revised: In fact, Theorem 2 says that, with high probability, the maximum coherence µ X over of any iterate X t will only be a mild constant factor of ⇡ 16 times worse than the coherence µ of the ground truth ZZ T .

Original: This is particularly surprising in view of the fact that every iteration of ScaledSGD involves inverting a potentially highly ill-conditioned matrix ( X T X )  1 .
Revised: This is particularly surprising in view of the fact that every iteration of ScaledSGD involves inverting a potentially highly ill-conditioned matrix ( X T X )  1 .

Original: In contrast, even without inverting matrices, Theorem 1 says that SGD is only able to keep µ X within a factor of  2 of µ , and only by shrinking the step-size ↵ by another factor of  .
Revised: In contrast, even without inverting matrices, Theorem 1 says that SGD is only able to keep µ X within a factor of  2 of µ , and only by shrinking the step-size ↵ by another factor of  .

Original: However, the price we pay for maintaining incoherence is that the quality of the initial point X 0 now gains a dependence on dimension d , in addition to the condition number  .
Revised: However, the price we pay for maintaining incoherence is that the quality of the initial point X 0 now gains a dependence on dimension d , in addition to the condition number  .

Original: In order to guarantee fast linear convergence independent of  , Theorem 2 requires X 0 to lie in the neighborhood k X 0 X T 

Revised: In order to guarantee fast linear convergence independent of  , Theorem 2 requires X 0 to lie in the neighborhood k X 0 X T

Original:  ZZ T k F < µr  min ( Z T Z ) /d
Revised:  ZZ T k F < µr  min ( Z T Z ) /d

Original: = (  d )  1 , so that ⇢ can be set to be the same order of magnitude as g max .
Revised: = (  d )  1 , so that ⇢ can be set to be the same order of magnitude as g max .

Original: In essence, the “effective” condition number of the ground truth has been worsened by another factor of d .
Revised: In essence, the “effective” condition number of the ground truth has been worsened by another factor of d .

Original: This shrinks the size of our local neighborhood by a factor of d , but has no impact on the convergence rate of the resulting iterations.
Revised: This shrinks the size of our local neighborhood by a factor of d , but has no impact on the convergence rate of the resulting iterations.

Original: In the limit that  !
Revised: In the limit that  !

Original: and the search rank r becomes overparameterized with respect to the true rank r ?
Revised: and the search rank r becomes overparameterized with respect to the true rank r ?

Original: <
Revised: <

Original: r of ZZ T , both full-batch gradient descent and SGD slows down to a sublinear convergence rate, in theory and in practice [5, 6].
Revised: r of ZZ T , both full-batch gradient descent and SGD slows down to a sublinear convergence rate, in theory and in practice [5, 6].

Original: While Theorem 2 is no longer applicable, we observe in our numerical experiments that ScaledSGD nevertheless maintains its fast linear convergence rate as if  = 1 .
Revised: While Theorem 2 is no longer applicable, we observe in our numerical experiments that ScaledSGD nevertheless maintains its fast linear convergence rate as if  = 1 .

Original: Following PrecGD [5], we believe that introducing a small identity perturbation to the scaling matrix of ScaledSGD, as in ( X T X + ⌘ I )  1 for some ⌘ ⇡ p f ( X ) , should be enough to rigorously extend Theorem 2 to the overparameterized regime.
Revised: Following PrecGD [5], we believe that introducing a small identity perturbation to the scaling matrix of ScaledSGD, as in ( X T X + ⌘ I )  1 for some ⌘ ⇡ p f ( X ) , should be enough to rigorously extend Theorem 2 to the overparameterized regime.

Original: We leave this extension as future work.
Revised: We leave this extension as future work.

Original:  Key ideas for the proof
Revised:  Key ideas for the proof

Original: We begin by explaining the mechanism by which SGD slows down when converging towards an ill-conditioned ground truth.
Revised: We begin by explaining the mechanism by which SGD slows down when converging towards an ill-conditioned ground truth.

Original: Recall that
Revised: Recall that

Original: [Equation]
Revised: [Equation]

Original: As XX T converges towards an ill-conditioned ground truth
Revised: As XX T converges towards an ill-conditioned ground truth

Original: [Equation]
Revised: [Equation]

Original: , the factor matrix X must become progressively ill-conditioned, with
Revised: , the factor matrix X must become progressively ill-conditioned, with

Original: [Equation]
Revised: [Equation]

Original: Therefore, it is possible for components of the error vector XX T  ZZ T to become “invisible” by aligning within the ill-conditioned subspaces of X .
Revised: Therefore, it is possible for components of the error vector XX T  ZZ T to become “invisible” by aligning within the ill-conditioned subspaces of X .

Original: As SGD progresses towards the solution, these ill-conditioned subspaces of X become the slowest components of the error vector to converge to zero.
Revised: As SGD progresses towards the solution, these ill-conditioned subspaces of X become the slowest components of the error vector to converge to zero.

Original: On the other hand, the maximum step-size that can be taken is controlled by the most wellconditioned subspaces of X .
Revised: On the other hand, the maximum step-size that can be taken is controlled by the most wellconditioned subspaces of X .

Original: A simple idea, therefore, is to rescale the ill-conditioned components of the gradient r f ( X ) in order to make the ill-conditioned subspaces of X more “visible”.
Revised: A simple idea, therefore, is to rescale the ill-conditioned components of the gradient r f ( X ) in order to make the ill-conditioned subspaces of X more “visible”.

Original: More concretely, deﬁne the local norm of the gradient as kr f ( X ) k X =
Revised: More concretely, deﬁne the local norm of the gradient as kr f ( X ) k X =

Original: kr f ( X )( X T X ) 1 / 2 k F and its corresponding dual norm as kr f ( X ) k
Revised: kr f ( X )( X T X ) 1 / 2 k F and its corresponding dual norm as kr f ( X ) k

Original: ⇤ X =
Revised: ⇤ X =

Original: kr f
Revised: kr f

Original: ( X )( X T X )  1 / 2 k F .
Revised: ( X )( X T X )  1 / 2 k F .

Original: It has long been known (see e.g. [23, 44]) that rescaling the gradient yields
Revised: It has long been known (see e.g. [23, 44]) that rescaling the gradient yields

Original: [Equation]
Revised: [Equation]

Original: where ✓ is the angle between the error vector XX T  ZZ T and the linear subspace { XY T + Y X T :
Revised: where ✓ is the angle between the error vector XX T  ZZ T and the linear subspace { XY T + Y X T :

Original: Y 2 R d ⇥ r } .
Revised: Y 2 R d ⇥ r } .

Original: This insight immediately suggests an iteration like X + = X  ↵ r f ( X )( X T X )  1 .
Revised: This insight immediately suggests an iteration like X + = X  ↵ r f ( X )( X T X )  1 .

Original: In fact, the gradients of f have some Lipschitz constant L , so
Revised: In fact, the gradients of f have some Lipschitz constant L , so

Original: [Equation]
Revised: [Equation]

Original: However, a naive analysis ﬁnds that L X = L/
Revised: However, a naive analysis ﬁnds that L X = L/

Original:  min ( X T X ) ⇡ L ·  , and this causes the step-size to shrink by a factor of  .
Revised:  min ( X T X ) ⇡ L ·  , and this causes the step-size to shrink by a factor of  .

Original: The main motivating insight behind ScaledGD [4, 9] and later PrecGD [5] is that, with a ﬁner analysis, it is possible to prove Lipschitz continuity under a local change of norm.
Revised: The main motivating insight behind ScaledGD [4, 9] and later PrecGD [5] is that, with a ﬁner analysis, it is possible to prove Lipschitz continuity under a local change of norm.

Original: Lemma 3 (Function descent) .
Revised: Lemma 3 (Function descent) .

Original: Let X, Z 2 R n ⇥ r satisfy k XX T  ZZ T k F  ⇢ min ( Z T Z ) where ⇢ < 1 / 2 .
Revised: Let X, Z 2 R n ⇥ r satisfy k XX T  ZZ T k F  ⇢ min ( Z T Z ) where ⇢ < 1 / 2 .

Original: Then, the function f ( X ) =
Revised: Then, the function f ( X ) =

Original: [Equation]
Revised: [Equation]

Original: satisﬁes
Revised: satisﬁes

Original: [Equation]
Revised: [Equation]

Original: for all k V k X  C · p f ( X ) with L X = 6 + 8 C + 2 C 2 =
Revised: for all k V k X  C · p f ( X ) with L X = 6 + 8 C + 2 C 2 =

Original: O (1 + C 2 ) .
Revised: O (1 + C 2 ) .

Original: This same idea can be “stochastiﬁed” in a straightforward manner.
Revised: This same idea can be “stochastiﬁed” in a straightforward manner.

Original: Conditioning on the current iterate X , then the new iterate X + = X
Revised: Conditioning on the current iterate X , then the new iterate X + = X

Original:  ↵ SG ( X )( X T X )  1 has expectation
Revised:  ↵ SG ( X )( X T X )  1 has expectation

Original: [Equation]
Revised: [Equation]

Original: The linear term evaluates as E [ SG ( X )( X T X )  1 ] =
Revised: The linear term evaluates as E [ SG ( X )( X T X )  1 ] =

Original: r f ( X )( X T X )  1 , while the quadratic term is
Revised: r f ( X )( X T X )  1 , while the quadratic term is

Original: [Equation]
Revised: [Equation]

Original: where g i ( X ) =
Revised: where g i ( X ) =

Original: e Ti X ( X T X )  1 X T
Revised: e Ti X ( X T X )  1 X T

Original: e
Revised: e

Original: i
Revised: i

Original: =
Revised: =

Original: ( k e Ti X k ⇤ X ) 2 .
Revised: ( k e Ti X k ⇤ X ) 2 .

Original: Combined, we obtain geometric convergence
Revised: Combined, we obtain geometric convergence

Original: [Equation]
Revised: [Equation]

Original: We see that the step-size depends crucially on the incoherence
Revised: We see that the step-size depends crucially on the incoherence

Original: g i ( X )  g max of the current iterate.
Revised: g i ( X )  g max of the current iterate.

Original: If the current iterate X is incoherent with g max
Revised: If the current iterate X is incoherent with g max

Original: = O (1 /d ) , then a step-size of ↵ =
Revised: = O (1 /d ) , then a step-size of ↵ =

Original: O (1 /d ) is possible, resulting in convergence in O ( dr log( d/ ✏ )) iterations, which can be shown using standard martingale techniques [10].
Revised: O (1 /d ) is possible, resulting in convergence in O ( dr log( d/ ✏ )) iterations, which can be shown using standard martingale techniques [10].

Original: But if the current iterate is g max
Revised: But if the current iterate is g max

Original: = O (1) , then only a step-size of ↵ =
Revised: = O (1) , then only a step-size of ↵ =

Original: O (1 /d 2 ) is possible, which forces us to compute d 2 iterations, thereby obviating the need to complete the matrix in the ﬁrst place.
Revised: O (1 /d 2 ) is possible, which forces us to compute d 2 iterations, thereby obviating the need to complete the matrix in the ﬁrst place.

Original: Therefore, in order for prove rapid linear convergence, we need to additionally show that with high probability, the coherence g k ( X )=(
Revised: Therefore, in order for prove rapid linear convergence, we need to additionally show that with high probability, the coherence g k ( X )=(

Original: k e Tk X k ⇤ X ) 2 remains O (1) throughout ScaledGD iterations.
Revised: k e Tk X k ⇤ X ) 2 remains O (1) throughout ScaledGD iterations.

Original: This is the most challenging part of our proof.
Revised: This is the most challenging part of our proof.

Original: Previous methods that applied a similar scaling to fullbatch GD [9] required an explicit projection onto the set of incoherent matrices at each iteration.
Revised: Previous methods that applied a similar scaling to fullbatch GD [9] required an explicit projection onto the set of incoherent matrices at each iteration.

Original: Applying a similar projection to ScaledSGD will take O ( d ) time, which destroys the scalability of our method.
Revised: Applying a similar projection to ScaledSGD will take O ( d ) time, which destroys the scalability of our method.

Original: On the other hand, Jin et al.
Revised: On the other hand, Jin et al.

Original: [10] showed that the randomness in SGD is enough to keep the coherence of the iterates within a factor of  2 times worse than the coherence of the ground truth, and only by a step-size of at most ↵
Revised: [10] showed that the randomness in SGD is enough to keep the coherence of the iterates within a factor of  2 times worse than the coherence of the ground truth, and only by a step-size of at most ↵

Original: =
Revised: =

Original: O (   1 ) .
Revised: O (   1 ) .

Original: Surprisingly, here we show that the randomness in ScaledSGD is enough to keep the coherence of the iterates with a constant factor of the coherence the ground truth, using a step-size with no dependence on  .
Revised: Surprisingly, here we show that the randomness in ScaledSGD is enough to keep the coherence of the iterates with a constant factor of the coherence the ground truth, using a step-size with no dependence on  .

Original: The following key lemma is the crucial insight of our proof.
Revised: The following key lemma is the crucial insight of our proof.

Original: First, it says that function g k ( X ) satisﬁes a “descent lemma” with respect to the local norm k · k ⇤ X .
Revised: First, it says that function g k ( X ) satisﬁes a “descent lemma” with respect to the local norm k · k ⇤ X .

Original: Second, and much more importantly, it says that descending g k ( X ) along the scaled gradient direction r f ( X )( X T X ) incurs a linear decrement 1  2 ⇢ 1  ⇢ g k ( X ) with no dependence of the condition number  .
Revised: Second, and much more importantly, it says that descending g k ( X ) along the scaled gradient direction r f ( X )( X T X ) incurs a linear decrement 1  2 ⇢ 1  ⇢ g k ( X ) with no dependence of the condition number  .

Original: This is in direct analogy to the function value decrement in (6), which has no dependence on  , and in direct contrast to the proof of Jin et al.
Revised: This is in direct analogy to the function value decrement in (6), which has no dependence on  , and in direct contrast to the proof of Jin et al.

Original: [10], which is only able to achieve a decrement of (8 /  ) g k
Revised: [10], which is only able to achieve a decrement of (8 /  ) g k

Original: ( X ) due to the lack of rescaling by ( X T X )  1 .
Revised: ( X ) due to the lack of rescaling by ( X T X )  1 .

Original: Lemma 4 (Coherence descent) .
Revised: Lemma 4 (Coherence descent) .

Original: Let g k ( X ) =
Revised: Let g k ( X ) =

Original: e Tk X
Revised: e Tk X

Original: ( X T X )  1 X T e k .
Revised: ( X T X )  1 X T e k .

Original: Under the same conditions as
Revised: Under the same conditions as

Original: Lemma 3, we have
Revised: Lemma 3, we have

Original: [Equation]
Revised: [Equation]

Original: r g k ( X ) i + 5( k V k
Revised: r g k ( X ) i + 5( k V k

Original: [Equation]
Revised: [Equation]

Original: X
Revised: X

Original: [Equation]
Revised: [Equation]

Original: 2 2 k V k ⇤ X , hr g k ( X ) , r f ( X )( X T X )  1 i   1  2 ⇢ 1  ⇢ g k ( X )  1  ⇢ p g k ( X ) g k ( Z )  .
Revised: 2 2 k V k ⇤ X , hr g k ( X ) , r f ( X )( X T X )  1 i   1  2 ⇢ 1  ⇢ g k ( X )  1  ⇢ p g k ( X ) g k ( Z )  .

Original: Conditioning on X , we have for the search direction V =
Revised: Conditioning on X , we have for the search direction V =

Original: SG ( X )( X T X )  1 and X + = X
Revised: SG ( X )( X T X )  1 and X + = X

Original: + V
Revised: + V

Original: E [ g k
Revised: E [ g k

Original: ( X + )]  g k ( X )  ↵ hr g k ( X ) , E
Revised: ( X + )]  g k ( X )  ↵ hr g k ( X ) , E

Original: [ V ] i + ↵ 2 · E
Revised: [ V ] i + ↵ 2 · E

Original: [Equation]
Revised: [Equation]

Original: 1 
Revised: 1

Original: [Equation]
Revised: [Equation]

Original: ) 1  O
Revised: ) 1  O

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: 2 max · ⇢ )  (1  ⇣↵ ) g k
Revised: 2 max · ⇢ )  (1  ⇣↵ ) g k

Original: ( X ) + ↵ · ⇣ 2 g max for ↵ =
Revised: ( X ) + ↵ · ⇣ 2 g max for ↵ =

Original: O ( ⇢  1 d  2 ) .
Revised: O ( ⇢  1 d  2 ) .

Original: It then follows that g k ( X + ) converges geometrically towards 12 g max in expectation, with a convergence rate (1  ⇣↵ ) that is independent of the condition number  :
Revised: It then follows that g k ( X + ) converges geometrically towards 12 g max in expectation, with a convergence rate (1  ⇣↵ ) that is independent of the condition number  :

Original: [Equation]
Revised: [Equation]

Original: The proof of Theorem 2 then follows from standard techniques, by making the two decrement conditions (6) and (7) into supermartingales and applying a standard concentration inequality.
Revised: The proof of Theorem 2 then follows from standard techniques, by making the two decrement conditions (6) and (7) into supermartingales and applying a standard concentration inequality.

Original: We defer the rigorous proof to appendix E.
Revised: We defer the rigorous proof to appendix E.

Original:  Experimental validation
Revised:  Experimental validation

Original: In this section we compare the practical performance of ScaledSGD and SGD for the RMSE loss function in Theorem 2 and two real-world loss functions: the pairwise RMSE loss used to complete
Revised: In this section we compare the practical performance of ScaledSGD and SGD for the RMSE loss function in Theorem 2 and two real-world loss functions: the pairwise RMSE loss used to complete

Original: Euclidean Distance Matrices (EDM) in wireless communication networks; and the Bayesian Personalized Ranking (BRP) loss used to generate personalized item recommendation in collaborative ﬁltering.
Revised: Euclidean Distance Matrices (EDM) in wireless communication networks; and the Bayesian Personalized Ranking (BRP) loss used to generate personalized item recommendation in collaborative ﬁltering.

Original: In each case, ScaledSGD remains highly efﬁcient since it only updates two or three rows at a time, and the preconditioner P can be computed through low-rank updates, for a per-iteration cost of O ( r 2 ) .
Revised: In each case, ScaledSGD remains highly efﬁcient since it only updates two or three rows at a time, and the preconditioner P can be computed through low-rank updates, for a per-iteration cost of O ( r 2 ) .

Original: All of our experiments use random Gaussian initializations and an initial P =  2 I .
Revised: All of our experiments use random Gaussian initializations and an initial P =  2 I .

Original: To be able to accurately measure and report the effects of ill-conditioning on ScaledSGD and SGD, we focus on small-scale synthetic datasets in the ﬁrst two experiments, for which the ground truth is explicitly known, and where the condition numbers can be ﬁnely controlled.
Revised: To be able to accurately measure and report the effects of ill-conditioning on ScaledSGD and SGD, we focus on small-scale synthetic datasets in the ﬁrst two experiments, for which the ground truth is explicitly known, and where the condition numbers can be ﬁnely controlled.

Original: In addition, to gauge the scalability of ScaledSGD on huge-scale real-world datasets, in the third experiment, we apply
Revised: In addition, to gauge the scalability of ScaledSGD on huge-scale real-world datasets, in the third experiment, we apply

Original: ScaledSGD to generate personalized item recommendation using MovieLens25M dataset [50], for which the underlying item-item matrix has more than 62,000 items and 100 million pairwise samples are used during training.
Revised: ScaledSGD to generate personalized item recommendation using MovieLens25M dataset [50], for which the underlying item-item matrix has more than 62,000 items and 100 million pairwise samples are used during training.

Original: (Due to space constraints, we defer the details on the experimental setup, mathematical formulations, and the actual update equations to Appendix A.)
Revised: (Due to space constraints, we defer the details on the experimental setup, mathematical formulations, and the actual update equations to Appendix A.)

Original: The code for all experiments are available at https://github.com/Hong-Ming/ScaledSGD .
Revised: The code for all experiments are available at https://github.com/Hong-Ming/ScaledSGD .

Original: Matrix completion with RMSE loss.
Revised: Matrix completion with RMSE loss.

Original: The problem formulation is discussed in Section 3.
Revised: The problem formulation is discussed in Section 3.

Original: Figureplots the error f ( X ) =
Revised: Figureplots the error f ( X ) =

Original: k XX T  M k 2 F as the number of epochs increases.
Revised: k XX T  M k 2 F as the number of epochs increases.

Original: As expected, in the well-conditioned case, both ScaledSGD and SGD converges to machine error at roughly the same linear rate.
Revised: As expected, in the well-conditioned case, both ScaledSGD and SGD converges to machine error at roughly the same linear rate.

Original: However, in the ill-conditioned case, SGD slows down signiﬁcantly while ScaledSGD converges at almost exactly the same rate as in the well-conditioned case.
Revised: However, in the ill-conditioned case, SGD slows down signiﬁcantly while ScaledSGD converges at almost exactly the same rate as in the well-conditioned case.

Original: [Figure]
Revised: [Figure]

Original: ScaledSGDSGD
Revised: ScaledSGDSGD

Original: Euclidean distance matrix (EDM) completion.
Revised: Euclidean distance matrix (EDM) completion.

Original: The Euclidean distance matrix (EDM) is a matrix of pairwise distance between d points in Euclidean space [51].
Revised: The Euclidean distance matrix (EDM) is a matrix of pairwise distance between d points in Euclidean space [51].

Original: In applications such as wireless sensor networks, estimation of unknown distances, i.e., completing the EDM is often required.
Revised: In applications such as wireless sensor networks, estimation of unknown distances, i.e., completing the EDM is often required.

Original: We emphasize that this loss function is a pairwise loss, meaning that each measurement indexes multiple elements of the ground truth matrix.
Revised: We emphasize that this loss function is a pairwise loss, meaning that each measurement indexes multiple elements of the ground truth matrix.

Original: To demonstrate the efﬁcacy of ScaledSGD, we conduct two experiments where D is wellconditioned and ill-conditioned respectively:
Revised: To demonstrate the efﬁcacy of ScaledSGD, we conduct two experiments where D is wellconditioned and ill-conditioned respectively:

Original: Experiment 1.
Revised: Experiment 1.

Original: We uniformly sample 30 points in a cube center at origin with side length 2 , and use them to compute the ground truth EDM D .
Revised: We uniformly sample 30 points in a cube center at origin with side length 2 , and use them to compute the ground truth EDM D .

Original: In this case, each row x i 2 R 3 corresponds to the coordinates of the i -th sample.
Revised: In this case, each row x i 2 R 3 corresponds to the coordinates of the i -th sample.

Original: The corresponding matrix
Revised: The corresponding matrix

Original: X 2 R 30 ⇥ 3 is well-conditioned because of the uniform sampling.
Revised: X 2 R 30 ⇥ 3 is well-conditioned because of the uniform sampling.

Original: Experiment 2.
Revised: Experiment 2.

Original: The ground truth
Revised: The ground truth

Original: EDM is generated with 25 samples lie in the same cube in experiment 1, and 5 samples lie far away from the the cube.
Revised: EDM is generated with 25 samples lie in the same cube in experiment 1, and 5 samples lie far away from the the cube.

Original: These ﬁve outliers make the corresponding X become ill-conditioned.
Revised: These ﬁve outliers make the corresponding X become ill-conditioned.

Original: [Figure]
Revised: [Figure]

Original: ScaledSGDSGD
Revised: ScaledSGDSGD

Original: Item-item collaborative ﬁltering (CF).
Revised: Item-item collaborative ﬁltering (CF).

Original: In the task of item-item collaborative ﬁltering (CF), the ground truth M is a d ⇥ d matrix where d is the number of items we wish to rank and the i, j -th of
Revised: In the task of item-item collaborative ﬁltering (CF), the ground truth M is a d ⇥ d matrix where d is the number of items we wish to rank and the i, j -th of

Original: M is a similarity measure between the items.
Revised: M is a similarity measure between the items.

Original: Our goal is to learn a low-rank matrix that preserves the ranking of similarity between the items.
Revised: Our goal is to learn a low-rank matrix that preserves the ranking of similarity between the items.

Original: For instance, given a pairwise sample ( i, j, k ) , if item i is more similar to item j than item k , then M ij > M ik .
Revised: For instance, given a pairwise sample ( i, j, k ) , if item i is more similar to item j than item k , then M ij > M ik .

Original: We want to learn a low-rank matrix that also has this property, i.e., the i, j -th entry is greater than the i, k -th entry.
Revised: We want to learn a low-rank matrix that also has this property, i.e., the i, j -th entry is greater than the i, k -th entry.

Original: To gauge the scalability of ScaledSGD on a huge-scale real-world dataset, we perform simulation on item-item collaborative ﬁltering using a 62 , 000 ⇥ 62 , 000 item-item matrix M obtained from
Revised: To gauge the scalability of ScaledSGD on a huge-scale real-world dataset, we perform simulation on item-item collaborative ﬁltering using a 62 , 000 ⇥ 62 , 000 item-item matrix M obtained from

Original: MovieLens25M dataset.
Revised: MovieLens25M dataset.

Original: The CF model is trained using Bayesian Personalized Ranking (BRP) loss [1] on a training set, which consists of 100 million pairwise samples in M .
Revised: The CF model is trained using Bayesian Personalized Ranking (BRP) loss [1] on a training set, which consists of 100 million pairwise samples in M .

Original: The performance of
Revised: The performance of

Original: CF model is evaluated using Area Under the ROC Curve (AUC) score [1] on a test set, which consists of 10 million pairwise samples in M .
Revised: CF model is evaluated using Area Under the ROC Curve (AUC) score [1] on a test set, which consists of 10 million pairwise samples in M .

Original: The BPR loss is a widely used loss function in the context of collaborative ﬁltering for the task of personalized recommendation, and the AUC score is a popular evaluation metric to measure the accuracy of the recommendation system.
Revised: The BPR loss is a widely used loss function in the context of collaborative ﬁltering for the task of personalized recommendation, and the AUC score is a popular evaluation metric to measure the accuracy of the recommendation system.

Original: We defer the detail deﬁnition of BPR loss and AUC score to Appendix A.4.
Revised: We defer the detail deﬁnition of BPR loss and AUC score to Appendix A.4.

Original: Figure 3 plots the training BPR loss and testing AUC score within the ﬁrst epoch (ﬁlled with red) and the second epoch (ﬁlled with blue).
Revised: Figure 3 plots the training BPR loss and testing AUC score within the ﬁrst epoch (ﬁlled with red) and the second epoch (ﬁlled with blue).

Original: In order to measure the efﬁcacy of ScaledSGD, we compare its testing AUC score against a standard baseline called the NP-Maximum [1], which is the best possible
Revised: In order to measure the efﬁcacy of ScaledSGD, we compare its testing AUC score against a standard baseline called the NP-Maximum [1], which is the best possible

Original: AUC score by non-personalized ranking methods.
Revised: AUC score by non-personalized ranking methods.

Original: For a rigorous deﬁnition, see Appendix A.4.
Revised: For a rigorous deﬁnition, see Appendix A.4.

Original: We emphasize two important points in the Figure 3.
Revised: We emphasize two important points in the Figure 3.

Original: First, the percentage of training samples needed for ScaledSGD to achieve the same testing AUC scores as NP-Maximum is roughly 4 times smaller than SGD.
Revised: First, the percentage of training samples needed for ScaledSGD to achieve the same testing AUC scores as NP-Maximum is roughly 4 times smaller than SGD.

Original: Though both ScaledSGD and SGD are able to achieve higher AUC score than
Revised: Though both ScaledSGD and SGD are able to achieve higher AUC score than

Original: NP-Maximum before ﬁnishing the ﬁrst epoch, ScaledSGD achieve the same AUC score as NPMaximum after training on 11% of training samples while SGD requires 46% of them.
Revised: NP-Maximum before ﬁnishing the ﬁrst epoch, ScaledSGD achieve the same AUC score as NPMaximum after training on 11% of training samples while SGD requires 46% of them.

Original: We note that in this experiment, the size of the training set is 100 million, this means that SGD would requiremillion more iterations than ScaledSGD before it can reach NP-Maximum.
Revised: We note that in this experiment, the size of the training set is 100 million, this means that SGD would requiremillion more iterations than ScaledSGD before it can reach NP-Maximum.

Original: Second, the percentage of training samples needed for ScaledSGD to converge after the ﬁrst epoch is roughly 5 times smaller than SGD.
Revised: Second, the percentage of training samples needed for ScaledSGD to converge after the ﬁrst epoch is roughly 5 times smaller than SGD.

Original: Given that both ScaledSGD and SGD converge to AUC score at around 0 .
Revised: Given that both ScaledSGD and SGD converge to AUC score at around 0 .

Original: 9 within the second epoch (area ﬁlled with blue), we indicate the percentage of training samples when both algorithms reach 0 .
Revised: 9 within the second epoch (area ﬁlled with blue), we indicate the percentage of training samples when both algorithms reach 0 .

Original: 9 AUC score in Figure 3.
Revised: 9 AUC score in Figure 3.

Original: As expected, ScaledSGD is able to converge using fewer samples than SGD, with only 16% of training samples.
Revised: As expected, ScaledSGD is able to converge using fewer samples than SGD, with only 16% of training samples.

Original: SGD, on the other hand, requires 81% training samples.
Revised: SGD, on the other hand, requires 81% training samples.

Original:  Conclusions
Revised:  Conclusions

Original: We propose an algorithm called ScaledSGD for huge scale online matrix completion.
Revised: We propose an algorithm called ScaledSGD for huge scale online matrix completion.

Original: For the nonconvex approach to solving matrix completion, ill-conditioning in the ground truth causes SGD to slow down signiﬁcantly.
Revised: For the nonconvex approach to solving matrix completion, ill-conditioning in the ground truth causes SGD to slow down signiﬁcantly.

Original: ScaledSGD preserves all the favorable qualities of SGD while making it immune to ill-conditioning.
Revised: ScaledSGD preserves all the favorable qualities of SGD while making it immune to ill-conditioning.

Original: For the RMSE loss, we prove that with an initial point close to the ground truth, ScaledSGD converges to an ✏ -accurate solution in O (log(1 / ✏ )) iterations, independent of the condition number  .
Revised: For the RMSE loss, we prove that with an initial point close to the ground truth, ScaledSGD converges to an ✏ -accurate solution in O (log(1 / ✏ )) iterations, independent of the condition number  .

Original: We also run numerical experiments on a wide range of other loss functions commonly used in applications such as collaborative ﬁltering, distance matrix recovery, etc.
Revised: We also run numerical experiments on a wide range of other loss functions commonly used in applications such as collaborative ﬁltering, distance matrix recovery, etc.

Original: We ﬁnd that ScaledSGD achieves similar acceleration on these losses, which means that it is widely applicable to many real problems.
Revised: We ﬁnd that ScaledSGD achieves similar acceleration on these losses, which means that it is widely applicable to many real problems.

Original: It remains future work to provide rigorous justiﬁcation for these observations.
Revised: It remains future work to provide rigorous justiﬁcation for these observations.

Revised: Acknowledgments

Revised: The authors thank Salar Fattahi for helpful discussions and feedback on an earlier draft.

Revised: Financial support for this work was provided in part by the NSF CAREER Award ECCS-2047462 and in part by C3.ai Inc. and the Microsoft Corporation via the C3.ai Digital Transformation Institute.