Title: Distributed Least Squares in Small Space via Sketching and Bias Reduction
Venue: NeurIPS 2024
Accepted: yes
Keywords: Matrix Sketching

Least squares

Randomized Linear Algebra

Random Matrix Theory
Abstract: Matrix sketching is a powerful tool for reducing the size of large data matrices. Yet there are fundamental limitations to this size reduction when we want to recover an accurate estimator for a task such as least square regression. We show that these limitations can be circumvented in the distributed setting by designing sketching methods that minimize the bias of the estimator, rather than its error. In particular, we give a sparse sketching method running in optimal space and current matrix multiplication time, which recovers a nearly-unbiased least squares estimator using two passes over the data. This leads to new communication-efficient distributed averaging algorithms for least squares and related tasks, which directly improve on several prior approaches. Our key novelty is a new bias analysis for sketched least squares, giving a sharp characterization of its dependence on the sketch sparsity. The techniques include new higher moment restricted Bai-Silverstein inequalities, which are of independent interest to the non-asymptotic analysis of deterministic equivalents for random matrices that arise from sketching.

Distributed Least Squares in Small Space via Sketching and Bias Reduction

Abstract

Matrix sketching is a powerful tool for reducing the size of large data matrices. Yet there are fundamental limitations to this size reduction when we want to recover an accurate estimator for a task such as least square regression. We show that these limitations can be circumvented in the distributed setting by designing sketching methods that minimize the bias of the estimator, rather than its error. In particular, we give a sparse sketching method running in optimal space and current matrix multiplication time, which recovers a nearly-unbiased least squares estimator using two passes over the data. This leads to new communication-efficient distributed averaging algorithms for least squares and related tasks, which directly improve on several prior approaches. Our key novelty is a new bias analysis for sketched least squares, giving a sharp characterization of its dependence on the sketch sparsity. The techniques include new higher-moment restricted Bai-Silverstein inequalities, which are of independent interest to the non-asymptotic analysis of deterministic equivalents for random matrices that arise from sketching.

1 Introduction

Matrix sketching is a powerful collection of randomized techniques for compressing large data matrices, developed over a long line of works as part of Randomized Numerical Linear Algebra [RandNLA, e.g., 45, 26, 37, 39, 21]. Sketching can be used to reduce the large dimension n of a data matrix A ∈ R n × d by applying a random sketching matrix (operator) S ∈ R m × n to obtain the sketch ˜ A = SA ∈ R m × d where m ≪ n . For example, sketching can be used to approximate the solution to the least squares problem, x ∗ = argmin x L ( x ) where L ( x ) = ∥ Ax -b ∥ 2 , by using a sketched estimator ˜ x = argmin x ∥ ˜ Ax -˜ b ∥ 2 , where ˜ A = SA and ˜ b = Sb .

Perhaps the simplest form of sketching is subsampling, where the sketching operator S selects a random sample of the rows of matrix A . However, the real advantage of sketching as a framework emerges as we consider more complex operators S , such as sub-Gaussian matrices [1], randomized Hadamard transforms [3], and sparse random matrices [12]. These approaches ensure higher quality and more robust compression of the data matrix, e.g., leading to provable ϵ -approximation guarantees for the estimate ˜ x in the least squares task, i.e., L (˜ x ) ≤ (1 + ϵ ) L ( x ∗ ) . Nevertheless, there are fundamental limitations to how far we can compress a data matrix using sketching while ensuring an ϵ -approximation. These limitations pose a challenge particularly in space-limited computing

Kevin Tan

A

Subsample

˜

O

(

d/ϵ

)

×

d

˜

O

(1

/ϵ

)

rows

˜

O

(1

/ϵ

)

rows

.

.

.

˜

O

(1

/ϵ

)

rows

Sketch

O

(

d

)

×

d

.

.

.

environments, such as for streaming algorithms where we observe the matrix A , say, one row at a time, and we have limited space for storing the sketch [11].

One strategy for overcoming the fundamental limitations of sketching as a compression tool is to look beyond the single approximation guarantee provided by a sketching-based estimator ˜ x , and consider how its broader statistical properties can be leveraged in a given computing environment. To that end, many recent works have demonstrated both theoretically and empirically that sketching-based estimators often exhibit not only approximation robustness but also statistical robustness, for instance enjoying sharp confidence intervals, effectiveness of statistical inference tools such as bootstrap and cross-validation, as well as accuracy boosting techniques such as distributed averaging [e.g., 35, 25, 32, 33]. Yet, these results have had limited impact on the traditional computational complexity analysis in RandNLA and sketching literature, as many of them either impose additional assumptions, focus on sharpening the constant factors, or require using more expensive sketching techniques. In this work, we demonstrate that the statistical properties of sketching-based estimators can in fact have a substantial impact on the computational trade-offs that arise in RandNLA.

Our key motivating example is the above mentioned least squares regression task. It is well understood that for an n × d least squares task, to recover an ϵ -approximate solution out of an m × d sketch, we need sketch size at least m = Ω( d/ϵ ) . This has been formalized in the streaming setting with a lower bound of Ω( ϵ -1 d 2 log( nd )) bits of space required, when all of the input numbers use O (log( nd )) bits of precision [11]. One setting where this can be circumvented is in the distributed computing model where the bits can be spread out across many machines, so that the per-machine space can be smaller. Here, one could for instance hope that we can maintain small O ( d ) × d sketches in q = O (1 /ϵ ) machines and then combine their estimates to recover an ϵ -approximate solution. A simple and attractive approach is to average the estimates ˜ x i produced by the individual machines, returning ˆ x = 1 q ∑ q i =1 ˜ x i , as this only requires each machine to communicate O ( d log( nd )) bits of information about its sketch. This approach requires the sketching-based estimates ˜ x i to have sufficiently small bias for the averaging scheme to be effective. While this has been demonstrated empirically in many cases, existing theoretical results still require relatively expensive sketching methods to recover low-bias estimators, leading to an unfortunate trade-off in the distributed averaging scheme between the time and space complexity required.

In this work, we address the time-space trade-off in distributed averaging of sketching-based estimators, by giving a sharp characterization of how their bias depends on the sparsity of the sketching matrix. Remarkably, we show that in the distributed streaming environment one can compress the data down to the minimum size of O ( d 2 log( nd )) bits at no extra computational cost, while still being able to recover an ϵ -approximate solution for least squares and related problems. Importantly, our results require the sketching matrix to be slightly denser than is necessary for obtaining approximation guarantees on a single estimate, and thus, cannot be recovered by standard RandNLA sampling methods such as approximate leverage score sampling [27].

Before we state the full result in the distributed setting, we give our main technical contribution. which is the following efficient construction of a low-bias least squares estimator in a single pass using only O ( d 2 log( nd )) bits of space, assuming all numbers use O (log( nd )) bits of precision. Below, γ > 0 denotes an arbitrarily small constant.

Theorem 1. Given streaming access to A ∈ R n × d and b ∈ R n , and direct access to a preconditioner matrix P ∈ R d × d such that κ ( AP ) ≤ α , within a single pass over ( A , b ) , in O ( γ -1 nnz( A ) + ϵ -1 αd 2+ γ polylog( d )) time and O ( d 2 log( nd )) bits of space, we can construct a randomized estimator ˜ x for the least squares solution x ∗ = argmin x ∥ Ax -b ∥ 2 such that:

Our estimator ˜ x is constructed at the end of the data pass from a sketch ( ˜ A , ˜ b ) where ˜ A = SA and ˜ b = Sb , by minimizing ∥ ˜ Ax -˜ b ∥ 2 using preconditioned conjugate gradient. Here, S is a carefully constructed sparse sketching matrix which is inspired by the so-called leverage score sparsified (LESS) embeddings [18]. Leverage scores represent the relative importances of the rows of A which are commonly used for subsampling in least squares (see Definition 1), and their estimates can be easily obtained in a single pass by using the preconditioner matrix P .

Our time complexity bound of ˜ O (nnz( A ) + d 2 /ϵ ) matches the time it would take (for a single machine) to subsample ˜ O ( d/ϵ ) rows of A according to the approximate leverage scores and produce an estimator ˜ x that achieves the ϵ -error bound ∥ A ˜ x -b ∥ 2 ≤ (1 + ϵ ) ∥ Ax ∗ -b ∥ 2 . However, this strategy requires either maintaining ˜ O ( d 2 /ϵ ) bits of space for the sketch, or computing ˜ x directly along the way, blowing up the runtime to ˜ O ( d ω /ϵ ) . Since approximate leverage score sampling leads to significant least squares bias, averaging can only improve this to ˜ O ( d ω / √ ϵ ) (see Table 1). An alternate strategy would be to combine leverage score sampling with a preconditioned mini-batch stochastic gradient descent (Weighted Mb-SGD), with mini-batches chosen so that they fit in ˜ O ( d 2 ) space. This achieves the same time and space complexity as our method, but due to the streaming access to A and the sequential nature of SGD, it requires O (1 /ϵ ) data passes.

Instead, our algorithm essentially mixes an ˜ O ( d/ϵ ) size leverage score sample into an O ( d ) size sketch, merging ˜ O (1 /ϵ ) rows of A into a single row of the sketch (see Figure 1). This results in better data compression compared to direct leverage score sampling, with only ˜ O ( d 2 ) bits of space, while retaining the same ˜ O (nnz( A ) + d 2 /ϵ ) runtime complexity as the above approaches and requiring only a single data pass. The resulting estimator ˜ x can no longer recover the ϵ -error bound, but remarkably, its expectation E [˜ x ] still does. To turn this into an improved estimator in a distributed model, we can simply average q = 1 /ϵ such estimators, i.e., ˆ x = 1 q ∑ q i =1 ˜ x i , obtaining E ∥ A ˆ x -b ∥ 2 ≤ (1 + 2 ϵ ) ∥ Ax ∗ -b ∥ 2 . As shown in Table 1, ours is the first result in this model to achieve ˜ O ( d 2 ) space in a single pass and faster than current matrix multiplication time O ( d ω ) .

Finally, by incorporating a preconditioning scheme, we illustrate how our construction can be used to design the first algorithm that solves least squares in current matrix multiplication time, constant parallel passes and O ( d 2 log( nd )) bits of space. We note that the O ( d ω ) cost comes only from the worst-case complexity of constructing the preconditioner P , which can often be accelerated in practice. The computational model used in Theorem 2 is described in detail in Section 3.

Theorem 2. Given A ∈ R n × d and b ∈ R n in the parallel computing model, using two parallel passes with q machines, we can compute ˜ x such that with probability 0 . 9

in O ( γ -1 nnz( A ) + d ω + ϵ -1 d 2+ γ polylog( d )) time, O ( d 2 log( nd )) bits of space and O ( d log( nd )) bits of communication. In particular, choosing q = 1 /ϵ we recover an O ( ϵ ) -approximation.

Remark 2. While the parallel computing model in Theorem 2 assumes that all machines have streaming access to the entire data matrix, this result can be easily extended to the setting where A has been randomly down-sampled or partitioned into separate size n chunks A 1 , ..., A q , and each machine constructs an estimate ˜ x i based on a sketch of its own chunk. Then, with the same computational guarantees as in Theorem 2, the averaged estimator ˜ x = 1 q ∑ q i =1 ˜ x i with probability 0 . 9 enjoys a guarantee of:

where B chunk is the bias that would be incurred if we solved each chunk exactly and averaged those solutions. Using existing guarantees for uniform down-sampling of least squares [Theorem 20, 43], one can bound this bias as B chunk = ˜ O ( µ qn +( µ n ) 2 ) , where µ is the coherence of the data matrix and n is the chunk size. For sufficiently large chunks, this bias is negligible compared to the error ϵ . We prove the above high probability guarantee in Theorem 6 in Appendix E.

Further applications. Our least squares analysis can be extended to other settings where prior works [e.g., 17, 18, 16, 22] have analyzed randomized estimators based on sparse sketching via techniques from asymptotic random matrix theory. The primary and most direct application involves correcting inversion bias in the so-called sketched inverse covariance estimate ( ˜ A ⊤ ˜ A ) -1 , which was the motivating task of [18], with applications including distributed second-order optimization and statistical uncertainty quantification, where quantities like ( ˜ A ⊤ ˜ A ) -1 x are approximated.

Theorem 3 (informal Theorem 5) . Given A ∈ R n × d and its LESS embedding S with sketch size m ≥ Cd and s non-zeros per row, the inverse covariance sketch ( m m -d A ⊤ S ⊤ SA ) -1 is an ( ϵ, δ ) -unbiased estimator of ( A ⊤ A ) -1 (see Definition 3) for ϵ = ˜ O ( (1 + √ d/s ) √ d m ) and δ = 1 / poly( d ) .

This result should be compared with ϵ = ˜ O ( (1 + d/s ) √ d m ) obtained by [18]. Thus, we get a direct improvement for very sparse sketches, i.e., s = o ( d ) . This can be immediately translated into an improved local convergence guarantee for Distributed Newton Sketch which is a second-order convex minimization algorithm used in settings where the Hessian matrix can be expressed as A ⊤ A for a tall matrix A , e.g., in generalized linear models like logistic regression. In this method, following Corollary 16 of [18] as well as related results [16, 44, 19], we use sketching and averaging to estimate a Newton step:

where g t is the gradient at x t and ˜ H 1 , ..., ˜ H q are Hessian sketches constructed by independent machines. The following corollary, which is a direct improvement over Corollary 16 of [18], shows that Distributed Newton Sketch on a generalized linear model task can achieve a fast local convergence rate of the form f ( x t +1 ) -f ( x ∗ ) f ( x t ) -f ( x ∗ ) = o (1) with O ( d ω ) time per iteration (see Appendix B.1 for details).

Corollary 1 (Distributed Newton Sketch) . Consider f ( x ) = 1 n ∑ n i =1 ℓ i ( x ⊤ ϕ i ) + λ 2 ∥ x ∥ 2 , where ℓ i are convex twice continuously differentiable functions, such that f has a Lipschitz Hessian, λ > 0 , and ϕ ⊤ i is the i th row of an n × d data matrix Φ . Given ϵ > 0 , there is a neighborhood U ϵ around the minimizer x ∗ = argmin x f ( x ) such that, for any x t ∈ U ϵ , using two parallel passes with ˜ O (1 /ϵ ) machines, we can compute a Distributed Newton Sketch update x t +1 such that

in O ( γ -1 nnz( Φ ) + d ω + ϵ -1 d 2+ γ polylog( d )) time, O ( d 2 log( nd )) bits of space and O ( d log( nd )) bits of communication.

Our Techniques. At the core of our analysis are techniques inspired by asymptotic random matrix theory (RMT) in the proportional limit [e.g., see 6]. Here, in order to establish the limiting spectral distribution (such as the Marchenko-Pastur law) of a random matrix ˜ A ⊤ ˜ A whose dimensions diverge to infinity, one aims to show the convergence of the Stieltjes transform of its resolvent matrix ( ˜ A ⊤ ˜ A -z I ) -1 . Recently, [18] showed that these techniques can be adapted to sparse sketching matrices (via leverage score sparsification) in order to characterize the bias of the sketched inverse covariance ( ˜ A ⊤ ˜ A ) -1 , where ˜ A = SA .

Our main contribution is two-fold. First, we show that a similar argument can also be applied to analyze the bias of the least squares estimator, ˜ x = ( ˜ A ⊤ ˜ A ) -1 ˜ A ⊤ ˜ b . Unlike the inverse covariance, this estimator no longer takes the form of a resolvent matrix, but its bias is also associated with the inverse, which means that we can use a leave-one-out argument to characterize the effect of removing a single row of the sketch on the estimation bias. Our second main contribution is to improve the sharpness of the bounds relative to the sparsity of the sketching matrix by combining a careful application of Hölder's inequality with a higher moments analysis of the restricted BaiSilverstein inequality for quadratic forms. Those improvements are not only applicable to the least squares analysis, but also to all existing RMT-style results for LESS embeddings, including the aforementioned inverse covariance estimation, as well as applications in stochastic optimization, resulting in the sketching cost of LESS embeddings dropping below matrix multiplication time.

2 Related Work

Randomized numerical linear algebra. RandNLA sketching techniques have been developed over a long line of works, starting from fast least squares approximations of [42]; for an overview, see [45, 26, 37, 39, 21] among others. Since then, these methods have been used in designing fast algorithms not only for least squares but also many other fundamental problems in numerical linear algebra and optimization including low-rank approximation [13, 34], l p regression [14], solving linear systems [28, 24] and more. Using sparse random matrices for matrix sketching also has a long history, including data-oblivious sketching methods such as CountSketch [12], OSNAP [40], and more [38]. Leverage score sparsification (LESS) was introduced by [18] as a data-dependent sparse sketching method to enable RMT-style analysis for sketching (see below).

Unbiased estimators for least squares. To put our results in a proper context, let us consider other approaches for producing near-unbiased estimators for least squares, see also Table 1. First, a well known folklore result states that the least squares estimator computed from a dense Gaussian sketching matrix is unbiased. The bias of other sketching methods, including leverage score sampling and OSNAP, has been studied by [43], showing that these methods need a √ ϵ -error guarantee to achieve an ϵ -bias which leads to little improvement unless ϵ is extremely small and the sketch size is sufficiently large. Another approach of constructing unbiased estimators for least squares, first proposed by [23], is based on subsampling with a non-i.i.d. importance sampling distribution based on Determinantal Point Processes [DPPs, 31, 20]. However, despite significant efforts [30, 7, 5], sampling from DPPs remains quite expensive: the fastest known algorithm requires running a Markov chain for polylog( n/ϵ ) many steps, each of which requires a separate data pass and takes O ( d ω ) time. Other approaches have also been considered which provide partial bias reduction for i.i.d. RandNLA subsampling schemes in various regimes that are are either much more expensive or not directly comparable to ours [2, 44].

Statistical and RMT analysis of sketching. Recently, there has been significant interest in statistical and random matrix theory (RMT) analysis of matrix sketching; see [21] for an overview. These approaches include both asymptotic analysis via limiting spectral distributions and deterministic equivalents [35, 25, 32, 33], and non-asymptotic analysis under statistical assumptions [36, 41, 4]. A number of works have shown that the RMT-style techniques based on deterministic equivalents can be made rigorously non-asymptotic for certain sketching methods such as dense sub-Gaussian [17], LESS matrices [16, 18], and other sparse matrices [9], which has been applied to low-rank approximation, fast subspace embeddings and stochastic optimization. Our new analysis can be viewed as a general strategy for directly improving the sparsity required by LESS embeddings (and thereby, the sketching time complexity) in many of these applications, specifically those that rely on analysis inspired by the calculus of deterministic equivalents via generalized Stieltjes transforms.

3 Preliminaries

Notations. In all our results, we use lowercase letters to denote scalars, lowercase boldface for vectors, and uppercase boldface for matrices. The norm ∥ · ∥ denotes the spectral norm for matrices and the Euclidean norm for vectors, whereas ∥ · ∥ F denotes the Frobenius norm for matrices. We use ⪯ to denote the p.s.d. ordering of matrices.

Computational model. We next clarify the computational model that is used in Theorem 2. We consider a central data server storing ( A , b ) , and q machines. The j th machine has a handle Stream( j ) , which can be used to open a stream and to read the next row/label pair ( a i , b i ) in the stream. After a full pass, the machine can re-open the handle and begin another pass over the data. The machines can operate their streams entirely asynchronously, and each has its own limited local storage space, e.g., in Theorem 2 we use O ( d 2 log( nd )) bits of space per machine. At the end, they can communicate some information back to the server, e.g., in Theorem 2, they communicate their final estimate vectors ˜ x j , using O ( d log( nd )) bits of communication. Then, the server computes the final estimate, in our case via averaging, ˜ x = 1 q ∑ q i =1 ˜ x i , which can be done either directly or via a map-reduce type architecture.

We define the parallel passes required by such an algorithm as the maximum number of times the stream is opened by any single machine. We analogously define time/space/communication costs by taking a maximum over the costs required by any single machine (for communication, this refers only to the number of bits sent from the machine back to the server).

Definitions and useful lemmas. In our framework, we construct a sparse sketching matrix S where sparsification is achieved using a probability distribution over rows of data matrix A , that is proportional to the leverage scores of A . The next definition [following, e.g., 9] provides the explicit definition of exact and approximate leverage scores for our setting.

Definition 1 ( ( β 1 , β 2 ) -approximate leverage scores) . Fix a matrix A ∈ R n × d and consider matrix U ∈ R n × d with orthonormal columns spanning the column space of A . Then, the leverage scores l i , 1 ≤ i ≤ n are defined as the row norms squared of U , i.e., l i = ∥ u i ∥ 2 , where u ⊤ i is the i th row of U . Furthermore, consider fixed β 1 , β 2 > 1 . Then ˜ l i are called ( β 1 , β 2 ) -approximate leverage scores for A if the following holds for all i

The approximate leverage scores can be computed by first constructing a preconditioner matrix P ∈ R d × d such that κ ( AP ) = O (1) , which takes O (nnz( A ) + d ω ) in a single pass, and then relying on the following norm approximation scheme.

Lemma 1 (Based on Lemma 7.2 from [10]) . Given A ∈ R n × d and P ∈ R d × d , using a single pass over A in time O ( γ -1 (nnz( A ) + d 2 )) for small constant γ > 0 , we can compute estimates ˜ l 1 , ..., ˜ l n such that with probability ≥ 0 . 95 :

In the next definition, we give the sparse sketching strategy used in our analysis. This approach is similar to the original leverage score sparsification proposed by [18], except: 1) we adapted it so that it can be implemented effectively in a single pass, and 2) we use it in a much sparser regime (fewer non-zeros per row).

Definition 2 ( ( s, β 1 , β 2 ) -LESS embedding) . Fix a matrix A ∈ R n × d and some s ≥ 0 . Let the tuple ( ˜ l 1 , · · · , ˜ l n ) denote ( β 1 , β 2 ) -approximate leverage scores for A . Let p i = min { 1 , sβ 1 ˜ l i d } . We define a ( s, β 1 , β 2 ) -approximate leverage score sparsifier ξ as follows.

Moreover, we define the ( s, β 1 , β 2 ) -leverage score sparsified (LESS) embedding of size m as matrix S ∈ R m × d with i.i.d. rows 1 √ m x i such that x i = diag( ξ i ) y i where ξ i denotes a randomly generated ( β 1 , β 2 ) -approximate leverage score sparsifier and y i ∈ R n consist of random ± 1 entries.

A key property of a sketching matrix is the subspace embedding property, defined below. It was recently shown by [9] that LESS embeddings require only polylogarithmically many non-zeros per row of S to prove that S is a subspace embedding for the data matrix A with the optimal m = O ( d ) sketching dimension. For our main results, it is sufficient to use η = O (1) below.

Lemma 2 (Subspace embedding for LESS, Theorem 1.3, [9]) . Fix η, δ > 0 . Consider β 1 , β 2 > 1 and a full rank matrix A ∈ R n × d . Then for a ( β 1 , β 2 ) -leverage score sparsified embedding S ∈ R m × n with s ≥ O (log 4 ( d/δ ) /η 4 ) and m = O (( d +log1 /δ ) /η 2 ) , with probability 1 -δ we have

4 Least Squares Bias Analysis

In this section we provide an outline of the bias analysis for the sketched least squares estimator constructed using a LESS embedding, leading to the proofs of our main results, Theorems 1 and 2. In particular, we prove the following main technical result (detailed proof in Appendix C).

Theorem 4 (Bias of LESS-sketched least squares) . Fix A ∈ R n × d and let S be an ( s, β 1 , β 2 ) -LESS embedding of size m for A . Let S satisfy (1) with η = 1 2 and probability 1 -δ where δ < 1 m 4 . Then there exists an event E with probability at least 1 -δ such that

Remark 3. Thus, the bias L ( E E [˜ x ]) -L ( x ∗ ) of the LESS estimator using O ( β 1 β 2 s ) = ˜ O ( s ) nonzeros per row of S is of the order ˜ O ( d 2 sm 2 + d m 2 ) · L ( x ∗ ) . By comparison, the standard expected loss bound which holds for sketched least squares (including this estimator) is E [ L (˜ x )] -L ( x ∗ ) ≤ ˜ O ( d m ) L ( x ∗ ) , and the best known bound on the bias of most standard sketched estimators (e.g., leverage score sampling) is ˜ O ( d 2 m 2 ) L ( x ∗ ) , given by [43]. So, our result recovers the standard bias bound for s = 1 and improves on it for s ≫ 1 by a factor of min { s, d } . At the end of the section, we discuss how to deal with the lower order term ˜ O ( d m 2 ) to reduce the bias further.

Proof sketch. Using a standard argument, we can replace the matrix A with the matrix U ∈ R n × d consisting of orthonormal columns spanning the column space of A , and assume that n = poly( d ) . Let S be an ( s, β 1 , β 2 ) -LESS embedding for U . Also, let b ∈ R n be a vector of responses/labels corresponding to n rows in U . Let ˜ x = argmin x ∥ SUx -Sb ∥ 2 . Furthermore for any x ∈ R d we can find the loss at x as L ( x ) = ∥ Ux -b ∥ 2 . Additionally, we use r to denote the residual b -Ux ∗ . We also define Q = ( γ U ⊤ S ⊤ SU ) -1 as the sketched inverse covariance matrix with scaling γ = m m -d representing the standard correction accounting for inversion bias. We condition on the high probability event E guaranteed in Lemma 2 and consider L ( E E (˜ x )) -L ( x ∗ ) . By Pythagorean theorem, we have L ( E E [˜ x ]) -L ( x ∗ ) = ∥ U ( E E [˜ x ]) -Ux ∗ ∥ 2 . Note that by the normal equations we have ˜ x = ( U ⊤ S ⊤ SU ) -1 U ⊤ S ⊤ Sb = γ QU ⊤ S ⊤ Sb , and also S ⊤ S = 1 m ∑ m i =1 x i x ⊤ i . These two facts lead to writing the bias as follows:

Using a leave-one-out technique , we replace Q with Q -i = ( γ U ⊤ S ⊤ -i S -i U ) -1 , where S -i denotes matrix S without the i th row, by noting that Q = ( γ U ⊤ S ⊤ -i S -i U + γ m U ⊤ x i x ⊤ i U ) -1 and applying the Sherman-Morrison formula. This leads to the following relation:

where γ i = 1 + γ m x ⊤ i UQ -i U ⊤ x i . Due to the subspace embedding assumption and assuming m large enough, we have ∥ Q ∥ = O (1) and also ∥ Q -i ∥ = O (1) . The first term ∥ Z 0 r ∥ 2 is quite straightforward to bound since, if not for the conditioning on the high probability event E , we would have E [ Q -i U ⊤ x i x ⊤ i r ] = E [ Q -i U ⊤ r ] = 0 , which follows from U ⊤ ( b -Ux ∗ ) = 0 . We get an upper bound on ∥ Z 0 r ∥ 2 as O ( d 2 log( d/δ ) sm 2 + d m 2 ) · ∥ r ∥ 2 , which is sufficient for us.

The central novelty of our analysis lies in bounding ∥ Z 2 r ∥ 2 for ( s, β 1 , β 2 ) -LESS embeddings, which is the dominant term. Our key observation is that, when examining a random variable of the form x ⊤ i v for some vector v , the dependence on the sparsity of row x i only arises when considering moments higher than 2 + 1 O (log( n )) , because otherwise we can simply rely on the fact that E [ x i x ⊤ i ] = I . Thus, when decomposing ∥ Z 2 r ∥ 2 , we must carefully separate the contribution of near-second moments vs the contribution of higher moments to the overall bound.

To obtain this separation, we start by applying Hölder's inequality on ∥ Z 2 r ∥ with p = O (log( n )) and q = 1 + 1 O (log( n )) to get

Furthermore applying Cauchy-Schwarz inequality on the second term leads to

Unlike [18], we exploit the fact that ∥ x i ∥ 1 /O (log( n )) = O (1) and get a constant upper bound on ( E E ∥ Q -i U ⊤ x i ∥ 2 q ) . However, this results in a much more careful argument, requiring now an upper bound on ( E E [ ∣ ∣ γ γ i -1 ∣ ∣ p ]) 1 /p for p = O (log( n )) . First, we observe that

where ¯ γ = 1 + γ m E E ( x ⊤ i UQ -i U ⊤ x i ) . In particular, for the second term, we have

To bound the first of these two terms, we prove a new version of the Restricted Bai-Silverstein inequality (Lemma 3) for ( s, β 1 , β 2 ) -LESS embeddings. Unlike [18], we provide a proof with any p and any ( β 1 , β 2 ) values. Furthermore, utilizing the subspace embedding guarantee from Lemma 2, we prove a much more general result where the number of non-zeros in the approximate leverage score sparsifier ξ can be much smaller than d (proof in Appendix D).

Lemma 3 (Restricted Bai-Silverstein for ( s, β 1 , β 2 ) -LESS embeddings) . Let p ∈ N be fixed and U ∈ R n × d be such that U ⊤ U = I . Let x i = diag( ξ ) y i where y i ∈ R n has independent ± 1 entries and ξ is an ( s, β 1 , β 2 ) -approximate leverage score sparsifier for U . Then for any matrix with 0 ⪯ C ⪯ O (1) · I and any δ > 0 we have for an absolute constant c > 0 .

Using Lemma 3, we upper bound the first term squared in (3) as ˜ O ( d m 2 ( 1 + d s )) . Moreover, also using Lemma 3, we get a matching upper bound on | γ -¯ γ | . Wethen design a martingale concentration argument to prove a high probability upper bound on the last remaining term, | tr( Q -i ) -E E tr( Q -i ) | , which implies the desired moment bound (proof in Appendix B), concluding the proof of Theorem 4.

Lemma 4. For given δ > 0 and matrix Q -i we have with probability 1 -δ :

Completing the proof of Theorem 1. First, suppose that ϵ ≥ O (polylog( d ) /d ) so that the bias bound can be achieved from Theorem 4. Our implementation is mainly based on the online construction of approximate leverage scores, given the preconditioner P , using Lemma 1. Briefly, this construction proceeds by first sketching P using a d × O (1 /γ ) Gaussian matrix G to produce the matrix ˜ P = PG , and then, for each observed row a i of A , we compute ˜ l i = ∥ a ⊤ i ˜ P ∥ 2 . Assuming without loss of generality that d = poly( n ) and adjusting γ , the estimates satisfy β 1 β 2 = O ( αd γ ) .

Next, we sample the non-zero entries of S corresponding to the observed row a i , i.e., the i -th column of S . Note that for this we only need to know the single leverage score estimate ˜ l i . Crucially for our analysis, the entries of this column need to be sampled i.i.d., which can be done in time proportional to the number of non-zeros in that column by first sampling a corresponding Binomial distribution to determine how many non-zeros we need, then picking a random subset of that size, and then sampling the random ± 1 values. Altogether, the cost of constructing the sketch is O ( γ -1 nnz( A )+ β 1 β 2 sd 2 ) = O ( γ -1 nnz( A )+ αϵ -1 d 2+ γ polylog( d )) by setting s = O (polylog( d ) /ϵ ) . Finally, once we construct the sketch, at the end of the pass we can run conjugate gradient preconditioned with P on the sketched problem, which takes ˜ O ( αd 2 ) .

We note that in the (somewhat artificial) regime where we require extremely small bias, i.e., ϵ = o (polylog( d ) /d ) , the bound claimed in Theorem 1 can still be obtained, since in this case for small enough γ we have d 2+ γ /ϵ = O ( d ω / √ ϵ ) with ω < 2 . 5 , so we can rely on direct leverage score sampling (which corresponds to s = 1 ), and instead of maintaining the sketch, we compute the estimator ˜ x = ( ˜ A ⊤ ˜ A ) -1 ˜ A ⊤ b directly along the way. This involves performing a separate d × d matrix multiplication after collecting each d leverage score samples, to gradually compute ˜ A ⊤ ˜ A , and then inverting the matrix at the end. From Theorem 4, we see that it suffices to set sketch size m = ˜ O ( d/ √ ϵ ) , which leads to the desired runtime.

Completing the proof of Theorem 2. Here, we use a slightly modified variant of Lemma 2, given as Theorem 1.4 in [9], which shows that using a single pass we can compute a sketch ˜ A in time O (nnz( A ) + d ω ) , which satisfies the subspace embedding property (1) with η = 1 2 . Then, we can perform the QR decomposition ˜ A = QR and set P = R -1 in additional time O ( d ω ) to obtain the desired preconditioner. Next, we use Theorem 1 to construct q i.i.d. estimators ˜ x i in a second parallel pass, and finally, the estimators are aggregated to compute ˆ x = 1 q ∑ q i =1 ˜ x i , which satisfies E ∥ A ˆ x -b ∥ 2 ≤ ( 1 + ϵ + O (1 /q ) ) ∥ Ax ∗ -b ∥ 2 . Applying Markov's inequality concludes the proof.

5 Experiments

In this section, we illustrate empirically how our results point to a practical free lunch phenomenon in distributed averaging of sketching-based estimators. As mentioned in Section 1, our construction from Theorem 1 essentially works by taking a subsample of the data and then mixing groups of those rows together to produce an even smaller sketch (see Figure 1). According to our theory, while the small sketch does not recover the same ϵ -small error as the larger subsample, it does recover an ϵ -small bias. Moreover, this happens without incurring any additional computational cost, as the cost of the sketching is proportional to the cost of simply reading the subsampled rows. This suggests that we can use sparse sketching to compress a data subsample down to a small size while retaining the least squares performance in a distributed averaging environment.

To verify this, we evaluate the effectiveness of distributed averaging of sketched least squares estimators on several benchmark datasets. Specifically, we visualize the relative error of the averaged sketch-and-solve estimator L (ˆ x ) -L ( x ∗ ) L ( x ∗ ) , against the number of machines q used to generate the estimate ˆ x = 1 q ∑ q i =1 ˜ x i . Each estimate ˜ x i is constructed with the same sparsification strategy used by LESS, except that instead of sparsifying the sketch with leverage scores, we instead sparsify them with uniform probabilities (which is often sufficient in practice). Following [16], we call the resulting method LESSUniform. Within each dataset, we perform four simulations, designed so that the total sketching cost stays the same for all four test cases, by simultaneously changing sketch size and sparsity. Concretely, we vary these so that the product (sketch size × nnz per row) stays the same in each case, so as to ensure that the total cost of sketching is fixed in each plot.

In Figure 2, on the X-axis we plot the number q of estimators being averaged, so that the bias of a single estimator appears on the right-hand side of the plot (large q ), whereas the variance (error) appears on the left-hand side ( q = 1 ). In each plot, the line with nnz per row = 1 (and large sketch size) corresponds to uniform subsampling, whereas the remaining ones are sketches produced by compressing that subsample. The plot shows that decreasing the sketch size (i.e., compressing the sample) does increase the error of a single estimator (as expected), however it also shows that the bias of these estimators remains essentially unchanged regardless of the sketch size (since all lines

MSD

Abalone

100

Sketch size=2000,nnz per row=1

100

Sketch size=128,nnzperrow=1

Sketch size=1000,nnz per row=2

Sketch size=64,nnz per row=2

1

Sketch size=400,nnz per row=5

1

Sketch size=32, nnz per row=4

Sketch size=200,nnz per row=10

Squares

Squares

10-1

10

Loss/Least

Loss/LeastS

10-2

10-2

100

101

102

103

104

100

101

102

103

104

Number of Machines

Number of Machines

meet as q →∞ ), confirming that suitable sparse sketches that 'compress' rows of data can preserve near-unbiasedness without increasing the cost.

This phenomenon may not occur for all sketching methods. Figure 4 within Appendix F showcases a few more interesting results. We first further demonstrate that suitable sketches that 'compress' rows of data A into a single row of ˜ A can preserve near-unbiasedness without increasing the cost. In particular, this desirable phenomenon that LESSUniform enjoys also extends to LESS proper, as well as the Gaussian and Subgaussian (Rademacher) sketches. In fact, we observe that LESS enjoys similar desirable performance as the Gaussian and Subgaussian sketches and virtually no least squares bias, while retaining the computational speedups of sparse sketching, suggesting that it attains the best of both worlds.

However, when we decrease the number of subsamples within leverage score subsampling, the bias introduced by subsampling increases as expected. This happens as the number of subsamples is reduced without increasing the amount of 'compression' as one would with LESS or LESSUniform. 1 We also show that the subsampled randomized Hadamard transform (SRHT) can exhibit some amount of least squares bias as the sketch size decreases. The numerical results shown for the bias introduced by leverage score subsampling and the SRHT complement the lower bounds established in [18].

6 Conclusions

We gave a new sparse sketching method that, using two passes over the data, produces a nearlyunbiased least squares estimator, which can be used to improve upon the space-time trade-offs of solving least squares in parallel or distributed environments via simple averaging. In particular, our algorithm is the first to require only O ( d 2 log( nd )) bits of space and current matrix multiplication time O ( d ω ) while obtaining an ϵ = o (1) approximation in few passes. Our techniques are of broader interest to sketching-based optimization algorithms, including Distributed Newton Sketch.

Acknowledgments

This work was partially supported by NSF CAREER CCF-2338655.

References

Dimitris Achlioptas. Database-friendly random projections: Johnson-Lindenstrauss with binary coins. Journal of computer and System Sciences , 66(4):671-687, 2003.

Naman Agarwal, Brian Bullins, and Elad Hazan. Second-order stochastic optimization for machine learning in linear time. The Journal of Machine Learning Research , 18(1):4148-4187, 2017.

Nir Ailon and Bernard Chazelle. The fast Johnson-Lindenstrauss transform and approximate nearest neighbors. SIAM Journal on computing , 39(1):302-322, 2009.

Ahmed El Alaoui and Michael W. Mahoney. Fast randomized kernel ridge regression with statistical guarantees. In Proceedings of the 28th International Conference on Neural Information Processing Systems , pages 775-783, 2015.

Nima Anari, Yang P Liu, and Thuy-Duong Vuong. Optimal sublinear sampling of spanning trees and determinantal point processes via average-case entropic independence. In 2022 IEEE 63rd Annual Symposium on Foundations of Computer Science (FOCS) , pages 123-134. IEEE, 2022.

Zhidong Bai and Jack W Silverstein. Spectral analysis of large dimensional random matrices , volume 20. Springer, 2010.

Daniele Calandriello, Michal Derezinski, and Michal Valko. Sampling from a k-dpp without looking at all items. Advances in Neural Information Processing Systems , 33:6889-6899, 2020.

Chih-Chung Chang and Chih-Jen Lin. LIBSVM: A library for support vector machines. ACM Transactions on Intelligent Systems and Technology , 2:27:1-27:27, 2011.

Shabarish Chenakkod, Michał Derezi´ nski, Xiaoyu Dong, and Mark Rudelson. Optimal embedding dimension for sparse subspace embeddings. In 56th Annual ACM Symposium on Theory of Computing , 2024.

Nadiia Chepurko, Kenneth L Clarkson, Praneeth Kacham, and David P Woodruff. Near-optimal algorithms for linear algebra in the current matrix multiplication time. In Proceedings of the 2022 Annual ACM-SIAM Symposium on Discrete Algorithms (SODA) , pages 3043-3068. SIAM, 2022.

Kenneth L Clarkson and David P Woodruff. Numerical linear algebra in the streaming model. In Proceedings of the forty-first annual ACM symposium on Theory of computing , pages 205-214, 2009.

Kenneth L Clarkson and David P Woodruff. Low rank approximation and regression in input sparsity time. In Proceedings of the forty-fifth annual ACM symposium on Theory of Computing , pages 81-90, 2013.

Michael B Cohen, Cameron Musco, and Christopher Musco. Input sparsity time low-rank approximation via ridge leverage score sampling. In Proceedings of the Twenty-Eighth Annual ACM-SIAM Symposium on Discrete Algorithms , pages 1758-1777. SIAM, 2017.

Michael B Cohen and Richard Peng. Lp row sampling by lewis weights. In Proceedings of the symposium on Theory of computing , pages 183-192, 2015.

Michał Derezi´ nski. Stochastic variance-reduced newton: Accelerating finite-sum minimization with large batches. arXiv preprint arXiv:2206.02702 , 2022.

Michał Derezi´ nski, Jonathan Lacotte, Mert Pilanci, and Michael W Mahoney. Newton-LESS: Sparsification without trade-offs for the sketched newton update. Advances in Neural Information Processing Systems , 34:2835-2847, 2021.

Michał Derezi´ nski, Feynman T Liang, Zhenyu Liao, and Michael W Mahoney. Precise expressions for random projections: Low-rank approximation and randomized newton. Advances in Neural Information Processing Systems , 33, 2020.

Michał Derezi´ nski, Zhenyu Liao, Edgar Dobriban, and Michael Mahoney. Sparse sketches with small inversion bias. In Conference on Learning Theory , pages 1467-1510. PMLR, 2021.

Michał Derezi´ nski and Michael W Mahoney. Distributed estimation of the inverse Hessian by determinantal averaging. In Advances in Neural Information Processing Systems 32 , pages 11401-11411. 2019.

Michał Derezi´ nski and Michael W Mahoney. Determinantal point processes in randomized numerical linear algebra. Notices of the American Mathematical Society , 68(1):34-45, 2021.

Michał Derezi´ nski and Michael W Mahoney. Recent and upcoming developments in randomized numerical linear algebra for machine learning. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pages 6470-6479, 2024.

Michał Derezi´ nski and Elizaveta Rebrova. Sharp analysis of sketch-and-project methods via a connection to randomized singular value decomposition. SIAM Journal on Mathematics of Data Science , 6(1):127-153, 2024.

Michał Derezi´ nski and Manfred K. Warmuth. Unbiased estimates for linear regression via volume sampling. In Advances in Neural Information Processing Systems 30 , pages 3087-3096, 2017.

Michał Derezi´ nski and Jiaming Yang. Solving dense linear systems faster than via preconditioning. In 56th Annual ACM Symposium on Theory of Computing , 2024.

Edgar Dobriban and Sifan Liu. Asymptotics for sketching in least squares regression. Advances in Neural Information Processing Systems , 32, 2019.

Petros Drineas and Michael W Mahoney. Randnla: randomized numerical linear algebra. Communications of the ACM , 59(6):80-90, 2016.

Petros Drineas, Michael W Mahoney, and S Muthukrishnan. Sampling algorithms for ℓ 2 regression and applications. In Proceedings of the seventeenth annual ACM-SIAM symposium on Discrete algorithm , pages 1127-1136, 2006.

Robert M Gower and Peter Richtárik. Randomized iterative methods for linear systems. SIAM Journal on Matrix Analysis and Applications , 36(4):1660-1690, 2015.

William B Johnson, Gideon Schechtman, and Joel Zinn. Best constants in moment inequalities for linear combinations of independent and exchangeable random variables. The Annals of Probability , pages 234-253, 1985.

Alex Kulesza and Ben Taskar. k-DPPs: Fixed-Size Determinantal Point Processes. In Proceedings of the 28th International Conference on Machine Learning , pages 1193-1200, June 2011.

Alex Kulesza and Ben Taskar. Determinantal Point Processes for Machine Learning . Now Publishers Inc., Hanover, MA, USA, 2012.

Jonathan Lacotte, Sifan Liu, Edgar Dobriban, and Mert Pilanci. Optimal iterative sketching methods with the subsampled randomized Hadamard transform. Advances in Neural Information Processing Systems , 33:9725-9735, 2020.

Daniel LeJeune, Pratik Patil, Hamid Javadi, Richard G Baraniuk, and Ryan J Tibshirani. Asymptotics of the sketched pseudoinverse. arXiv preprint arXiv:2211.03751 , 2022.

Yi Li and David Woodruff. Input-sparsity low rank approximation in schatten norm. In International Conference on Machine Learning , pages 6001-6009. PMLR, 2020.

Miles E Lopes, Shusen Wang, and Michael W Mahoney. A bootstrap method for error estimation in randomized matrix multiplication. The Journal of Machine Learning Research , 20(1):14341473, 2019.

Ping Ma, Yongkai Chen, Xinlian Zhang, Xin Xing, Jingyi Ma, and Michael W Mahoney. Asymptotic analysis of sampling estimators for randomized numerical linear algebra algorithms. The Journal of Machine Learning Research , 23(1):7970-8014, 2022.

Per-Gunnar Martinsson and Joel A Tropp. Randomized numerical linear algebra: Foundations and algorithms. Acta Numerica , 29:403-572, 2020.

Xiangrui Meng and Michael W. Mahoney. Low-distortion subspace embeddings in inputsparsity time and applications to robust linear regression. In Proceedings of the Symposium on Theory of Computing , STOC '13, pages 91-100, 2013.

R. Murray, J. Demmel, M. W. Mahoney, N. B. Erichson, M. Melnichenko, O. A. Malik, L. Grigori, M. Derezi´ nski, M. E. Lopes, T. Liang, and H. Luo. Randomized Numerical Linear Algebra - a perspective on the field with an eye to software. Technical Report arXiv preprint arXiv:2302.11474, 2023.

Jelani Nelson and Huy L Nguyên. Osnap: Faster numerical linear algebra algorithms via sparser subspace embeddings. In 2013 ieee 54th annual symposium on foundations of computer science , pages 117-126. IEEE, 2013.

G. Raskutti and M. W. Mahoney. A statistical perspective on randomized sketching for ordinary least-squares. Journal of Machine Learning Research , 17(214):1-31, 2016.

Tamas Sarlos. Improved approximation algorithms for large matrices via random projections. In 2006 47th annual IEEE symposium on foundations of computer science (FOCS'06) , pages 143-152. IEEE, 2006.

S. Wang, A. Gittens, and M. W. Mahoney. Sketched ridge regression: Optimization perspective, statistical perspective, and model averaging. Journal of Machine Learning Research , 18(218):150, 2018.

Shusen Wang, Fred Roosta, Peng Xu, and Michael W Mahoney. GIANT: globally improved approximate newton method for distributed optimization. Advances in Neural Information Processing Systems , 31:2332-2342, 2018.

David P Woodruff. Sketching as a tool for numerical linear algebra. Foundations and Trends® in Theoretical Computer Science , 10(1-2):1-157, 2014.

A Detailed preliminaries

We start by providing several classical results, used in our analysis. The following formula provides a way to compute the inverse of matrix A after a rank1 update, given the inverse before the update.

̸

Lemma 5 (Sherman-Morrison formula) . For an invertible matrix A ∈ R d × d and vector u , v ∈ R d , A + uv ⊤ is invertible if and only if 1 + v ⊤ A -1 u = 0 . If this holds then,

The following inequality provides a crucial tool for writing expectation of the product of two random variables as the product of higher individual moments.

Lemma 6 (Hölder's inequality) . For real-valued random variables X and Y ,

where p, q > 0 are Hölder's conjugates, i.e. 1 p + 1 q = 1 .

The following technical lemmas provide concentration results for the sum of random quantities. We collect these results here and then refer to them while using in our analysis.

Lemma 7 (Matrix Chernoff Inequality) . For i = 1 , 2 , · · · , n consider a sequence Z i of d × d positive semi-definite random matrices such that E [ 1 n ∑ i Z i ] = I d and ∥ Z i ∥ ≤ R . Then for any ϵ > 0 , we have

where λ max denotes taking the maximum eigenvalue.

Lemma 8 (Azuma's inequality) . If { Y 0 , Y 1 , Y 2 , · · · } is a martingale with | Y j -Y j -1 | ≤ c j then for any m,λ > 0 we have

Lemma 9 (Rosenthal's inequality ([29], Theorem 2.5 and Corollary 2.6)) . Let 1 ≤ p < ∞ and X 1 , X 2 , · · · , X n are nonnegative, independent random variables with finite p th moments then,

Furthermore, for mean-zero independent and symmetric random variables we have

Lemma 10 (Bai-Silverstein's Inequality Lemma B.26 from [6]) . Let B be a d × d be a fixed matrix and x be a random vector of independent entries. Let E [ x i ] = 0 and E [ x 2 i ] = 1 ,and E | x j | l ≤ ν l . Then for any p ≥ 1 ,

In particular,

B Inversion bias analysis

We use the notion of an ( ϵ, δ ) unbiased estimator of A as defined in [18].

Definition 3 ( ( ϵ, δ ) -unbiased estimator) . For ϵ, δ > 0 , a random positive definite matrix B ∈ R d × d is called an ( ϵ, δ ) unbiased estimator of A if there exists an event E with Pr( E ) ≥ 1 -δ such that,

when conditioned on the event E .

In this section, we give a formal statement and proof for Theorem 3, which is then used in the proof of Theorem 4. We replace A with U such that U consists of d orthonormal columns spanning the column space of A . Here S m ∈ R m × n denotes a LESS sketching matrix with independent rows 1 √ m x ⊤ , x ⊤ = y ⊤ · diag( ξ ) where y consists of ± 1 Rademacher entries and ξ is an ( s, β 1 , β 2 ) -approximate leverage score sparsifier. Note that E [ xx ⊤ ] = I n . We assume that the sketching matrix S m consists of m ≥ 10 d i.i.d rows and 3 divides m . Also, we assume that the S m satisfies the subspace embedding condition for U (Theorem 2) with η = 1 2 . Let Q = ( γ U ⊤ S ⊤ m S m U ) -1 where γ = m m -d .

Theorem 5 (Small inversion bias for ( s, β 1 , β 2 ) -LESS embeddings) . Let δ > 0 satisfy δ < 1 m 4 and m ≥ O ( d ) . Let S m ∈ R m × n be an ( s, β 1 , β 2 ) -LESS embedding for data matrix U ∈ R n × d such that U ⊤ U = I . Then there exists an event E with Pr( E ) ≥ 1 -δ such that,

Proof. Let S -i denote S m without the i th row, and S -ij denote S m with the i th and j th rows removed. Let Q -i = ( γ U ⊤ S ⊤ -i S -i U ) -1 and Q -ij = ( γ U ⊤ S ⊤ -ij S -ij U ) -1 . We proceed with the same proof strategy as adopted in [18]. We define the events E j as follows:

Note that event E j means that the sketching matrix with just (1 / 3) rd rows (scaled to maintain unbiasedness of the sketch) from S m satisfies a lower spectral approximation of U ⊤ U = I . Also we notice that events E 1 , E 2 , E 3 are independent, and for any pair ( i, j ) there exists at least one event E k , k ∈ { 1 , 2 , 3 } such that E k is independent of both x i and x j . Furthermore conditioned on E k we have

Note that as guaranteed in Theorem 2, we have Pr( E k ) ≥ 1 -δ ′ for all k and therefore Pr( E ) ≥ 1 -δ with δ ′ = δ/ 3 . Let E E denote the expectation conditioned on the event E .

where γ i = 1 + γ m x ⊤ i UQ -i U ⊤ x i . The second equality follows by noting that S ⊤ m S m = 1 m ∑ m i =1 x i x ⊤ i and using linearity of expectation. The third equality holds due to the application of Sherman Morrison's (Lemma 5) formula on Q = ( γ U ⊤ S ⊤ -i S -i U + γ m U ⊤ x i x ⊤ i U ) -1 . We start by upper bounding Z 0 and use the following from [18].

Lemma 11 (Upper bound on ∥ Z 0 ∥ ) . For any k > 0 we have,

Using Chebyshev's inequality we have, Pr( x ⊤ i UU ⊤ x i ≥ x ) ≤ Var ( x ⊤ i UU ⊤ x i ) x 2 . Using Restricted Bai-Silverstein inequality for ( s, β 1 , β 2 ) -approximate LESS embeddings i.e., Lemma 3 (proved in Lemma 16) with p = 2 , we have Var ( x ⊤ i UU ⊤ x i ) ≤ cd · ( 1 + d log( d/δ ) s ) for some absolute constant c . Let k = m 2 and δ ′ < 1 m 4 we get,

We use the bound on the term ∥ Z 1 ∥ directly from [18], provided below as a Lemma.

Lemma 12 (Upper bound on ∥ Z 1 ∥ , [18]) .

It remains to upper bound ∥ Z 2 ∥ .

Applying Hölder's inequality with p = O (log( n )) and q = p p -1 = 1 + ∆ for ∆ = 1 O (log( n )) , we get,

where we used Cauchy-Schwarz inequality on E E [ | v ⊤ Q -i U ⊤ x i x ⊤ i Uz | q ] . Now note that x i = diag( ξ ) · y i , we have ∥ x i ∥ = poly( n ) and therefore ∥ x i ∥ ∆ = O (1) . We now show the terms involving exponents depending on q are O (1) as highlighted in the inequality (6).

Here E ′ is an event independent of x i . Without loss of generality, we can assume that E ′ = E 1 ∧ E 2 . We first condition on Q -i and take expectation over x i . Also note that Q -i and x i are independent and event E ′ is independent of x i , and furthermore E E ′ [ x ⊤ i x i ] = E [ x ⊤ i x i ] = 1 . We get,

Now we use that conditioned on E ′ , ∥ Q -i ∥ ≤ 6 and ∥ U ⊤ ∥ = 1 , we get,

Similarly, using ∥ x i ∥ ∆ = O (1) and E ′ [ x ⊤ i x i ] = E [ x ⊤ i x i ] = 1 ,

Now we prove an upper bound on ( E E [( γ γ i -1 ) p ]) 1 /p . Without loss of generality, we assume that p is even. We have,

where E ′ is event independent of x i . The above can be upper bounded as,

where ¯ γ = 1 + γ m E E ′ ( x ⊤ i UQ -i U ⊤ x i ) . As E ′ is independent of x i and Q -i is independent of x i , we get E E ′ ( x ⊤ i UQ -i U ⊤ x i ) = E E ′ tr( Q -i ) . Therefore, ¯ γ = 1 + γ m E E ′ tr( Q -i ) . We now aim to upper bound ( E E ′ [( γ i -¯ γ ) p ]) 1 /p as,

Using our new Restricted Bai-Silverstein inequality from Lemma 3 (restated as Lemma 16 and proven in Appendix D), we have

Wenowconsider ( E E ′ [tr( Q -i ) -E E ′ tr( Q -i )] p ) 1 /p . In Lemma 4 (restated below as Lemma 13), we show that | tr( Q -i ) -E E ′ tr( Q -i ) | ≤ c ′ γ √ m · d log 4 . 5 ( m/δ ) with probability at least 1 -δ . Conditioned on this high-probability event we have,

for an absolute constant c ′ > 0 . Therefore we get with probability at least 1 -δ ,

As m>d , we get ( E E ′ [( γ i -¯ γ ) p ]) 1 /p = O ( √ d log 4 . 5 ( n/δ ) m · ( 1 + √ d s )) . Also using the analysis in [18] for upper bounding | γ -¯ γ | , we get a matching upper bound on | γ -¯ γ | as follows:

Substituting these bounds in (7) and then in (6) we get,

Combining the upper bounds for Z 0 , Z 1 and Z 2 using relations (4,5,8) we conclude our proof.

We now provide the proof of Lemma 4, which we restate in the following Lemma.

Lemma 13. For given δ > 0 and matrix Q -i we have with probability 1 -δ :

for an absolute constant c ′ > 0 .

Proof. Writing tr( Q -i ) -E E ′ tr( Q -i ) as a finite sum, we have,

Denoting X j = E E ′ ,j tr( Q -i ) with X 0 = E E ′ tr( Q -i ) , we have the following formulation

with E E ′ ,j -1 [ X j ] = X j -1 . The random sequence X j forms a martingale and tr( Q -i ) -E E ′ tr( Q -i ) = X m -X 0 . We find an upper bound on | X j -X j -1 | . To achieve that we note,

Therefore with ψ j = ( E E ′ ,j -E E ′ ,j -1 )tr( Q -ij -Q -i ) and χ j = -( E E ′ ,j -E E ′ ,j -1 )tr( Q -ij ) , we have,

From [18], we have | χ j | ≤ 1 m . We now prove an upper bound on ψ j .

Now look at the term x ⊤ j UQ 2 -ij U ⊤ x j . For any a > 0 and any k > 0 , by Markov's inequality we have,

Let E 1 be an event independent of both x i and x j and have probability at least 1 -δ ′ . Therefore we have E [ · ] ≤ 2 · E E 1 [ · ] . We get,

We now upper bound both terms on the right-hand side separately. Considering the term E E 1 [ | x ⊤ j UQ 2 -ij U ⊤ x j -tr( Q 2 -ij ) | k and using Lemma 3 we get,

Now considering the second term in (9), i.e., E E 1 [(tr( Q 2 -ij )) k ] . We use that conditioned on E 1 , we have Q -ij ⪯ 6 I d . Therefore,

Substituting (10) and (11) in (9),

for some potentially different constant c . Consider k = ⌈ log( m/δ ) log(2) ⌉ and a = 4 · c · k 3 · d · √ k log( d/δ ) and we have with probability at least 1 -δ/m ,

This implies that for an absolute constant c ′ we have,

Therefore we now have an upper bound for | ψ j |

This means for all j , we have,

with probability at least 1 -δ . Consider c j = 4 c ′ · γ m · d · log 3 . 5 ( m/δ ) . Then ∑ m j =1 c 2 j = γ 2 m · 16 c ′ 2 · d 2 log 7 ( m/δ ) . Applying Azuma's inequality (Lemma 8) with λ = γ √ m · 4 c ′ · d · log 3 . 5 ( m/δ ) . We get with probability at least 1 -δ and for potentially different absolute constant c ′ > 0 :

This concludes our proof.

B.1 Proof of Corollary 1

We consider the following Distributed Newton Sketch method:

where g t = ∇ f ( x t ) and ˜ H i = m m -d A ⊤ t S ⊤ i S i A t is a sketched estimate of the Hessian H t = ∇ 2 f ( x t ) , when expressed as the matrix product H t = A ⊤ t A t for the appropriately chosen n × d matrix A t (as can be done for any generalized linear model). This update can be viewed as an approximate Newton step with the Hessian inverse estimate ˆ H -1 t = 1 q ∑ q i =1 ˜ H -1 i . As long as each ˜ H -1 t is an ( √ ϵ, δ/ 2 q ) -unbiased estimator of H -1 t , then using Lemma 34 from [18] we get that the averaged Hessian inverse with q = ˜ O (1 /ϵ ) satisfies:

Using standard approximate Newton analysis (e.g., see Lemmas 1 and 3, and the related discussion in [15]), this implies that when x t is in a sufficiently small neighborhood around x ∗ (determined solely by the strong convexity and Lipschitz constants of the Hessian of f ), we have:

Adjusting the constants appropriately, relying on Theorem 3 for constructing ˜ H i , and following the complexity analysis from Theorem 2, we recover the claim.

C Least squares bias analysis: Proof of Theorem 4

In this section, we aim to prove Theorem 4. Let x ∗ = argmin x ∥ Ux -b ∥ 2 where U ∈ R n × d is the data matrix containing n data points and b ∈ R n is a vector containing labels corresponding to n data points. We adopt the same notations as used in the proof of Theorem 5. Let ˜ x = argmin x ∥ S m Ux -S m b ∥ 2 . Furthermore for any x ∈ R d we can find the loss at x as L ( x ) = ∥ Ux -b ∥ 2 . Additionally, we use r to denote the residual b -Ux ∗ . We aim to provide an upper bound on the bias introduced due to this sketch and solve paradigm, i.e. L ( E (˜ x )) -L ( x ∗ ) . Similar to Theorem 4 we condition on the high probability event E and consider L ( E E (˜ x )) -L ( x ∗ ) . By Pythagorean theorem, we have L ( E E [˜ x ]) -L ( x ∗ ) = ∥ U ( E E [˜ x ]) -Ux ∗ ∥ 2 . Also,

Note that b -UU ⊤ b = b -Ux ∗ = r . We get,

where we used linearity of expectation in the last line combined with S ⊤ m S m = 1 m ∑ m i =1 x i x ⊤ i . Using Sherman-Morrison formula (Lemma 5) we have QU ⊤ x i = Q -i U ⊤ x i 1+ γ m x ⊤ i UQ -i U ⊤ x i . Denote γ i = 1 + γ m x ⊤ i UQ -i U ⊤ x i and substitute we get,

So we get the following decomposition:

Note that a similar decomposition was considered in the proof of Theorem 5 (see Appendix B) with slightly different Z 0 and Z 2 . We first bound ∥ Z 0 r ∥ 2 in the following argument. Without loss of generality, we assume that events E 1 and E 2 are independent of x i and E ′ = E 1 ∧ E 2 .

Now note that in the first term Q -i is independent of x i and x i is also independent of the event E ′ . Using this with the fact that E [ x i x ⊤ i ] = I we get E E ′ [ Q -i U ⊤ x i x ⊤ i r ] = E E ′ [ Q -i U ⊤ r ] = 0 , since U ⊤ r = 0 . Therefore,

The last inequality holds because conditioned on E ′ , we know that ∥ Q -i ∥ ≤ 6 . Using CauchySchwarz inequality we have,

Note that E E ′ [ r ⊤ x i x ⊤ i r ] = ∥ r ∥ 2 . Also, we can bound E E ′ [ x ⊤ i UU ⊤ x i · 1 ¬E 3 ] as following:

Using Chebyshev's inequality, Pr( x ⊤ i UU ⊤ x i ≥ x ) ≤ Var ( x ⊤ i UU ⊤ x i ) x 2 . By Restricted Bai-Silverstein, Lemma 3 with p = 2 , we have Var ( x ⊤ i UU ⊤ x i ) ≤ c · ( d + d 2 log( d/δ ) s ) for some absolute constant c . Let y = m 2 and δ ′ < 1 m 4 we get,

This finishes upper bounding ∥ Z 0 r ∥ 2 as:

Now we proceed with ∥ Z 2 r ∥ 2 ,

Applying Hölder's inequality with p = O (log( n )) and q = 1 + ∆ where ∆ = 1 O (log( n )) , we get,

Since ∥ x i ∥ = poly( n ) , we have ∥ x i ∥ 2∆ = O (1) and therefore we have ∥ Q -i U ⊤ x i ∥ 2∆ = O (1) ∥ Q -i U ⊤ ∥ 2∆ . Also E E ′ ∥ Q -i U ⊤ x i ∥ 2 = E E ′ ∥ Q -i ∥ 2 . Using that conditioned on E ′ we have ∥ Q -i ∥ ≤ 6 , we get ( E E ∥ Q -i U ⊤ x i ∥ 2 q ) 1 / 2 q = O (1) . Similarly using ∥ x i ∥ 2∆ = O (1) we get ( E E ∥ x ⊤ i r ∥ 2 q ) 1 / 2 q = O (1) ∥ r ∥ . This gives us:

Now using (8) from the proof of Theorem 5, we get,

Finally the bound for ∥ Z 2 r ∥ 2 follows as:

Combining (13) and (14) we conclude our proof.

D Higher-Moment Restricted Bai-Silvestein (Lemma 3)

In this section, we prove Lemma 3. We need the following two auxiliary lemmas to derive the main theorem of this section. The first lemma uses Matrix-Chernoff concentration inequality to upper bound the spectral norm of U ⊤ ξ U ξ where ξ is a ( s, β 1 , β 2 ) -approximate leverage score sparsifier.

Lemma 14 (Spectral norm bound with leverage score sparsifier) . Let U ∈ R n × d has orthonormal columns. Let ξ be a ( s, β 1 , β 2 ) -approximate leverage score sparsifier for U , and denote U ξ = diag( ξ ) U . Then for any δ > 0 we have,

Proof. Writing U ⊤ ξ U ξ as a sum of matrices we have,

where p i = min { 1 , sβ 1 ˜ l i d } and Z i = b i p i u i u ⊤ i . Note that Z ′ i s are independent random variables and E [ Z i ] = u i u ⊤ i . Also ∑ n i =1 E Z i = U ⊤ U = I d . If p i = 1 then Z i = u i u ⊤ i and therefore ∥ Z i ∥ = ∥ u i ∥ 2 ≤ 1 . If p i < 1 , we have ∥ Z i ∥ ≤ 1 p i ∥ u i ∥ 2 = d sβ 1 ˜ l i · l i . As l i ≤ β 1 ˜ l i , we get ∥ Z i ∥ ≤ d s . Therefore ∥ Z i ∥ ≤ max { 1 , d s } for all i . Denote R = max { 1 , d s } . We use Matrix Chernoff (Lemma 7) to upper bound the largest eigenvalue of U ⊤ ξ U ξ . For any ϵ > 0 , we have,

With R = max { 1 , d s } and depending on the case whether s ≤ d or s > d we get

Let A ξ denote the event ∥ ∥ ∥ U ⊤ ξ U ξ ∥ ∥ ∥ ≤ 1+ 3 d log( d/δ ) s , holding with probability at least 1 -δ , for small δ > 0 . In the next result we upper bound the higher moments of the trace of U ξ CU ⊤ ξ for any matrix C ⪯ O (1) · I . We first prove the upper bound in the case when the high probability event A ξ does not occur.

Lemma 15 (Trace moment bound over small probability event) . Let k ∈ N be fixed. Let U ∈ R n × d have orthonormal columns. Let ξ be a ( s, β 1 , β 2 ) -approximate leverage score sparsifier for U . Let U ξ = diag( ξ ) U . Also let Pr( A ξ ) ≥ 1 -1 (12 d ) 4 k and event E ′ be independent of the sparsifier ξ . Then we have,

for any fixed matrix C such that 0 ⪯ C ⪯ 6 I .

Proof.

The last equality holds because E ′ is independent of ξ . Consider some fixed y > 0 .

The last inequality holds because by Chebyshev's inequality Pr((tr( U ξ CU ⊤ ξ )) 2 k ≥ x ) ≤ E [(tr( U ξ CU ⊤ ξ )) 4 k ] x 2 . Also note that,

For 1 ≤ i ≤ n , let R i be random variables denoting b i p i u ⊤ i Cu i . Then note that R i are independent random variables with E [ R i ] = u ⊤ i Cu i . Also R i are non-negative random variables with finite (4 k ) th moment. Using Rosenthal's inequality (Lemma 9) we get,

Now ∑ n i =1 E [ R i ] = tr( UCU ⊤ ) and E [ R 4 k i ] can be found as follows: if p i = 1 then R i = u ⊤ i Cu i and therefore R 4 k i = ( u ⊤ i Cu i ) 4 k ≤ ∥ u i ∥ 2(4 k -1) u ⊤ i C 4 k u i ≤ u ⊤ i C 4 k u i , if p i < 1 , we have,

Now using l i ≤ β 1 ˜ l i we get,

Therefore,

Using the above inequality along with using C ⪯ 6 I we get,

Substituting the above bound in (15), it follows that:

For y > (12 d ) 4 k and δ < 1 y

We are now ready to prove the main result of this section. The following result, which is central to our analysis, upper bounds the high moments of a deviation of a quadratic form from its mean.

Lemma 16 (Restricted Bai-Silverstein for ( s, β 1 , β 2 ) -LESS embedding) . Let p ∈ N be fixed and U ∈ R n × d have orthonormal columns. Let x i = diag( ξ ) y i where y i ∈ R n has independent ± 1 entries and ξ is a ( s, β 1 , β 2 ) -approximate leverage score sparsifier for U . Let U ξ = diag( ξ ) U . Then for any matrix 0 ⪯ C ⪯ 6 I and any δ > 0 we have,

for an absolute constant c > 0 .

First, consider T 1 , substitute ξ 2 i = b i p i , and assume exponent p to be even.

For 1 ≤ i ≤ n consider random variables R i where R i = b i p i u ⊤ i Cu i . Furthermore let Y i = R i -u ⊤ i Cu i . Then ∑ n i =1 Y i = ∑ n i =1 ( b i p i -1) u ⊤ i Cu i and E [ Y i ] = 0 . Y i are independent mean zero random variables with finite p th moments and therefore we can use Rosenthal's inequality (for symmetric random variables, Lemma 9) to get,

where A ( p ) is a constant depending on p . We bound T 1 1 and T 2 1 separately, starting with T 1 1 as,

Recall that Y i = R i -u ⊤ i Cu i . R i is always non-negative because C is a positive semi-definite matrix. Therefore we have,

We find the p th moment of R i , E ( R i ) p = ( u ⊤ i Cu i ) p if p i = 1 . If p i < 1 then,

where in the last inequality we use ∥ u i ∥ 2( p -1) = l p -1 i ≤ ( β 1 ˜ l i ) p -1 . Summing over i from 1 to n we get an upper bound for T 1 1 as the following:

Now we upper bound T 2 1 . Note that E ( Y i ) 2 = E [ R i ] 2 = ( u ⊤ i Cu i ) 2 if p i = 1 and if p i < 1 we have E [ R i ] 2 = 1 p i · ( u ⊤ i Cu i ) 2 ≤ d s · 1 β 1 ˜ l i ∥ u i ∥ 2 u ⊤ i C 2 u i ≤ d s u ⊤ i C 2 u i . Summing over i from 1 to n we get an upper bound for T 2 1 as,

Substituting (18) and (19) in (17) we have,

Now we aim to upper bound the term T 2 in (16) i.e., E [ tr( U ξ CU ⊤ ξ ) -y ⊤ i U ξ CU ⊤ ξ y i ] p . First, we condition over ξ and take expectation over y i . This requires using standard Bai-Silverstein inequality (Lemma 10), we get,

where B ( p ) is a constant depending on p . Since y i consists of ± 1 entries we have ν 4 , ν 2 p ≤ 1 . Also using tr( AB ) ≤ tr( A )tr( B ) and considering the high probability event A ξ capturing ∥ ∥ ∥ U ⊤ ξ U ξ ∥ ∥ ∥ ≤ ( 1 + max ( 3 d log( d/δ ) s , 3 log( d/δ ) )) . We get the following,

The first term in the last inequality follows from the Matrix-Chernoff (Lemma 7) and the second term follows from Lemma 15 by considering k = p/ 2 (assuming that δ is small enough so that Lemma 15 is satisfied). We now upper-bound E ( tr( U ξ C 2 U ⊤ ξ ) ) p/ 2 using Rosenthal's inequality for uncentered (non-symmetric) random variables,

For 1 ≤ i ≤ n consider independent random variables R ′ i = b i p i u ⊤ i C 2 u i . We have,

Here R ′ i are positive random variables with finite ( p/ 2) th moment. Using Rosenthal's inequality we get,

It is straightforward to upper bound E ( R ′ i ) p/ 2 as,

Summing over i from 1 to n we get upper bound for T 1 2 as,

Now we consider T 2 2 . It is simply given as,

Combining (23) and (24) and substituting in (22) we get,

Substituting (25) in (21) and let w = max(1 , d/s ) ,

Combining the bounds for T 1 (20) and T 2 (26) substituting in (16), and noting that tr( UC k U ⊤ ) = tr( C k ) for any k , we get,

Now we specify the various constants depending on p . We have A ( p ) ≤ (2 p ) p , B ( p ) ≤ (2 p ) p , C ( p/ 2) ≤ p p/ 2 . Also we use tr( C k ) ≤ 6 k · d since C ⪯ 6 I . This implies for an absolute constant c we have,

where δ > 0 is now arbitrary.

E Application to distributed settings with partitioned data

In this section, we extend our main result to settings where the dataset is partitioned across multiple machines. We illustrate this in the distributed setting considered in [43]. Let A ∈ R N × d and b ∈ R N with rank( A ) = d . In practice, we randomly shuffle the rows of A and b , and partition the data uniformly across q machines, with every machine getting chunks of size n denoted by ( A i , b i ) . In this distributed setup, A i can be considered a uniformly subsampled sketch of A with sketch size n . However, similar to the analysis in [43], in the following analysis we assume that each machine constructs its sketch ( A i , b i ) by uniformly sampling rows from A and b with replacement and independently from other machines. This setting is not ideally a partition but we believe similar results can also be shown for the setting where we partition the dataset after reshuffling the rows. Let µ = N max n i =1 ℓ i ( A ) denote the matrix coherence of A . Let S ∈ R n × N be a uniform sampling matrix with size n , where we will assume that S also scales the samples by √ N/n , which implies that E [ S ⊤ S ] = I . Then the following holds:

Lemma 17 ( Subspace embedding using uniform subsampling; based on Theorem 12 in [43] ) . Let S ∈ R n × N be a uniform sampling matrix of size n where n ≥ O ( µ log( d/δ )) . Then with probability 1 -δ we have,

Let x ∗ = A † b denote the minimizer of least squares loss ∥ Ax -b ∥ 2 and x ∗ i = A † i b i denote the solution to the sketched least squares problem at i th machine. Here A i = S i A and b i = S i b , where S i denotes the uniform subsampling matrix for i th machine. We can use the model averaging result from [43] to prove the following result.

Lemma 18 ( Model averaging result; adapted from Theorem 20 in [43] ) . Let there be q machines and each machine constructs its sketch ( A i , b i ) independently by uniformly sampling n rows from A and b , where n = O ( µ log( dq/δ )) . Then with probability 0 . 99 ,

where ¯ x ∗ = 1 q ∑ q i =1 x ∗ i and c > 0 is an absolute constant 2 .

Since computing x ∗ i at any machine could be potentially expensive due to large n , we instead find ˜ x i at i th machine using the sketching procedure from Theorem 2. Note that this requires applying our debiasing techniques to the least squares problem min ∥ A i x -b i ∥ 2 at every machine. To that end, let ¯ x = 1 q ∑ q i =1 ˜ x i and consider E P ∥ A ¯ x -b ∥ 2 , where E P means conditional expectation given the sketches ( A i , b i ) at every machine. The following theorem states that if the bias from partitioning is much smaller than the sketching error in each machine, then our approach can be successfully used to reduce the bias of the sketching estimate so that averaging will work as desired.

Theorem 6. Let there be q machines and each machine constructs its sketch ( A i , b i ) independently by uniformly sampling n rows from A and b , where n ≥ O ( µ log( dq/δ )) . Also, let i th machine construct ˜ x i via locally sketching A i and b i such that ˜ x i satisfies Theorem 1 (with A i and b i ) at i th machine. Then, the estimator ¯ x = 1 q ∑ q i =1 ˜ x i averaged across the machines satisfies

with probability at least 0 . 9 and an absolute constant c ′ > 0 .

Proof. By the Pythagorean theorem we have:

We proceed to upper bound E ∥ A (¯ x -x ∗ ) ∥ 2

We use Lemma 18 to upper bound the second term and get,

On the other hand,

We upper bound both terms in (30) separately. Proceeding with the first term we have,

The equality (31) holds as all ˜ x ′ i s are independent. The inequality (32) holds due to the subspace embedding property (Lemma 17). The inequality (33) holds due to the variance bound in Theorem 1. The last inequality holds because x ∗ i minimizes the least squares loss ∥ A i x -b i ∥ 2 at i th machine. We now upper bound the second term in (30) as

The inequality (35) holds due to the subspace embedding property (Lemma 17) and inequality (36) holds due to the bias bound in Theorem 1. Therefore we get,

Combining (28, 29, 34, 37) and assuming η < 1 we get,

As E ( 1 q ∑ q i =1 ∥ A i x ∗ -b i ∥ 2 ) = ∥ Ax ∗ -b ∥ 2 , we just use Markov's inequality to upperbound the second term in (38) by 160 q ∥ Ax ∗ -b ∥ 2 with probability 0 . 95 . For the last term, we have

̸

Again 1 q ∑ q i =1 ∥ A i x ∗ -b i ∥ 2 ≤ 20 ∥ Ax ∗ -b ∥ 2 with probability 0 . 95 . The last term is an average of q 2 -q random variables and we can provide uniform bound for the expectation of all of them using the independence among different machines. We have E [ ∥ A i x ∗ -b i ∥ · ∥ A j x ∗ -b j ∥ ] ≤ √ E [ ∥ A i x ∗ -b i ∥ 2 ] · √ E [ ∥ A j x ∗ -b j ∥ 2 ] = ∥ Ax ∗ -b ∥ 2 . Again using Markov's inequality we can upper bound the average of these q 2 -q random variables with 20 times their expectation with probability 0 . 95 , finishing the proof of the theorem.

F Further experimental details

Here, we provide a small set of numerical experiments to empirically examine the relative error of distributed averaging estimates from individual machines to return an estimator ˆ x = 1 q ∑ q i =1 ˜ x i . First, we show that when the sketching-based estimates ˜ x i do have a non-negligible bias, so that the distributed averaging estimator ˆ x remains inconsistent in the number of machines - even with an unlimited number of machines, as long as the space on each machine is limited, the averaged estimator's performance will be limited by the bias of the individual estimates. Second, we show that one can use sketching to compress a data subsample at no extra computational cost, without increasing its bias, which we refer to as the free lunch in distributed averaging via sketching.

We examine three benchmark regression datasets, Abalone (4177 rows, 8 features), Boston (506 rows, 13 features), and YearPredictionMSD (truncated to the first 2500 rows, with 90 features), from the libsvm repository from [8]. All the experiments were run on an i9-13900k processor with 128GB of RAM and an NVIDIA RTX 3090 GPU. We repeat the experiment 100 times, with the shaded region representing the standard error. We visualize the relative error of the averaged sketchand-solve estimator L (ˆ x ) -L ( x ∗ ) L ( x ∗ ) , against the number of machines q used to generate the estimate ˆ x = 1 q ∑ q i =1 ˜ x i .

Each estimate ˜ x i was constructed with the same sparsification strategy used by LESS, except that instead of sparsifying the sketch with leverage scores, we instead sparsify them with uniform probabilities. Following [16], we call the resulting method LESSUniform. Within each dataset, we perform four simulations, each with different sketch sizes and different numbers of nonzero entries per row. We vary these so that the product (sketch size × nnz per row) stays the same, so as to ensure that the total cost of sketching is fixed in each plot.

As expected, decreasing the sketch size while increasing the number of nonzeros per row (effectively increasing the amount of 'compression' occurring here by sparse sketching) increases the error in all three datasets. However, remarkably, increasing the amount of 'compression' does not seem to increase the amount of bias. We can therefore conclude that sparse sketches preserve nearunbiasedness, while enabling us to reduce the sketch size from subsampling without incurring any additional computational cost. The increase in error can be mitigated in a distributed setting by increasing the number of estimates/machines.

F.1 Comparison with other sketching estimators

To further illustrate the phenomenon that suitable sketched least squares estimators enjoy small bias, we provide a further set of numerical experiments below in Figure 4 for the Boston and Abalone

Boston

Sketchsize=340,nnzperrow=1

100

1

Loss

1

Sketchsize=30,nnzperrow=13

10-1

Squares

10-2

Loss/Least

10-3

100

101

102

103

104

Number of Machines

datasets. Recall that in Figures 2 and 3, we show that increasing the amount of 'compression' of the sample, by reducing the sketch size when using sparse sketching matrices with uniform probabilities, does not increase the bias of the sketched estimator. Figure 4 further demonstrates that Gaussian and Subgaussian sketches, which are computationally much more expensive than our proposed sparse sketches, exhibit virtually no bias in sketched least squares (as expected). A similar conclusion can be made for the proper LESS method, which uses leverage score estimates, and is therefore more expensive than the LESSUniform method we use in other experiments, but still potentially much cheaper than Subgaussian sketches.

However, the unbiasedness does not hold for all sparse sketches, as we demonstrate below in Figure 4. At a high level, we show that sketches constructed using leverage score subsampling and the SRHT can exhibit a non-negligible level of least squares bias.

Gaussian and Subgaussian sketches, as we see in the figure, enjoy near-unbiasedness while not introducing additional bias as the sketch size decreases. This does not hold for leverage score subsampling. When we decrease the number of subsamples within leverage score subsampling, the bias introduced by subsampling increases as suggested by the lower bound in Theorem 10 of [18]. Intuitively, this happens as the number of subsamples is reduced without increasing the amount of 'compression' as one would with LESS or LESSUniform.

We also show that the subsampled randomized Hadamard transform (SRHT) also introduces increased bias as the sketch size decreases. This complements the lower bound established for the failure of SRHT and other data-oblivious sparse embeddings like CountSketch to satisfy the restricted Bai-Silverstein property that is a structural condition for near-unbiasedness sketches [18].

For completeness, we also demonstrate the desirable performance of LESS proper. Figure 4 demonstrates that the desirable phenomenon that LESSUniform enjoys also extends to LESS proper. In fact, we observe that LESS enjoys similar desirable performance as the Gaussian and Subgaussian sketches and does not increase (in fact, it has minimal) least squares bias, while enjoying the computational speedups that sparse sketches also enjoy, suggesting a best-of-both-worlds property.

Boston - Leverage Score Subsampling

Abalone-LeverageScore Subsampling

Loss/Least Squares Loss

-1

10-1

10°

-·Subsamples=170

Subsamples=340

Subsamples=70

Loss/Least Squares Loss - 1

10-1

10°

Subsamples=60

Subsamples=30

Subsamples=130

Subsamples=20

Subsamples=30

10-2

10-2

10-

10-3.

10-4

10-5

10-5

10-6.

10-6

100

101

102

103

104

105

100

101

102

103

104

105

Number of Machines

Number of Machines

BostonDataset-SRHTSketch

AbaloneDataset-SRHTSketch

10-1

Sketch Size=340

10-1

Sketch Size= 70

T-sson sasenbs isean/sson

Sketch Size= 32

Sketch Size=16

10

10-2

10-3

10-3

j1111 1

10-4

10-5

101

102

103

104

101

102

103

104

Number of machines

Number of machines

BostonDataset-GaussianSketch

AbaloneDataset-GaussianSketch

10-1

Sketch Size=340

10-1

Sketch Size= 70

Sketch Size= 32

10

10-2

Sketch Size=16

10

-3

10

-3

10-4

10-4

i......

101

102

103

104

101

102

103

104

Number of machines

Number of machines

BostonDataset-SubgaussianSketch

AbaloneDataset-SubgaussianSketch

10-1

Sketch Size=340

10-1

T-ssot sasenbs asean/sso

Sketch Size= 70

-sson sasenbs isean/sso

Sketch Size= 32

10

-2

Sketch Size=16

10-²

10-3

10-4

10-5

10-5

101

102

103

104

101

102

103

104

Number of machines

Number of machines

Boston-LESS

Abalone-LESS

100

Sketch size=340, nnz per row=1

100

Sketch size=130, nnz per row=1

Squares Loss-

Sketch size=60, nnz per row=2

10-1

10

-1

Sketch size=30, nnz per row=13

Sketch size=20,nnz per row=8

10-2

10

-2

10-3

10-3

Loss/Least S

10-4

10-5

10-5

100

101

102

103

104

100

101

102

103

104

Number of Machines

Number of Machines

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The paper provides concise mathematical proofs for the claims made in the abstract.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [NA]

Justification: The results proved in the paper hold in a very general setting, without any limitations or requiring strong assumptions.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All theoretical results have been proved. Most of the formal proofs are part of the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The experiments demonstrate the results for solving a few least-squares problems using the Matrix sketching technique described in the paper. All the experimental results are easily reproducible.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: A supplementary zipfile containing the code is provided to reproduce the experimental results. One of the experiments requires the file 'YearPredictionMSD.txt' file, which can be obtained at ( https://archive.ics.uci.edu/dataset/ 203/yearpredictionmsd ).

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The paper describes the parameters e.g. sketch size, sparsity of sketch (nnz per row), and number of estimates(number of machines) along with dataset details to understand the results.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: The shaded areas in our plots denote the standard error obtained after averaging the result over 100 trials.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All experiments were run on an i9-13900k with 128GB of RAM and a NVIDIA RTX 3090 GPU.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This work aligns with the NeurIPS code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our theoretical work provides a Matrix sketching algorithm, recovering a nearly unbiased least squares estimator. Solving linear systems approximately is an important component of a large number of algorithms from varied domains. Our work can thus provide a way to recover approximation solutions with a smaller bias. This does not lead to any negative societal impact.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This work does not release any data or model, and carries no risk of misuse.

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We test our theoretical results on a few publicly available datasets that have been cited in the paper as well.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: No new assets are released in this work.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This work does not involve crowdsourcing or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This work does not involve crowdsourcing or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSketching is a technique from randomized numerical linear algebra which compresses an input matrix $A$ by multiplication with a random matrix $S$. Recent works [18] have shown how to characterize the bias in the sketched inverse covariance $(\\tilde A^\\top\\tilde A)^{-1}$, where $\\tilde A = SA$ for an input matrix $A$ and a sketching matrix $S$. This work extends this machinery to construct approximately unbiased estimators $\\tilde x$ for the least squares regression problem $\\min_x \\|Ax-b\\|$. As the authors show in Theorem 2, this gives an efficient distributed algorithm by averaging the unbiased estimator over $q = O(1/\\epsilon)$ servers. The techniques introduced by the authors also improve prior results on bias-free estimation of the inverse covariance.\n\nSTRENGTHS:\nThe authors introduce a new way to exploit unbiased estimators, which appears to be different from prior works, which incorporate debiasing in a second order optimization framework (e.g. https://arxiv.org/abs/2007.01327). The new technique simply averages local solutions, and conceptually seems simpler than prior techniques.\n\nTechnically, the authors improve prior results on bias-free estimation of the inverse covariance by showing an improved moment bound on a certain random variable by carefully exploiting better bounds for smaller moments via a Holder’s inequality. Other new tools from random matrix theory are additionally needed to handle this change. These ideas, as well as their improvements to bias-free estimation of the inverse covariance, are interesting and may be useful for future results in this area.\n\nWEAKNESSES:\nThe result and techniques seem to be a bit niche and incremental. The notion of distributed computation is different from ones that I see in other works in the sketching literature (e.g. https://arxiv.org/abs/1408.5823, https://arxiv.org/abs/1504.06729) where the input is partitioned across multiple servers. Instead, this work considers a setting where there is just one stream which contains the entire input, and the central server has access to $q$ machines that can each access this stream (see the Computational model section), which seems nonstandard. The number of servers needed is also rather large (e.g. $q = O(1/\\epsilon)$), which also seems restrictive. I encourage the authors to include references that prove results in this setting if there are any others. The main technical novelty (line 275) also seems to be an improvement to a bound in the analysis of LESS embeddings from prior work, while the overall sketching framework is largely unchanged.\n\nQUESTIONS:\n- The authors may consider discussion the related work https://arxiv.org/abs/2203.09755 on distributed least squares and https://arxiv.org/abs/2007.01327 on using debiasing approaches\n\nOther comments\n\n- Lemma 2 should have a probability statement (e.g., with probability at least $1-\\delta$)\n\nLIMITATIONS:\nThe authors claim that the work does not have limitations in the checklist. Some of the points I raised in the weaknesses section could be discussed in the work, such as restrictions on the computational model and the number of servers.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSketched least squares involve estimating the term $(X^TX)^{-1}$ which has a high bias when the sketch matrix $S$ is not sub-Gaussian. This paper gives a sparse sketching method using a LESS embedding which runs in optimal space and current matrix multiplication time, where $S$ is sparse, and constructed based on the leverage scores of the data matrix $A$ (Definition 2 $(s,\\beta_1,\\beta_2)$-LESS embedding). The paper also improves the sharpness of the probability bounds which is applicable to similar problems in RMT using LESS embeddings. For $s=1$, nothing is different, but when $s > 1$, the bias bound is reduced.\n\nSTRENGTHS:\nInversion bias for estimates of $(X^TX)^{-1}$ is a challenging problem in least-squares sketching. Sub-gaussian sketches have high computational cost but low bias ; other sketches have low computational cost but high bias. The sketch proposed in this paper minimizes the bias of the estimator only requires 1 parallel pass over the data, and has runtime of nnz($A$) + $\\tilde{O}(d^2/\\epsilon)$ which is much faster. Moreover, computing the estimate can be done in parallel, with only the final result needing to be averaged. \n\nThe authors rigorously prove that their sketch fulfills the above criteria by using techniques from random matrix theory, and this is not trivial at all. I find the key contribution in the paper is proving these statements, and the authors did an excellent job here. The paper is extremely well written and easy to follow. Of particular note is the bias analysis for the least squares estimator (Section 4) which is concisely written, where the main ideas (proof sketch) given in the main paper, and technical details in the appendix. The ideas given in the proof sketch can be used for similar problems, and are straightforward to understand. Moreover, the technical details come with sufficient exposition such that it is straightforward for a reader to understand the direction the proof is going (which is certainly much appreciated).\n\nWEAKNESSES:\n1. It would be nice to have experiments with other sketching methods to (empirically) justify some statements on the bias, variance (although not stated) and computational time, e.g. estimators mentioned in Table 1. For example, are there sketching estimators with higher bias, but less variance? \n\nThere are some minor typos, e.g. line 212 reference missing, lines 282 to Equation (2) at bottom of the page is missing a bracket for the numerous expectations (in contrast to lines 542 onwards in the appendix).\n\nIn Appendix A, notation for concentration inequalities should be looked at and made consistent, e.g. Lemma 6 / H{\\\"o}lder's inequality should have a $\\frac{1}{q}$, Lemma 7 ($\\lambda$ max isn't defined), Lemma 8 / Azuma's inequality ($\\lambda$, $m$ should be consistent), Lemma 10 is missing a bracket for $\\mathbb E[x_i^2]$. \n\nThe presentation of the proof for Theorem 5 was slightly jarring (due to Lemma 11, Lemma 12 appearing in the proof), but there also doesn't seem to be a good way to include them (since referring to the two lemmas requires the upper bounds, and flipping back a page is also inconvenient). Maybe a solution is to indent the Lemmas, or box them up?\n\nThe dot before line 626 (after 72) should be removed.\n\nQUESTIONS:\n1. It would be nice to have experiments with other sketching methods to (empirically) justify some statements on the bias, variance (although not stated) and computational time, e.g. estimators mentioned in Table 1. For example, are there sketching estimators with higher bias, but less variance? \n\n2. Despite being clear to read, I had to go back and forth a bit to find out what the novelty is. I appreciate the clarity and thoroughness of explaining the bounds on the bias and variance, runtime, but I would like it if Definition 1 & 2 came much earlier (or at least maybe an informal Definition 2 after line 85?)\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-02"}

comments: SUMMARY:
The paper studies the least squares regression task and improves the space and communication amount in a distributed setting to be independent of $\epsilon$. The key to achieve this is by sketching the data in blocks that have an $\epsilon$ dependence which can be reduced to only d-dependencies by aggregating them. For the covariance A^TA this becomes only a d x d matrix and for the other required term A^Tb it can be handled by communicating and aggregating just the solutions which have a dimension of d. Although the idea seems very simple described this way, the analysis seems highly non-trivial and has very interesting aspects and novel techniques (or at least new to me). It analyzes in this setting not only the standard least squares error but also a bias term, which allows for smaller sketch dimensions, and only 1/sqrt(eps) dependencies.

STRENGTHS:
* gives least squares sketching results with lower time, space, and communication complexities
* interesting techniques
* some further applications are given, though details are fully in the appendix
* very good writing

WEAKNESSES:
* motivation of the model, seems to be a niche
* very low improvements but there seems no big gap that can be leveraged

QUESTIONS:
* I am slightly confused by the definition of the bias. Is $||AE(\tilde x)-b||$ standard in some literature? I would say this is the variance of the expected sketched estimator. Why would I be interested in this expected estimator, instead of the actual outcome after sketching?
* I would rather define the bias as $||\tilde x -x^*||$ which I think does not allow for any improvements over standard sketching results, right?

LIMITATIONS:
f
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents new techniques for distributed least squares regression using matrix sketching. The key contributions are:\n\n1. A sparse sketching method that produces a nearly-unbiased least squares estimator in two passes over the data, using optimal space and current matrix multiplication time.\n2. Improved communication-efficient distributed averaging algorithms for least squares and related tasks.\n3. A novel bias analysis for sketched least squares, characterizing its dependence on sketch sparsity. This includes new higher-moment restricted Bai-Silverstein inequalities.\n\nThe theoretical results are backed by experiments on real datasets showing the practical benefits of the approach.\n\nSTRENGTHS:\n1. Provides a sparse sketching method that achieves near-unbiased least squares estimation in optimal space and current matrix multiplication time.\n2. Achieves O(d^2 log(nd)) bits of space, which is optimal. Matches current matrix multiplication time O(d^ω), improving over previous approaches.\n3. Introduces a new bias analysis for sketched least squares that sharply characterizes dependence on sketch sparsity. The techniques developed may be applicable to other sketching and randomized linear algebra problems.\n\nWEAKNESSES:\n1. Experiments are conducted on only a few datasets. And it does not explore a wide range of problem sizes or distributed computing scenarios.\n2. Primarily focused on least squares regression, with limited discussion of extensions to other problems.\n3. Some of the theoretical results rely on assumptions (e.g., about leverage score approximation) that may not always hold in practice.\n\nQUESTIONS:\n1. Are there natural extensions of this work to other loss functions beyond least squares?\n2. Do you expect the techniques developed here to be applicable to other sketching problems beyond least squares? If so, which ones?\n\nMinor Comments:\n\n1. The abstract could more clearly state the key theoretical results/bounds achieved\n2. Some additional discussion of practical implications and potential applications would be valuable\n3. A few typos noted (e.g. line 330)\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}