Distributed Least Squares in Small Space via Sketching and Bias Reduction

Summary

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

Strengths

The 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. Technically, 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.

Weaknesses

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

Questions

- 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 Other comments - Lemma 2 should have a probability statement (e.g., with probability at least $1-\delta$)

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

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

Summary

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

Strengths

Inversion 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. The 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).

Weaknesses

1. 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? There 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). In 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]$. The 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? The dot before line 626 (after 72) should be removed.

Questions

1. 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? 2. 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?)

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

N/A

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?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

f

Summary

This paper presents new techniques for distributed least squares regression using matrix sketching. The key contributions are: 1. 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. 2. Improved communication-efficient distributed averaging algorithms for least squares and related tasks. 3. A novel bias analysis for sketched least squares, characterizing its dependence on sketch sparsity. This includes new higher-moment restricted Bai-Silverstein inequalities. The theoretical results are backed by experiments on real datasets showing the practical benefits of the approach.

Strengths

1. Provides a sparse sketching method that achieves near-unbiased least squares estimation in optimal space and current matrix multiplication time. 2. Achieves O(d^2 log(nd)) bits of space, which is optimal. Matches current matrix multiplication time O(d^ω), improving over previous approaches. 3. 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.

Weaknesses

1. Experiments are conducted on only a few datasets. And it does not explore a wide range of problem sizes or distributed computing scenarios. 2. Primarily focused on least squares regression, with limited discussion of extensions to other problems. 3. Some of the theoretical results rely on assumptions (e.g., about leverage score approximation) that may not always hold in practice.

Questions

1. Are there natural extensions of this work to other loss functions beyond least squares? 2. Do you expect the techniques developed here to be applicable to other sketching problems beyond least squares? If so, which ones? Minor Comments: 1. The abstract could more clearly state the key theoretical results/bounds achieved 2. Some additional discussion of practical implications and potential applications would be valuable 3. A few typos noted (e.g. line 330)

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

yes