Original: L EARNING THE P OSITIONS IN C OUNT S KETCH
Revised: L EARNING THE P OSITIONS IN C OUNT S KETCH

Original: A BSTRACT
Revised: A BSTRACT

Original: We consider sketching algorithms which first compress data by multiplication with a random sketch matrix, and then apply the sketch to quickly solve an optimization problem, e.g., low-rank approximation and regression.
Revised: We consider sketching algorithms which first compress data by multiplication with a random sketch matrix, and then apply the sketch to quickly solve an optimization problem, e.g., low-rank approximation and regression.

Original: In the learning-based sketching paradigm proposed by Indyk, Vakilian, and Yuan (2019), the sketch matrix is found by choosing a random sparse matrix, e.g., CountSketch, and then the values of its non-zero entries are updated by running gradient descent on a training data set.
Revised: In the learning-based sketching paradigm proposed by Indyk, Vakilian, and Yuan (2019), the sketch matrix is found by choosing a random sparse matrix, e.g., CountSketch, and then the values of its non-zero entries are updated by running gradient descent on a training data set.

Original: Despite the growing body of work on this paradigm, a noticeable omission is that the locations of the non-zero entries of previous algorithms were fixed, and only their values were learned.
Revised: Despite the growing body of work on this paradigm, a noticeable omission is that the locations of the non-zero entries of previous algorithms were fixed, and only their values were learned.

Original: In this work, we propose the first learning-based algorithms that also optimize the locations of the non-zero entries.
Revised: In this work, we propose the first learning-based algorithms that also optimize the locations of the non-zero entries.

Original: Our first proposed algorithm is based on a greedy algorithm.
Revised: Our first proposed algorithm is based on a greedy algorithm.

Original: However, one drawback of the greedy algorithm is its slower training time.
Revised: However, one drawback of the greedy algorithm is its slower training time.

Original: We fix this issue and propose approaches for learning a sketching matrix for both low-rank approximation and Hessian approximation for second order optimization.
Revised: We fix this issue and propose approaches for learning a sketching matrix for both low-rank approximation and Hessian approximation for second order optimization.

Original: The latter is helpful for a range of constrained optimization problems, such as LASSO and matrix estimation with a nuclear norm constraint.
Revised: The latter is helpful for a range of constrained optimization problems, such as LASSO and matrix estimation with a nuclear norm constraint.

Original: Both approaches achieve good accuracy with a fast running time.
Revised: Both approaches achieve good accuracy with a fast running time.

Original: Moreover, our experiments suggest that our algorithm can still reduce the error significantly even if we only have a very limited number of training matrices.
Revised: Moreover, our experiments suggest that our algorithm can still reduce the error significantly even if we only have a very limited number of training matrices.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Indyk et al.
Revised: Indyk et al.

Original: (2019) investigated learning-based sketching algorithms for low-rank approximation.
Revised: (2019) investigated learning-based sketching algorithms for low-rank approximation.

Original: A sketching algorithm is a method of constructing approximate solutions for optimization problems via summarizing the data.
Revised: A sketching algorithm is a method of constructing approximate solutions for optimization problems via summarizing the data.

Original: In particular, linear sketching algorithms compress data by multiplication with a sparse “sketch matrix” and then use just the compressed data to find an approximate solution.
Revised: In particular, linear sketching algorithms compress data by multiplication with a sparse “sketch matrix” and then use just the compressed data to find an approximate solution.

Original: Generally, this technique results in much faster or more space-efficient algorithms for a fixed approximation error.
Revised: Generally, this technique results in much faster or more space-efficient algorithms for a fixed approximation error.

Original: The pioneering work of Indyk et al.
Revised: The pioneering work of Indyk et al.

Original:  shows it is possible to learn sketch matrices for low-rank approximation (LRA) with better average performance than classical sketches.
Revised: (2019) shows it is possible to learn sketch matrices for low-rank approximation (LRA) with better average performance than classical sketches.

Original: In this model, we assume inputs come from an unknown distribution and learn a sketch matrix with strong expected performance over the distribution.
Revised: In this model, we assume inputs come from an unknown distribution and learn a sketch matrix with strong expected performance over the distribution.

Original: This distributional assumption is often realistic – there are many situations where a sketching algorithm is applied to a large batch of related data.
Revised: This distributional assumption is often realistic – there are many situations where a sketching algorithm is applied to a large batch of related data.

Original: For example, genomics researchers might sketch DNA from different individuals, which is known to exhibit strong commonalities.
Revised: For example, genomics researchers might sketch DNA from different individuals, which is known to exhibit strong commonalities.

Original: The high-performance computing industry also uses sketching, e.g., researchers at NVIDIA have created standard implementations of sketching algorithms for CUDA, a widely used GPU library.
Revised: The high-performance computing industry also uses sketching, e.g., researchers at NVIDIA have created standard implementations of sketching algorithms for CUDA, a widely used GPU library.

Original: They investigated the (classical) sketched singular value decomposition (SVD), but found that the solutions were not accurate enough across a spectrum of inputs (Chien &
Revised: They investigated the (classical) sketched singular value decomposition (SVD), but found that the solutions were not accurate enough across a spectrum of inputs (Chien &

Original: Bernabeu, 2019).
Revised: Bernabeu, 2019).

Original: This is precisely the issue addressed by the learned sketch paradigm where we optimize for “good” average performance across a range of inputs.
Revised: This is precisely the issue addressed by the learned sketch paradigm where we optimize for “good” average performance across a range of inputs.

Original: While promising results have been shown using previous learned sketching techniques, notable gaps remain.
Revised: While promising results have been shown using previous learned sketching techniques, notable gaps remain.

Original: In particular, all previous methods work by initializing the sketching matrix with a random sparse matrix, e.g., each column of the sketching matrix has a single non-zero value chosen at a uniformly random position.
Revised: In particular, all previous methods work by initializing the sketching matrix with a random sparse matrix, e.g., each column of the sketching matrix has a single non-zero value chosen at a uniformly random position.

Original: Then, the values of the non-zero entries are updated by running gradient descent on a training data set, or via other methods.
Revised: Then, the values of the non-zero entries are updated by running gradient descent on a training data set, or via other methods.

Original: However, the locations of the non-zero entries are held fixed throughout the entire training process.
Revised: However, the locations of the non-zero entries are held fixed throughout the entire training process.

Original: Clearly this is sub-optimal.
Revised: Clearly this is sub-optimal.

Original: Indeed, suppose the input matrix A is an n × d matrix with first
Revised: Indeed, suppose the input matrix A is an n × d matrix with first

Original: d rows equal to the d × d identity matrix, and remaining rows equal to 0 .
Revised: d rows equal to the d × d identity matrix, and remaining rows equal to 0 .

Original: A random sketching matrix S with a single non-zero per column is known to require m = Ω( d 2 ) rows in order for S
Revised: A random sketching matrix S with a single non-zero per column is known to require m = Ω( d 2 ) rows in order for S

Original: · A to preserve the rank of A (Nelson & Nguyên, 2014); this follows by a birthday paradox argument.
Revised: · A to preserve the rank of A (Nelson & Nguyên, 2014); this follows by a birthday paradox argument.

Original: On the other hand, it is clear that if S is a d × n matrix with first d rows equal to the identity matrix, then ∥ S · Ax ∥ 2 = ∥
Revised: On the other hand, it is clear that if S is a d × n matrix with first d rows equal to the identity matrix, then ∥ S · Ax ∥ 2 = ∥

Original: Ax ∥for all vectors x , and so S preserves not only the rank of A but all important spectral properties.
Revised: Ax ∥for all vectors x , and so S preserves not only the rank of A but all important spectral properties.

Original: A random matrix would be very unlikely to choose the non-zero entries in the first d columns of S so perfectly, whereas an algorithm trained to optimize the locations of the non-zero entries would notice and correct for this.
Revised: A random matrix would be very unlikely to choose the non-zero entries in the first d columns of S so perfectly, whereas an algorithm trained to optimize the locations of the non-zero entries would notice and correct for this.

Original: This is precisely the gap in our understanding that we seek to fill.
Revised: This is precisely the gap in our understanding that we seek to fill.

Original: Learned CountSketch Paradigm of Indyk et al.
Revised: Learned CountSketch Paradigm of Indyk et al.

Original: Throughout the paper, we assume our data A ∈ R n × d is sampled from an unknown distribution D .
Revised: Throughout the paper, we assume our data A ∈ R n × d is sampled from an unknown distribution D .

Original: Specifically, we have a training set
Revised: Specifically, we have a training set

Original: Tr = { A 1 , . . .
Revised: Tr = { A 1 , . . .

Original: , A N } ∈ D .
Revised: , A N } ∈ D .

Original: The generic form of our optimization problems is min X f ( A, X ) , where
Revised: The generic form of our optimization problems is min X f ( A, X ) , where

Original: A ∈ R n × d is the input matrix.
Revised: A ∈ R n × d is the input matrix.

Original: For a given optimization problem and a set S of sketching matrices, define ALG ( S , A ) to be the output of the classical sketching algorithm resulting from using S ; this uses the sketching matrices in S to map the given input A and construct an approximate solution ˆ X .
Revised: For a given optimization problem and a set S of sketching matrices, define ALG ( S , A ) to be the output of the classical sketching algorithm resulting from using S ; this uses the sketching matrices in S to map the given input A and construct an approximate solution ˆ X .

Original: We remark that the number of sketches used by an algorithm can vary and in its simplest case, S is a single sketch, but in more complicated sketching approaches we may need to apply sketching more than once—hence S may also denote a set of more than one sketching matrix.
Revised: We remark that the number of sketches used by an algorithm can vary and in its simplest case, S is a single sketch, but in more complicated sketching approaches we may need to apply sketching more than once—hence S may also denote a set of more than one sketching matrix.

Original: The learned sketch framework has two parts: (1) offline sketch learning and (2) “online” sketching (i.e., applying the learned sketch and some sketching algorithm to possibly unseen data).
Revised: The learned sketch framework has two parts: (1) offline sketch learning and (2) “online” sketching (i.e., applying the learned sketch and some sketching algorithm to possibly unseen data).

Original: In offline sketch learning, the goal is to construct a CountSketch matrix (abbreviated as CS matrix) with the minimum expected error for the problem of interest.
Revised: In offline sketch learning, the goal is to construct a CountSketch matrix (abbreviated as CS matrix) with the minimum expected error for the problem of interest.

Original: Formally, that is,
Revised: Formally, that is,

Original: [Equation]
Revised: [Equation]

Original: where X ∗ denotes the optimal solution.
Revised: where X ∗ denotes the optimal solution.

Original: Moreover, the minimum is taken over all possible constructions of CS .
Revised: Moreover, the minimum is taken over all possible constructions of CS .

Original: We remark that when ALG needs more than one CS to be learned (e.g., in the sketching algorithm we consider for LRA), we optimize each CS independently using a surrogate loss function.
Revised: We remark that when ALG needs more than one CS to be learned (e.g., in the sketching algorithm we consider for LRA), we optimize each CS independently using a surrogate loss function.

Original: In the second part of the learned sketch paradigm, we take the sketch from part one and use it within a sketching algorithm.
Revised: In the second part of the learned sketch paradigm, we take the sketch from part one and use it within a sketching algorithm.

Original: This learned sketch and sketching algorithm can be applied, again and again, to different inputs.
Revised: This learned sketch and sketching algorithm can be applied, again and again, to different inputs.

Original: Finally, we augment the sketching algorithm to provide worst-case guarantees when used with learned sketches.
Revised: Finally, we augment the sketching algorithm to provide worst-case guarantees when used with learned sketches.

Original: The goal is to have good performance on A ∈ D while the worst-case performance on A ̸∈ D remains comparable to the guarantees of classical sketches.
Revised: The goal is to have good performance on A ∈ D while the worst-case performance on A ̸∈ D remains comparable to the guarantees of classical sketches.

Original: We remark that the learned matrix S is trained offline only once using the training data.
Revised: We remark that the learned matrix S is trained offline only once using the training data.

Original: Hence, no additional computational cost is incurred when solving the optimization problem on the test data.
Revised: Hence, no additional computational cost is incurred when solving the optimization problem on the test data.

Original: Our Results.
Revised: Our Results.

Original: In this work, in addition to learning the values of the non-zero entries, we learn the locations of the non-zero entries.
Revised: In this work, in addition to learning the values of the non-zero entries, we learn the locations of the non-zero entries.

Original: Namely, we propose three algorithms that learn the locations of the non-zero entries in CountSketch.
Revised: Namely, we propose three algorithms that learn the locations of the non-zero entries in CountSketch.

Original: Our first algorithm (Section 4) is based on a greedy search.
Revised: Our first algorithm (Section 4) is based on a greedy search.

Original: The empirical result shows that this approach can achieve a good performance.
Revised: The empirical result shows that this approach can achieve a good performance.

Original: Further, we show that the greedy algorithm is provably beneficial for LRA when inputs follow a certain input distribution (Section F).
Revised: Further, we show that the greedy algorithm is provably beneficial for LRA when inputs follow a certain input distribution (Section F).

Original: However, one drawback of the greedy algorithm is its much slower training time.
Revised: However, one drawback of the greedy algorithm is its much slower training time.

Original: We then fix this issue and propose two specific approaches for optimizing the positions for the sketches for low-rank approximation and second-order optimization, which run much faster than all previous algorithms while achieving better performance.
Revised: We then fix this issue and propose two specific approaches for optimizing the positions for the sketches for low-rank approximation and second-order optimization, which run much faster than all previous algorithms while achieving better performance.

Original: For low-rank approximation, our approach is based on first sampling a small set of rows based on their ridge leverage scores, assigning each of these sampled rows to a unique hash bucket, and then placing each non-sampled remaining row in the hash bucket containing the sampled row for which it is most similar to, i.e., for which it has the largest dot product with.
Revised: For low-rank approximation, our approach is based on first sampling a small set of rows based on their ridge leverage scores, assigning each of these sampled rows to a unique hash bucket, and then placing each non-sampled remaining row in the hash bucket containing the sampled row for which it is most similar to, i.e., for which it has the largest dot product with.

Original: We also show that the worst-case guarantee of this approach is strictly better than that of the classical Count-Sketch (see Section 5).
Revised: We also show that the worst-case guarantee of this approach is strictly better than that of the classical Count-Sketch (see Section 5).

Original: For sketch-based second-order optimization where we focus on the case that n ≫ d , we observe that the actual property of the sketch matrix we need is the subspace embedding property .
Revised: For sketch-based second-order optimization where we focus on the case that n ≫ d , we observe that the actual property of the sketch matrix we need is the subspace embedding property .

Original: We next optimize this property of the sketch matrix.
Revised: We next optimize this property of the sketch matrix.

Original: We provably show that the sketch matrix S needs fewer rows, with optimized positions of the non-zero entries, when the input matrix A has a small number of rows with a heavy leverage score.
Revised: We provably show that the sketch matrix S needs fewer rows, with optimized positions of the non-zero entries, when the input matrix A has a small number of rows with a heavy leverage score.

Original: More precisely, while CountSketch takes O ( d 2 / ( δϵ 2 )) rows with failure probability δ , in our construction, S requires only O (( d polylog(1 /ϵ )
Revised: More precisely, while CountSketch takes O ( d 2 / ( δϵ 2 )) rows with failure probability δ , in our construction, S requires only O (( d polylog(1 /ϵ )

Original: + log(1 /δ ))
Revised: + log(1 /δ ))

Original: /ϵ 2 ) rows if A has at most d polylog(1 /ϵ )
Revised: /ϵ 2 ) rows if A has at most d polylog(1 /ϵ )

Original: /ϵ
Revised: /ϵ

Original: 2 rows with leverage score at least ϵ/d .
Revised: 2 rows with leverage score at least ϵ/d .

Original: This is a quadratic improvement in d and an exponential improvement in δ .
Revised: This is a quadratic improvement in d and an exponential improvement in δ .

Original: In practice, it is not necessary to calculate the leverage scores.
Revised: In practice, it is not necessary to calculate the leverage scores.

Original: Instead, we show in our experiments that the indices of the rows of heavy leverage score can be learned and the induced S is still accurate.
Revised: Instead, we show in our experiments that the indices of the rows of heavy leverage score can be learned and the induced S is still accurate.

Original: We also consider a new learning objective, that is, we directly optimize the subspace embedding property of the sketching matrix instead of optimizing the error in the objective function of the optimization problem in hand.
Revised: We also consider a new learning objective, that is, we directly optimize the subspace embedding property of the sketching matrix instead of optimizing the error in the objective function of the optimization problem in hand.

Original: This demonstrates a significant advantage over non-learned sketches, and has a fast training time (Section 6).
Revised: This demonstrates a significant advantage over non-learned sketches, and has a fast training time (Section 6).

Original: We show strong empirical results for real-world datasets.
Revised: We show strong empirical results for real-world datasets.

Original: For low-rank approximation, our methods reduce the errors by 70% than classical sketches under the same sketch size, while we reduce the errors by 30% than previous learning-based sketches.
Revised: For low-rank approximation, our methods reduce the errors by 70% than classical sketches under the same sketch size, while we reduce the errors by 30% than previous learning-based sketches.

Original: For second-order optimization, we show that the convergence rate can be reduced by 87% over the non-learned CountSketch for the LASSO problem on a real-world dataset.
Revised: For second-order optimization, we show that the convergence rate can be reduced by 87% over the non-learned CountSketch for the LASSO problem on a real-world dataset.

Original: We also evaluate our approaches in the few-shot learning setting where we only have a limited amount of training data (Indyk et al., 2021).
Revised: We also evaluate our approaches in the few-shot learning setting where we only have a limited amount of training data (Indyk et al., 2021).

Original: We show our approach reduces the error significantly even if we only have one training matrix (Sections 7 and 8).
Revised: We show our approach reduces the error significantly even if we only have one training matrix (Sections 7 and 8).

Original: This approach clearly runs faster than all previous methods.
Revised: This approach clearly runs faster than all previous methods.

Original: Additional Related Work.
Revised: Additional Related Work.

Original: In the last few years, there has been much work on leveraging machine learning techniques to improve classical algorithms.
Revised: In the last few years, there has been much work on leveraging machine learning techniques to improve classical algorithms.

Original: We only mention a few examples here which are based on learned sketches.
Revised: We only mention a few examples here which are based on learned sketches.

Original: One related body of work is data-dependent dimensionality reduction, such as an approach for pair-wise/multi-wise similarity preservation for indexing big data (Wang et al., 2017), learned sketching for streaming problems (Indyk et al., 2019; Aamand et al., 2019;
Revised: One related body of work is data-dependent dimensionality reduction, such as an approach for pair-wise/multi-wise similarity preservation for indexing big data (Wang et al., 2017), learned sketching for streaming problems (Indyk et al., 2019; Aamand et al., 2019;

Original: Jiang et al., 2020; Cohen et al., 2020; Eden et al., 2021; Indyk et al., 2021), learned algorithms for nearest neighbor search (Dong et al., 2020), and a method for learning linear projections for general applications (Hegde et al., 2015).
Revised: Jiang et al., 2020; Cohen et al., 2020; Eden et al., 2021; Indyk et al., 2021), learned algorithms for nearest neighbor search (Dong et al., 2020), and a method for learning linear projections for general applications (Hegde et al., 2015).

Original: While we also learn linear embeddings, our embeddings are optimized for the specific application of low rank approximation.
Revised: While we also learn linear embeddings, our embeddings are optimized for the specific application of low rank approximation.

Original: In fact, one of our central challenges is that the theory and practice of learned sketches generally needs to be tailored to each application.
Revised: In fact, one of our central challenges is that the theory and practice of learned sketches generally needs to be tailored to each application.

Original: Our work builds off of (Indyk et al., 2019), which introduced gradient descent optimization for LRA, but a major difference is that we also optimize the locations of the non-zero entries.
Revised: Our work builds off of (Indyk et al., 2019), which introduced gradient descent optimization for LRA, but a major difference is that we also optimize the locations of the non-zero entries.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: Notation.
Revised: Notation.

Original: Denote the canonical basis vectors of R n by e 1 , . .
Revised: Denote the canonical basis vectors of R n by e 1 , . .

Original: e n .
Revised: e n .

Original: Suppose that A has singular value decomposition (SVD)
Revised: Suppose that A has singular value decomposition (SVD)

Original: A = U Σ V ⊤ .
Revised: A = U Σ V ⊤ .

Original: Define [ A ] k = U k Σ k V ⊤ k to be the optimal rank- k approximation to A , computed by the truncated SVD.
Revised: Define [ A ] k = U k Σ k V ⊤ k to be the optimal rank- k approximation to A , computed by the truncated SVD.

Original: Also, define the Moore-Penrose pseudo-inverse of A to be A † = V Σ − 1 U ⊤ , where Σ − 1 is constructed by inverting the non-zero diagonal entries.
Revised: Also, define the Moore-Penrose pseudo-inverse of A to be A † = V Σ − 1 U ⊤ , where Σ − 1 is constructed by inverting the non-zero diagonal entries.

Original: Let row ( A ) and col ( A ) be the row space and the column space of A , respectively.
Revised: Let row ( A ) and col ( A ) be the row space and the column space of A , respectively.

Original: CountSketch.
Revised: CountSketch.

Original: We define S C ∈
Revised: We define S C ∈

Original: R m × n as a classical CountSketch (abbreviated as CS ).
Revised: R m × n as a classical CountSketch (abbreviated as CS ).

Original: It is a sparse matrix with one nonzero entry from {± 1 } per column.
Revised: It is a sparse matrix with one nonzero entry from {± 1 } per column.

Original: The position and value of this nonzero entry are chosen uniformly at random.
Revised: The position and value of this nonzero entry are chosen uniformly at random.

Original: CountSketch matrices can be succinctly represented by two vectors.
Revised: CountSketch matrices can be succinctly represented by two vectors.

Original: We define p ∈ [ m ] n , v ∈ R n as the positions and values of the nonzero entries, respectively.
Revised: We define p ∈ [ m ] n , v ∈ R n as the positions and values of the nonzero entries, respectively.

Original: Further, we let CS ( p, v ) be the CountSketch constructed from vectors p and v .
Revised: Further, we let CS ( p, v ) be the CountSketch constructed from vectors p and v .

Original: Below we define the objective function f ( · , · ) and a classical sketching algorithm ALG ( S , A ) for each individual problem.
Revised: Below we define the objective function f ( · , · ) and a classical sketching algorithm ALG ( S , A ) for each individual problem.

Original: Low-rank approximation (LRA).
Revised: Low-rank approximation (LRA).

Original: In LRA, we find a rank- k approximation of our data that minimizes the Frobenius norm of the approximation error.
Revised: In LRA, we find a rank- k approximation of our data that minimizes the Frobenius norm of the approximation error.

Original: For A ∈ R n × d , min rank- k B f LRA ( A, B ) =
Revised: For A ∈ R n × d , min rank- k B f LRA ( A, B ) =

Original: min rank-
Revised: min rank-

Original: k X ∥
Revised: k X ∥

Original: A − B ∥ 2 F .
Revised: A − B ∥ 2 F .

Original: Usually, instead of outputting the a whole B ∈ R n × d , the algorithm outputs two factors Y ∈ R n × k and X ∈ R k ×
Revised: Usually, instead of outputting the a whole B ∈ R n × d , the algorithm outputs two factors Y ∈ R n × k and X ∈ R k ×

Original: d
Revised: d

Original: such that B = Y X for efficiency.
Revised: such that B = Y X for efficiency.

Original: Indyk et al.
Revised: Indyk et al.

Original: (2019) considered Algorithm 1, which only compresses one side of the input matrix
Revised: (2019) considered Algorithm 1, which only compresses one side of the input matrix

Original: A .
Revised: A .

Original: However, in practice often both dimensions of the matrix
Revised: However, in practice often both dimensions of the matrix

Original: A are large.
Revised: A are large.

Original: Hence, in this work we consider Algorithm 2 that compresses both sides of A .
Revised: Hence, in this work we consider Algorithm 2 that compresses both sides of A .

Original: Constrained regression.
Revised: Constrained regression.

Original: Given a vector b ∈ R n , a matrix A ∈ R n × d ( n ≫ d ) and a convex set C , we want to find x to minimize the squared error
Revised: Given a vector b ∈ R n , a matrix A ∈ R n × d ( n ≫ d ) and a convex set C , we want to find x to minimize the squared error

Original: [Equation]
Revised: [Equation]

Original: Iterative Hessian Sketch.
Revised: Iterative Hessian Sketch.

Original: The Iterative Hessian Sketching (IHS) method (Pilanci & Wainwright, 2016) solves the constrained least-squares problem by iteratively performing the update
Revised: The Iterative Hessian Sketching (IHS) method (Pilanci & Wainwright, 2016) solves the constrained least-squares problem by iteratively performing the update

Original: [Equation]
Revised: [Equation]

Original: where S t +1 is a sketching matrix.
Revised: where S t +1 is a sketching matrix.

Original: It is not difficult to see that for the unsketched version ( S t +1 is the identity matrix) of (2.2), the optimal solution x t +1 coincides with the optimal solution to the original constrained regression problem (2.1).
Revised: It is not difficult to see that for the unsketched version ( S t +1 is the identity matrix) of (2.2), the optimal solution x t +1 coincides with the optimal solution to the original constrained regression problem (2.1).

Original: The IHS approximates the Hessian A ⊤ A by a sketched version ( S t +1 A ) ⊤ ( S t +1 A ) to improve runtime, as S t +

Revised: The IHS approximates the Hessian A ⊤ A by a sketched version ( S t +1 A ) ⊤ ( S t +1 A ) to improve runtime, as S t +

Original: A typically has very few rows.
Revised: A typically has very few rows.

Original: Algorithm a sketch S (see (Clarkson & Woodruff, 2009,
Revised: Algorithm a sketch S (see (Clarkson & Woodruff, 2009,

Original: Sec. 4.1.1))
Revised: Sec. 4.1.1))

Original: Input: A ∈ R n × d , S ∈ R m × n 1: U, Σ , V ⊤ ← C OMPACT SVD ( SA ) ▷ { r = rank( SA ) , U ∈ R m × r , V ∈ R d × r } 2: Return: [ AV ] k V ⊤ los (2006); Clarkson & Woodruff (2017)
Revised: Input: A ∈ R n × d , S ∈ R m × n 1: U, Σ , V ⊤ ← C OMPACT SVD ( SA ) ▷ { r = rank( SA ) , U ∈ R m × r , V ∈ R d × r } 2: Return: [ AV ] k V ⊤ los (2006); Clarkson & Woodruff (2017)

Original: Avron et al.
Revised: Avron et al.

Original: Input: A ∈ R n × d , S ∈ R m S × n , R ∈
Revised: Input: A ∈ R n × d , S ∈ R m S × n , R ∈

Original: R
Revised: R

Original: m R × d ,
Revised: m R × d ,

Original: V ∈ R m V × n , W ∈ R m W ×
Revised: V ∈ R m V × n , W ∈ R m W ×

Original: d 1: U C [ T C
Revised: d 1: U C [ T C

Original: T ′ C ] ← V AR ⊤ , (cid:20) T ⊤ D T ′⊤
Revised: T ′ C ] ← V AR ⊤ , (cid:20) T ⊤ D T ′⊤

Original: D (cid:21)
Revised: D (cid:21)

Original: U ⊤ D ←
Revised: U ⊤ D ←

Original: SAW ⊤ with U C , U D orthogonal 2: G ← V AW ⊤ , Z ′
Revised: SAW ⊤ with U C , U D orthogonal 2: G ← V AW ⊤ , Z ′

Original: L Z ′ R ← [ U ⊤ C GU D ] k 3: Z L ← (cid:2)
Revised: L Z ′ R ← [ U ⊤ C GU D ] k 3: Z L ← (cid:2)

Original: Z ′
Revised: Z ′

Original: L ( T − 1 D ) ⊤ 0 (cid:3) , Z R ← (cid:20) T − 1 C Z ′ R(cid:21) 4: Z ← Z L Z R 5: return: AR ⊤ ZSA in form P n × k , Q k × d
Revised: L ( T − 1 D ) ⊤ 0 (cid:3) , Z R ← (cid:20) T − 1 C Z ′ R(cid:21) 4: Z ← Z L Z R 5: return: AR ⊤ ZSA in form P n × k , Q k × d

Original: Learning-Based Algorithms in the Few-Shot
Revised: Learning-Based Algorithms in the Few-Shot

Original: Setting.
Revised: Setting.

Original: Recently, Indyk et al.
Revised: Recently, Indyk et al.

Original: (2021) studied learning-based algorithms for LRA in the setting where we have access to limited data or computing resources.
Revised: (2021) studied learning-based algorithms for LRA in the setting where we have access to limited data or computing resources.

Original: We provide a brief explanation of learning-based algorithms in the Few-Shot setting in Appendix A.3.
Revised: We provide a brief explanation of learning-based algorithms in the Few-Shot setting in Appendix A.3.

Original: Leverage Scores and Ridge Leverage Scores.
Revised: Leverage Scores and Ridge Leverage Scores.

Original: Given a matrix A , the leverage score of the i -th row a i of A is defined to be τ i := a i ( A ⊤ A ) † a ⊤ i , which is the squared ℓ 2 -norm of the i -th row of U , where A = U Σ V T is the singular value decomposition of A .
Revised: Given a matrix A , the leverage score of the i -th row a i of A is defined to be τ i := a i ( A ⊤ A ) † a ⊤ i , which is the squared ℓ 2 -norm of the i -th row of U , where A = U Σ V T is the singular value decomposition of A .

Original: Given a regularization parameter λ , the ridge leverage score of the i -th row a i of A is defined to be τ i := a i ( A ⊤ A + λI ) † a ⊤ i .
Revised: Given a regularization parameter λ , the ridge leverage score of the i -th row a i of A is defined to be τ i := a i ( A ⊤ A + λI ) † a ⊤ i .

Original: Our learning-based algorithms employs the ridge leverage score sampling technique proposed in (Cohen et al., 2017), which shows that sampling proportional to ridge leverage scores gives a good solution to LRA.
Revised: Our learning-based algorithms employs the ridge leverage score sampling technique proposed in (Cohen et al., 2017), which shows that sampling proportional to ridge leverage scores gives a good solution to LRA.

Original:  D ESCRIPTION OF O UR A PPROACH
Revised:  D ESCRIPTION OF O UR A PPROACH

Original: We describe our contributions to the learning-based sketching paradigm which, as mentioned, is to learn the locations of the non-zero values in the sketch matrix.
Revised: We describe our contributions to the learning-based sketching paradigm which, as mentioned, is to learn the locations of the non-zero values in the sketch matrix.

Original: To learn a CountSketch for the given training data set, we locally optimize the following in two stages:
Revised: To learn a CountSketch for the given training data set, we locally optimize the following in two stages:

Original: [Equation]
Revised: [Equation]

Original: 3.1) (1) compute the positions of the non-zero entries, then (2) fix the positions and optimize their values.
Revised: 3.1) (1) compute the positions of the non-zero entries, then (2) fix the positions and optimize their values.

Original: Stage 1: Optimizing Positions.
Revised: Stage 1: Optimizing Positions.

Original: In Section 4, we provide a greedy search algorithm for this stage, as our starting point.
Revised: In Section 4, we provide a greedy search algorithm for this stage, as our starting point.

Original: In Section 5 and 6, we provide our specific approaches for optimizing the positions for the sketches for low-rank approximation and second-order optimization.
Revised: In Section 5 and 6, we provide our specific approaches for optimizing the positions for the sketches for low-rank approximation and second-order optimization.

Original: Stage 2: Optimizing Values.
Revised: Stage 2: Optimizing Values.

Original: This stage is similar to the approach of Indyk et al.
Revised: This stage is similar to the approach of Indyk et al.

Original: However, instead of the power method, we use an automatic differentiation package, PyTorch (Paszke et al., 2019), and we pass it our objective min v ∈ R n E A ∈D [ f ( A, ALG ( CS ( p, v ) , A ))] , implemented as a chain of differentiable operations.
Revised: However, instead of the power method, we use an automatic differentiation package, PyTorch (Paszke et al., 2019), and we pass it our objective min v ∈ R n E A ∈D [ f ( A, ALG ( CS ( p, v ) , A ))] , implemented as a chain of differentiable operations.

Original: It will automatically compute the gradient using the chain rule.
Revised: It will automatically compute the gradient using the chain rule.

Original: We also consider new approaches to optimize the values for LRA (proposed in Indyk et al.
Revised: We also consider new approaches to optimize the values for LRA (proposed in Indyk et al.

Original: (2021), see
Revised: (2021), see

Original: Appendix A.3 for details) and second-order optimization (proposed in Section 6).
Revised: Appendix A.3 for details) and second-order optimization (proposed in Section 6).

Original: Worst-Cases Guarantees.
Revised: Worst-Cases Guarantees.

Original: In Appendix D, we show that both of our approaches for the above two problems can perform no worse than a classical sketching matrix when A does not follow the distribution D .
Revised: In Appendix D, we show that both of our approaches for the above two problems can perform no worse than a classical sketching matrix when A does not follow the distribution D .

Original: In particular, for LRA, we show that the sketch monotonicity property holds for the time-optimal sketching algorithm for low rank approximation.
Revised: In particular, for LRA, we show that the sketch monotonicity property holds for the time-optimal sketching algorithm for low rank approximation.

Original: For second-order optimization, we propose an algorithm which runs in input-sparsity time and can test for and use the better of a random sketch and a learned sketch.
Revised: For second-order optimization, we propose an algorithm which runs in input-sparsity time and can test for and use the better of a random sketch and a learned sketch.

Original:  S KETCH L EARNING : G REEDY S EARCH
Revised:  S KETCH L EARNING : G REEDY S EARCH

Original: REEDY
Revised: REEDY

Original: EARCH
Revised: EARCH

Original: Input: f, ALG , Tr = { A 1 , ..., A N ∈ R n × d } ; sketch dimension m 1: initialize S L
Revised: Input: f, ALG , Tr = { A 1 , ..., A N ∈ R n × d } ; sketch dimension m 1: initialize S L

Original: = O m × n 2: for i = 1 to n do 3: j ← arg min j ∈ [ m ] (cid:88)
Revised: = O m × n 2: for i = 1 to n do 3: j ← arg min j ∈ [ m ] (cid:88)

Original: A ∈ Tr f ( A, ALG ( S L ±
Revised: A ∈ Tr f ( A, ALG ( S L ±

Original: e j e ⊤ i , A )) 4:
Revised: e j e ⊤ i , A )) 4:

Original: S L ← S L ± ( e j e ⊤ i ) 5: end for 6: return p for S L =
Revised: S L ← S L ± ( e j e ⊤ i ) 5: end for 6: return p for S L =

Original: CS ( p, v )
Revised: CS ( p, v )

Original: When S is a CountSketch, computing SA amounts to hashing the n rows of A into the m ≪ n rows of SA .
Revised: When S is a CountSketch, computing SA amounts to hashing the n rows of A into the m ≪ n rows of SA .

Original: The optimization is a combinatorial optimization problem with an empirical risk minimization (ERM) objective.
Revised: The optimization is a combinatorial optimization problem with an empirical risk minimization (ERM) objective.

Original: The naïve solution is to compute the objective value of the exponentially many ( m n ) possible placements, but this is clearly intractable.
Revised: The naïve solution is to compute the objective value of the exponentially many ( m n ) possible placements, but this is clearly intractable.

Original: Instead, we iteratively construct a full placement in a greedy fashion.
Revised: Instead, we iteratively construct a full placement in a greedy fashion.

Original: We start with S as a zero matrix.
Revised: We start with S as a zero matrix.

Original: Then, we iterate through the columns of S in an order determined by the algorithm, adding a nonzero entry to each.
Revised: Then, we iterate through the columns of S in an order determined by the algorithm, adding a nonzero entry to each.

Original: The best position in each column is the one that minimizes Eq.
Revised: The best position in each column is the one that minimizes Eq.

Original: (3.1) if an entry were to be added there.
Revised: (3.1) if an entry were to be added there.

Original: For each column, we evaluate
Revised: For each column, we evaluate

Original: Eq.
Revised: Eq.

Original: (3.1) O ( m ) times, once for each prospective half-built sketch.
Revised: (3.1) O ( m ) times, once for each prospective half-built sketch.

Original: While this greedy strategy is simple to state, additional tactics are required for each problem to make it more tractable.
Revised: While this greedy strategy is simple to state, additional tactics are required for each problem to make it more tractable.

Original: Usually the objective evaluation (Algorithm 3, line 3) is too slow, so we must leverage our insight into their sketching algorithms to pick a proxy objective.
Revised: Usually the objective evaluation (Algorithm 3, line 3) is too slow, so we must leverage our insight into their sketching algorithms to pick a proxy objective.

Original: Note that we can reuse these proxies for value optimization, since they may make gradient computation faster too.
Revised: Note that we can reuse these proxies for value optimization, since they may make gradient computation faster too.

Original: Proxy objective for LRA.
Revised: Proxy objective for LRA.

Original: For the two-sided sketching algorithm, we can assume that the two factors
Revised: For the two-sided sketching algorithm, we can assume that the two factors

Original: X, Y has the form Y =
Revised: X, Y has the form Y =

Original: AR ⊤ ˜ Y and X = ˜ XSA , where S and R are both CS matrices, so we optimize the positions in both S and R .
Revised: AR ⊤ ˜ Y and X = ˜ XSA , where S and R are both CS matrices, so we optimize the positions in both S and R .

Original: We cannot use f ( A, ALG ( S, R, A )) as our objective because then we would have to consider combinations of placements between S and R .
Revised: We cannot use f ( A, ALG ( S, R, A )) as our objective because then we would have to consider combinations of placements between S and R .

Original: To find a proxy, we note that a prerequisite for good performance is for row ( SA ) and col ( AR ⊤ ) to both contain a good rank- k approximation to A (see proof of Lemma C.5).
Revised: To find a proxy, we note that a prerequisite for good performance is for row ( SA ) and col ( AR ⊤ ) to both contain a good rank- k approximation to A (see proof of Lemma C.5).

Original: Thus, we can decouple the optimization of S and R .
Revised: Thus, we can decouple the optimization of S and R .

Original: The proxy objective for S is (cid:13)(cid:13)
Revised: The proxy objective for S is (cid:13)(cid:13)

Original: [ AV ] k V ⊤ − A (cid:13)(cid:13) 2 F where SA = U Σ V ⊤ .
Revised: [ AV ] k V ⊤ − A (cid:13)(cid:13) 2 F where SA = U Σ V ⊤ .

Original: In this expression, ˆ X
Revised: In this expression, ˆ X

Original: =
Revised: =

Original: [ AV ] k V ⊤ is the best rank- k approximation to A in row ( SA ) .
Revised: [ AV ] k V ⊤ is the best rank- k approximation to A in row ( SA ) .

Original: The proxy objective for R is defined analogously.
Revised: The proxy objective for R is defined analogously.

Original: In Appendix F, we show the greedy algorithm is provably beneficial for LRA when inputs follow the spiked covariance or the Zipfian distribution.
Revised: In Appendix F, we show the greedy algorithm is provably beneficial for LRA when inputs follow the spiked covariance or the Zipfian distribution.

Original: Despite the good empirical performance we present in
Revised: Despite the good empirical performance we present in

Original: Section 7, one drawback is its much slower training time.
Revised: Section 7, one drawback is its much slower training time.

Original: Also, for the iterative sketching method for second-order optimization, it is non-trivial to find a proxy objective because the input of the i -th iteration depends on the solution to the ( i − 1) -th iteration, for which the greedy approach sometimes does not give a good solution.
Revised: Also, for the iterative sketching method for second-order optimization, it is non-trivial to find a proxy objective because the input of the i -th iteration depends on the solution to the ( i − 1) -th iteration, for which the greedy approach sometimes does not give a good solution.

Original: In the next section, we will propose our specific approach for optimizing the positions of the sketches for low-rank approximation and second-order optimization, both of which achieve a very high accuracy and can finish in a very short amount of time.
Revised: In the next section, we will propose our specific approach for optimizing the positions of the sketches for low-rank approximation and second-order optimization, both of which achieve a very high accuracy and can finish in a very short amount of time.

Original:  S KETCH L EARNING : L OW -R ANK A PPROXIMATION
Revised:  S KETCH L EARNING : L OW -R ANK A PPROXIMATION

Original: Now we present a conceptually new algorithm which runs much faster and empirically achieves similar error bounds as the greedy search approach.
Revised: Now we present a conceptually new algorithm which runs much faster and empirically achieves similar error bounds as the greedy search approach.

Original: Moreover, we show that this algorithm has strictly better guarantees than the classical Count-Sketch.
Revised: Moreover, we show that this algorithm has strictly better guarantees than the classical Count-Sketch.

Original: To achieve this, we need a more careful analysis.
Revised: To achieve this, we need a more careful analysis.

Original: To provide some intuition, if rank( SA )
Revised: To provide some intuition, if rank( SA )

Original: =
Revised: =

Original: k and
Revised: k and

Original: SA = U Σ V ⊤ , then the rank- k approximation cost is exactly (cid:13)(cid:13)
Revised: SA = U Σ V ⊤ , then the rank- k approximation cost is exactly (cid:13)(cid:13)

Original: AV V ⊤ − A (cid:13)(cid:13) 2 F , the projection cost onto col ( V ) .
Revised: AV V ⊤ − A (cid:13)(cid:13) 2 F , the projection cost onto col ( V ) .

Original: Minimizing it is equivalent to maximizing the sum of squared projection coefficients:
Revised: Minimizing it is equivalent to maximizing the sum of squared projection coefficients:

Original: [Equation]
Revised: [Equation]

Original: As mentioned, computing SA actually amounts to hashing the n rows of A to the m rows of SA .
Revised: As mentioned, computing SA actually amounts to hashing the n rows of A to the m rows of SA .

Original: Hence, intuitively, if we can put similar rows into the same bucket, we may get a smaller error.
Revised: Hence, intuitively, if we can put similar rows into the same bucket, we may get a smaller error.

Original: Our algorithm is given in Algorithm 4.
Revised: Our algorithm is given in Algorithm 4.

Original: Suppose that we want to form matrix S with m rows.
Revised: Suppose that we want to form matrix S with m rows.

Original: At the beginning of the algorithm, we sample m rows according to the ridge leverage scores of A .
Revised: At the beginning of the algorithm, we sample m rows according to the ridge leverage scores of A .

Original: By the property of the ridge leverage score, the subspace spanned by this set of sampled rows contains an approximately optimal solution to the low rank approximation problem.
Revised: By the property of the ridge leverage score, the subspace spanned by this set of sampled rows contains an approximately optimal solution to the low rank approximation problem.

Original: Hence, we map these rows to separate “buckets” of SA .
Revised: Hence, we map these rows to separate “buckets” of SA .

Original: Then, we need to decide the locations of the remaining rows (i.e., the non-sampled rows).
Revised: Then, we need to decide the locations of the remaining rows (i.e., the non-sampled rows).

Original: Ideally, we want similar rows to be mapped into the same bucket.
Revised: Ideally, we want similar rows to be mapped into the same bucket.

Original: To achieve this, we use the m sampled rows as reference points and assign each (nonsampled) row A i to the p -th bucket in SA if the normalized row A i and C p have the largest inner product (among all possible buckets).
Revised: To achieve this, we use the m sampled rows as reference points and assign each (nonsampled) row A i to the p -th bucket in SA if the normalized row A i and C p have the largest inner product (among all possible buckets).

Original: Once the locations of the non-zero entries are fixed, the next step is to determine the values of these entries.
Revised: Once the locations of the non-zero entries are fixed, the next step is to determine the values of these entries.

Original: We follow the same idea proposed in (Indyk et al., 2021): for each block
Revised: We follow the same idea proposed in (Indyk et al., 2021): for each block

Original: A ( i ) , one natural approach is to choose the unit vector s i ∈ R
Revised: A ( i ) , one natural approach is to choose the unit vector s i ∈ R

Original: | I i | that preserves as much of the Frobenius norm of A ( i ) as possible, i.e., to maximize (cid:13)(cid:13)
Revised: | I i | that preserves as much of the Frobenius norm of A ( i ) as possible, i.e., to maximize (cid:13)(cid:13)

Original: s
Revised: s

Original: ⊤ i
Revised: ⊤ i

Original: A ( i ) (cid:13)(cid:13) 2 2 .
Revised: A ( i ) (cid:13)(cid:13) 2 2 .

Original: Hence, we set s i to be the top left singular vector of A ( i ) .
Revised: Hence, we set s i to be the top left singular vector of A ( i ) .

Original: In our experiments, we observe that this step reduces the error of downstream value optimizations performed by SGD.
Revised: In our experiments, we observe that this step reduces the error of downstream value optimizations performed by SGD.

Original: To obtain a worst-case guarantee, we show that w.h.p., the row span of the sampled rows C i is a good subspace.
Revised: To obtain a worst-case guarantee, we show that w.h.p., the row span of the sampled rows C i is a good subspace.

Original: We set the matrix S 2 to be the sampling matrix that samples C i .
Revised: We set the matrix S 2 to be the sampling matrix that samples C i .

Original: The final output of our algorithm is the vertical concatenation of S 1 and S 2 .
Revised: The final output of our algorithm is the vertical concatenation of S 1 and S 2 .

Original: Here S 1 performs well empirically, while Shas a worst-case guarantee for any input.
Revised: Here S 1 performs well empirically, while Shas a worst-case guarantee for any input.

Original: Combining Lemma E.2 and the sketch monotonicity for low rank approximation in Section D, we get that O
Revised: Combining Lemma E.2 and the sketch monotonicity for low rank approximation in Section D, we get that O

Original: ( k log k + k/ϵ ) rows is enough for a (1 ± ϵ ) - approximation for the input matrix A induced by Tr , which is better than the Ω( k 2 ) rows required of a non-learned Count-Sketch, even if its non-zero values have been further improved by the previous learning-based algorithms in (Indyk et al., 2019; 2021).
Revised: ( k log k + k/ϵ ) rows is enough for a (1 ± ϵ ) - approximation for the input matrix A induced by Tr , which is better than the Ω( k 2 ) rows required of a non-learned Count-Sketch, even if its non-zero values have been further improved by the previous learning-based algorithms in (Indyk et al., 2019; 2021).

Original: As a result, under the assumption of the input data, we may expect that S will still be good for the test data.
Revised: As a result, under the assumption of the input data, we may expect that S will still be good for the test data.

Original: We defer the proof to Appendix E.1.
Revised: We defer the proof to Appendix E.1.

Original: In Appendix A, we shall show that the assumptions we make in Theorem 5.1 are reasonable.
Revised: In Appendix A, we shall show that the assumptions we make in Theorem 5.1 are reasonable.

Original: We also provide an empirical comparison between Algorithm 4 and some of its variants, as well as some adaptive sketching methods on the training sample.
Revised: We also provide an empirical comparison between Algorithm 4 and some of its variants, as well as some adaptive sketching methods on the training sample.

Original: The evaluation result shows that only our algorithm has a significant improvement for the test data, which suggests that both ridge leverage score sampling and row bucketing are essential.
Revised: The evaluation result shows that only our algorithm has a significant improvement for the test data, which suggests that both ridge leverage score sampling and row bucketing are essential.

Original: Theorem 5.1.
Revised: Theorem 5.1.

Original: Let S ∈ R 2 m × n be given by concatenating the sketching matrices S 1 , S 2 computed by
Revised: Let S ∈ R 2 m × n be given by concatenating the sketching matrices S 1 , S 2 computed by

Original: Algorithm 4 with input A induced by Tr and let B ∈ R n × d .
Revised: Algorithm 4 with input A induced by Tr and let B ∈ R n × d .

Original: Then with probability at least 1 − δ , we have min rank- k X : row ( X ) ⊆ row ( SB )
Revised: Then with probability at least 1 − δ , we have min rank- k X : row ( X ) ⊆ row ( SB )

Original: ∥
Revised: ∥

Original: B − X ∥ 2 F ≤ (1 + ϵ )
Revised: B − X ∥ 2 F ≤ (1 + ϵ )

Original: ∥
Revised: ∥

Original: B − B k ∥ 2 F if one of the following holds:.
Revised: B − B k ∥ 2 F if one of the following holds:.

Original: 1. m
Revised: 1. m

Original: =
Revised: =

Original: O ( β · ( k log k + k/ϵ )) , δ = 0 . 1 , and τ i ( B ) ≥ 

Revised: O ( β · ( k log k + k/ϵ )) , δ = 0 . 1 , and τ i ( B ) ≥

Original: β τ i ( A ) for all i ∈ [ n ] .
Revised: β τ i ( A ) for all i ∈ [ n ] .

Original: 2. m
Revised: 2. m

Original: =
Revised: =

Original: O ( k log k + k/ϵ ) , δ = 0 .
Revised: O ( k log k + k/ϵ ) , δ = 0 .

Original: 1 + 1 . 1 β , and the total variation distance d tv ( p, q ) ≤ β , where p, q are sampling probabilities defined as p i
Revised: 1 + 1 . 1 β , and the total variation distance d tv ( p, q ) ≤ β , where p, q are sampling probabilities defined as p i

Original: =
Revised: =

Original: τ i ( A ) (cid:80) i τ i
Revised: τ i ( A ) (cid:80) i τ i

Original: ( A ) and q i =
Revised: ( A ) and q i =

Original: τ i ( B ) (cid:80) i τ i ( B ) .
Revised: τ i ( B ) (cid:80) i τ i ( B ) .

Original: Time Complexity.
Revised: Time Complexity.

Original: As mentioned, an advantage of our second approach is that it significantly reduces the training time.
Revised: As mentioned, an advantage of our second approach is that it significantly reduces the training time.

Original: We now discuss the training times of different algorithms.
Revised: We now discuss the training times of different algorithms.

Original: For the value-learning algorithms in (Indyk et al., 2019), each iteration requires computing a differentiable SVD to perform gradient descent, hence the runtime is at least Ω( n it · T ) , where n it is the number of iterations (usually set > 500 ) and T is the time to compute an SVD.
Revised: For the value-learning algorithms in (Indyk et al., 2019), each iteration requires computing a differentiable SVD to perform gradient descent, hence the runtime is at least Ω( n it · T ) , where n it is the number of iterations (usually set > 500 ) and T is the time to compute an SVD.

Original: For the greedy algorithm, there are m choices for each column, hence the runtime is at least Ω( m n · T ) .
Revised: For the greedy algorithm, there are m choices for each column, hence the runtime is at least Ω( m n · T ) .

Original: For our second approach, the most complicated step is to compute the ridge leverage scores of A and then the SVD of each submatrix.
Revised: For our second approach, the most complicated step is to compute the ridge leverage scores of A and then the SVD of each submatrix.

Original: Hence, the total runtime is at most O ( T ) .
Revised: Hence, the total runtime is at most O ( T ) .

Original: We note that the time complexities discussed here are all for training time.
Revised: We note that the time complexities discussed here are all for training time.

Original: There is no additional runtime cost for the test data.
Revised: There is no additional runtime cost for the test data.

Original:  S KETCH L EARNING : S ECOND -O RDER O PTIMIZATION
Revised:  S KETCH L EARNING : S ECOND -O RDER O PTIMIZATION

Original: In this section, we consider optimizing the sketch matrix in the context of second-order methods.
Revised: In this section, we consider optimizing the sketch matrix in the context of second-order methods.

Original: The key observation is that for many sketching-based second-order methods, the crucial property of the sketching matrix is the so-called subspace embedding property: for a matrix A ∈ R n × d , we say a matrix S ∈ R m × n is a (1 ± ϵ ) -subspace embedding for the column space of A if (1 − ϵ ) ∥
Revised: The key observation is that for many sketching-based second-order methods, the crucial property of the sketching matrix is the so-called subspace embedding property: for a matrix A ∈ R n × d , we say a matrix S ∈ R m × n is a (1 ± ϵ ) -subspace embedding for the column space of A if (1 − ϵ ) ∥

Original: Ax ∥ 2 ≤ ∥ SAx ∥ 2 ≤ (1 + ϵ ) ∥
Revised: Ax ∥ 2 ≤ ∥ SAx ∥ 2 ≤ (1 + ϵ ) ∥

Original: Ax ∥ 2 for all x ∈ R d .
Revised: Ax ∥ 2 for all x ∈ R d .

Original: For example, consider the iterative Hessian sketch, which performs the update (2.2) to compute { x t } t .
Revised: For example, consider the iterative Hessian sketch, which performs the update (2.2) to compute { x t } t .

Original: Pilanci & Wainwright (2016) showed that if
Revised: Pilanci & Wainwright (2016) showed that if

Original: S 1 , . . .
Revised: S 1 , . . .

Original: , S t +1 are (1 + O ( ρ )) -subspace embeddings of A , then ∥
Revised: , S t +1 are (1 + O ( ρ )) -subspace embeddings of A , then ∥

Original: A ( x t − x ∗ ) ∥ 2 ≤ ρ t ∥
Revised: A ( x t − x ∗ ) ∥ 2 ≤ ρ t ∥

Original: Ax ∗ ∥ 2 .
Revised: Ax ∗ ∥ 2 .

Original: Thus, if S i is a good subspace embedding of A and we will have a good convergence guarantee.
Revised: Thus, if S i is a good subspace embedding of A and we will have a good convergence guarantee.

Original: Therefore, unlike (Indyk et al., 2019), which treats the training objective in a black-box manner, we shall optimize the subspace embedding property of the matrix A .
Revised: Therefore, unlike (Indyk et al., 2019), which treats the training objective in a black-box manner, we shall optimize the subspace embedding property of the matrix A .

Original: Optimizing positions.
Revised: Optimizing positions.

Original: We consider the case that A has a few rows of large leverage score, as well as access to an oracle which reveals a superset of the indices of such rows.
Revised: We consider the case that A has a few rows of large leverage score, as well as access to an oracle which reveals a superset of the indices of such rows.

Original: Formally, let τ i ( A ) be the leverage score of the i -th row of A and I ∗
Revised: Formally, let τ i ( A ) be the leverage score of the i -th row of A and I ∗

Original: = { i : τ i ( A ) ≥ ν } be the set of rows with large leverage score.
Revised: = { i : τ i ( A ) ≥ ν } be the set of rows with large leverage score.

Original: Suppose that a superset I ⊇ I ∗ is known to the algorithm.
Revised: Suppose that a superset I ⊇ I ∗ is known to the algorithm.

Original: In the experiments we train an oracle to predict such rows.
Revised: In the experiments we train an oracle to predict such rows.

Original: We can maintain all rows in I explicitly and apply a Count-Sketch to the remaining rows, i.e., the rows in [ n ] \ I .
Revised: We can maintain all rows in I explicitly and apply a Count-Sketch to the remaining rows, i.e., the rows in [ n ] \ I .

Original: Up to permutation of the rows, we can write
Revised: Up to permutation of the rows, we can write

Original: [Equation]
Revised: [Equation]

Original: S =
Revised: S =

Original: [Equation]
Revised: [Equation]

Original: where S ′ is a random Count-Sketch matrix of m rows.
Revised: where S ′ is a random Count-Sketch matrix of m rows.

Original: Clearly S has a single non-zero entry per column.
Revised: Clearly S has a single non-zero entry per column.

Original: We have the following theorem, whose proof is postponed to Section E.2.
Revised: We have the following theorem, whose proof is postponed to Section E.2.

Original: Intuitively, the proof for Count-Sketch in (Clarkson & Woodruff, 2017) handles rows of large leverage score and rows of small leverage score separately.
Revised: Intuitively, the proof for Count-Sketch in (Clarkson & Woodruff, 2017) handles rows of large leverage score and rows of small leverage score separately.

Original: The rows of large leverage score are to be perfectly hashed while the rows of small leverage score will concentrate in the sketch by the Hanson-Wright inequality.
Revised: The rows of large leverage score are to be perfectly hashed while the rows of small leverage score will concentrate in the sketch by the Hanson-Wright inequality.

Original: Theorem 6.1.
Revised: Theorem 6.1.

Original: Let ν = ϵ/d .
Revised: Let ν = ϵ/d .

Original: Suppose that m =
Revised: Suppose that m =

Original: O (( d/ϵ 2 )(polylog(1 /ϵ )
Revised: O (( d/ϵ 2 )(polylog(1 /ϵ )

Original: + log(1 /δ ))) , δ ∈ (0 , 1 /m ] and d = Ω((1 /ϵ )
Revised: + log(1 /δ ))) , δ ∈ (0 , 1 /m ] and d = Ω((1 /ϵ )

Original: polylog(1 /ϵ ) log 2 (1 /δ )) .
Revised: polylog(1 /ϵ ) log 2 (1 /δ )) .

Original: Then, there exists a distribution on S of the form in (6.1) with m + | I | rows such that Pr (cid:8) ∀ x ∈ col
Revised: Then, there exists a distribution on S of the form in (6.1) with m + | I | rows such that Pr (cid:8) ∀ x ∈ col

Original: ( A ) , | ∥ Sx ∥ 22 − ∥ x ∥ 22 | > ϵ ∥ x ∥ 22 (cid:9) ≤ δ .
Revised: ( A ) , | ∥ Sx ∥ 22 − ∥ x ∥ 22 | > ϵ ∥ x ∥ 22 (cid:9) ≤ δ .

Original: In particular, when δ = 1 /m , the sketching matrix S has O (( d/ϵ 2 ) polylog( d/ϵ )) rows.
Revised: In particular, when δ = 1 /m , the sketching matrix S has O (( d/ϵ 2 ) polylog( d/ϵ )) rows.

Original: Hence, if there happen to be
Revised: Hence, if there happen to be

Original: at most d polylog(1 /ϵ )
Revised: at most d polylog(1 /ϵ )

Original: /ϵ 2 rows of leverage score at least ϵ/d , the overall sketch length for embedding colsp( A ) can be reduced to O (( d polylog(1 /ϵ )
Revised: /ϵ 2 rows of leverage score at least ϵ/d , the overall sketch length for embedding colsp( A ) can be reduced to O (( d polylog(1 /ϵ )

Original: +log(1 /δ ))
Revised: +log(1 /δ ))

Original: /ϵ 2 )
Revised: /ϵ 2 )

Original: , a quadratic improvement in d and an exponential improvement in δ over the original sketch length of O ( d 2 / ( ϵ 2 δ )) for Count-Sketch.
Revised: , a quadratic improvement in d and an exponential improvement in δ over the original sketch length of O ( d 2 / ( ϵ 2 δ )) for Count-Sketch.

Original: In the worst case there could be O
Revised: In the worst case there could be O

Original: ( d 2 /ϵ ) such rows, though empirically we do not observe this.
Revised: ( d 2 /ϵ ) such rows, though empirically we do not observe this.

Original: In Section 8, we shall show it is possible to learn the indices of the heavy rows for real-world data.
Revised: In Section 8, we shall show it is possible to learn the indices of the heavy rows for real-world data.

Original: Optimizing values.
Revised: Optimizing values.

Original: When we fix the positions of the non-zero entries, we aim to optimize the values by gradient descent.
Revised: When we fix the positions of the non-zero entries, we aim to optimize the values by gradient descent.

Original: Rather than the previous black-box way in (Indyk et al., 2019) that minimizes (cid:80)
Revised: Rather than the previous black-box way in (Indyk et al., 2019) that minimizes (cid:80)

Original: i f
Revised: i f

Original: ( A, ALG ( S, A )) , we propose the following objective loss function for the learning algorithm
Revised: ( A, ALG ( S, A )) , we propose the following objective loss function for the learning algorithm

Original: L ( S, A ) =
Revised: L ( S, A ) =

Original: (cid:80) A i ∈A ∥
Revised: (cid:80) A i ∈A ∥

Original: ( A i R i )
Revised: ( A i R i )

Original: ⊤ A i R i − I ∥ F , over all the training data, where R i comes from the QR decomposition of SA
Revised: ⊤ A i R i − I ∥ F , over all the training data, where R i comes from the QR decomposition of SA

Original: i =
Revised: i =

Original: Q i R − 1 i .
Revised: Q i R − 1 i .

Original: The intuition for this loss function is given by the lemma below, whose proof is deferred to Section E.3.
Revised: The intuition for this loss function is given by the lemma below, whose proof is deferred to Section E.3.

Original: Lemma 6.2.
Revised: Lemma 6.2.

Original: Suppose that ϵ ∈ (0 , 12 ) , S ∈ R m × n , A ∈ R n × d of full column rank, and SA = QR is the QR-decomposition of SA .
Revised: Suppose that ϵ ∈ (0 , 12 ) , S ∈ R m × n , A ∈ R n × d of full column rank, and SA = QR is the QR-decomposition of SA .

Original: If ∥ ( AR − 1 ) ⊤ AR − 1 − I ∥ op ≤ ϵ , then S is a (1 ± ϵ ) -subspace embedding of col ( A ) .
Revised: If ∥ ( AR − 1 ) ⊤ AR − 1 − I ∥ op ≤ ϵ , then S is a (1 ± ϵ ) -subspace embedding of col ( A ) .

Original: Lemma 6.2 implies that if the loss function over A train is small and the distribution of A test is similar to A train , it is reasonable to expect that S is a good subspace embedding of A test .
Revised: Lemma 6.2 implies that if the loss function over A train is small and the distribution of A test is similar to A train , it is reasonable to expect that S is a good subspace embedding of A test .

Original: Here we use the
Revised: Here we use the

Original: Frobenius norm rather than operator norm in the loss function because it will make the optimization problem easier to solve, and our empirical results also show that the performance of the Frobenius norm is better than that of the operator norm.
Revised: Frobenius norm rather than operator norm in the loss function because it will make the optimization problem easier to solve, and our empirical results also show that the performance of the Frobenius norm is better than that of the operator norm.

Original:  E XPERIMENTS : L OW -R ANK A PPROXIMATION
Revised:  E XPERIMENTS : L OW -R ANK A PPROXIMATION

Original: In this section, we evaluate the empirical performance of our learning-based approach for LRA on three datasets.
Revised: In this section, we evaluate the empirical performance of our learning-based approach for LRA on three datasets.

Original: For each, we fix the sketch size and compare the approximation error ∥
Revised: For each, we fix the sketch size and compare the approximation error ∥

Original: A − X ∥ F − ∥
Revised: A − X ∥ F − ∥

Original: A − A k ∥ F averaged over 10 trials.
Revised: A − A k ∥ F averaged over 10 trials.

Original: In order to make position optimization more efficient, in lineof Algorithm 3), instead of computing many rank- 1 SVD updates, we use formulas for fast rankSVD updates (Brand, 2006).
Revised: In order to make position optimization more efficient, in lineof Algorithm 3), instead of computing many rank- 1 SVD updates, we use formulas for fast rankSVD updates (Brand, 2006).

Original: For the greedy method, we used several Nvidia GeForce GTX 1080 Ti machines.
Revised: For the greedy method, we used several Nvidia GeForce GTX 1080 Ti machines.

Original: For the maximum inner product method, the experiments are conducted on a laptop with a 1.90GHz CPU and 16GB RAM.
Revised: For the maximum inner product method, the experiments are conducted on a laptop with a 1.90GHz CPU and 16GB RAM.

Original: Datasets.
Revised: Datasets.

Original: We use the three datasets from (Indyk et al., 2019): (1, 2) Friends, Logo (image): frames from a short video of the TV show Friends and of a logo being painted; (3) Hyper (image): hyperspectral images from natural scenes.
Revised: We use the three datasets from (Indyk et al., 2019): (1, 2) Friends, Logo (image): frames from a short video of the TV show Friends and of a logo being painted; (3) Hyper (image): hyperspectral images from natural scenes.

Original: Additional details are in Table A.1.
Revised: Additional details are in Table A.1.

Original: Baselines.
Revised: Baselines.

Original: We compare our approach to the following baselines.
Revised: We compare our approach to the following baselines.

Original: Classical CS : a random CountSketch.
Revised: Classical CS : a random CountSketch.

Original: IVY19 : a sparse sketch with learned values, and random positions for the non-zero entries.
Revised: IVY19 : a sparse sketch with learned values, and random positions for the non-zero entries.

Original: Ours (greedy) : a sparse sketch where both the values and positions of the non-zero entries are learned.
Revised: Ours (greedy) : a sparse sketch where both the values and positions of the non-zero entries are learned.

Original: The positions are learned by Algorithm 3.
Revised: The positions are learned by Algorithm 3.

Original: The values are learned similarly to (Indyk et al., 2019).
Revised: The values are learned similarly to (Indyk et al., 2019).

Original: Ours (inner product) : a sparse sketch where both the values and the positions of the non-zero entries are learned.
Revised: Ours (inner product) : a sparse sketch where both the values and the positions of the non-zero entries are learned.

Original: The positions are learned by S 1 in Algorithm 4.
Revised: The positions are learned by S 1 in Algorithm 4.

Original: IVY19 and greedy algorithm use the full training set and our Algorithm 4 takes the input as the average over the entire training matrix.
Revised: IVY19 and greedy algorithm use the full training set and our Algorithm 4 takes the input as the average over the entire training matrix.

Original: We also give a sensitivity analysis for our algorithm, where we compare our algorithm with the following variants: Only row sampling (perform projection by ridge leverage score sampling), ℓsampling (Replace leverage score sampling with ℓ 2 -norm row sampling and maintain the same downstream step), and Randomly Grouping (Use ridge leverage score sampling but randomly distribute the remaining rows).
Revised: We also give a sensitivity analysis for our algorithm, where we compare our algorithm with the following variants: Only row sampling (perform projection by ridge leverage score sampling), ℓsampling (Replace leverage score sampling with ℓ 2 -norm row sampling and maintain the same downstream step), and Randomly Grouping (Use ridge leverage score sampling but randomly distribute the remaining rows).

Original: The result shows none of these variants outperforms non-learned sketching.
Revised: The result shows none of these variants outperforms non-learned sketching.

Original: We defer the results of this part to Appendix A.1.
Revised: We defer the results of this part to Appendix A.1.

Original: Result Summary.
Revised: Result Summary.

Original: Our empirical results are provided in Table 7.1 for both Algorithm 2 and Algorithm 1, where the errors take an average over 10 trials.
Revised: Our empirical results are provided in Table 7.1 for both Algorithm 2 and Algorithm 1, where the errors take an average over 10 trials.

Original: We use the average of all training matrices from Tr , as the input to the algorithm 4.
Revised: We use the average of all training matrices from Tr , as the input to the algorithm 4.

Original: We note that all the steps of our training algorithms are done on the training data.
Revised: We note that all the steps of our training algorithms are done on the training data.

Original: Hence, no additional computational cost is incurred for the sketching algorithm on the test data.
Revised: Hence, no additional computational cost is incurred for the sketching algorithm on the test data.

Original: Experimental parameters (i.e., learning rate for gradient descent) can be found in Appendix G.
Revised: Experimental parameters (i.e., learning rate for gradient descent) can be found in Appendix G.

Original: For both sketching algorithms, Ours are always the best of the four sketches.
Revised: For both sketching algorithms, Ours are always the best of the four sketches.

Original: It is significantly better than Classical CS , obtaining improvements of around 70% .
Revised: It is significantly better than Classical CS , obtaining improvements of around 70% .

Original: It also obtains a roughly 30% improvement over IVY19 .
Revised: It also obtains a roughly 30% improvement over IVY19 .

Original: [Figure]
Revised: [Figure]

Original: Wall-Clock Times.
Revised: Wall-Clock Times.

Original: The offline learning runtime is in Table 7.2, which is the time to train a sketch on A train .
Revised: The offline learning runtime is in Table 7.2, which is the time to train a sketch on A train .

Original: We can see that although the greedy method will take much longer (1h 45min), our second approach is much faster ( seconds) than the previous algorithm in (Indyk et al., 2019) (3 min) and can still achieve a similar error as the greedy algorithm.
Revised: We can see that although the greedy method will take much longer (1h 45min), our second approach is much faster ( seconds) than the previous algorithm in (Indyk et al., 2019) (3 min) and can still achieve a similar error as the greedy algorithm.

Original: The reason is that Algorithm 4 only needs to compute the ridge leverage scores on the training matrix once, which is actually much cheaper than
Revised: The reason is that Algorithm 4 only needs to compute the ridge leverage scores on the training matrix once, which is actually much cheaper than

Original: IVY19 which needs to compute a differentiable SVD many times during gradient descent.
Revised: IVY19 which needs to compute a differentiable SVD many times during gradient descent.

Original: In Section A.4, we also study the performance of our approach in the few-shot learning setting, which has been studied in Indyk et al.
Revised: In Section A.4, we also study the performance of our approach in the few-shot learning setting, which has been studied in Indyk et al.

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original:  E XPERIMENTS : S ECOND -O RDER O PTIMIZATION
Revised:  E XPERIMENTS : S ECOND -O RDER O PTIMIZATION

Original: In this section, we consider the IHS on the following instance of LASSO regression:
Revised: In this section, we consider the IHS on the following instance of LASSO regression:

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

Original: where λ is a parameter.
Revised: where λ is a parameter.

Original: We also study the performance of the sketches on the matrix estimation with a nuclear norm constraint problem, the fast regression solver (van den Brand et al.
Revised: We also study the performance of the sketches on the matrix estimation with a nuclear norm constraint problem, the fast regression solver (van den Brand et al.

Original: (2021)), as well as the use of sketches for first-order methods.
Revised: (2021)), as well as the use of sketches for first-order methods.

Original: The results can be found in Appendix B. All of our experiments are conducted on a laptop with a 1.90GHz CPU and 16GB RAM.
Revised: The results can be found in Appendix B. All of our experiments are conducted on a laptop with a 1.90GHz CPU and 16GB RAM.

Original: The offline training is done separately using a single GPU.
Revised: The offline training is done separately using a single GPU.

Original: The details of the implementation are deferred to Appendix G.
Revised: The details of the implementation are deferred to Appendix G.

Original: Dataset.
Revised: Dataset.

Original: We use the Electric 1 dataset of residential electric load measurements.
Revised: We use the Electric 1 dataset of residential electric load measurements.

Original: Each row of the matrix corresponds to a different residence.
Revised: Each row of the matrix corresponds to a different residence.

Original: Matrix columns are consecutive measurements at different times.
Revised: Matrix columns are consecutive measurements at different times.

Original: Here A i ∈ R 370 × 9 , b i ∈ R 370 × 1 , and | ( A, b ) train | = 320 , | ( A, b ) test | = 80 .
Revised: Here A i ∈ R 370 × 9 , b i ∈ R 370 × 1 , and | ( A, b ) train | = 320 , | ( A, b ) test | = 80 .

Original: We set λ = 15 .
Revised: We set λ = 15 .

Original: Experiment Setting.
Revised: Experiment Setting.

Original: We compare the learned sketch against the classical Count-Sketch 2 .
Revised: We compare the learned sketch against the classical Count-Sketch 2 .

Original: We choose m = 6 d, 8 d, 10 d and consider the error f
Revised: We choose m = 6 d, 8 d, 10 d and consider the error f

Original: ( x ) − f ( x ∗ ) .
Revised: ( x ) − f ( x ∗ ) .

Original: For the heavy-row Count-Sketch, we allocate 30% of the sketch space to the rows of the heavy row candidates.
Revised: For the heavy-row Count-Sketch, we allocate 30% of the sketch space to the rows of the heavy row candidates.

Original: For this dataset, each row represents a specific residence and hence there is a strong pattern of the distribution of the heavy rows.
Revised: For this dataset, each row represents a specific residence and hence there is a strong pattern of the distribution of the heavy rows.

Original: We select the heavy rows according to the number of times each row is heavy in the training data.
Revised: We select the heavy rows according to the number of times each row is heavy in the training data.

Original: We give a detailed discussion about this in Appendix B.1.
Revised: We give a detailed discussion about this in Appendix B.1.

Original: We highlight that it is still possible to recognize the pattern of the rows even if the row orders of the test data are permuted.
Revised: We highlight that it is still possible to recognize the pattern of the rows even if the row orders of the test data are permuted.

Original: We also consider optimizing the non-zero values after identifying the heavy rows, using our new approach in Section 6.
Revised: We also consider optimizing the non-zero values after identifying the heavy rows, using our new approach in Section 6.

Original: Results.
Revised: Results.

Original: We plot in Figures 7.1 the mean errors on a logarithmic scale.
Revised: We plot in Figures 7.1 the mean errors on a logarithmic scale.

Original: The average offline training time is 3 .
Revised: The average offline training time is 3 .

Original: 67 s to find a superset of the heavy rows over the training data and 66 s to optimize the values when m = 10 d , which are both faster than the runtime of Indyk et al.
Revised: 67 s to find a superset of the heavy rows over the training data and 66 s to optimize the values when m = 10 d , which are both faster than the runtime of Indyk et al.

Original: with the same parameters.
Revised: with the same parameters.

Original: Note that the learned matrix S is trained offline only once using the training data.
Revised: Note that the learned matrix S is trained offline only once using the training data.

Original: Hence, no additional computational cost is incurred when solving the optimization problem on the test data.
Revised: Hence, no additional computational cost is incurred when solving the optimization problem on the test data.

Original: We see all methods display linear convergence, that is, letting e k denote the error in the k -th iteration, we have e k ≈ ρ k e 1 for some convergence rate ρ .
Revised: We see all methods display linear convergence, that is, letting e k denote the error in the k -th iteration, we have e k ≈ ρ k e 1 for some convergence rate ρ .

Original: A smaller convergence rate implies a faster convergence.
Revised: A smaller convergence rate implies a faster convergence.

Original: We calculate an estimated rate of convergence ρ = ( e k /e 1 ) 1 /k with k = 7 .
Revised: We calculate an estimated rate of convergence ρ = ( e k /e 1 ) 1 /k with k = 7 .

Original: We can see both sketches, especially the sketch that optimizes both the positions and values, show significant improvements.
Revised: We can see both sketches, especially the sketch that optimizes both the positions and values, show significant improvements.

Original: When the sketch size is small ( 6 d ), this sketch has a convergence rate that is just 13.2% of that of the classical Count-Sketch, and when the sketch size is large ( 10 d ), this sketch has a smaller convergence rate that is just 12.1%.
Revised: When the sketch size is small ( 6 d ), this sketch has a convergence rate that is just 13.2% of that of the classical Count-Sketch, and when the sketch size is large ( 10 d ), this sketch has a smaller convergence rate that is just 12.1%.

Original: A CKNOWLEDGEMENTS
Revised: A CKNOWLEDGEMENTS

Original: Yi Li would like to thank for the partial support from the Singapore Ministry of Education under Tier 1 grant RG75/21.
Revised: Yi Li would like to thank for the partial support from the Singapore Ministry of Education under Tier 1 grant RG75/21.

Original: Honghao Lin and David Woodruff were supported in part by an Office of Naval
Revised: Honghao Lin and David Woodruff were supported in part by an Office of Naval

Original: Research (ONR) grant N00014-18-1-2562.
Revised: Research (ONR) grant N00014-18-1-2562.

Original: Ali Vakilian was supported by NSF award CCF-1934843.
Revised: Ali Vakilian was supported by NSF award CCF-1934843.

Original: A A DDITIONAL E XPERIMENTS : L OW -R ANK A PPROXIMATION
Revised: A A DDITIONAL E XPERIMENTS : L OW -R ANK A PPROXIMATION

Original: The details (data dimension, N train , etc.) are presented in Table A.1.
Revised: The details (data dimension, N train , etc.) are presented in Table A.1.

Original: [Table]
Revised: [Table]

Original: A. S ENSITIVITY A NALYSIS OF A LGORITHM
Revised: A. S ENSITIVITY A NALYSIS OF A LGORITHM

Original: [Table]
Revised: [Table]

Original: In this section we explore how sensitive the performance of our Algorithm 4 is to the ridge leverage score sampling and maximum inner product grouping process.
Revised: In this section we explore how sensitive the performance of our Algorithm 4 is to the ridge leverage score sampling and maximum inner product grouping process.

Original: We consider the following baselines:
Revised: We consider the following baselines:

Original: • ℓ 2 norm sampling: we sample the rows according to their squared length instead of doing ridge leverage score sampling.
Revised: • ℓ 2 norm sampling: we sample the rows according to their squared length instead of doing ridge leverage score sampling.

Original: • Only ridge leverage score sampling: the subspace spanned by only the sampled rows from ridge leverage score sampling.
Revised: • Only ridge leverage score sampling: the subspace spanned by only the sampled rows from ridge leverage score sampling.

Original: • Randomly grouping: we put the sampled rows into different buckets as before, but randomly divide the non-sampled rows into buckets.
Revised: • Randomly grouping: we put the sampled rows into different buckets as before, but randomly divide the non-sampled rows into buckets.

Original: The results are shown in Table A.2.
Revised: The results are shown in Table A.2.

Original: Here we set k = 30 , m = 60 as an example.
Revised: Here we set k = 30 , m = 60 as an example.

Original: To show the difference of the initialization method more clearly, we compare the error using the one-sided sketching Algorithm 1 and do not further optimize the non-zeros values.
Revised: To show the difference of the initialization method more clearly, we compare the error using the one-sided sketching Algorithm 1 and do not further optimize the non-zeros values.

Original: From the table we can see both that ridge leverage score sampling and the downstream grouping process are necessary, otherwise the error will be similar or even worse than that of the classical Count-Sketch.
Revised: From the table we can see both that ridge leverage score sampling and the downstream grouping process are necessary, otherwise the error will be similar or even worse than that of the classical Count-Sketch.

Original: A. T OTAL V ARIATION D ISTANCE
Revised: A. T OTAL V ARIATION D ISTANCE

Original: As we have shown in Theorem 5.1, if the total variation distance between the row sampling probability distributions p and q is O (1) , we have a worst-case guarantee of O ( k log k + k/ϵ ) , which is strictly better than the Ω( k 2 ) lower bound for the random CountSketch, even when its non-zero values have been optimized.
Revised: As we have shown in Theorem 5.1, if the total variation distance between the row sampling probability distributions p and q is O (1) , we have a worst-case guarantee of O ( k log k + k/ϵ ) , which is strictly better than the Ω( k 2 ) lower bound for the random CountSketch, even when its non-zero values have been optimized.

Original: We now study the total variation distance between the train and test matrix in our dataset.
Revised: We now study the total variation distance between the train and test matrix in our dataset.

Original: The result is shown in Figure A.1.
Revised: The result is shown in Figure A.1.

Original: From the figure we can see that for all the three dataset, the total variation distance is bounded by a constant, which suggests that the assumptions are reasonable for real-world data.
Revised: From the figure we can see that for all the three dataset, the total variation distance is bounded by a constant, which suggests that the assumptions are reasonable for real-world data.

Original: [Figure]
Revised: [Figure]

Original: A. L EARNING -B ASED A LGORITHMS FOR L OW -R ANK A PPROXIMATION IN THE F EW -S HOT S ETTING
Revised: A. L EARNING -B ASED A LGORITHMS FOR L OW -R ANK A PPROXIMATION IN THE F EW -S HOT S ETTING

Original: In this section, we will give a brief explanation of the two algorithms proposed in Indyk et al.
Revised: In this section, we will give a brief explanation of the two algorithms proposed in Indyk et al.

Original: Both algorithms aim to optimize the non-zero values of a Count-Sketch matrix under fixed locations of the non-zero entries.
Revised: Both algorithms aim to optimize the non-zero values of a Count-Sketch matrix under fixed locations of the non-zero entries.

Original: One-shot closed-form algorithm.
Revised: One-shot closed-form algorithm.

Original: Given a sparsity pattern of a Count-Sketch matrix S ∈ R m × n , it partitions the rows of A into m blocks
Revised: Given a sparsity pattern of a Count-Sketch matrix S ∈ R m × n , it partitions the rows of A into m blocks

Original: A (1) , ..., A ( m ) as follows: let I i
Revised: A (1) , ..., A ( m ) as follows: let I i

Original: = { j : S ij = 1 } .
Revised: = { j : S ij = 1 } .

Original: The block A ( i ) ∈ R | I i |× d is the sub-matrix of A that contains the rows whose indices are in I i .
Revised: The block A ( i ) ∈ R | I i |× d is the sub-matrix of A that contains the rows whose indices are in I i .

Original: The goal here is for each block
Revised: The goal here is for each block

Original: A ( i ) , to choose a (non-sparse) one-dimensional sketching vector s i ∈ R
Revised: A ( i ) , to choose a (non-sparse) one-dimensional sketching vector s i ∈ R

Original: | I i | .
Revised: | I i | .

Original: The first approach is to set s i to be the top left singular vector of A ( i ) , which is the algorithm 1Shot2Vec .
Revised: The first approach is to set s i to be the top left singular vector of A ( i ) , which is the algorithm 1Shot2Vec .

Original: Another approach is to set s i to be a left singular vector of A ( i ) chosen randomly and proportional to its squared singular value.
Revised: Another approach is to set s i to be a left singular vector of A ( i ) chosen randomly and proportional to its squared singular value.

Original: The main advantage of the latter approach over the previous one is that it endows the algorithm with provable guarantees on the LRA error.
Revised: The main advantage of the latter approach over the previous one is that it endows the algorithm with provable guarantees on the LRA error.

Original: The 1Shot2Vec algorithm combines both ways, obtaining the benefits of both approaches.
Revised: The 1Shot2Vec algorithm combines both ways, obtaining the benefits of both approaches.

Original: The advantage of these two algorithms is that they extract a sketching matrix by an analytic computation, requiring neither
Revised: The advantage of these two algorithms is that they extract a sketching matrix by an analytic computation, requiring neither

Original: GPU access nor auto-gradient functionality.
Revised: GPU access nor auto-gradient functionality.

Original: Few-shot SGD algorithm.
Revised: Few-shot SGD algorithm.

Original: In this algorithm, the authors propose a new loss function for LRA, namely,
Revised: In this algorithm, the authors propose a new loss function for LRA, namely,

Original: [Equation]
Revised: [Equation]

Original: where A = U Σ V ⊤ is the SVD-decomposition of A and U k ∈ R n × k denotes the submatrix of U that contains its first k columns.
Revised: where A = U Σ V ⊤ is the SVD-decomposition of A and U k ∈ R n × k denotes the submatrix of U that contains its first k columns.

Original: I 

Revised: I

Original: ∈ R k × d denotes the result of augmenting the identity matrix of order k with d − k additional zero columns on the right.
Revised: ∈ R k × d denotes the result of augmenting the identity matrix of order k with d − k additional zero columns on the right.

Original: This loss function is motivated by the analysis of prior LRA algorithms that use random sketching matrices.
Revised: This loss function is motivated by the analysis of prior LRA algorithms that use random sketching matrices.

Original: It is faster to compute and differentiate than the previous empirical loss in Indyk et al.
Revised: It is faster to compute and differentiate than the previous empirical loss in Indyk et al.

Original: In the experiments the authors also show that this loss function can achieve a smaller error in a shorter amount of time, using a small number of randomly sampled training matrices, though the final error will be larger than that of the previous algorithm in Indyk et al.
Revised: In the experiments the authors also show that this loss function can achieve a smaller error in a shorter amount of time, using a small number of randomly sampled training matrices, though the final error will be larger than that of the previous algorithm in Indyk et al.

Original: (2019) if we allow a longer training time and access to the whole training set
Revised: (2019) if we allow a longer training time and access to the whole training set

Original: Tr .
Revised: Tr .

Original: A. E XPERIMENTS : LRA IN THE FEW - SHOT SETTING
Revised: A. E XPERIMENTS : LRA IN THE FEW - SHOT SETTING

Original: In the rest of this section, we study the performance of our second approach in the few-shot learning setting.
Revised: In the rest of this section, we study the performance of our second approach in the few-shot learning setting.

Original: We first consider the case where we only have one training matrix randomly sampled from Tr .
Revised: We first consider the case where we only have one training matrix randomly sampled from Tr .

Original: Here, we compare our method with the 1Shot2Vec method proposed in (Indyk et al., 2021) in the same setting ( k = 10 , m = 40 ) as in their empirical evaluation.
Revised: Here, we compare our method with the 1Shot2Vec method proposed in (Indyk et al., 2021) in the same setting ( k = 10 , m = 40 ) as in their empirical evaluation.

Original: The result is shown in Table A.3.
Revised: The result is shown in Table A.3.

Original: Compared to 1Shot2Vec , our method reduces the error by around 50% , and has an even slightly faster runtime.
Revised: Compared to 1Shot2Vec , our method reduces the error by around 50% , and has an even slightly faster runtime.

Original: Indyk et al.
Revised: Indyk et al.

Original: (2021) also proposed a FewShotSGD algorithm which further improves the non-zero values of the sketches after different initialization methods.
Revised: (2021) also proposed a FewShotSGD algorithm which further improves the non-zero values of the sketches after different initialization methods.

Original: We compare the performance of this approach for different initialization methods: in all initialization methods, we only use one training matrix and we use three training matrices for the FewShotSGD step.
Revised: We compare the performance of this approach for different initialization methods: in all initialization methods, we only use one training matrix and we use three training matrices for the FewShotSGD step.

Original: The results are shown in
Revised: The results are shown in

Original: Table A.4.
Revised: Table A.4.

Original: We report the minimum error of 50 iterations of the FewShotSGD because we aim to compare the computational efficiency for different methods.
Revised: We report the minimum error of 50 iterations of the FewShotSGD because we aim to compare the computational efficiency for different methods.

Original: From the table we see that our approach plus the FewShotSGD method can achieve a much smaller error, with around a 50% improvement upon (Indyk et al., 2021).
Revised: From the table we see that our approach plus the FewShotSGD method can achieve a much smaller error, with around a 50% improvement upon (Indyk et al., 2021).

Original: Moreover, even without further optimization by FewShotSGD , our initialization method for learning the non-zero locations in CountSketch obtains a smaller error than other methods (even when they are optimized with 1ShotSGD or FewShotSGD learning).
Revised: Moreover, even without further optimization by FewShotSGD , our initialization method for learning the non-zero locations in CountSketch obtains a smaller error than other methods (even when they are optimized with 1ShotSGD or FewShotSGD learning).

Original: B A DDITIONAL E XPERIMENTS : S ECOND -O RDER O PTIMIZATION
Revised: B A DDITIONAL E XPERIMENTS : S ECOND -O RDER O PTIMIZATION

Original: As we mentioned in Section 8, despite the number of problems that learned sketches have been applied to, they have not been applied to convex optimization, or say, iterative sketching algorithms in general.
Revised: As we mentioned in Section 8, despite the number of problems that learned sketches have been applied to, they have not been applied to convex optimization, or say, iterative sketching algorithms in general.

Original: To demonstrate the difficulty, we consider the Iterative Hessian Sketch (IHS) as an example.
Revised: To demonstrate the difficulty, we consider the Iterative Hessian Sketch (IHS) as an example.

Original: In that scheme, suppose that we have k iterations of the algorithm.
Revised: In that scheme, suppose that we have k iterations of the algorithm.

Original: Then we need k independent
Revised: Then we need k independent

Original: [Table] [Table]
Revised: [Table] [Table]

Original: sketching matrices (otherwise the solution may diverge).
Revised: sketching matrices (otherwise the solution may diverge).

Original: A natural way is to follow the method in (Indyk et al., 2019)
Revised: A natural way is to follow the method in (Indyk et al., 2019)

Original: [Table]
Revised: [Table]

Original: which is to minimize the following quantity
Revised: which is to minimize the following quantity

Original: [Equation]
Revised: [Equation]

Original: where the minimization is taken over k Count-Sketch matrices S 1 , . .
Revised: where the minimization is taken over k Count-Sketch matrices S 1 , . .

Original: , S k .
Revised: , S k .

Original: In this case, however, calculating the gradient with respect to S 1 would involve all iterations and in each iteration we need to solve a constrained optimization problem.
Revised: In this case, however, calculating the gradient with respect to S 1 would involve all iterations and in each iteration we need to solve a constrained optimization problem.

Original: Hence, it would be difficult and intractable to compute the gradients.
Revised: Hence, it would be difficult and intractable to compute the gradients.

Original: An alternative way is to train k sketching matrices sequentially, that is, learn the sketching matrix for the i -th iteration using a local loss function for the i -th iteration
Revised: An alternative way is to train k sketching matrices sequentially, that is, learn the sketching matrix for the i -th iteration using a local loss function for the i -th iteration

Original: [Table]
Revised: [Table]

Original: and then using the learned matrix in the i -th iteration to generate the training data for the ( i + 1) -st iteration.
Revised: and then using the learned matrix in the i -th iteration to generate the training data for the ( i + 1) -st iteration.

Original: However, the empirical results suggest that it works for the first iteration only, because in this case the training data for the ( i + 1) -st iteration depends on the solution to the i -th iteration and may become farther away from the test data in later iterations.
Revised: However, the empirical results suggest that it works for the first iteration only, because in this case the training data for the ( i + 1) -st iteration depends on the solution to the i -th iteration and may become farther away from the test data in later iterations.

Original: The core problem here is that the method proposed in Indyk et al.
Revised: The core problem here is that the method proposed in Indyk et al.

Original: (2019) treats the training process in a black-box way, which is difficult to extend to iterative methods.
Revised: (2019) treats the training process in a black-box way, which is difficult to extend to iterative methods.

Original: B. T HE D ISTRIBUTION OF THE H EAVY R OWS
Revised: B. T HE D ISTRIBUTION OF THE H EAVY R OWS

Original: In our experiments, we hypothesize that in real-world data there may be an underlying pattern which can help us identify the heavy rows.
Revised: In our experiments, we hypothesize that in real-world data there may be an underlying pattern which can help us identify the heavy rows.

Original: In the Electric dataset, each row of the matrix corresponds to a specific residence and the heavy rows are concentrated on some specific rows.
Revised: In the Electric dataset, each row of the matrix corresponds to a specific residence and the heavy rows are concentrated on some specific rows.

Original: To exemplify this, we study the heavy leverage score rows distribution over the Electric dataset.
Revised: To exemplify this, we study the heavy leverage score rows distribution over the Electric dataset.

Original: For a row i ∈ [370] , let f i denote the number of times that row
Revised: For a row i ∈ [370] , let f i denote the number of times that row

Original: i is heavy out of 320 training data points from the Electric dataset, where we say row i is heavy if ℓ i ≥ 5 d/n .
Revised: i is heavy out of 320 training data points from the Electric dataset, where we say row i is heavy if ℓ i ≥ 5 d/n .

Original: Below we list all 74 pairs ( i, f i ) with f i > 0 .
Revised: Below we list all 74 pairs ( i, f i ) with f i > 0 .

Original: Observe that the heavy rows are concentrated on a set of specific row indices.
Revised: Observe that the heavy rows are concentrated on a set of specific row indices.

Original: There are only 30 rows i with f i ≥ 50 .
Revised: There are only 30 rows i with f i ≥ 50 .

Original: We view this as strong evidence for our hypothesis.
Revised: We view this as strong evidence for our hypothesis.

Original: Heavy Rows Distribution Under Permutation.
Revised: Heavy Rows Distribution Under Permutation.

Original: We note that even though the order of the rows has been changed, we can still recognize the patterns of the rows.
Revised: We note that even though the order of the rows has been changed, we can still recognize the patterns of the rows.

Original: We continue to use the Electric dataset as an example.
Revised: We continue to use the Electric dataset as an example.

Original: To address the concern that a permutation may break the sketch, we can measure the similarity between vectors, that is, after processing the training data, we can instead test similarity on the rows of the test matrix and use this to select the heavy rows, rather than an index which may simply be permuted.
Revised: To address the concern that a permutation may break the sketch, we can measure the similarity between vectors, that is, after processing the training data, we can instead test similarity on the rows of the test matrix and use this to select the heavy rows, rather than an index which may simply be permuted.

Original: To illustrate this method, we use the following example on the Electric dataset, using locality sensitive hashing.
Revised: To illustrate this method, we use the following example on the Electric dataset, using locality sensitive hashing.

Original: After processing the training data, we obtain a set I of indices of heavy rows.
Revised: After processing the training data, we obtain a set I of indices of heavy rows.

Original: For each i ∈
Revised: For each i ∈

Original: I , we pick q = 3 independent standard Gaussian vectors g 1 , g 2 , g 3 ∈ R d , and compute f ( r i ) =
Revised: I , we pick q = 3 independent standard Gaussian vectors g 1 , g 2 , g 3 ∈ R d , and compute f ( r i ) =

Original: ( g T 1 r i , g T 2 r i , g T 3 r i ) ∈ R 3 , where r i takes an average of the i -th rows over all training sets.
Revised: ( g T 1 r i , g T 2 r i , g T 3 r i ) ∈ R 3 , where r i takes an average of the i -th rows over all training sets.

Original: Let A be the test matrix.
Revised: Let A be the test matrix.

Original: For each i ∈ I , let j i = argmin
Revised: For each i ∈ I , let j i = argmin

Original: j ∥ f ( A j ) − f
Revised: j ∥ f ( A j ) − f

Original: ( r i ) ∥ 2 .
Revised: ( r i ) ∥ 2 .

Original: We take the j i -th row to be a heavy row in our learned sketch.
Revised: We take the j i -th row to be a heavy row in our learned sketch.

Original: This method only needs an additional O (1) passes over the entries of A and hence, the extra time cost is negligible.
Revised: This method only needs an additional O (1) passes over the entries of A and hence, the extra time cost is negligible.

Original: To test the performance of the method, we randomly pick a matrix from the test set and permute its rows.
Revised: To test the performance of the method, we randomly pick a matrix from the test set and permute its rows.

Original: The result shows that when k is small, we can roughly recover 70% of the top- k heavy rows, and we plot below the regression error using the learned Count-Sketch matrix generated this way, where we set m = 90 and k = 0 .
Revised: The result shows that when k is small, we can roughly recover 70% of the top- k heavy rows, and we plot below the regression error using the learned Count-Sketch matrix generated this way, where we set m = 90 and k = 0 .

Original: 3 m = 27 .
Revised: 3 m = 27 .

Original: We can see that the learned method still obtains a significant improvement.
Revised: We can see that the learned method still obtains a significant improvement.

Original: [Figure]
Revised: [Figure]

Original: In many applications, for the problem
Revised: In many applications, for the problem

Original: [Figure]
Revised: [Figure]

Original: it is reasonable to model the matrix X ∗ as having low rank.
Revised: it is reasonable to model the matrix X ∗ as having low rank.

Original: Similar to ℓ 1 -minimization for compressive sensing, a standard relaxation of the rank constraint
Revised: Similar to ℓ 1 -minimization for compressive sensing, a standard relaxation of the rank constraint

Original: is to minimize the nuclear norm of X , defined as ∥ X ∥ ∗ := (cid:80) min { d 1 ,d 2 } j = σ j ( X ) , where σ j ( X ) is the j -th largest singular value of X .
Revised: is to minimize the nuclear norm of X , defined as ∥ X ∥ ∗ := (cid:80) min { d 1 ,d 2 } j = σ j ( X ) , where σ j ( X ) is the j -th largest singular value of X .

Original: Hence, the matrix estimation problem we consider here is
Revised: Hence, the matrix estimation problem we consider here is

Original: [Equation]
Revised: [Equation]

Original: such that
Revised: such that

Original: [Equation]
Revised: [Equation]

Original: where ρ > 0 is a user defined radius as a regularization parameter.
Revised: where ρ > 0 is a user defined radius as a regularization parameter.

Original: We conduct Iterative Hessian Sketch (IHS) experiments on the following dataset:
Revised: We conduct Iterative Hessian Sketch (IHS) experiments on the following dataset:

Original: • Tunnel 3 : The data set is a time series of gas concentrations measured by eight sensors in a wind tunnel.
Revised: • Tunnel 3 : The data set is a time series of gas concentrations measured by eight sensors in a wind tunnel.

Original: Each ( A, B ) corresponds to a different data collection trial.
Revised: Each ( A, B ) corresponds to a different data collection trial.

Original: A i ∈ R 13530 × 5 , B i ∈ R 13530 ×, | ( A, B ) | train =
Revised: A i ∈ R 13530 × 5 , B i ∈ R 13530 ×, | ( A, B ) | train =

Original: 144 , | ( A, B ) | test = 36 .
Revised: 144 , | ( A, B ) | test = 36 .

Original: In our nuclear norm constraint, we set ρ = 10 .
Revised: In our nuclear norm constraint, we set ρ = 10 .

Original: Experiment Setting : We choose m = 7 d, 10 d for the Tunnel dataset.
Revised: Experiment Setting : We choose m = 7 d, 10 d for the Tunnel dataset.

Original: We consider the error 12 ∥ AX − B ∥ 22 − 12 ∥
Revised: We consider the error 12 ∥ AX − B ∥ 22 − 12 ∥

Original: AX ∗ − B ∥ 22 .
Revised: AX ∗ − B ∥ 22 .

Original: The leverage scores of this dataset are very uniform.
Revised: The leverage scores of this dataset are very uniform.

Original: Hence, for this experiment we only consider optimizing the values of the non-zero entries.
Revised: Hence, for this experiment we only consider optimizing the values of the non-zero entries.

Original: Results of Our Experiments : We plot on a logarithmic scale the mean errors of the dataset in
Revised: Results of Our Experiments : We plot on a logarithmic scale the mean errors of the dataset in

Original: Figures B.2.
Revised: Figures B.2.

Original: We can see that when m = 7 d , the gradient-based sketch, based on the first 6 iterations, has a rate of convergence that is 48% of the random sketch, and when m = 10 d , the gradient-based sketch has a rate of convergence that is 29% of the random sketch.
Revised: We can see that when m = 7 d , the gradient-based sketch, based on the first 6 iterations, has a rate of convergence that is 48% of the random sketch, and when m = 10 d , the gradient-based sketch has a rate of convergence that is 29% of the random sketch.

Original: B. F AST R EGRESSION S OLVER
Revised: B. F AST R EGRESSION S OLVER

Original: Consider an unconstrained convex optimization problem min x f ( x ) , where f is smooth and strongly convex, and its Hessian ∇ 2 f is Lipschitz continuous.
Revised: Consider an unconstrained convex optimization problem min x f ( x ) , where f is smooth and strongly convex, and its Hessian ∇ 2 f is Lipschitz continuous.

Original: This problem can be solved by Newton’s
Revised: This problem can be solved by Newton’s

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: method, which iteratively performs the update
Revised: method, which iteratively performs the update

Original: [Equation]
Revised: [Equation]

Original: provided it is given a good initial point x 0 .
Revised: provided it is given a good initial point x 0 .

Original: In each step, it requires solving a regression problem of the form min z (cid:13)(cid:13)
Revised: In each step, it requires solving a regression problem of the form min z (cid:13)(cid:13)

Original: A ⊤ Az − y (cid:13)(cid:13) 2 , which, with access to A , can be solved with a fast regression solver in (van den Brand et al., 2021).
Revised: A ⊤ Az − y (cid:13)(cid:13) 2 , which, with access to A , can be solved with a fast regression solver in (van den Brand et al., 2021).

Original: The regression solver first computes a preconditioner R via a QR decomposition such that SAR has orthonormal columns, where S is a sketching matrix, then solves ˆ
Revised: The regression solver first computes a preconditioner R via a QR decomposition such that SAR has orthonormal columns, where S is a sketching matrix, then solves ˆ

Original: z = arg min z ′ (cid:13)(cid:13) ( AR ) ⊤ ( AR ) z ′
Revised: z = arg min z ′ (cid:13)(cid:13) ( AR ) ⊤ ( AR ) z ′

Original: − y (cid:13)(cid:13) 2 by gradient descent and returns R ˆ z in the end.
Revised: − y (cid:13)(cid:13) 2 by gradient descent and returns R ˆ z in the end.

Original: Here, the point of sketching is that the QR decomposition of SA can be computed much more efficiently than the
Revised: Here, the point of sketching is that the QR decomposition of SA can be computed much more efficiently than the

Original: QR decomposition of A , since S has only a small number of rows.
Revised: QR decomposition of A , since S has only a small number of rows.

Original: In this section, We consider the unconstrained least squares problem min x f ( x ) with f ( x ) = 12 ∥
Revised: In this section, We consider the unconstrained least squares problem min x f ( x ) with f ( x ) = 12 ∥

Original: Ax − b ∥ 22 using the Electric dataset, using the above fast regression solver.
Revised: Ax − b ∥ 22 using the Electric dataset, using the above fast regression solver.

Original: Training : Note that ∇ 2 f ( x ) =
Revised: Training : Note that ∇ 2 f ( x ) =

Original: A ⊤ A , independent of x .
Revised: A ⊤ A , independent of x .

Original: In the t -th round of Newton’s method, by (B.1), we need to solve a regression problem
Revised: In the t -th round of Newton’s method, by (B.1), we need to solve a regression problem

Original: min z (cid:13)(cid:13)
Revised: min z (cid:13)(cid:13)

Original: A ⊤ Az − y (cid:13)(cid:13) 2 2 with y = ∇ f ( x t ) .
Revised: A ⊤ Az − y (cid:13)(cid:13) 2 2 with y = ∇ f ( x t ) .

Original: Hence, we can use the same methods in the preceding subsection to optimize the learned sketch S i .
Revised: Hence, we can use the same methods in the preceding subsection to optimize the learned sketch S i .

Original: For a general problem where ∇ 2 f ( x ) depends on x , one can take x t to be the solution obtained from the learned sketch S t to generate A and y for the ( t + 1)
Revised: For a general problem where ∇ 2 f ( x ) depends on x , one can take x t to be the solution obtained from the learned sketch S t to generate A and y for the ( t + 1)

Original: -st round, train a learned sketch S t +1 , and repeat this process.
Revised: -st round, train a learned sketch S t +1 , and repeat this process.

Original: Experiment Setting : For the Electric dataset, we set m = 10 d = 90 .
Revised: Experiment Setting : For the Electric dataset, we set m = 10 d = 90 .

Original: We observe that the classical
Revised: We observe that the classical

Original: Count-Sketch matrix makes the solution diverge terribly in this setting.
Revised: Count-Sketch matrix makes the solution diverge terribly in this setting.

Original: To make a clearer comparison, we consider the following sketch matrix:
Revised: To make a clearer comparison, we consider the following sketch matrix:

Original: • Gaussian sketch: S = 1 √
Revised: • Gaussian sketch: S = 1 √

Original: m G , where G ∈ R m × n with i.i
Revised: m G , where G ∈ R m × n with i.i

Original: d. N (0 , 1) entries.
Revised: d. N (0 , 1) entries.

Original: • Sparse Johnson-Lindenstrauss Transform (SJLT): S is the vertical concatenation of s independent Count-Sketch matrices, each of dimension m/s × n .
Revised: • Sparse Johnson-Lindenstrauss Transform (SJLT): S is the vertical concatenation of s independent Count-Sketch matrices, each of dimension m/s × n .

Original: We note that the above sketching matrices require more time to compute SA but need fewer rows to be a subspace embedding than the classical Count-Sketch matrix.
Revised: We note that the above sketching matrices require more time to compute SA but need fewer rows to be a subspace embedding than the classical Count-Sketch matrix.

Original: For the step length η in gradient descent, we set η = 1 in all iterations of the learned sketches.
Revised: For the step length η in gradient descent, we set η = 1 in all iterations of the learned sketches.

Original: For classical random sketches, we set η in the following two ways: (a) η = 1 in all iterations and (b) η = 1 in the first iteration and η = 0 .
Revised: For classical random sketches, we set η in the following two ways: (a) η = 1 in all iterations and (b) η = 1 in the first iteration and η = 0 .

Original: in all subsequent iterations.
Revised: in all subsequent iterations.

Original: Experimental Results : We examine the accuracy of the subproblem min z (cid:13)(cid:13)
Revised: Experimental Results : We examine the accuracy of the subproblem min z (cid:13)(cid:13)

Original: A ⊤ Az − y (cid:13)(cid:13) 2 2 and define the error to be (cid:13)(cid:13)
Revised: A ⊤ Az − y (cid:13)(cid:13) 2 2 and define the error to be (cid:13)(cid:13)

Original: A ⊤ ARz t − y (cid:13)(cid:13) 2 / ∥
Revised: A ⊤ ARz t − y (cid:13)(cid:13) 2 / ∥

Original: y ∥ 2 .
Revised: y ∥ 2 .

Original: We consider the subproblems in the first three iterations of the global Newton method.
Revised: We consider the subproblems in the first three iterations of the global Newton method.

Original: The results are plotted in Figure B.3.
Revised: The results are plotted in Figure B.3.

Original: Note that Count-Sketch causes a terrible divergence of the subroutine and is thus omitted in the plots.
Revised: Note that Count-Sketch causes a terrible divergence of the subroutine and is thus omitted in the plots.

Original: Still, we observe that in setting (a) of η , the other two classical sketches cause the subroutine to diverge.
Revised: Still, we observe that in setting (a) of η , the other two classical sketches cause the subroutine to diverge.

Original: In setting (b) of η , the other two classical sketches lead to convergence but their error is significantly larger than that of the learned sketches, in each of the first three calls to the subroutine.
Revised: In setting (b) of η , the other two classical sketches lead to convergence but their error is significantly larger than that of the learned sketches, in each of the first three calls to the subroutine.

Original: The error of the learned sketch is less than 0 .
Revised: The error of the learned sketch is less than 0 .

Original: 01 in all iterations of all three subroutine calls, in both settings (a) and (b) of η .
Revised: 01 in all iterations of all three subroutine calls, in both settings (a) and (b) of η .

Original: We also plot a figure on the convergence of the global Newton method.
Revised: We also plot a figure on the convergence of the global Newton method.

Original: Here, for each subroutine, we only run one iteration, and plot the error of the original least squares problem.
Revised: Here, for each subroutine, we only run one iteration, and plot the error of the original least squares problem.

Original: The result is shown in Figure B.4, which clearly displays a significantly faster decay with learned sketches.
Revised: The result is shown in Figure B.4, which clearly displays a significantly faster decay with learned sketches.

Original: The rate of convergence using heavy-rows sketches is 80 .
Revised: The rate of convergence using heavy-rows sketches is 80 .

Original: 6% of that using Gaussian or sparse JL sketches.
Revised: 6% of that using Gaussian or sparse JL sketches.

Original: B. 
Revised: B.

Original: [Equation]
Revised: [Equation]

Original: O RDER O PTIMIZATION
Revised: O RDER O PTIMIZATION

Original: In this section, we study the use of the sketch in first-order methods.
Revised: In this section, we study the use of the sketch in first-order methods.

Original: Particularly, let QR − 1 =
Revised: Particularly, let QR − 1 =

Original: SA be the QR-decomposition for SA , where S is a sketch matrix.
Revised: SA be the QR-decomposition for SA , where S is a sketch matrix.

Original: We use R as an (approximate) preconditioner and use gradient descent to solve the problem min ∥ ARx − b ∥ 2 .
Revised: We use R as an (approximate) preconditioner and use gradient descent to solve the problem min ∥ ARx − b ∥ 2 .

Original: Here we use the Electric dataset where A is 370 × 9 and we set S to have 90 rows.
Revised: Here we use the Electric dataset where A is 370 × 9 and we set S to have 90 rows.

Original: The result is shown in the following table, where the time includes the time to compute R .
Revised: The result is shown in the following table, where the time includes the time to compute R .

Original: We can see that if we use a learned sketch matrix, the error converges very fast when we set the learning rate to be 1 and 0 .
Revised: We can see that if we use a learned sketch matrix, the error converges very fast when we set the learning rate to be 1 and 0 .

Original: 1 , while the classical
Revised: 1 , while the classical

Original: Count-Sketch will lead to divergence.
Revised: Count-Sketch will lead to divergence.

Original: [Table]
Revised: [Table]

Original: C P RELIMINARIES : T HEOREMS AND A DDITIONAL A LGORITHMS
Revised: C P RELIMINARIES : T HEOREMS AND A DDITIONAL A LGORITHMS

Original: In this section, we provide the full description of the time-optimal sketching algorithm for LRA in
Revised: In this section, we provide the full description of the time-optimal sketching algorithm for LRA in

Original: Algorithm 2.
Revised: Algorithm 2.

Original: We also provide several definitions and lemmas that are used in the proofs of our results for LRA.
Revised: We also provide several definitions and lemmas that are used in the proofs of our results for LRA.

Original: Definition C.1 (Affine Embedding) .
Revised: Definition C.1 (Affine Embedding) .

Original: Given a pair of matrices A and B , a matrix S is an affine ϵ -embedding if for all X of the appropriate shape, ∥ S ( AX − B ) ∥ 2 F = (1 ± ϵ )
Revised: Given a pair of matrices A and B , a matrix S is an affine ϵ -embedding if for all X of the appropriate shape, ∥ S ( AX − B ) ∥ 2 F = (1 ± ϵ )

Original: ∥ AX − B ∥ 2 F .
Revised: ∥ AX − B ∥ 2 F .

Original: Lemma C.2 (Clarkson & Woodruff (2017); Lemma 40) .
Revised: Lemma C.2 (Clarkson & Woodruff (2017); Lemma 40) .

Original: Let A be an n × d matrix and let S ∈
Revised: Let A be an n × d matrix and let S ∈

Original: R O (1 /ϵ 2 ) × n be a CountSketch matrix.
Revised: R O (1 /ϵ 2 ) × n be a CountSketch matrix.

Original: Then with constant probability, ∥ SA ∥ 2 F = (1 ± ϵ ) ∥
Revised: Then with constant probability, ∥ SA ∥ 2 F = (1 ± ϵ ) ∥

Original: A ∥ 2 F .
Revised: A ∥ 2 F .

Original: The following result is shown in Clarkson & Woodruff (2017) and sharpened with Nelson & Nguyên (2013); Meng & Mahoney (2013).
Revised: The following result is shown in Clarkson & Woodruff (2017) and sharpened with Nelson & Nguyên (2013); Meng & Mahoney (2013).

Original: Lemma C.3.
Revised: Lemma C.3.

Original: Given matrices A, B with n rows, a CountSketch with O (rank( A ) 2 /ϵ 2 ) rows is an affine ϵ -embedding matrix with constant probability.
Revised: Given matrices A, B with n rows, a CountSketch with O (rank( A ) 2 /ϵ 2 ) rows is an affine ϵ -embedding matrix with constant probability.

Original: Moreover, the matrix product SA can be computed in O (nnz( A )) time, where nnz( A ) denotes the number of non-zero entries of matrix A .
Revised: Moreover, the matrix product SA can be computed in O (nnz( A )) time, where nnz( A ) denotes the number of non-zero entries of matrix A .

Original: Lemma C.4 (Sarlos (2006); Clarkson & Woodruff (2017)) .
Revised: Lemma C.4 (Sarlos (2006); Clarkson & Woodruff (2017)) .

Original: Suppose that A ∈ R n × d and B ∈ R n × d ′ .
Revised: Suppose that A ∈ R n × d and B ∈ R n × d ′ .

Original: Let S ∈ R m × n be a CountSketch with m = rank( A ) 2 ϵ. Let ˜ X = arg min rank-
Revised: Let S ∈ R m × n be a CountSketch with m = rank( A ) 2 ϵ. Let ˜ X = arg min rank-

Original: k X ∥
Revised: k X ∥

Original: SAX − SB ∥ 2 F .
Revised: SAX − SB ∥ 2 F .

Original: Then,
Revised: Then,

Original: With constant probability, (cid:13)(cid:13)(cid:13)
Revised: With constant probability, (cid:13)(cid:13)(cid:13)

Original: A ˜ X − B (cid:13)(cid:13)(cid:13) 2 F ≤ (1 + ϵ ) min rank-
Revised: A ˜ X − B (cid:13)(cid:13)(cid:13) 2 F ≤ (1 + ϵ ) min rank-

Original: k X ∥ AX − B ∥ 2 F .
Revised: k X ∥ AX − B ∥ 2 F .

Original: In other words, in O (nnz( A ) + nnz( B ) + m ( d + d ′ )) time, we can reduce the problem to a smaller (multi-response regression) problem with m rows whose optimal solution is a (1 + ϵ ) -approximate solution to the original instance.
Revised: In other words, in O (nnz( A ) + nnz( B ) + m ( d + d ′ )) time, we can reduce the problem to a smaller (multi-response regression) problem with m rows whose optimal solution is a (1 + ϵ ) -approximate solution to the original instance.

Original: The (1 + ϵ ) -approximate solution ˜ X can be computed in time O
Revised: The (1 + ϵ ) -approximate solution ˜ X can be computed in time O

Original: (nnz( A ) + nnz( B ) + mdd ′
Revised: (nnz( A ) + nnz( B ) + mdd ′

Original: + min( m 2 d, md 2 )) .
Revised: + min( m 2 d, md 2 )) .

Original: Now we turn our attention to the time-optimal sketching algorithm for LRA.
Revised: Now we turn our attention to the time-optimal sketching algorithm for LRA.

Original: The next lemma is known, though we include it for completeness Avron et al.
Revised: The next lemma is known, though we include it for completeness Avron et al.

Original: Lemma C.5.
Revised: Lemma C.5.

Original: Suppose that S ∈ R m S × n and R ∈ R m R ×
Revised: Suppose that S ∈ R m S × n and R ∈ R m R ×

Original: d are sparse affine ϵ -embedding matrices for ( A k , A ) and (( SA ) ⊤ , A ⊤ ) , respectively.
Revised: d are sparse affine ϵ -embedding matrices for ( A k , A ) and (( SA ) ⊤ , A ⊤ ) , respectively.

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Consider the following multiple-response regression problem:
Revised: Consider the following multiple-response regression problem:

Original: [Equation]
Revised: [Equation]

Original: Note that since X =
Revised: Note that since X =

Original: I k is a feasible solution to Eq.
Revised: I k is a feasible solution to Eq.

Original: (C.1), min rank-
Revised: (C.1), min rank-

Original: k X ∥
Revised: k X ∥

Original: A k X − A ∥ 2 F = ∥
Revised: A k X − A ∥ 2 F = ∥

Original: A k − A ∥ 2 F .
Revised: A k − A ∥ 2 F .

Original: Let S ∈ R m S × n be a sketching matrix that satisfies the condition of Lemma C.4 (Item 1) for A :
Revised: Let S ∈ R m S × n be a sketching matrix that satisfies the condition of Lemma C.4 (Item 1) for A :

Original: = A k and B := A .
Revised: = A k and B := A .

Original: By the normal equations, the rank- k minimizer of ∥ SA k X − SA ∥ 2 F is ( SA k ) + SA .
Revised: By the normal equations, the rank- k minimizer of ∥ SA k X − SA ∥ 2 F is ( SA k ) + SA .

Original: Hence,
Revised: Hence,

Original: [Equation]
Revised: [Equation]

Original: which in particular shows that a (1 + ϵ ) -approximate
Revised: which in particular shows that a (1 + ϵ ) -approximate

Original: rank-
Revised: rank-

Original: k approximation of A exists in the row space of SA .
Revised: k approximation of A exists in the row space of SA .

Original: In other words,
Revised: In other words,

Original: [Equation]
Revised: [Equation]

Original: Next, let
Revised: Next, let

Original: R ∈ R m R ×
Revised: R ∈ R m R ×

Original: d be a sketching matrix which satisfies the condition of Lemma C.4 (Item 1) for A := ( SA ) ⊤ and B :=
Revised: d be a sketching matrix which satisfies the condition of Lemma C.4 (Item 1) for A := ( SA ) ⊤ and B :=

Original: A ⊤ .
Revised: A ⊤ .

Original: Let Y denote the rank- k minimizer of (cid:13)(cid:13) R ( SA ) ⊤ X ⊤ − RA ⊤
Revised: Let Y denote the rank- k minimizer of (cid:13)(cid:13) R ( SA ) ⊤ X ⊤ − RA ⊤

Original: [Equation]
Revised: [Equation]

Original: 2 F .
Revised: 2 F .

Original: Hence,
Revised: Hence,

Original: [Equation]
Revised: [Equation]

Original: Note that by the normal equations, again rowsp( Y ⊤ ) ⊆ rowsp( RA ⊤ ) and we can write Y =
Revised: Note that by the normal equations, again rowsp( Y ⊤ ) ⊆ rowsp( RA ⊤ ) and we can write Y =

Original: AR ⊤ Z where rank( Z ) =
Revised: AR ⊤ Z where rank( Z ) =

Original: k .
Revised: k .

Original: Thus,
Revised: Thus,

Original: [Equation]
Revised: [Equation]

Original: Lemma C.6 (Avron et al.
Revised: Lemma C.6 (Avron et al.

Original: Lemma 27) .
Revised: Lemma 27) .

Original: For C ∈ R p × m ′ , D ∈ R m × p ′ , G ∈ R p × p ′ , the following problem
Revised: For C ∈ R p × m ′ , D ∈ R m × p ′ , G ∈ R p × p ′ , the following problem

Original: [Equation]
Revised: [Equation]

Original: can be solved in O ( pm ′ r C + p ′
Revised: can be solved in O ( pm ′ r C + p ′

Original: mr D + pp ′ ( r D + r C )) time, where r C = rank( C ) ≤ min { m ′ , p } and r D = rank( D ) ≤ min { m, p ′ } .
Revised: mr D + pp ′ ( r D + r C )) time, where r C = rank( C ) ≤ min { m ′ , p } and r D = rank( D ) ≤ min { m, p ′ } .

Original: Lemma C.7.
Revised: Lemma C.7.

Original: Let S ∈ R m S
Revised: Let S ∈ R m S

Original: × d , R ∈
Revised: × d , R ∈

Original: R m R ×
Revised: R m R ×

Original: d be CountSketch (CS) matrices such that
Revised: d be CountSketch (CS) matrices such that

Original: [Equation]
Revised: [Equation]

Original: Let V ∈ R m 2 Rβ 2 × n , and W ∈ R m 2 Sβ 2 × d be CS matrices.
Revised: Let V ∈ R m 2 Rβ 2 × n , and W ∈ R m 2 Sβ 2 × d be CS matrices.

Original: Then, Algorithm 2 gives a (1 + O ( β + γ )) - approximation in time nnz
Revised: Then, Algorithm 2 gives a (1 + O ( β + γ )) - approximation in time nnz

Original: ( A )
Revised: ( A )

Original: + O (
Revised: + O (

Original: m 4 S β 2 + m 4 R β 2 + m 2 S m 2 R ( m S + m R ) β
Revised: m 4 S β 2 + m 4 R β 2 + m 2 S m 2 R ( m S + m R ) β

Original: [Equation]
Revised: [Equation]

Original: with constant probability.
Revised: with constant probability.

Original: Proof.
Revised: Proof.

Original: The approximation guarantee follows from Eq.
Revised: The approximation guarantee follows from Eq.

Original: (C.6) and the fact that V and W are affine β -embedding matrices of AR ⊤ and SA , respectively (see Lemma C.3).
Revised: (C.6) and the fact that V and W are affine β -embedding matrices of AR ⊤ and SA , respectively (see Lemma C.3).

Original: The algorithm first computes C =
Revised: The algorithm first computes C =

Original: V AR ⊤ , D = SAW ⊤ , G = V AW ⊤ which can be done in time O
Revised: V AR ⊤ , D = SAW ⊤ , G = V AW ⊤ which can be done in time O

Original: (nnz( A )) .
Revised: (nnz( A )) .

Original: As an example, we bound the time to compute C = V AR .
Revised: As an example, we bound the time to compute C = V AR .

Original: Note that since
Revised: Note that since

Original: V is a CS, V A can be computed in O (nnz( A )) time and the number of non-zero entries in the resulting matrix is at most nnz( A ) .
Revised: V is a CS, V A can be computed in O (nnz( A )) time and the number of non-zero entries in the resulting matrix is at most nnz( A ) .

Original: Hence, since R is a CS as well, C can be computed in time
Revised: Hence, since R is a CS as well, C can be computed in time

Original: O (nnz( A ) + nnz( V A ))
Revised: O (nnz( A ) + nnz( V A ))

Original: =
Revised: =

Original: O (nnz( A )) .
Revised: O (nnz( A )) .

Original: Then, it takes an extra O (( m 3 S
Revised: Then, it takes an extra O (( m 3 S

Original: + m 3 R
Revised: + m 3 R

Original: + m 2 S m 2 R ) /β 2 ) time to store C, D and G in matrix form.
Revised: + m 2 S m 2 R ) /β 2 ) time to store C, D and G in matrix form.

Original: Next, as we showed in Lemma C.6, the time to compute Z in
Revised: Next, as we showed in Lemma C.6, the time to compute Z in

Original: Algorithm 2 is O ( m 4 S β 2 + m 4 R β 2 + m 2 S m 2 R
Revised: Algorithm 2 is O ( m 4 S β 2 + m 4 R β 2 + m 2 S m 2 R

Original: m S + m R ) β) .
Revised: m S + m R ) β) .

Original: Finally, it takes O (nnz( A )
Revised: Finally, it takes O (nnz( A )

Original: + k ( nm S + dm R )) time to compute Q =
Revised: + k ( nm S + dm R )) time to compute Q =

Original: AR ⊤ Z L and P =
Revised: AR ⊤ Z L and P =

Original: Z R SA and to return the solution in the form of P n × k Q k × d .
Revised: Z R SA and to return the solution in the form of P n × k Q k × d .

Original: Hence, the total runtime is
Revised: Hence, the total runtime is

Original: [Equation]
Revised: [Equation]

Original: D A TTAINING W ORST -C ASE G UARANTEES D. L OW -R ANK A PPROXIMATION
Revised: D A TTAINING W ORST -C ASE G UARANTEES D. L OW -R ANK A PPROXIMATION

Original: We shall provide the following two methods to achieve worst case guarantees: MixedSketch—whose guarantee is via the sketch monotonicity property, and approximate comparison method (a.k.a.
Revised: We shall provide the following two methods to achieve worst case guarantees: MixedSketch—whose guarantee is via the sketch monotonicity property, and approximate comparison method (a.k.a.

Original: ApproxCheck), which just approximately evaluates the cost of two solutions and takes the better one.
Revised: ApproxCheck), which just approximately evaluates the cost of two solutions and takes the better one.

Original: These methods asymptotically achieve the same worst-case guarantee.
Revised: These methods asymptotically achieve the same worst-case guarantee.

Original: However, for any input matrix A and any pair of sketches S, T , the performance of the MixedSketch method on ( A, S, T ) is never worse than the performance of its corresponding ApproxCheck method on ( A, S, T ) , and can be much better.
Revised: However, for any input matrix A and any pair of sketches S, T , the performance of the MixedSketch method on ( A, S, T ) is never worse than the performance of its corresponding ApproxCheck method on ( A, S, T ) , and can be much better.

Original: Remark D.1 .
Revised: Remark D.1 .

Original: Let A = diag(2 , 2 , √ 2 , √ 2) , and suppose the goal is to find a rank- 2 approximation of A .
Revised: Let A = diag(2 , 2 , √ 2 , √ 2) , and suppose the goal is to find a rank- 2 approximation of A .

Original: Consider two sketches S and T such that SA and TA capture span ( e 1 , e 3 ) and span ( e 2 , e 4 ) , respectively.
Revised: Consider two sketches S and T such that SA and TA capture span ( e 1 , e 3 ) and span ( e 2 , e 4 ) , respectively.

Original: Then for both SA and TA , the best solution in the subspace of one of these two spaces is a ( 32 ) -approximation:
Revised: Then for both SA and TA , the best solution in the subspace of one of these two spaces is a ( 32 ) -approximation:

Original: ∥ A − A 2 ∥ 2 F = 4 and ∥
Revised: ∥ A − A 2 ∥ 2 F = 4 and ∥

Original: A − P SA ∥ 2 F = ∥
Revised: A − P SA ∥ 2 F = ∥

Original: A − P TA ∥ 2 F = 6 where P SA and
Revised: A − P TA ∥ 2 F = 6 where P SA and

Original: P TA respectively denote the best approximation of A in the space spanned by SA and TA .
Revised: P TA respectively denote the best approximation of A in the space spanned by SA and TA .

Original: However, if we find the best rank- 2 approximation of A , Z , inside the span of the union of SA and
Revised: However, if we find the best rank- 2 approximation of A , Z , inside the span of the union of SA and

Original: TA , then ∥ A − Z ∥ 2 F = 4 .
Revised: TA , then ∥ A − Z ∥ 2 F = 4 .

Original: Since ApproxCheck just chooses the better of SA and TA by evaluating their costs, it misses out on the opportunity to do as well as MixedSketch.
Revised: Since ApproxCheck just chooses the better of SA and TA by evaluating their costs, it misses out on the opportunity to do as well as MixedSketch.

Original: Here, we show the sketch monotonicity property for LRA.
Revised: Here, we show the sketch monotonicity property for LRA.

Original: Theorem D.2.
Revised: Theorem D.2.

Original: Let A ∈ R n ×
Revised: Let A ∈ R n ×

Original: d be an input matrix, V and W be η -affine embeddings, and S 1 ∈ R m S × n , R 1 ∈
Revised: d be an input matrix, V and W be η -affine embeddings, and S 1 ∈ R m S × n , R 1 ∈

Original: R m R × n be arbitrary matrices.
Revised: R m R × n be arbitrary matrices.

Original: Consider arbitrary extensions to S 1 , R 1 : S, R (e.g.,
Revised: Consider arbitrary extensions to S 1 , R 1 : S, R (e.g.,

Original: S is a concatenation of S 1 with an arbitrary matrix with the same number of columns).
Revised: S is a concatenation of S 1 with an arbitrary matrix with the same number of columns).

Original: Then, (cid:13)(cid:13)
Revised: Then, (cid:13)(cid:13)

Original: A − ALG LRA (( S, R, V, W ) , A ))
Revised: A − ALG LRA (( S, R, V, W ) , A ))

Original: (cid:13)(cid:13) 2 F ≤ (1 + η ) 2 ∥
Revised: (cid:13)(cid:13) 2 F ≤ (1 + η ) 2 ∥

Original: A − ALG LRA (( S 1 , R 1 , V, W ) , A ) ∥ 2 F
Revised: A − ALG LRA (( S 1 , R 1 , V, W ) , A ) ∥ 2 F

Original: Proof.
Revised: Proof.

Original: We have (cid:13)(cid:13)
Revised: We have (cid:13)(cid:13)

Original: A − ALG LRA (( S, R, V, W ) , A ) (cid:13)(cid:13) 2 F ≤ (1 + η ) min rank-
Revised: A − ALG LRA (( S, R, V, W ) , A ) (cid:13)(cid:13) 2 F ≤ (1 + η ) min rank-

Original: k X (cid:13)(cid:13)
Revised: k X (cid:13)(cid:13)

Original: ARXSA − A (cid:13)(cid:13) 2 F =
Revised: ARXSA − A (cid:13)(cid:13) 2 F =

Original: (1 + η ) min rank-
Revised: (1 + η ) min rank-

Original: k X : X ∈ row ( SA ) ∩ col ( AR ) ∥
Revised: k X : X ∈ row ( SA ) ∩ col ( AR ) ∥

Original: X − A ∥ 2 F , which is in turn at most (1 + η ) min rank-
Revised: X − A ∥ 2 F , which is in turn at most (1 + η ) min rank-

Original: k X : X ∈ row ( S 1 A ) ∩ col ( AR 1 )
Revised: k X : X ∈ row ( S 1 A ) ∩ col ( AR 1 )

Original: ∥ X − A ∥ 2 F = (1 + η ) min rank-
Revised: ∥ X − A ∥ 2 F = (1 + η ) min rank-

Original: k X ∥ AR 1 XS 1 A − A ∥ 2 F ≤ (1 + η ) 2 ∥
Revised: k X ∥ AR 1 XS 1 A − A ∥ 2 F ≤ (1 + η ) 2 ∥

Original: A − ALG LRA (( S 1 , R 1 , V, W ) , A ) ∥ 2 F , where we use the fact the V, W are affine η -embeddings
Revised: A − ALG LRA (( S 1 , R 1 , V, W ) , A ) ∥ 2 F , where we use the fact the V, W are affine η -embeddings

Original: (Definition C.1), as well as the fact that ( col ( AR 1 ) ∩ row ( S 1 A )) ⊆ (cid:0)
Revised: (Definition C.1), as well as the fact that ( col ( AR 1 ) ∩ row ( S 1 A )) ⊆ (cid:0)

Original: col ( AR ) ∩ row ( SA ) (cid:1) .
Revised: col ( AR ) ∩ row ( SA ) (cid:1) .

Original: ApproxCheck for LRA.
Revised: ApproxCheck for LRA.

Original: We give the pseudocode for the ApproxCheck method and prove that the runtime of this method for LRA is of the same order as the classical time-optimal sketching algorithm of LRA.
Revised: We give the pseudocode for the ApproxCheck method and prove that the runtime of this method for LRA is of the same order as the classical time-optimal sketching algorithm of LRA.

Original: Input: learned sketches S L , R L , V L , W L ; classical sketches S C , R C , V C , W C ; β ; A ∈ R n × d 1: P L , Q L ← ALG LRA ( S L , R L , V L , W L , A ) , P C Q C ← ALG LRA ( S C , R C , V C , W C , A ) 2: Let S ∈ R O
Revised: Input: learned sketches S L , R L , V L , W L ; classical sketches S C , R C , V C , W C ; β ; A ∈ R n × d 1: P L , Q L ← ALG LRA ( S L , R L , V L , W L , A ) , P C Q C ← ALG LRA ( S C , R C , V C , W C , A ) 2: Let S ∈ R O

Original: (1 /β 2 ) × n , R ∈ R
Revised: (1 /β 2 ) × n , R ∈ R

Original: O
Revised: O

Original: (1 /β 2 ) ×
Revised: (1 /β 2 ) ×

Original: d be classical CountSketch matrices 3: ∆ L ← (cid:13)(cid:13) S ( P L Q L − A ) R ⊤ (cid:13)(cid:13) 2 F , ∆ C ← (cid:13)(cid:13) S ( P C Q C − A ) R ⊤ (cid:13)(cid:13) 2 F 4: if ∆ L ≤ ∆ C then 5: return P L Q L 6: end if 7: return P C Q C
Revised: d be classical CountSketch matrices 3: ∆ L ← (cid:13)(cid:13) S ( P L Q L − A ) R ⊤ (cid:13)(cid:13) 2 F , ∆ C ← (cid:13)(cid:13) S ( P C Q C − A ) R ⊤ (cid:13)(cid:13) 2 F 4: if ∆ L ≤ ∆ C then 5: return P L Q L 6: end if 7: return P C Q C

Original: Theorem D.3.
Revised: Theorem D.3.

Original: Assume we have data A ∈ R n × d , learned sketches S L ∈ R poly( kϵ ) × n , R L ∈
Revised: Assume we have data A ∈ R n × d , learned sketches S L ∈ R poly( kϵ ) × n , R L ∈

Original: R poly( k ϵ ) ×
Revised: R poly( k ϵ ) ×

Original: d , V L ∈ R poly( kϵ ) × n , W L ∈ R poly( kϵ ) ×
Revised: d , V L ∈ R poly( kϵ ) × n , W L ∈ R poly( kϵ ) ×

Original: d which attain a (1 + O ( γ )) -approximation, classical sketches of the same size, S C , R C , V C , W C , which attain a (1 + O ( ϵ ))
Revised: d which attain a (1 + O ( γ )) -approximation, classical sketches of the same size, S C , R C , V C , W C , which attain a (1 + O ( ϵ ))

Original: -approximation, and a tradeoff parameter β .
Revised: -approximation, and a tradeoff parameter β .

Original: Then, Algorithm 5 attains a (1 + β + min( γ, ϵ )) -approximation in
Revised: Then, Algorithm 5 attains a (1 + β + min( γ, ϵ )) -approximation in

Original: O (nnz( A ) + ( n + d ) poly( kϵ ) + k 4 β 4 · poly( kϵ )) time.
Revised: O (nnz( A ) + ( n + d ) poly( kϵ ) + k 4 β 4 · poly( kϵ )) time.

Original: Proof.
Revised: Proof.

Original: Let ( P L , Q L ) , ( P C , Q C ) be the approximate rank- k approximations of A in factored form using ( S L , R L ) and ( S O , R O ) .
Revised: Let ( P L , Q L ) , ( P C , Q C ) be the approximate rank- k approximations of A in factored form using ( S L , R L ) and ( S O , R O ) .

Original: Then, clearly,
Revised: Then, clearly,

Original: [Equation]
Revised: [Equation]

Original: Let Γ L = P L Q L − A , Γ C = P C Q C − A and Γ M = arg min( ∥ S Γ L R ∥ F , ∥ S Γ C R ∥ F ) .
Revised: Let Γ L = P L Q L − A , Γ C = P C Q C − A and Γ M = arg min( ∥ S Γ L R ∥ F , ∥ S Γ C R ∥ F ) .

Original: Then, ∥
Revised: Then, ∥

Original: [Equation]
Revised: [Equation]

Original: ▷ by Lemma C.
Revised: ▷ by Lemma C.

Original: [Equation]
Revised: [Equation]

Original: ▷ by Eq.
Revised: ▷ by Eq.

Original: (D
Revised: (D

Original: [Equation]
Revised: [Equation]

Original: Runtime analysis.
Revised: Runtime analysis.

Original: By Lemma C.7, Algorithm 2 computes P L , Q L and P C , Q C in time O
Revised: By Lemma C.7, Algorithm 2 computes P L , Q L and P C , Q C in time O

Original: (nnz( A ) + k
Revised: (nnz( A ) + k

Original: [Equation]
Revised: [Equation]

Original: + k
Revised: + k

Original: [Equation]
Revised: [Equation]

Original: ϵ 2 ( n + dk 2 ϵ 4 )) .
Revised: ϵ 2 ( n + dk 2 ϵ 4 )) .

Original: Next, once we have P L , Q L and P C , Q C , it takes O (nnz( A ) + kβ 4 ) time to compute ∆ L and ∆ C .
Revised: Next, once we have P L , Q L and P C , Q C , it takes O (nnz( A ) + kβ 4 ) time to compute ∆ L and ∆ C .

Original: [Equation]
Revised: [Equation]

Original: To interpret the above theorem, note that when ϵ ≫ k ( n + d ) − 4 , we can set β − 

Revised: To interpret the above theorem, note that when ϵ ≫ k ( n + d ) − 4 , we can set β −

Original: =
Revised: =

Original: O ( k ( n + d ) − 4 ) so that Algorithm 5 has the same asymptotic runtime as the best (1 + ϵ ) -approximation algorithm for
Revised: O ( k ( n + d ) − 4 ) so that Algorithm 5 has the same asymptotic runtime as the best (1 + ϵ ) -approximation algorithm for

Original: LRA with the classical CountSketch.
Revised: LRA with the classical CountSketch.

Original: Moreover, Algorithm 5 is a (1 + o ( ϵ )) -approximation when the learned sketch outperforms classical sketches,
Revised: Moreover, Algorithm 5 is a (1 + o ( ϵ )) -approximation when the learned sketch outperforms classical sketches,

Original: γ =
Revised: γ =

Original: o ( ϵ ) .
Revised: o ( ϵ ) .

Original: On the other hand, when the learned sketches perform poorly, γ = Ω( ϵ ) , the worst-case guarantee of Algorithm 5 remains (1 + O ( ϵ )) .
Revised: On the other hand, when the learned sketches perform poorly, γ = Ω( ϵ ) , the worst-case guarantee of Algorithm 5 remains (1 + O ( ϵ )) .

Original: D. S ECOND -O RDER O PTIMIZATION
Revised: D. S ECOND -O RDER O PTIMIZATION

Original: For the sketches for second-order optimization, the monotonicity property does not hold.
Revised: For the sketches for second-order optimization, the monotonicity property does not hold.

Original: Below we provide an input-sparsity algorithm which can test for and use the better of a random sketch and a learned sketch.
Revised: Below we provide an input-sparsity algorithm which can test for and use the better of a random sketch and a learned sketch.

Original: Our theorem is as follows.
Revised: Our theorem is as follows.

Original: Theorem D.4.
Revised: Theorem D.4.

Original: Let ϵ ∈
Revised: Let ϵ ∈

Original: (0 , 0 . 09) be a constant and S 1 a learned Count-Sketch matrix.
Revised: (0 , 0 . 09) be a constant and S 1 a learned Count-Sketch matrix.

Original: Suppose that A is of full rank.
Revised: Suppose that A is of full rank.

Original: There is an algorithm whose output is a solution ˆ x which, with probability at least 0 . 98 , satisfies that
Revised: There is an algorithm whose output is a solution ˆ x which, with probability at least 0 . 98 , satisfies that

Original: ∥ A (ˆ x − x ∗ ) ∥ 2 ≤ O (cid:0) min (cid:8)
Revised: ∥ A (ˆ x − x ∗ ) ∥ 2 ≤ O (cid:0) min (cid:8)

Original: Z 2 ( S 1 )
Revised: Z 2 ( S 1 )

Original: Z 1 ( S 1 ) , ϵ (cid:9)(cid:1) ∥
Revised: Z 1 ( S 1 ) , ϵ (cid:9)(cid:1) ∥

Original: Ax ∗ ∥ 2 , where x ∗ = arg min x ∈C ∥ Ax − b ∥ 2 is the least-squares solution.
Revised: Ax ∗ ∥ 2 , where x ∗ = arg min x ∈C ∥ Ax − b ∥ 2 is the least-squares solution.

Original: Furthermore, the algorithm runs in O (nnz( A ) log( 1 ϵ ) + poly( d ϵ )) time.
Revised: Furthermore, the algorithm runs in O (nnz( A ) log( 1 ϵ ) + poly( d ϵ )) time.

Original: Consider the minimization problem
Revised: Consider the minimization problem

Original: [Equation]
Revised: [Equation]

Original: which is used as a subroutine for the IHS (cf.
Revised: which is used as a subroutine for the IHS (cf.

Original: We note that in this subroutine if we let x ← x − x i − 1 , b ← b −
Revised: We note that in this subroutine if we let x ← x − x i − 1 , b ← b −

Original: Ax i − 1 , C ← C − x i − 1 , we would get the guarantee of the i -th iteration of the original IHS.
Revised: Ax i − 1 , C ← C − x i − 1 , we would get the guarantee of the i -th iteration of the original IHS.

Original: To analyze the performance of the learned sketch, we define the following quantities (corresponding exactly to the unconstrained case in (Pilanci & Wainwright, 2016))
Revised: To analyze the performance of the learned sketch, we define the following quantities (corresponding exactly to the unconstrained case in (Pilanci & Wainwright, 2016))

Original: [Equation]
Revised: [Equation]

Original: When S is a (1 + ϵ ) -subspace embedding of colsp( A ) , we have Z 1 ( S ) ≥ 1 − ϵ and Z 2 ( S ) ≤ 2 ϵ .
Revised: When S is a (1 + ϵ ) -subspace embedding of colsp( A ) , we have Z 1 ( S ) ≥ 1 − ϵ and Z 2 ( S ) ≤ 2 ϵ .

Original: For a general sketching matrix S , the following is the approximation guarantee of ˆ Z 1 and ˆ Z 2 , which are estimates of Z 1 ( S ) and Z 2 ( S ) , respectively.
Revised: For a general sketching matrix S , the following is the approximation guarantee of ˆ Z 1 and ˆ Z 2 , which are estimates of Z 1 ( S ) and Z 2 ( S ) , respectively.

Original: The main idea is that AR − 1 is well-conditioned, where R is as calculated in Algorithm 6.
Revised: The main idea is that AR − 1 is well-conditioned, where R is as calculated in Algorithm 6.

Original: Lemma D.5.
Revised: Lemma D.5.

Original: Suppose that η ∈ (0 , 13 ) is a small constant, A is of full rank and S has poly( d/η ) rows.
Revised: Suppose that η ∈ (0 , 13 ) is a small constant, A is of full rank and S has poly( d/η ) rows.

Original: The function E STIMATE ( S, A ) returns in O ((nnz( A ) log 1 η +poly( dη )) time
Revised: The function E STIMATE ( S, A ) returns in O ((nnz( A ) log 1 η +poly( dη )) time

Original: ˆ Z 1 , ˆ Z 2 which with probability at least 0 .
Revised: ˆ Z 1 , ˆ Z 2 which with probability at least 0 .

Original: 99 satisfy that Z 1 ( S ) 1+ η ≤ ˆ Z 1 ≤ Z 1 ( S ) 1 − η and Z 2 ( S ) (1+ η ) 2 − 3 η ≤ ˆ Z 2 ≤ Z 2 ( S ) (1 − η ) 2 + 3 η .
Revised: 99 satisfy that Z 1 ( S ) 1+ η ≤ ˆ Z 1 ≤ Z 1 ( S ) 1 − η and Z 2 ( S ) (1+ η ) 2 − 3 η ≤ ˆ Z 2 ≤ Z 2 ( S ) (1 − η ) 2 + 3 η .

Original: Proof.
Revised: Proof.

Original: Suppose that AR − 1 = UW , where U ∈ R n × d has orthonormal columns, which form an orthonormal basis of the column space of A .
Revised: Suppose that AR − 1 = UW , where U ∈ R n × d has orthonormal columns, which form an orthonormal basis of the column space of A .

Original: Since T is a subspace embedding of the column space of A with probability 0 .
Revised: Since T is a subspace embedding of the column space of A with probability 0 .

Original: 99 , it holds for all x ∈ R d that
Revised: 99 , it holds for all x ∈ R d that

Original: [Equation]
Revised: [Equation]

Original: Since
Revised: Since

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: we have that
Revised: we have that

Original: [Equation]
Revised: [Equation]

Original: It is easy to see that
Revised: It is easy to see that

Original: [Equation]
Revised: [Equation]

Original: and thus, min
Revised: and thus, min

Original: [Equation]
Revised: [Equation]

Original: Recall that SUW = SAR − 1 .
Revised: Recall that SUW = SAR − 1 .

Original: We see that
Revised: We see that

Original: [Equation]
Revised: [Equation]

Original: By definition,
Revised: By definition,

Original: [Equation]
Revised: [Equation]

Original: It follows from (D.4)
Revised: It follows from (D.4)

Original: that
Revised: that

Original: [Equation]
Revised: [Equation]

Original: op
Revised: op

Original: [Equation]
Revised: [Equation]

Original: and from (D.4), (D.3) and Lemma 5.36 of Vershynin (2012) that
Revised: and from (D.4), (D.3) and Lemma 5.36 of Vershynin (2012) that

Original: [Equation]
Revised: [Equation]

Original: Since
Revised: Since

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: op
Revised: op

Original: [Equation]
Revised: [Equation]

Original: ≤
Revised: ≤

Original: [Equation]
Revised: [Equation]

Original: it follows that
Revised: it follows that

Original: [Equation]
Revised: [Equation]

Original: We have so far proved the correctness of the approximation and we next analyze the runtime below.
Revised: We have so far proved the correctness of the approximation and we next analyze the runtime below.

Original: Since S and T are sparse, computing SA and TA takes O (nnz( A )) time.
Revised: Since S and T are sparse, computing SA and TA takes O (nnz( A )) time.

Original: The QR decomposition of TA , which is a matrix of size poly( d/η ) × d , can be computed in poly( d/η ) time.
Revised: The QR decomposition of TA , which is a matrix of size poly( d/η ) × d , can be computed in poly( d/η ) time.

Original: The matrix
Revised: The matrix

Original: SAR − 1 can be computed in poly( d ) time.
Revised: SAR − 1 can be computed in poly( d ) time.

Original: Since it has size poly( d/η ) × d , its smallest singular value can be computed in poly( d/η ) time.
Revised: Since it has size poly( d/η ) × d , its smallest singular value can be computed in poly( d/η ) time.

Original: To approximate Z 2 ( S ) , we can use the power method to estimate (cid:13)(cid:13) ( SAR − 1 ) T SAR − 1 − I (cid:13)(cid:13) op up to a (1 ± η ) -factor in O ((nnz( A ) + poly( d/η )) log(1 /η )) time.
Revised: To approximate Z 2 ( S ) , we can use the power method to estimate (cid:13)(cid:13) ( SAR − 1 ) T SAR − 1 − I (cid:13)(cid:13) op up to a (1 ± η ) -factor in O ((nnz( A ) + poly( d/η )) log(1 /η )) time.

Original: Now we are ready to prove Theorem D.4.
Revised: Now we are ready to prove Theorem D.4.

Original: Proof of Theorem D.4.
Revised: Proof of Theorem D.4.

Original: In Lemma D.5, we have with probability at least 0 .
Revised: In Lemma D.5, we have with probability at least 0 .

Original: 99 that
Revised: 99 that

Original: [Equation]
Revised: [Equation]

Original: and similarly,
Revised: and similarly,

Original: [Equation]
Revised: [Equation]

Original: Note that since S 2 is an ϵ -subspace embedding with probability at least 0 .
Revised: Note that since S 2 is an ϵ -subspace embedding with probability at least 0 .

Original: 99 , we have that Z 1 ( S 2 ) ≥ 1 − ϵ and Z 2 ( S 2 ) /Z
Revised: 99 , we have that Z 1 ( S 2 ) ≥ 1 − ϵ and Z 2 ( S 2 ) /Z

Original: 1 ( S 2 ) ≤ 2 .
Revised: 1 ( S 2 ) ≤ 2 .

Original: 2 ϵ .
Revised: 2 ϵ .

Original: Consider Z 1 ( S 1 ) .
Revised: Consider Z 1 ( S 1 ) .

Original: First, we consider the case where Z 1 ( S 1 ) < 1 / 2 .
Revised: First, we consider the case where Z 1 ( S 1 ) < 1 / 2 .

Original: Observe that Z 2 ( S ) ≥ 1 − Z 1 ( S ) .
Revised: Observe that Z 2 ( S ) ≥ 1 − Z 1 ( S ) .

Original: We have in this case (cid:98) Z 1 , 2 / (cid:98)
Revised: We have in this case (cid:98) Z 1 , 2 / (cid:98)

Original: Z 1 , 1 > 1 / 5 ≥ 2 .
Revised: Z 1 , 1 > 1 / 5 ≥ 2 .

Original: 2 ϵ ≥ Z 2 ( S 2 ) /Z
Revised: 2 ϵ ≥ Z 2 ( S 2 ) /Z

Original: 1 ( S 2 ) .
Revised: 1 ( S 2 ) .

Original: In this case our algorithm will choose Scorrectly.
Revised: In this case our algorithm will choose Scorrectly.

Original: Next, assume that Z 1 ( S 1 ) ≥ 1 / 2 .
Revised: Next, assume that Z 1 ( S 1 ) ≥ 1 / 2 .

Original: Now we have with probability at least 0 .
Revised: Now we have with probability at least 0 .

Original: 98 that
Revised: 98 that

Original: [Equation]
Revised: [Equation]

Original: Therefore, when Z 2 ( S 1 ) /Z 1 ( S 1 ) ≤ c 1 Z 2 ( S 2 ) /Z
Revised: Therefore, when Z 2 ( S 1 ) /Z 1 ( S 1 ) ≤ c 1 Z 2 ( S 2 ) /Z

Original: ( S 2 ) for some small absolute constant c 1 > 0 , we will have (cid:98) Z 1 , 2 / (cid:98)
Revised: ( S 2 ) for some small absolute constant c 1 > 0 , we will have (cid:98) Z 1 , 2 / (cid:98)

Original: Z 1 , 1 < (cid:98) Z 2 , 2 / (cid:98) Z 2 , 1 , and our algorithm will choose S 1 correctly.
Revised: Z 1 , 1 < (cid:98) Z 2 , 2 / (cid:98) Z 2 , 1 , and our algorithm will choose S 1 correctly.

Original: If Z 2 ( S 1 ) /Z 1 ( S 1 ) ≥ C 1 ϵ for some absolute constant C 1 > 0 , our algorithm will choose S 2 correctly.
Revised: If Z 2 ( S 1 ) /Z 1 ( S 1 ) ≥ C 1 ϵ for some absolute constant C 1 > 0 , our algorithm will choose S 2 correctly.

Original: In the remaining case, both ratios Z 2 ( S 2 ) /Z
Revised: In the remaining case, both ratios Z 2 ( S 2 ) /Z

Original: 1 ( S 2 ) and Z 2 ( S 1 ) /Z
Revised: 1 ( S 2 ) and Z 2 ( S 1 ) /Z

Original: ( S 1 ) are at most max { C 2 , 3 } ϵ , and the guarantee of the theorem holds automatically.
Revised: ( S 1 ) are at most max { C 2 , 3 } ϵ , and the guarantee of the theorem holds automatically.

Original: The correctness of our claim then follows from Proposition 1 of Pilanci & Wainwright (2016), together with the fact that S 2 is a random subspace embedding.
Revised: The correctness of our claim then follows from Proposition 1 of Pilanci & Wainwright (2016), together with the fact that S 2 is a random subspace embedding.

Original: The runtime follows from Lemma D.5 and
Revised: The runtime follows from Lemma D.5 and

Original: Theorem 2.2 of Cormode & Dickens (2019).
Revised: Theorem 2.2 of Cormode & Dickens (2019).

Original: E S KETCH L EARNING : O MITTED P ROOFS E. P ROOF OF T HEOREM 5. 
Revised: E S KETCH L EARNING : O MITTED P ROOFS E. P ROOF OF T HEOREM 5.

Original: We need the following lemmas for the ridge leverage score sampling in (Cohen et al., 2017).
Revised: We need the following lemmas for the ridge leverage score sampling in (Cohen et al., 2017).

Original: Lemma E.1 ((Cohen et al., 2017, Lemma 4)) .
Revised: Lemma E.1 ((Cohen et al., 2017, Lemma 4)) .

Original: Let λ =
Revised: Let λ =

Original: ∥
Revised: ∥

Original: A − A k ∥ 2 F /k .
Revised: A − A k ∥ 2 F /k .

Original: Then we have (cid:80) i τ
Revised: Then we have (cid:80) i τ

Original: i ( A ) ≤ 2 k .
Revised: i ( A ) ≤ 2 k .

Original: Lemma E.2 ((Cohen et al., 2017, Theorem 7)) .
Revised: Lemma E.2 ((Cohen et al., 2017, Theorem 7)) .

Original: Let λ =
Revised: Let λ =

Original: ∥
Revised: ∥

Original: A − A k ∥ 2 F /k and ˜ τ
Revised: A − A k ∥ 2 F /k and ˜ τ

Original: i ≥ τ i ( A ) be an overestimate to the i -th ridge leverage score of A .
Revised: i ≥ τ i ( A ) be an overestimate to the i -th ridge leverage score of A .

Original: Let p i =
Revised: Let p i =

Original: ˜
Revised: ˜

Original: τ i / (cid:80) i ˜ τ i .
Revised: τ i / (cid:80) i ˜ τ i .

Original: If C is a matrix that is constructed by sampling t
Revised: If C is a matrix that is constructed by sampling t

Original: = O ((log k + log(1 /δ ) ϵ ) · (cid:80) i ˜ τ i ) rows of A , each set to a i with probability p
Revised: = O ((log k + log(1 /δ ) ϵ ) · (cid:80) i ˜ τ i ) rows of A , each set to a i with probability p

Original: i , then with probability at least 1 − δ we have
Revised: i , then with probability at least 1 − δ we have

Original: [Equation]
Revised: [Equation]

Original: Recall that the sketch monotonicity for low-rank approximation says that concatenating two sketching matrices S 1 and S 2 will not increase the error compared to the single sketch matrix S 1 or S 2 , Now matter how S 1 and S 2 are constructed.
Revised: Recall that the sketch monotonicity for low-rank approximation says that concatenating two sketching matrices S 1 and S 2 will not increase the error compared to the single sketch matrix S 1 or S 2 , Now matter how S 1 and S 2 are constructed.

Original: (see Section D.1 and Section 4 in (Indyk et al., 2019))
Revised: (see Section D.1 and Section 4 in (Indyk et al., 2019))

Original: Proof.
Revised: Proof.

Original: We first consider the first condition.
Revised: We first consider the first condition.

Original: From the condition that τ i ( B ) ≥ 1 β
Revised: From the condition that τ i ( B ) ≥ 1 β

Original: τ i ( A ) we know that if we sample m =
Revised: τ i ( A ) we know that if we sample m =

Original: O ( β · ( k log k + k/ϵ )) rows according to τ i ( A ) .
Revised: O ( β · ( k log k + k/ϵ )) rows according to τ i ( A ) .

Original: The actual probability that the i -th row of B gets sampled is
Revised: The actual probability that the i -th row of B gets sampled is

Original: [Equation]
Revised: [Equation]

Original: From (cid:80) i τ i ( B ) ≤ 2 k and Lemma E.2 (recall the sketch monotonicity property for LRA), we have that with probability at least 9 / 10 , S 2 is a matrix such that
Revised: From (cid:80) i τ i ( B ) ≤ 2 k and Lemma E.2 (recall the sketch monotonicity property for LRA), we have that with probability at least 9 / 10 , S 2 is a matrix such that

Original: [Equation]
Revised: [Equation]

Original: Hence, since S =
Revised: Hence, since S =

Original: [ S 1 S 2 ] , from the the sketch monotonicity property for LRA we have that
Revised: [ S 1 S 2 ] , from the the sketch monotonicity property for LRA we have that

Original: [Equation]
Revised: [Equation]

Original: Now we consider the second condition.
Revised: Now we consider the second condition.

Original: Suppose that { X i } i ≤ m and { Y i } i ≤ m are a sequence of m =
Revised: Suppose that { X i } i ≤ m and { Y i } i ≤ m are a sequence of m =

Original: O ( k log k + k/ϵ ) samples from [ n ] according to the sampling probability distribution p and q , where p i
Revised: O ( k log k + k/ϵ ) samples from [ n ] according to the sampling probability distribution p and q , where p i

Original: =
Revised: =

Original: τ i ( A ) (cid:80) i τ i
Revised: τ i ( A ) (cid:80) i τ i

Original: ( A ) and q i =
Revised: ( A ) and q i =

Original: τ i ( B ) (cid:80) i τ i ( B ) .
Revised: τ i ( B ) (cid:80) i τ i ( B ) .

Original: Let S be the set of index i such that X
Revised: Let S be the set of index i such that X

Original: i ̸
Revised: i ̸

Original: = Y i .
Revised: = Y i .

Original: From the property of the total variation distance, we get that
Revised: From the property of the total variation distance, we get that

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: From Markov’s inequality we get that with probability at least 1 − 1 .
Revised: From Markov’s inequality we get that with probability at least 1 − 1 .

Original: 1 β , | S | ≤ 1 / (1 . 1 β
Revised: 1 β , | S | ≤ 1 / (1 . 1 β

Original: [Equation]
Revised: [Equation]

Original: · βm = 1011 m .
Revised: · βm = 1011 m .

Original: Let T be the set of index i such that X i
Revised: Let T be the set of index i such that X i

Original: = Y i .
Revised: = Y i .

Original: We have that with probability at least 1 − 1 .
Revised: We have that with probability at least 1 − 1 .

Original: 1 β , | T | ≥ m − 10 11 m = Ω
Revised: 1 β , | T | ≥ m − 10 11 m = Ω

Original: ( k log k + k/ϵ ) .
Revised: ( k log k + k/ϵ ) .

Original: Because that { Y i } i ∈ T is i.i.d samples according to q and the actual sample we take is { X i } i ∈ T .
Revised: Because that { Y i } i ∈ T is i.i.d samples according to q and the actual sample we take is { X i } i ∈ T .

Original: From Lemma E.2 we get that with probability at least 9 / 10 , the row space of B T satisfies
Revised: From Lemma E.2 we get that with probability at least 9 / 10 , the row space of B T satisfies

Original: [Equation]
Revised: [Equation]

Original: Similarly, from the the sketch monotonicity property we have that with probability at least 0 .
Revised: Similarly, from the the sketch monotonicity property we have that with probability at least 0 .

Original: 9 − 1 .
Revised: 9 − 1 .

Original: 1 β
Revised: 1 β

Original: [Equation]
Revised: [Equation]

Original: X : row ( X )
Revised: X : row ( X )

Original: [Equation]
Revised: [Equation]

Original: row (
Revised: row (

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

Original: E
Revised: E

Original: [Equation]
Revised: [Equation]

Original: P ROOF OF T HEOREM 6. 
Revised: P ROOF OF T HEOREM 6.

Original: First we prove the following lemma.
Revised: First we prove the following lemma.

Original: Lemma E.3.
Revised: Lemma E.3.

Original: Let δ ∈ (0 , 1 /m ] .
Revised: Let δ ∈ (0 , 1 /m ] .

Original: It holds with probability at least 1 − δ that
Revised: It holds with probability at least 1 − δ that

Original: [Equation]
Revised: [Equation]

Original: provided that
Revised: provided that

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We shall adapt the proof of Theorem 5 in (Bourgain et al., 2015) to our setting.
Revised: We shall adapt the proof of Theorem 5 in (Bourgain et al., 2015) to our setting.

Original: Let T denote the unit sphere in colsp( A ) and set the sparsity parameter s = 1 .
Revised: Let T denote the unit sphere in colsp( A ) and set the sparsity parameter s = 1 .

Original: Observe that ∥ Sx ∥ 22 = ∥ x I ∥ 22 + ∥
Revised: Observe that ∥ Sx ∥ 22 = ∥ x I ∥ 22 + ∥

Original: Sx I c ∥ 22 , and so it suffices to show that
Revised: Sx I c ∥ 22 , and so it suffices to show that

Original: [Equation]
Revised: [Equation]

Original: for x ∈ T .
Revised: for x ∈ T .

Original: We make the following definition, as in (2.6) of (Bourgain et al., 2015):
Revised: We make the following definition, as in (2.6) of (Bourgain et al., 2015):

Original: [Equation]
Revised: [Equation]

Original: and thus, S ′ x I c =
Revised: and thus, S ′ x I c =

Original: A δ,x σ .
Revised: A δ,x σ .

Original: Also by E ∥ S
Revised: Also by E ∥ S

Original: ′ x I c ∥ 22 = ∥ x
Revised: ′ x I c ∥ 22 = ∥ x

Original: I c ∥ 22 , one has sup
Revised: I c ∥ 22 , one has sup

Original: [Equation]
Revised: [Equation]

Original: Now, in (2.7) of (Bourgain et al., 2015) we instead define a semi-norm
Revised: Now, in (2.7) of (Bourgain et al., 2015) we instead define a semi-norm

Original: [Equation]
Revised: [Equation]

Original: Then (2.8) continues to hold, and (2.9) as well as (2.10) continue to hold if the supremum in the left-hand side is replaced with the left-hand side of (E.1).
Revised: Then (2.8) continues to hold, and (2.9) as well as (2.10) continue to hold if the supremum in the left-hand side is replaced with the left-hand side of (E.1).

Original: At the beginning of Theorem 5, we define
Revised: At the beginning of Theorem 5, we define

Original: U ( i ) to be U , but each row j ∈
Revised: U ( i ) to be U , but each row j ∈

Original: I c is multiplied by δ ij and each row j ∈
Revised: I c is multiplied by δ ij and each row j ∈

Original: I is zeroed out.
Revised: I is zeroed out.

Original: Then we have in the first step of (4.5) that
Revised: Then we have in the first step of (4.5) that

Original: [Equation]
Revised: [Equation]

Original: instead of equality.
Revised: instead of equality.

Original: One can verify that the rest of (4.5) goes through.
Revised: One can verify that the rest of (4.5) goes through.

Original: It remains true that ∥·∥ δ ≤ (1 / √ s ) ∥·∥ 2 , and thus (4.6) holds.
Revised: It remains true that ∥·∥ δ ≤ (1 / √ s ) ∥·∥ 2 , and thus (4.6) holds.

Original: One can verify that the rest of the proof of Theorem 5 in (Bourgain et al., 2015) continues to hold if we replace (cid:80) nj =1 with (cid:80) j ∈
Revised: One can verify that the rest of the proof of Theorem 5 in (Bourgain et al., 2015) continues to hold if we replace (cid:80) nj =1 with (cid:80) j ∈

Original: I
Revised: I

Original: c and max 1 ≤ j ≤ n with max j ∈ I c , noting that
Revised: c and max 1 ≤ j ≤ n with max j ∈ I c , noting that

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: Thus, the symmetrization inequalities on
Revised: Thus, the symmetrization inequalities on

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: L
Revised: L

Original: [Equation]
Revised: [Equation]

Original: continue to hold.
Revised: continue to hold.

Original: The result then follows, observing that max j ∈ I c
Revised: The result then follows, observing that max j ∈ I c

Original: ∥ P E e j ∥ 2 ≤ ν .
Revised: ∥ P E e j ∥ 2 ≤ ν .

Original: The subspace embedding guarantee now follows as a corollary.
Revised: The subspace embedding guarantee now follows as a corollary.

Original: Theorem 6.1.
Revised: Theorem 6.1.

Original: Let ν = ϵ/d .
Revised: Let ν = ϵ/d .

Original: Suppose that m =
Revised: Suppose that m =

Original: Ω
Revised: Ω

Original: (( d/ϵ 2 )(polylog(1 /ϵ )
Revised: (( d/ϵ 2 )(polylog(1 /ϵ )

Original: + log(1 /δ ))) , δ ∈ (0 , 1 /m ) and d = Ω((1 /ϵ )
Revised: + log(1 /δ ))) , δ ∈ (0 , 1 /m ) and d = Ω((1 /ϵ )

Original: polylog(1 /ϵ ) log 2 (1 /δ )) .
Revised: polylog(1 /ϵ ) log 2 (1 /δ )) .

Original: Then, there exists a distribution on S with m + | I | rows such that
Revised: Then, there exists a distribution on S with m + | I | rows such that

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: One can verify that the two conditions in Lemma E.3 are satisfied if
Revised: One can verify that the two conditions in Lemma E.3 are satisfied if

Original: [Equation]
Revised: [Equation]

Original: The last condition is satisfied if
Revised: The last condition is satisfied if

Original: [Equation]
Revised: [Equation]

Original: E. P ROOF OF L EMMA 6. 
Revised: E. P ROOF OF L EMMA 6.

Original: Proof.
Revised: Proof.

Original: On the one hand, since Q = SAR is an orthogonal matrix, we have
Revised: On the one hand, since Q = SAR is an orthogonal matrix, we have

Original: [Equation]
Revised: [Equation]

Original: On the other hand, the assumption implies that
Revised: On the other hand, the assumption implies that

Original: [Equation]
Revised: [Equation]

Original: that is,
Revised: that is,

Original: [Equation]
Revised: [Equation]

Original: Combining both (E.2) and (E.3) leads to
Revised: Combining both (E.2) and (E.3) leads to

Original: [Equation]
Revised: [Equation]

Original: Equivalently, it can be written as
Revised: Equivalently, it can be written as

Original: [Equation]
Revised: [Equation]

Original: The claimed result follows from the fact that 1 / √ 1 + ϵ ≥ 1 − ϵ and 1 / √ 1 − ϵ ≤ 1 + ϵ whenever ϵ ∈ (0 , √ 5 − 1 2 ] .
Revised: The claimed result follows from the fact that 1 / √ 1 + ϵ ≥ 1 − ϵ and 1 / √ 1 − ϵ ≤ 1 + ϵ whenever ϵ ∈ (0 , √ 5 − 1 2 ] .

Original: F L OCATION O PTIMIZATION IN C OUNT S KETCH : G REEDY S EARCH
Revised: F L OCATION O PTIMIZATION IN C OUNT S KETCH : G REEDY S EARCH

Original: While the position optimization idea is simple, one particularly interesting aspect is that it is provably better than a random placement in some scenarios (Theorem. F.1).
Revised: While the position optimization idea is simple, one particularly interesting aspect is that it is provably better than a random placement in some scenarios (Theorem. F.1).

Original: Specifically, it is provably beneficial for LRA when inputs follow the spiked covariance model or Zipfian distributions, which are common for real data.
Revised: Specifically, it is provably beneficial for LRA when inputs follow the spiked covariance model or Zipfian distributions, which are common for real data.

Original: Spiked covariance model.
Revised: Spiked covariance model.

Original: Every matrix A ∈ R n × d from the distribution A sp ( s, ℓ ) has s < k “heavy” rows A r 1 , · · · , A r s of norm ℓ > 1 .
Revised: Every matrix A ∈ R n × d from the distribution A sp ( s, ℓ ) has s < k “heavy” rows A r 1 , · · · , A r s of norm ℓ > 1 .

Original: The indices of the heavy rows can be arbitrary, but must be the same for all members of A sp ( s, ℓ ) and are unknown to the algorithm.
Revised: The indices of the heavy rows can be arbitrary, but must be the same for all members of A sp ( s, ℓ ) and are unknown to the algorithm.

Original: The remaining (“light”) rows have unit norm.
Revised: The remaining (“light”) rows have unit norm.

Original: In other words, let R = { r 1 , . . .
Revised: In other words, let R = { r 1 , . . .

Original: , r s } .
Revised: , r s } .

Original: For all rows A i , i ∈ [ n ] , A i = ℓ · v i if i ∈ R and A i = v i otherwise, where v i is a uniformly random unit vector.
Revised: For all rows A i , i ∈ [ n ] , A i = ℓ · v i if i ∈ R and A i = v i otherwise, where v i is a uniformly random unit vector.

Original: Zipfian on squared row norms.
Revised: Zipfian on squared row norms.

Original: Every A ∈ R n × d ∼ A zipf has rows which are uniformly random and orthogonal.
Revised: Every A ∈ R n × d ∼ A zipf has rows which are uniformly random and orthogonal.

Original: Each A has 2 i +1 rows of squared norm n 2 / 2 2 i for i ∈ [1 , . . .
Revised: Each A has 2 i +1 rows of squared norm n 2 / 2 2 i for i ∈ [1 , . . .

Original: , O (log( n ))] .
Revised: , O (log( n ))] .

Original: We also assume that each row has the same squared norm for all members of A zipf .
Revised: We also assume that each row has the same squared norm for all members of A zipf .

Original: Theorem F.1.
Revised: Theorem F.1.

Original: Consider a matrix A from either the spiked covariance model or a Zipfian distribution.
Revised: Consider a matrix A from either the spiked covariance model or a Zipfian distribution.

Original: Let S L denote a CountSketch constructed by Algorithm 3 that optimizes the positions of the non-zero values with respect to A .
Revised: Let S L denote a CountSketch constructed by Algorithm 3 that optimizes the positions of the non-zero values with respect to A .

Original: Let S C denote a CountSketch matrix.
Revised: Let S C denote a CountSketch matrix.

Original: Then there is a fixed η > 0 such that, min
Revised: Then there is a fixed η > 0 such that, min

Original: [Equation]
Revised: [Equation]

Original: Remark F.2 .
Revised: Remark F.2 .

Original: Note that the above theorem implicitly provides an upper bound on the generalization error of the greedy placement method on the two distributions that we considered in this paper.
Revised: Note that the above theorem implicitly provides an upper bound on the generalization error of the greedy placement method on the two distributions that we considered in this paper.

Original: More precisely, for each of these two distributions, if Π is learned via our greedy approach over a set of sampled training matrices, the solution returned by the sketching algorithm using Π over any (test) matrix
Revised: More precisely, for each of these two distributions, if Π is learned via our greedy approach over a set of sampled training matrices, the solution returned by the sketching algorithm using Π over any (test) matrix

Original: A sampled from the distribution has error at most (1 − η ) min rank- k X ∈ rowsp( S C A )
Revised: A sampled from the distribution has error at most (1 − η ) min rank- k X ∈ rowsp( S C A )

Original: ∥ X − A ∥ 2 F .
Revised: ∥ X − A ∥ 2 F .

Original: A key structural property of the matrices from these two distributions that is crucial in our analysis is the ϵ -almost orthogonality of their rows (i.e., (normalized) pairwise inner products are at most ϵ ).
Revised: A key structural property of the matrices from these two distributions that is crucial in our analysis is the ϵ -almost orthogonality of their rows (i.e., (normalized) pairwise inner products are at most ϵ ).

Original: Hence, we can find a QR-factorization of the matrix of such vectors where the upper diagonal matrix
Revised: Hence, we can find a QR-factorization of the matrix of such vectors where the upper diagonal matrix

Original: R has diagonal entries close to 

Revised: R has diagonal entries close to

Original: and entries above the diagonal are close to 0 .
Revised: and entries above the diagonal are close to 0 .

Original: To state our result, we first provide an interpretation of the location optimization task as a selection of hash function for the rows of A .
Revised: To state our result, we first provide an interpretation of the location optimization task as a selection of hash function for the rows of A .

Original: Note that left-multiplying A by CountSketch S ∈ R m × n is equivalent to hashing the rows of A to m bins with coefficients in {± 1 } .
Revised: Note that left-multiplying A by CountSketch S ∈ R m × n is equivalent to hashing the rows of A to m bins with coefficients in {± 1 } .

Original: The greedy algorithm proceeds through the rows of A (in some order) and decides which bin to hash to, denoting this by adding an entry to S .
Revised: The greedy algorithm proceeds through the rows of A (in some order) and decides which bin to hash to, denoting this by adding an entry to S .

Original: The intuition is that our greedy approach separates heavy-norm rows (which are important “directions” in the row space) into different bins.
Revised: The intuition is that our greedy approach separates heavy-norm rows (which are important “directions” in the row space) into different bins.

Original: Proof Sketch of Theorem F.

Revised: Proof Sketch of Theorem F.

Original: The first step is to observe that in the greedy algorithm, when rows are examined according to a non-decreasing order of squared norms, the algorithm will isolate rows into their singleton bins until all bins are filled.
Revised: The first step is to observe that in the greedy algorithm, when rows are examined according to a non-decreasing order of squared norms, the algorithm will isolate rows into their singleton bins until all bins are filled.

Original: In particular, this means that the heavy norm rows will all be isolated—e.g., for the spiked covariance model, Lemma F.8 presents the formal statement.
Revised: In particular, this means that the heavy norm rows will all be isolated—e.g., for the spiked covariance model, Lemma F.8 presents the formal statement.

Original: Next, we show that none of the rows left to be processed (all light rows) will be assigned to the same bin as a heavy row.
Revised: Next, we show that none of the rows left to be processed (all light rows) will be assigned to the same bin as a heavy row.

Original: The main proof idea is to compare the cost of “colliding” with a heavy row to the cost of “avoiding” the heavy rows.
Revised: The main proof idea is to compare the cost of “colliding” with a heavy row to the cost of “avoiding” the heavy rows.

Original: This is the main place we use the properties of the aforementioned distributions and the fact that each heavy row is already mapped to a singleton bin.
Revised: This is the main place we use the properties of the aforementioned distributions and the fact that each heavy row is already mapped to a singleton bin.

Original: Overall, we show that at the end of the algorithm no light row will be assigned to the bins that contain heavy rows—the formal statement and proof for the spiked covariance model is in Lemma F.12.
Revised: Overall, we show that at the end of the algorithm no light row will be assigned to the bins that contain heavy rows—the formal statement and proof for the spiked covariance model is in Lemma F.12.

Original: Finally, we can interpret the randomized construction of CountSketch as a “balls and bins” experiment.
Revised: Finally, we can interpret the randomized construction of CountSketch as a “balls and bins” experiment.

Original: In particular, considering the heavy rows, we compute the expected number of bins (i.e., rows in
Revised: In particular, considering the heavy rows, we compute the expected number of bins (i.e., rows in

Original: S C A ) that contain a heavy row.
Revised: S C A ) that contain a heavy row.

Original: Note that the expected number of rows in S C A that do not contain any heavy row is k · (1 − 1 k ) s ≥ k · e − sk − 1 .
Revised: Note that the expected number of rows in S C A that do not contain any heavy row is k · (1 − 1 k ) s ≥ k · e − sk − 1 .

Original: Hence, the number of rows in S C A that contain a heavy row of A is at most k (1 − e − sk − 1 ) .
Revised: Hence, the number of rows in S C A that contain a heavy row of A is at most k (1 − e − sk − 1 ) .

Original: Thus, at least s − k (1 − e − sk − 1 ) heavy rows are not mapped to an isolated bin (i.e., they collide with some other heavy rows).
Revised: Thus, at least s − k (1 − e − sk − 1 ) heavy rows are not mapped to an isolated bin (i.e., they collide with some other heavy rows).

Original: Then, it is straightforward to show that the squared loss of the solution corresponding to S C is larger than the squared loss of the solution corresponding to S L , the CountSketch constructed by Algorithm 3—please see Lemma F.14 for the formal statement of its proof.
Revised: Then, it is straightforward to show that the squared loss of the solution corresponding to S C is larger than the squared loss of the solution corresponding to S L , the CountSketch constructed by Algorithm 3—please see Lemma F.14 for the formal statement of its proof.

Original: Preliminaries and notation.
Revised: Preliminaries and notation.

Original: Left-multiplying A by a CountSketch S ∈ R m × n is equivalent to hashing the rows of A to m bins with coefficients in {± 1 } .
Revised: Left-multiplying A by a CountSketch S ∈ R m × n is equivalent to hashing the rows of A to m bins with coefficients in {± 1 } .

Original: The greedy algorithm proceeds through the rows of A (in some order) and decides which bin to hash to, which we can think of as adding an entry to S .
Revised: The greedy algorithm proceeds through the rows of A (in some order) and decides which bin to hash to, which we can think of as adding an entry to S .

Original: We will denote the bins by b i and their summed contents by w i .
Revised: We will denote the bins by b i and their summed contents by w i .

Original: F. S PIKED COVARIANCE MODEL WITH SPARSE LEFT SINGULAR VECTORS .
Revised: F. S PIKED COVARIANCE MODEL WITH SPARSE LEFT SINGULAR VECTORS .

Original: To recap, every matrix A ∈ R n × d from the distribution A sp ( s, ℓ ) has s < k “heavy” rows ( A
Revised: To recap, every matrix A ∈ R n × d from the distribution A sp ( s, ℓ ) has s < k “heavy” rows ( A

Original: r 1 , · · · , A r s ) of norm ℓ > 1 .
Revised: r 1 , · · · , A r s ) of norm ℓ > 1 .

Original: The indices of the heavy rows can be arbitrary, but must be the same for all members of the distribution and are unknown to the algorithm.
Revised: The indices of the heavy rows can be arbitrary, but must be the same for all members of the distribution and are unknown to the algorithm.

Original: The remaining rows (called “light” rows) have unit norm.
Revised: The remaining rows (called “light” rows) have unit norm.

Original: In other words: let R = { r 1 , . . .
Revised: In other words: let R = { r 1 , . . .

Original: , r s } .
Revised: , r s } .

Original: For all rows A i , i ∈ [ n ] :
Revised: For all rows A i , i ∈ [ n ] :

Original: [Equation]
Revised: [Equation]

Original: where v i is a uniformly random unit vector.
Revised: where v i is a uniformly random unit vector.

Original: We also assume that S r , S g ∈ R k × n and that the greedy algorithm proceeds in a non-increasing row norm order.
Revised: We also assume that S r , S g ∈ R k × n and that the greedy algorithm proceeds in a non-increasing row norm order.

Original: Proof sketch.
Revised: Proof sketch.

Original: First, we show that the greedy algorithm using a non-increasing row norm ordering will isolate heavy rows (i.e., each is alone in a bin).
Revised: First, we show that the greedy algorithm using a non-increasing row norm ordering will isolate heavy rows (i.e., each is alone in a bin).

Original: Then, we conclude by showing that this yields a better k -rank approximation error when d is sufficiently large compared to n .
Revised: Then, we conclude by showing that this yields a better k -rank approximation error when d is sufficiently large compared to n .

Original: We begin with some preliminary observations that will be of use later.
Revised: We begin with some preliminary observations that will be of use later.

Original: It is well-known that a set of uniformly random vectors is ϵ -almost orthogonal (i.e., the magnitudes of their pairwise inner products are at most ϵ ).
Revised: It is well-known that a set of uniformly random vectors is ϵ -almost orthogonal (i.e., the magnitudes of their pairwise inner products are at most ϵ ).

Original: Observation F.3 .
Revised: Observation F.3 .

Original: Let v 1 , · · · , v n ∈ R
Revised: Let v 1 , · · · , v n ∈ R

Original: d be a set of random unit vectors.
Revised: d be a set of random unit vectors.

Original: Then with probability 1 − 1 / poly( n ) , we have |⟨ v i , v j ⟩| ≤ 2 (cid:113) log nd , ∀ i < j ≤ n .
Revised: Then with probability 1 − 1 / poly( n ) , we have |⟨ v i , v j ⟩| ≤ 2 (cid:113) log nd , ∀ i < j ≤ n .

Original: We define ϵ = 2 (cid:113) log nd .
Revised: We define ϵ = 2 (cid:113) log nd .

Original: Observation F.4 .
Revised: Observation F.4 .

Original: Let u 1 , · · · , u t be a set of vectors such that for each pair of i < j ≤ t , |⟨ u i ∥ u i ∥ , u j ∥
Revised: Let u 1 , · · · , u t be a set of vectors such that for each pair of i < j ≤ t , |⟨ u i ∥ u i ∥ , u j ∥

Original: u j ∥
Revised: u j ∥

Original: ⟩| ≤ ϵ , and g i , · · · , g j ∈ {− 1 , 1 } .
Revised: ⟩| ≤ ϵ , and g i , · · · , g j ∈ {− 1 , 1 } .

Original: Then, t
Revised: Then, t

Original: [Equation]
Revised: [Equation]

Original: Next, a straightforward consequence of ϵ -almost orthogonality is that we can find a QR-factorization of the matrix of such vectors where R (an upper diagonal matrix) has diagonal entries close to 1 and entries above the diagonal are close to 0 .
Revised: Next, a straightforward consequence of ϵ -almost orthogonality is that we can find a QR-factorization of the matrix of such vectors where R (an upper diagonal matrix) has diagonal entries close to 1 and entries above the diagonal are close to 0 .

Original: Lemma F.5.
Revised: Lemma F.5.

Original: Let u 1 , · · · , u t ∈ R
Revised: Let u 1 , · · · , u t ∈ R

Original: d be a set of unit vectors such that for any pair of i < j ≤ t , |⟨ u i , u j ⟩| ≤ ϵ where
Revised: d be a set of unit vectors such that for any pair of i < j ≤ t , |⟨ u i , u j ⟩| ≤ ϵ where

Original: ϵ
Revised: ϵ

Original: = O ( t − 2 ) .
Revised: = O ( t − 2 ) .

Original: There exists an orthonormal basis e 1 , · · · , e t for the subspace spanned by u 1 , · · · , u t such that for each i ≤ t , u i = (cid:80) ij =1 a i,j e j where a 2 i,i ≥ 1 − (cid:80) i − 1 j =1 j 2 · ϵand for each j <
Revised: There exists an orthonormal basis e 1 , · · · , e t for the subspace spanned by u 1 , · · · , u t such that for each i ≤ t , u i = (cid:80) ij =1 a i,j e j where a 2 i,i ≥ 1 − (cid:80) i − 1 j =1 j 2 · ϵand for each j <

Original: i , a 2 i,j ≤ j 2 ϵ 2 .
Revised: i , a 2 i,j ≤ j 2 ϵ 2 .

Original: Proof.
Revised: Proof.

Original: We follow the Gram-Schmidt process to construct the orthonormal basis e 1 , · · · , e t of the space spanned by u 1 , · · · , u t , by first setting e 1 =
Revised: We follow the Gram-Schmidt process to construct the orthonormal basis e 1 , · · · , e t of the space spanned by u 1 , · · · , u t , by first setting e 1 =

Original: u
Revised: u

Original: and then processing u 2 , · · · , u t , one-by-one.
Revised: and then processing u 2 , · · · , u t , one-by-one.

Original: The proof is by induction.
Revised: The proof is by induction.

Original: We show that once the first j vectors u 1 , · · · , u j are processed, the statement of the lemma holds for these vectors.
Revised: We show that once the first j vectors u 1 , · · · , u j are processed, the statement of the lemma holds for these vectors.

Original: Note that the base case of the induction trivially holds as u 1 = e 1 .
Revised: Note that the base case of the induction trivially holds as u 1 = e 1 .

Original: Next, suppose that the induction hypothesis holds for the first ℓ vectors u 1 , · · · , u ℓ .
Revised: Next, suppose that the induction hypothesis holds for the first ℓ vectors u 1 , · · · , u ℓ .

Original: Claim F.6 .
Revised: Claim F.6 .

Original: For each j ≤ ℓ , a 2 ℓ +1 ,j ≤ j 2 ϵ 2 .
Revised: For each j ≤ ℓ , a 2 ℓ +1 ,j ≤ j 2 ϵ 2 .

Original: Proof.
Revised: Proof.

Original: The proof of the claim is itself by induction.
Revised: The proof of the claim is itself by induction.

Original: Note that, for j = 1 and using the fact that |⟨ u 1 , u ℓ +1 ⟩| ≤ ϵ , the statement holds and a 2 ℓ +1 , 1 ≤ ϵ 2 .
Revised: Note that, for j = 1 and using the fact that |⟨ u 1 , u ℓ +1 ⟩| ≤ ϵ , the statement holds and a 2 ℓ +1 , 1 ≤ ϵ 2 .

Original: Next, suppose that the statement holds for all j ≤ i < ℓ .
Revised: Next, suppose that the statement holds for all j ≤ i < ℓ .

Original: Then using that |⟨
Revised: Then using that |⟨

Original: u i +1 , u ℓ +1 ⟩| ≤ ϵ , | a ℓ +1 ,i +1 | ≤ ( |⟨ u ℓ +1 , u
Revised: u i +1 , u ℓ +1 ⟩| ≤ ϵ , | a ℓ +1 ,i +1 | ≤ ( |⟨ u ℓ +1 , u

Original: i +

Revised: i +

Original: |
Revised: |

Original: + i (cid:88)
Revised: + i (cid:88)

Original: j = | a ℓ +1 ,j | · | a i +1 ,j | ) / | a
Revised: j = | a ℓ +1 ,j | · | a i +1 ,j | ) / | a

Original: i +1 ,i +1 | ≤ ( ϵ + i (cid:88)
Revised: i +1 ,i +1 | ≤ ( ϵ + i (cid:88)

Original: j = j 2 ϵ 2 ) / | a
Revised: j = j 2 ϵ 2 ) / | a

Original: i +1 ,i +1 |
Revised: i +1 ,i +1 |

Original: ▷ by the induction hypothesis on a ℓ +1 ,j for j ≤ i ≤ ( ϵ + i (cid:88)
Revised: ▷ by the induction hypothesis on a ℓ +1 ,j for j ≤ i ≤ ( ϵ + i (cid:88)

Original: j = j 2 ϵ 2 ) / (1 − i (cid:88)
Revised: j = j 2 ϵ 2 ) / (1 − i (cid:88)

Original: j = j 2 · ϵ 2 ) 1 /▷ by the induction hypothesis on a i +1 ,i + ≤ ( ϵ + i (cid:88)
Revised: j = j 2 · ϵ 2 ) 1 /▷ by the induction hypothesis on a i +1 ,i + ≤ ( ϵ + i (cid:88)

Original: j = j 2 ϵ 2 ) · (1 − i (cid:88)
Revised: j = j 2 ϵ 2 ) · (1 − i (cid:88)

Original: j = j 2 · ϵ 2 ) 1 / 2 · (1 + 2 · i (cid:88)
Revised: j = j 2 · ϵ 2 ) 1 / 2 · (1 + 2 · i (cid:88)

Original: j = j 2 ϵ 2 ) ≤ ( ϵ + i (cid:88)
Revised: j = j 2 ϵ 2 ) ≤ ( ϵ + i (cid:88)

Original: j = j 2 ϵ 2 ) · (1 + 2 · i (cid:88)
Revised: j = j 2 ϵ 2 ) · (1 + 2 · i (cid:88)

Original: j = j 2 ϵ 2 ) ≤ ϵ (( i (cid:88)
Revised: j = j 2 ϵ 2 ) ≤ ϵ (( i (cid:88)

Original: j = j 2 ϵ ) · (1 + 4 ϵ · i (cid:88)
Revised: j = j 2 ϵ ) · (1 + 4 ϵ · i (cid:88)

Original: j = j 2 ϵ ) + 1) ≤
Revised: j = j 2 ϵ ) + 1) ≤

Original: ϵ ( i + 1)
Revised: ϵ ( i + 1)

Original: ▷ by ϵ =
Revised: ▷ by ϵ =

Original: O ( t − 2 )
Revised: O ( t − 2 )

Original: Finally, since ∥ u ℓ +1 ∥
Revised: Finally, since ∥ u ℓ +1 ∥

Original: = 1 , a 2 ℓ +1 ,ℓ +1 ≥ 1 − (cid:80) ℓj =

Revised: = 1 , a 2 ℓ +1 ,ℓ +1 ≥ 1 − (cid:80) ℓj =

Original: j 2 ϵ 2 .
Revised: j 2 ϵ 2 .

Original: Corollary F.7.
Revised: Corollary F.7.

Original: Suppose that
Revised: Suppose that

Original: ϵ
Revised: ϵ

Original: =
Revised: =

Original: O ( t − 2 ) .
Revised: O ( t − 2 ) .

Original: There exists an orthonormal basis e 1 , · · · , e t for the space spanned by the randomly picked vectors v 1 , · · · , v t , of unit norm, so that for each i , v i = (cid:80) ij =1 a i,j e j where a 2 i,i ≥ 1 − (cid:80) i − 1 j =1 j 2 · ϵ 2 and for each j < i , a 2 i,j ≤ j 2 · ϵ 2 .
Revised: There exists an orthonormal basis e 1 , · · · , e t for the space spanned by the randomly picked vectors v 1 , · · · , v t , of unit norm, so that for each i , v i = (cid:80) ij =1 a i,j e j where a 2 i,i ≥ 1 − (cid:80) i − 1 j =1 j 2 · ϵ 2 and for each j < i , a 2 i,j ≤ j 2 · ϵ 2 .

Original: Proof.
Revised: Proof.

Original: The proof follows from Lemma F.5 and the fact that the set of vectors v 1 , · · · , v t is ϵ -almost orthogonal (by Observation F.3).
Revised: The proof follows from Lemma F.5 and the fact that the set of vectors v 1 , · · · , v t is ϵ -almost orthogonal (by Observation F.3).

Original: The first main step is to show that the greedy algorithm (with non-increasing row norm ordering) will isolate rows into their own bins until all bins are filled.
Revised: The first main step is to show that the greedy algorithm (with non-increasing row norm ordering) will isolate rows into their own bins until all bins are filled.

Original: In particular, this means that the heavy rows (the first to be processed) will all be isolated.
Revised: In particular, this means that the heavy rows (the first to be processed) will all be isolated.

Original: We note that because we set rank( SA )
Revised: We note that because we set rank( SA )

Original: =
Revised: =

Original: k , the k -rank approximation cost is the simplified expression (cid:13)(cid:13) AV V ⊤ − A (cid:13)(cid:13) 2 F , where U Σ V ⊤ =
Revised: k , the k -rank approximation cost is the simplified expression (cid:13)(cid:13) AV V ⊤ − A (cid:13)(cid:13) 2 F , where U Σ V ⊤ =

Original: SA , rather than (cid:13)(cid:13)
Revised: SA , rather than (cid:13)(cid:13)

Original: [ AV ] k V ⊤ − A (cid:13)(cid:13) 2 F .
Revised: [ AV ] k V ⊤ − A (cid:13)(cid:13) 2 F .

Original: This is just the projection cost onto row ( SA ) .
Revised: This is just the projection cost onto row ( SA ) .

Original: Also, we observe that minimizing this projection cost is the same as maximizing the sum of squared projection coefficients:
Revised: Also, we observe that minimizing this projection cost is the same as maximizing the sum of squared projection coefficients:

Original: [Equation]
Revised: [Equation]

Original: In the following sections, we will prove that our greedy algorithm makes certain choices by showing that these choices maximize the sum of squared projection coefficients.
Revised: In the following sections, we will prove that our greedy algorithm makes certain choices by showing that these choices maximize the sum of squared projection coefficients.

Original: Lemma F.8.
Revised: Lemma F.8.

Original: For any matrix A or batch of matrices A , at the end of iteration k , the learned
Revised: For any matrix A or batch of matrices A , at the end of iteration k , the learned

Original: CountSketch matrix S maps each row to an isolated bin.
Revised: CountSketch matrix S maps each row to an isolated bin.

Original: In particular, heavy rows are mapped to isolated bins.
Revised: In particular, heavy rows are mapped to isolated bins.

Original: Proof.
Revised: Proof.

Original: For any iteration i ≤ k , we consider the choice of assigning A i to an empty bin versus an occupied bin.
Revised: For any iteration i ≤ k , we consider the choice of assigning A i to an empty bin versus an occupied bin.

Original: Without loss of generality, let this occupied bin be b i − 1 , which already contains A i − 1 .
Revised: Without loss of generality, let this occupied bin be b i − 1 , which already contains A i − 1 .

Original: We consider the difference in cost for empty versus occupied.
Revised: We consider the difference in cost for empty versus occupied.

Original: We will do this cost comparison for A j with j ≤ i − 2 , j ≥ i + 1 , and finally, j ∈ { i − 1 , i } .
Revised: We will do this cost comparison for A j with j ≤ i − 2 , j ≥ i + 1 , and finally, j ∈ { i − 1 , i } .

Original: First, we let { e 1 , . . .
Revised: First, we let { e 1 , . . .

Original: , e i } be an orthonormal basis for { A 1 , . . .
Revised: , e i } be an orthonormal basis for { A 1 , . . .

Original: , A i } such that for each r ≤ i ,
Revised: , A i } such that for each r ≤ i ,

Original: A r = (cid:80) rj =1 a r,j e j where a r,r > 0 .
Revised: A r = (cid:80) rj =1 a r,j e j where a r,r > 0 .

Original: This exists by Lemma F.5.
Revised: This exists by Lemma F.5.

Original: Let { e 1 , . . .
Revised: Let { e 1 , . . .

Original: , e i − 2 , e } be an orthonormal basis for { A 1 , . . .
Revised: , e i − 2 , e } be an orthonormal basis for { A 1 , . . .

Original: A
Revised: A

Original: i +2 , A
Revised: i +2 , A

Original: i − 1 ±
Revised: i − 1 ±

Original: A i } .
Revised: A i } .

Original: Now, e = c 

Revised: Now, e = c

Original: e
Revised: e

Original: i − 1 + c 1 e i for some c 0 , c 1 because ( A i − 1 ±
Revised: i − 1 + c 1 e i for some c 0 , c 1 because ( A i − 1 ±

Original: A i ) − proj { e 1 ,...,
Revised: A i ) − proj { e 1 ,...,

Original: e i − 2 } ( A i − 1 ±
Revised: e i − 2 } ( A i − 1 ±

Original: A i ) ∈ span ( e i − 1 , e i ) .
Revised: A i ) ∈ span ( e i − 1 , e i ) .

Original: We note that c 20 + c 21 = 1 because we let e be a unit vector.
Revised: We note that c 20 + c 21 = 1 because we let e be a unit vector.

Original: We can find c 0 , c 1 to be:
Revised: We can find c 0 , c 1 to be:

Original: [Equation]
Revised: [Equation]

Original: j ≤ i − 2 : The cost is zero for both cases because A j ∈ span ( { e 1 , . . .
Revised: j ≤ i − 2 : The cost is zero for both cases because A j ∈ span ( { e 1 , . . .

Original: , e i − 2 } ) .
Revised: , e i − 2 } ) .

Original: j ≥ i + 1 : We compare the rewards (sum of squared projection coefficients) and find that { e 1 , . . .
Revised: j ≥ i + 1 : We compare the rewards (sum of squared projection coefficients) and find that { e 1 , . . .

Original: , e i − 2 , e } is no better than { e 1 , . . .
Revised: , e i − 2 , e } is no better than { e 1 , . . .

Original: , e i } .
Revised: , e i } .

Original: [Equation]
Revised: [Equation]

Original: ▷ Cauchy-Schwarz inequality
Revised: ▷ Cauchy-Schwarz inequality

Original: [Equation]
Revised: [Equation]

Original: j ∈ { i − 1 ,
Revised: j ∈ { i − 1 ,

Original: i
Revised: i

Original: We compute the sum of squared projection coefficients of A i − 1 and A i onto e : (
Revised: We compute the sum of squared projection coefficients of A i − 1 and A i onto e : (

Original: [Equation]
Revised: [Equation]

Original: On the other hand, the sum of squared projection coefficients of A i − 1 and A i onto e i − 1 ∪ e i is:
Revised: On the other hand, the sum of squared projection coefficients of A i − 1 and A i onto e i − 1 ∪ e i is:

Original: [Equation]
Revised: [Equation]

Original: Hence, the difference between the sum of squared projections of A
Revised: Hence, the difference between the sum of squared projections of A

Original: i − 1 and A i onto e and e i − 1 ∪ e
Revised: i − 1 and A i onto e and e i − 1 ∪ e

Original: i is ((F.3) - (F.2))
Revised: i is ((F.3) - (F.2))

Original: [Equation]
Revised: [Equation]

Original: Thus, we find that { e 1 , . . .
Revised: Thus, we find that { e 1 , . . .

Original: , e i } is a strictly better basis than { e 1 , . . .
Revised: , e i } is a strictly better basis than { e 1 , . . .

Original: , e i − 2 , e } .
Revised: , e i − 2 , e } .

Original: This means the greedy algorithm will choose to place A i in an empty bin.
Revised: This means the greedy algorithm will choose to place A i in an empty bin.

Original: Next, we show that none of the rows left to be processed (all light rows) will be assigned to the same bin as a heavy row.
Revised: Next, we show that none of the rows left to be processed (all light rows) will be assigned to the same bin as a heavy row.

Original: The main proof idea is to compare the cost of “colliding” with a heavy row to the cost of “avoiding” the heavy rows.
Revised: The main proof idea is to compare the cost of “colliding” with a heavy row to the cost of “avoiding” the heavy rows.

Original: Specifically, we compare the decrease (before and after bin assignment of a light row) in sum of squared projection coefficients, lower-bounding it in the former case and upper-bounding it in the latter.
Revised: Specifically, we compare the decrease (before and after bin assignment of a light row) in sum of squared projection coefficients, lower-bounding it in the former case and upper-bounding it in the latter.

Original: We introduce some results that will be used in Lemma F.12.
Revised: We introduce some results that will be used in Lemma F.12.

Original: Claim F.9 .
Revised: Claim F.9 .

Original: Let A k + r , r ∈
Revised: Let A k + r , r ∈

Original: [1 , . . . , n − k ] be a light row not yet processed by the greedy algorithm.
Revised: [1 , . . . , n − k ] be a light row not yet processed by the greedy algorithm.

Original: Let { e 1 , . . .
Revised: Let { e 1 , . . .

Original: , e k } be the Gram-Schmidt basis for the current { w 1 , . . .
Revised: , e k } be the Gram-Schmidt basis for the current { w 1 , . . .

Original: , w k } .
Revised: , w k } .

Original: Let β
Revised: Let β

Original: = O ( n − 1 k − 3 ) upper bound the inner products of the normalized A k + r , w 1 , . . .
Revised: = O ( n − 1 k − 3 ) upper bound the inner products of the normalized A k + r , w 1 , . . .

Original: , w k .
Revised: , w k .

Original: Then, for any bin i , ⟨ e i , A k + r ⟩ 2 ≤ β 2 · k 2 .
Revised: Then, for any bin i , ⟨ e i , A k + r ⟩ 2 ≤ β 2 · k 2 .

Original: Proof.
Revised: Proof.

Original: This is a straightforward application of Lemma F.5.
Revised: This is a straightforward application of Lemma F.5.

Original: From that, we have ⟨ A k + r , e i ⟩ 2 ≤ i 2 β 2 , for i ∈ [1 , . . .
Revised: From that, we have ⟨ A k + r , e i ⟩ 2 ≤ i 2 β 2 , for i ∈ [1 , . . .

Original: , k ] , which means ⟨ A k + r , e i ⟩ 2 ≤ k 

Revised: , k ] , which means ⟨ A k + r , e i ⟩ 2 ≤ k

Original: β 2 .
Revised: β 2 .

Original: Claim F.10 .
Revised: Claim F.10 .

Original: Let A k
Revised: Let A k

Original: + r be a light row that has been processed by the greedy algorithm.
Revised: + r be a light row that has been processed by the greedy algorithm.

Original: Let { e 1 , . . .
Revised: Let { e 1 , . . .

Original: , e k } be the Gram-Schmidt basis for the current { w 1 , . . .
Revised: , e k } be the Gram-Schmidt basis for the current { w 1 , . . .

Original: , w k } .
Revised: , w k } .

Original: If A k + r is assigned to bin b k − 1 (w.l.o.g.), the squared projection coefficient of A k + r onto e i , i ̸
Revised: If A k + r is assigned to bin b k − 1 (w.l.o.g.), the squared projection coefficient of A k + r onto e i , i ̸

Original: = k − 1 is at most 4 β 2 · k 2 , where β
Revised: = k − 1 is at most 4 β 2 · k 2 , where β

Original: = O ( n − 1 k − 3 ) upper bounds the inner products of normalized A k + r , w 1 , · · · , w k .
Revised: = O ( n − 1 k − 3 ) upper bounds the inner products of normalized A k + r , w 1 , · · · , w k .

Original: Proof.
Revised: Proof.

Original: Without loss of generality, it suffices to bound the squared projection of A k + r onto the direction of w k that is orthogonal to the subspace spanned by w 1 , · · · , w k − 1 .
Revised: Without loss of generality, it suffices to bound the squared projection of A k + r onto the direction of w k that is orthogonal to the subspace spanned by w 1 , · · · , w k − 1 .

Original: Let e 1 , · · · , e k be an orthonormal basis of w 1 , · · · , w k guaranteed by Lemma F.5.
Revised: Let e 1 , · · · , e k be an orthonormal basis of w 1 , · · · , w k guaranteed by Lemma F.5.

Original: Next, we expand the orthonormal basis to include e k +1 so that we can write the normalized vector of A k + r as v k + r = (cid:80) k
Revised: Next, we expand the orthonormal basis to include e k +1 so that we can write the normalized vector of A k + r as v k + r = (cid:80) k

Original: +1 j =1 b j e j .
Revised: +1 j =1 b j e j .

Original: By a similar approach to the proof of Lemma F.5, for each j ≤ k − 2 , b j ≤ β 2 j 2 .
Revised: By a similar approach to the proof of Lemma F.5, for each j ≤ k − 2 , b j ≤ β 2 j 2 .

Original: Next, since |⟨ w k , v k + r ⟩| ≤ β ,
Revised: Next, since |⟨ w k , v k + r ⟩| ≤ β ,

Original: [Equation]
Revised: [Equation]

Original: ▷ similar to the proof of Lemma F.
Revised: ▷ similar to the proof of Lemma F.

Original: < 2 β · k
Revised: < 2 β · k

Original: Hence, the squared projection of A k + r onto e k is at most 4 β 2 · k 2 ·∥ A k + r ∥ 22 .
Revised: Hence, the squared projection of A k + r onto e k is at most 4 β 2 · k 2 ·∥ A k + r ∥ 22 .

Original: We assumed ∥ A k + r ∥ = 1 ; hence, the squared projection of A k + r onto e k is at most 4 β 2 · k 2 .
Revised: We assumed ∥ A k + r ∥ = 1 ; hence, the squared projection of A k + r onto e k is at most 4 β 2 · k 2 .

Original: Claim F.11 .
Revised: Claim F.11 .