Distributed Least Squares in Small Space via Sketching and Bias Reduction

We thank the reviewer for the comments and questions. We are glad that the reviewer appreciates the simplicity of our algorithmic approach, and that the reviewer finds our technical contributions interesting and useful. Below, we provide clarifications on the reviewer's remarks, including on the distributed computation framework and the number of machines, as well as providing a comparison with the four references provided by the reviewer. We will be sure to include all of those details in the final version. If you think the responses adequately address your concerns, we encourage you to consider increasing your score.

• Novelty in our techniques. Our main contribution is in the theoretical analysis, which includes several new ideas (including the higher-moments version of Restricted Bai-Silverstein inequality, as well as the careful use of H"older's inequality in the analysis of the dominant term). In fact, these ideas are not niche, since they are relevant to many instances of RMT-style analysis (i.e., analysis relying on the Stieltjes transform of the resolvent matrix) for sparse sketching operators, which has been used for Newton Sketch [16], Randomized SVD [17], and Sketch-and-Project [21]. We chose to focus on least squares, as this gives the clearest computational improvements.
• Number of machines required. The use of $O(1/\epsilon)$ machines in our main results comes from the small space complexity, namely $O(d^2\log nd)$ bits per machine, that we impose in our setting. However, if we relax this constraint, then there is a natural trade-off between the space complexity and number of machines required by our methods: If we allow $((1/\theta)d^2\log nd)$ bits per machine for some $\theta\in [\epsilon,1]$, then we only need $O(\theta/\epsilon)$ machines. Setting $\theta=\epsilon$, we recover the standard sketched least squares on a single machine, while $\theta=1$ recovers our main results, but of course we can freely interpolate between those two extremes.
• Comparison with [BP22, arXiv:2203.09755]. This work indeed considers essentially the same problem setting as we do. In fact their Algorithm 1 is the same basic distributed averaging procedure that we use, and their computational model described in Remark 2.3 (Option 2) matches our single-server multiple-machines setup (except that they allow random access instead of streaming access to the data). Our results can be viewed as a direct improvement over their guarantees for Gaussian sketches (e.g., their Theorem 2.2), since we obtain analogous guarantees for extremely sparse sketches which are far more computationally efficient. In fact, we already mention these Gaussian sketch guarantees as a baseline for our work (see Table 1 for a comparison), and we will certainly add the reference to this work in that context.
• Comparison with [DBPM20, arXiv:2007.01327]. This work also considers a very similar problem setting to us, except they focus on regularized least squares, and how the regularization parameter affects the bias of the sketched least squares estimator, with distributed averaging as a main motivation (their main results also assume that all workers have access to the centralized data for the purpose of sketching). Similarly as for the earlier reference, their theoretical results require expensive sketching methods (called "surrogate sketches"), which are based on Determinantal Point Processes. We use these sketches as one of our theoretical baselines (see Table 1 for a comparison), and we will certainly add this work as a reference. It is worth noting that our random matrix theory techniques can likely be used to extend the regularization-based debiasing techniques developed in that work to fast sparse sketching. We leave this as a promising direction for future work.
• Comparison with [BKLW14, arXiv:1408.5823] and [BWZ16, arXiv:1504.06729]. The distributed computation framework considered in those works partitions the data across multiple servers. Importantly, note that these works solve the task of Principal Component Analysis, which is different from our task of least squares regression, so these works are not directly comparable. Nevertheless, our methods and results can be naturally extended to the multiple-server setting, as we discuss below. The main reason we focused the paper on a more centralized computational model, with one server and multiple machines, is because this allowed us to obtain worst-case results which are independent of condition-number type quantities (such as the data coherence defined below).
• Results for data partitioned into multiple servers. As mentioned above, our results can be extended to the setting where a dataset is partitioned into multiple servers, so that each one of the $q$ machines is accessing a separate chunk of the data. Suppose that a dataset $A,b$ of size $N\times d$ is partitioned uniformly at random into smaller chunks of size $n=N/q$, and each machine constructs an estimate $\tilde x_i$ based on a sketch of its own chunk. Then with the same computational guarantees on each machine as in Theorem 2, the averaged estimator $\tilde x=\frac1q\sum_i\tilde x_i$ will with high probability enjoy a guarantee of $||A\tilde x-b||\leq (1+\epsilon + \tilde O(\mu/n))||Ax^*-b||$, where $\mu=N\max l_i(A)$ is the coherence of the dataset (here, $l_i$ denotes the $i$th leverage score). The additional error term $\tilde O(\mu/n)$, which arises from the partitioning (e.g., see [43]), is often negligible compared to the sketching error $\epsilon$ when the chunk size $n$ is sufficiently larger than the sketch size. An analogous result can be obtained for Distributed Newton Sketch. We will add these claims to the final version.

We thank the reviewer for a careful read and detailed comments. We will address all of the typos and presentation suggestions in the final version.

• Additional experiments. We include additional experiments with other sketching methods (all of which are more computationally expensive than the fast sparse LESSUniform method we used in the paper): Leverage Score Sampling, Gaussian and Subgaussian sketches, as well as Subsampled Randomized Hadamard Transform (see our general response and the PDF). We note that the fact that sketching methods tend to yield smaller least squares bias than uniform subsampling has been empirically observed in prior works [42], which is why we did not focus on this here. Our main contribution is to provide the first sharp theoretical characterization of this phenomenon for extremely sparse sketches.
• Are there sketching estimators with higher bias, but less variance? By definition, the bias has to be no larger than the variance, but there can be cases where the two quantities are comparable in size (which implies that distributed averaging will not be effective). The main example here is i.i.d. sub-sampling, including leverage score sampling. There are theoretical lower bounds [18] which show that (in some cases) the bias of leverage score sampling may not be much smaller than the variance.

We thank the reviewer for comments and feedback. We will revise the abstract, and also expand our discussion of applications beyond least squares as outlined below.

• Extensions to other loss functions beyond least squares. In addition to least squares, we provide a more broadly applicable result in Theorem 3, which in particular applies to minimizing general convex loss functions via the framework of Distributed Newton Sketch (see Corollary 1). Thus, our sketching methods and proof techniques are of broader interest to sketching-based optimization algorithms for a variety of loss functions, including the logistic loss, other generalized linear model losses such as the hinge loss, etc.
• Applications to other sketching problems. Our theoretical analysis, which includes several new ideas (including the higher-moment version of Restricted Bai-Silverstein inequality, as well as the careful use of H"older's inequality in the analysis) is relevant to many instances of RMT-style analysis (i.e., analysis relying on the Stieltjes transform of the resolvent matrix) for sparse sketching operators. This RMT-style analysis has been used in Newton Sketch [16], Randomized SVD [17], and Sketch-and-Project [21]. We chose to focus on distributed least squares, as this gives the clearest worst-case computational improvements.
• Distributed computing scenarios. In fact, our methods and results are applicable much more generally than the single-server multiple-machine computation model used in the paper. For example, they can be naturally extended to the multiple-server model [43], where the data is randomly partitioned into multiple chunks stored on separate servers, which is common in the literature (see response to Reviewer h7iG for details). The main reason we focused the paper on the single-server multiple-machine model is because this allowed us to obtain worst-case results that are independent of condition-number type quantities (those are unavoidable in the multiple-server model, given our other computational constraints).
• Experiments. Our main contribution is to provide the first sharp theoretical characterization of the least squares bias for extremely sparse sketches, and this is also where we focused our experiments. Nevertheless, we include additional experiments (see the general response and PDF) with other sketching methods (all of which are more computationally expensive than the fast sparse LESSUniform method we used in the paper): Leverage Score Sampling, Gaussian and Subgaussian sketches, as well as Subsampled Randomized Hadamard Transform.
• Assumptions about leverage score approximation. Our main results, Theorems 2 and 3, do not require any assumptions related to leverage score approximation, as they include leverage score approximation as part of the algorithmic procedure.

We thank the reviewer for the positive feedback, as well as the comments and questions. We will address them in the final version.

• Definition of bias. At a high level, we rely on the statistical definition of the bias of an estimator, which is (informally) the difference between the expectation of that estimator and its target quantity (the estimand). In the case of a least squares estimator $\tilde x$, the most standard statistical notion of bias would be $E[\tilde x]-x^*$. Since the estimator is multivariate, one will typically measure the amount of bias by taking the norm, i.e., $||E[\tilde x]-x^*||$ (as opposed to $||\tilde x-x^*||$ which is typically referred to as the estimation error). In our case, since we are interested in the least squares prediction vector, i.e. $A\tilde x$, as an estimator of the vector $b$, we compute the bias as $||E[A\tilde x]-b||=||A E[\tilde x]-b||$. This turns out to be the exact right notion for the purpose of distributed averaging and bounding the least squares loss.
• Motivation of the model. Our model is motivated by the general distributed averaging framework (also known as model averaging, or bagging), which is widely used in many settings beyond least squares. In fact, our methods and results are applicable much more generally than the single-server multiple-machine computation model used in the paper. For example, they can be naturally extended to the multiple-server model [43], where the data is randomly partitioned into multiple chunks stored on separate servers, which is common in the literature (see response to Reviewer h7iG for details). The main reason we focused the paper on the single-server multiple-machine model is because this allowed us to obtain worst-case results that are independent of condition-number type quantities (those are unavoidable in the multiple-server model, given our other computational constraints).