Linear Regression using Heterogeneous Data Batches

We thank the reviewer for many positive comments about our paper. Below we address the remaining comments of the reviewer one by one.

1. Note that even when $k$ is arbitrarily large, the number of batches and batch sizes required in Theorem 2.1 and Corollary 2.2 are reasonable as in that case they depend on $\alpha$ instead of $k$, implying that batches can be recovered even for arbitrary large $k$. We also provide numerical experiments to further confirm this fact (see Figure 4 in the appendix.) Our main observation that helps in such a recovery for large $k$ is that even when $k=\infty$, a small subspace can be estimated that contains gradients for all components with a significant fraction present.

2. As discussed in line 112 of section 1.2, our sample complexity is lower than Kong et al. (2020). Our time complexity $\tilde O(|B_s| d^2/\alpha_m)$ is at most $\tilde O(1/\alpha_m)$ times the complexity of Kong et al. (2020). But as is clear from the value and our numerical experiments show, this complexity is still reasonable in practice.

3. Note that of all previous works, only Kong et al. (2020) take advantage of the batch structure. All other works mentioned in the paper do not, and hence they will naturally fare worse than the baseline we considered. For that reason we did not include them in our results.

4. We note that in Figures 1 and 2, $k$ is kept constant at 16. The small gap between the two algorithms for small values of medium batch size is because if it is too small then it falls below the requirement of Theorem 2.1. In that case, both our algorithm and baselines fail to recover the regression vector.

We thank the reviewer for the careful review and positive comments. The reviewer’s comments are addressed herein.

1. We thank the reviewer for pointing out the typos, and will correct them along with some others we found.

2. The condition part of Theorem 2.3 can be more clearly stated as follows: Let index $i$ be in set $I$ and let list $L$ be a list of $d$ dimensional vectors. Suppose for a given $\beta>0$ at least one of the vectors in $L$ is within distance $\beta$ from the regression vector of distribution $\mathcal D_i$, namely $\min_{w\in L} ||w-w_i||\le \beta$.

3. We thank the reviewer for the reference to the recent work on mixed linear regression. It is clearly relevant, and we will include it in our paper.

4. (see next point)

5. Addressing 4 and 5 together, note that, as stated in corollary 2, the $L$ identified by our algorithm is at most $\tilde O(1/\alpha_m)$, even for arbitrary large $k$. Furthermore, if all $k$ sub-populations follow the assumption in Theorem 2.1 the argument preceding corollary 2.2 can be modified as follows to show that the list $L$ identified by our algorithm can be trimmed down to have at most $k$ elements. In such a case, for any choice of $b^*$ from $B_m$ to run algorithm 1, the regression vector estimates $\hat w$ will satisfy $||\hat w-w_i||\le \epsilon \sigma$ for some $i\in [k]$. A simple argument using the triangle inequality shows that removing the estimates from $L$ until no two estimates are within distance $2\epsilon\sigma$ from the list $L$ will result in a list $L$ of size $\le k$ such that for all $i$ we have $||w_i-\hat w||\le 3\epsilon \sigma$. Also, we note that we already included the literature on list decoding. We referred to a recent work of list decodable linear regression using batches [DJKS22] in section 1.2 and other references in the appendix. We can discuss these connections in more detail as per the reviewer’s suggestion.

6. The input distributions need not be known. We will mention it clearly in the contributions following your suggestion.

7. If for two distributions the regression vectors have no separation then when the algorithm is run using a medium-sized batch as $b^*$ from either of two distributions we will be estimating the same regression vector. The upper bound on the size of the output list is independent of the separation.

8. As shown by the synthetic experiments in Figures 1-4, the algorithm is practical even for reasonable-size datasets and dimension $d$. As such there are no barriers to trying it out on real data. We leave it to future work to find interesting and relevant datasets for this setting.

We thank the reviewer for many positive comments. The reviewer’s main concerns involve the number of parameters and distributional assumptions, both addressed below.

While the number of parameters may seem large, each has a natural and important contribution.

Two accuracy parameters $\epsilon$ and $\delta$ that are part and parcel of any PAC learning algorithm. They specify the accuracy we are looking for and determine the number of samples needed to achieve it.

Two regression parameters: the length $M$ of the regression vector, and the additive noise variance $\sigma$, reflect the complexity of the input-output relation.

Two population parameters: the number $k$ of sub-populations and the fraction of batches $\alpha$ present from the relevant sub-population. Note that the algorithm combines these parameters to set one parameter $\ell$.

Finally, two distribution parameters: the L4-L2 hypercontractivity parameter $C$ and the condition number of input covariance matrix $C_1$. These parameters reflect the complexity/richness of the distribution class and substantially generalize input distributions considered in the past for mixed linear regression, as they only considered sub-gaussian distributions with the identity covariance matrix.

The L4-L2 hypercontractivity parameter significantly generalizes the sub-gaussian distribution assumption and allows for heavier-tailed distributions. The condition-number $C_1$ of the input covariance matrices significantly generalizes the identity covariance matrix assumption for input distribution that corresponds to $C_1=1$.

Though the number of parameters seems large, the algorithm needs only an upper bound on their values. If the upper bounds are tight, they will result in a good estimate. If for any parameter we do not have a tight upper bound, we can run the algorithm for different values of the parameters, dividing their possible range logarithmically, and choose the parameters that achieve a small error on a holdout set. We can add this approach in more detail in the final version of the paper and formally define error computation and selection for our setting.

Furthermore, it is easy to see that these upper bounds are needed. For example, even for the simplest case when all samples are from the same sub-population, some condition on the anti-concentration of $XX^T$ as above is needed to obtain finite sample complexities [Lecué and Mendelson, 2016, Oliveira, 2016] in heavy-tailed linear regression, and this corresponds to a low condition number.

The reviewer also asks about the intuition behind Remark 2.1. Our algorithm and its analysis are based on gradient descent, and can handle stochastic noise. Therefore we require only the distribution of gradients to be similar for batches in the same sub-population and since the analysis can tolerate statistical noise it can also tolerate deviation in the distribution of gradients across batches as long as it is small.

We thank the reviewer for the positive review, the appreciation of our contribution, and the two suggestions that we discuss next.

The first suggestion concerns rearranging Sections 1.1 and 1.2. Section 1.1 discusses the paper’s main results, and section 1.2 addresses the results’ improvement over prior work. We agree with the reviewer that the comparison in Section 1.2 results in some duplication of results outlined in Section 1.1, and will rearrange some of the material to minimize the repetition.

The second comment concerns R-way partitioning of the data. More concretely, our algorithm runs in R rounds, and our theoretical analysis requires the use of independent samples in each round.

We thank the reviewer for mentioning that Duchi’s work might help eliminate the need for this R-way partitioning in our analysis either entirely or under appropriate constraints.

In searching papers on Ergodic Gradient Descent by Duchi et. al. we found the paper “Ergodic Mirror Descent” that we believe the reviewer referred to. While this paper performs gradient descent using ergodic samples, and hence generalizes the independent samples we use, it still does not eliminate the need of new samples in each round of gradient descent. In the limited rebuttal period, we could not find a way to extend the analysis in the paper to allow using the same samples in each algorithm round, and we leave that for future work.

Note however that to achieve estimation accuracy $\epsilon \sigma$ with noise variance $\le \sigma^2$, and regression vectors of length at most $M$, the number of rounds $R$ required by the algorithm is just $\log(M/\sigma)$. Hence the factor $R$ increase in the sample complexity implied by our theoretical analysis is at most a logarithmic factor. Furthermore, in our numerical experiments, the algorithm recovers the regression vectors even when the same samples are used in each round. However, we agree with the reviewer that without a theoretical proof it is unclear if it will not be necessary for some extreme datasets.

Thank you for the careful review and constructive feedback, and for appreciating the contribution, novelty, presentation clarity, and intuition.

Regarding your comments and questions:

Line 123: Section 1.3 summarizes the algorithm’s innovations, and we will elaborate on how they help overcome assumptions in prior works.

Line 332: We will break it into several sentences so it is easier to follow.

Line 836: We will add the missing reference.

Next, we answer the two questions raised in the review.

For an example motivating the number of small batches, observe that in many recommendation systems, most users rate only a few items. The ratings by any such user may serve as a small batch. Since these systems have many users, the number of small batches will also be large. Observe also that typically relatively few users provide a fair number of ratings, and they can serve as medium-sized batches.

For a test justifying that $k$ can be large, Figure 4 in the appendix shows an experiment where the number $k$ of subpopulations is large (100), while the number of subpopulations with sufficiently many batches for enabling recovery is small (4). Our theoretical analysis shows that this experiment can be repeated with even larger $k$, without impacting the algorithm’s performance.