Linear Regression using Heterogeneous Data Batches

In many learning applications, data are collected from multiple sources, each providing a \emph{batch} of samples that by itself is insufficient to learn its input-output relationship. A common approach assumes that the sources fall in one of several unknown subgroups, each with an unknown input distribution and input-output relationship. We consider one of this setup's most fundamental and important manifestations where the output is a noisy linear combination of the inputs, and there are $k$ subgroups, each with its own regression vector. Prior work~\cite{kong2020meta} showed that with abundant small-batches, the regression vectors can be learned with only few, $\tildeΩ( k^{3/2})$, batches of medium-size with $\tildeΩ(\sqrt k)$ samples each. However, the paper requires that the input distribution for all $k$ subgroups be isotropic Gaussian, and states that removing this assumption is an ``interesting and challenging problem". We propose a novel gradient-based algorithm that improves on the existing results in several ways. It extends the applicability of the algorithm by: (1) allowing the subgroups' underlying input distributions to be different, unknown, and heavy-tailed; (2) recovering all subgroups followed by a significant proportion of batches even for infinite $k$; (3) removing the separation requirement between the regression vectors; (4) reducing the number of batches and allowing smaller batch sizes.

Paper

Similar papers

Peer review

Reviewer D4Ej7/10 · confidence 4/52024-07-05

Summary

This paper addresses the regression problem in scenarios where heterogeneous data is collected from multiple sources, necessitating learning tasks on small batches of samples. The approach involves dividing the data sources into k subgroups with unknown distributions. This study advances the work of Kong et al. (2020) by introducing a gradient-based algorithm, which does not require the restrictive assumptions present in previous works, such as the assumption of isotropic Gaussian distributions for all k subgroups. The properties of the proposed solution are examined through both theoretical and numerical analyses.

Strengths

The general quality of the paper is commendable, and it addresses all pertinent issues effectively. It adequately covers related works and provides a comprehensive perspective on the topic. The main topic is an important research area with applications in a wide range of real-world projects. Existing works in this area often rely on several assumptions regarding distribution and batch sizes, which can limit their efficiency and accuracy. This paper offers a more general solution that does not depend on restrictive assumptions, thereby enhancing its applicability and robustness.

Weaknesses

The main concerns about the paper include: It has not been properly discussed how this solution allows the recovery of models for sub-populations with a significant fraction of batches when there are a large number of subgroups. This aspect is mentioned several times but not explained clearly. In particular, the paper lacks a detailed explanation of how this work can handle large subgroups more effectively than the reference work by Kong et al. (2020). While it can improve upon previous related works, a detailed complexity analysis is needed. The proposed work is more complex than other works, such as Kong et al. (2020), in terms of both theory and computations. Eliminating the effect of restrictive assumptions can raise the complexity. Additionally, the experiments are not sufficient for this work. Many aspects should be checked, and only a comparison with one baseline cannot support the claims made in the paper. Several other works in this area can be used for comparison, even if they have been proposed only for specific cases. In Figures 1 and 2, the differences between the errors of the baselines when k is small are not significant. It would be beneficial for the authors to provide a theoretical explanation for this observation. Understanding why the error rates converge or show minimal variance under these conditions can provide deeper insights into the behavior of the algorithms and the impact of the number of subgroups on their performance.

Questions

See Weaknesses.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The limitations have been partially addressed in the paper.

Reviewer pdPW7/10 · confidence 3/52024-07-08

Summary

The paper considers the problem of linear regression with heterogeneous data batches, where a user receives a collection of batches of data with not necessarily the same underlying signal distribution and they must attempt to infer a list of heterogeneous signals present in this data. This paper proposes a novel gradient-based method to approach this problem that attains provable and favourable guarantees with less assumptions than previous work.

Strengths

Originality and significance: The paper considers a novel way to tackle the problem of linear regression with heterogeneous batches. The work is significant as their algorithm removes a number of assumptions in prior literature, and seems more performant than previous algorithms. The theory provided is solid and sound. The idea is good, and the description of the proof method and the grad clipping idea was good. Quality and clarity: The problem is well motivated, and the contextualization with respect to prior works is made clear.

Weaknesses

Clarity: - There are multiple spelling mistakes: a) “liner” above section header 1.1, b) “which follows [a] linear regression model” in the middle of page 2, c) “liner” on the last paragraph of page 3, and many others. Please re-read the paper carefully for spelling mistakes. - Theorem 2.3: The first sentence does not make sense to me. What does it mean? More comments: - Missing reference to recent novel developments in MLR: https://proceedings.mlr.press/v206/tan23a.html. The sample complexity in this work is not super-polynomial, but is indeed linear, and this should be mentioned, despite it still having to satisfy assumptions 2 and 4. - I think one drawback of the work has not been mentioned when compared to MLR: you output a list of size L, which is potentially much bigger than the number of signals in the original data, and this should be mentioned. This relates to the problem of list-decoding, and it would be good to also mention that literature and how this relates: https://arxiv.org/abs/1905.05679. - I appreciate that multiple assumptions have been removed, but also importantly one property of the solution has changed from MLR: you are content now with a list of possible signals. From this, it is less surprising that assumptions can be removed, as now the assumptions required and the problem formulation is closer overall to list-decodable linear regression. Could you comment on this? - The input distributions can be different, but they must be known? Whether yes or no, I would mention this clearly in the contribution. - I think the work could benefit from a better discussion on what happens to the algorithm when to regression vectors have no separation? How does the algorithm deal with that? Do we expect the size of the outputted list to change? - A real data experiment would be nice, to evidence this work as that of potential practical significance, but this is not necessary and I still find the work strong and a significant contribution without it.

Questions

What are the barriers to implementing/trying this on real data?

Rating

7

Confidence

3

Soundness

4

Presentation

2

Contribution

3

Limitations

The authors have addressed most limitations in the work, I acknowledge the last paragraph of the discussions. Further addressing the comments above would help as well.

Reviewer oZju7/10 · confidence 4/52024-07-11

Summary

This paper explores the concurrent learning of multiple linear models from various small batches of input-output pairs, each derived from a distinct, small-sized source dataset. Although there can be a large number of sources, it is assumed that each dataset belongs to one of several fixed but unknown subgroups, where each subgroup follows a specific input-output rule. These rules may not be linear for all subgroups; however, a significant number of subgroups, representing a substantial portion of the datasets, follow separate linear rules. The paper provides theoretical guarantees that the linear models for most datasets can be recovered, even with minimal data in each dataset. Moreover, the algorithm's computational complexity is polynomial. The paper claims considerable improvements over existing works, with a detailed comparison with prior research in Section 1.2. The authors have detailed their algorithm, which (similar to some prior works) is based on "subspace estimation," in Section 3. Additionally, the paper includes a series of computational experiments, although I have not thoroughly examined them. I have not fully reviewed the proofs regarding the nature and efficiency of the proposed algorithm. I would like to see other reviewers' comments. However, my current assessment is that this paper is a good candidate for NeurIPS. I vote for acceptance.

Strengths

The paper is well-written and well-motivated. The problem addressed is natural and aligns with numerous modern applications in machine learning and data science. The review of prior works is thorough, and the authors have effectively positioned their work within the existing literature. The authors claim significant advancements in their bounds concerning the relationship to certain parameters and dimensions compared to previous works. For example, their bounds require at least $\Omega\left(k^{3/2}\right)$ batches of data, where each batch contains at least $\Omega\left(k^{1/2}/\epsilon^2\right)$ samples. As can be seen, the number of batches is independent of $\epsilon$ (the estimation error). In comparison, some of the most notable prior works require the number of batches to grow significantly with $\epsilon$. The authors have mitigated several significant constraints present in prior works, including eliminating the assumption of a Gaussian distribution for feature vectors. Additionally, this work removes the assumption that all subgroups must follow a linear model. In fact, the authors extend the concept to accommodate scenarios where the number of subgroups can even approach infinity, ensuring the recovery of the "majority" of subgroups with linear models.

Weaknesses

There are too many hyper-parameters in Algorithm 1, which are provided to the algorithm as input. However, I found the discussion on how to tune these hyper-parameters in practice to be lacking. It would be beneficial to include more detailed guidelines or heuristics for selecting appropriate values for these hyper-parameters, potentially based on empirical studies or theoretical insights. Additionally, some of the assumptions in Section 2.1, such as the distributional properties of the input data and the specific conditions under which the theoretical guarantees hold, may need more clarification. Providing concrete examples or scenarios illustrating these assumptions would help in understanding their practical implications and ensuring they are not overly restrictive.

Questions

What are the main intuitions behind the assumptions made in Remark 2.1? (L. 196)

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

-

Reviewer Xj2G7/10 · confidence 3/52024-07-13

Summary

The paper presents an algorithm that handles a situation of linear regression from many heterogeneous batches of data, and simultaneously learns all linear separators for which there are large enough batches of data witnessing them. The actual guarantees are more nuanced, covering the common case of a heavy-tailed batch size distribution with completely unknown and potentially difficult-to-distinguish batches.

Strengths

There's a lot to like about the overall approach to the problem. I would compliment it as being "applied theory," in contrast to a lot of prior work which takes a more purely theoretical assumption-laden mindset (as commented upon also by the authors, in describing the many assumptions that they are able to remove). Algorithmic techniques like clipping and median-finding are coupled with appropriate theoretical choices, to achieve a sample complexity that's nearly optimal for these regimes. The applicability of the problem itself is more than meets the eye; theory can motivate usage of certain algorithmic techniques. In this case, working in the dual gradient domain to do clipping has a very helpful effect, as the authors prove -- and this is currently not standard for the problem.

Weaknesses

Sec. 1.2 and 1.1 cover overlapping material, and in addition they could be switched with each other. I found 1.2 much clearer about the top-level contributions (e.g. removal of assumptions), and 1.1 more of a detailed theoretical comparison. Overall, some rearrangements of the initial section 1 would improve presentation (e.g. "Our work" line 34 unclear reference at that point for the reader). Would suggest moving forward "improvement over prior work" subsection even within Sec. 1.2. It takes 3.5 pages to formally define the setting with batches, continuing the reader's ambiguity somewhat. For an paper that uses some nice applied algorithmic techniques (see Strengths), it would be nice to run it in a less contrived manner; the R-way partitioning of the data seriously affects applicability (see Questions).

Questions

There remains a significant gap between the scenario analyzed here and a practical one. The number of GD steps R is a real bottleneck because it requires partitioning the data R ways. Though I understand the authors saying "this division may not be necessary for practical implementation," there are certainly extreme datasets in which it would and would not be necessary. It seems to me like this may be just an analysis consideration, and a more refined analysis might highlight under what stability-type conditions the partitioning would be needed. This is significant; a large R makes it possible to get some of the nice adaptivity properties e.g. to \Delta . Duchi et al. have work on Ergodic GD which might be of interest for its techniques in this matter, and its parametrizations; doubtless there is more recent follow-up work in this line. Would appreciate comment from the authors on this issue of R.

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes.

Reviewer ntyP7/10 · confidence 3/52024-07-14

Summary

The paper proposes an algorithm to solve linear regression problems from batched data. The regression coefficients for each batch are assumed to be heterogeneous but come from $k$ possible values. The algorithm proposes a few novel approaches to identify which medium-sized batches are close to a given batch. It also improves the statistical error by estimating a low-rank subspace from small batches. The knowledge-integration techniques developed in the paper are intriguing, and the numerical performance is better than [KSS+20].

Strengths

The introduction of the algorithm is mostly clear, despite its complexity. The intuitions are well-explained. The gradient clipping technique is novel compared with [KSS+20]. Also, the use of gradient descent in the algorithm presents an interesting contribution.

Weaknesses

In line 123, the authors claim that the model does not rely on the previous assumptions. Authors can also briefly discuss how their algorithm innovations circumvents these assumptions. Some minor issues about writing: The sentence in line 332 is very long and hard to follow. In line 836, there is an error.

Questions

The number of small batches required scales with $d$, which is relatively large in high dimensions. Can authors give specific real-life examples to justify the assumptions on $|B_s|$ and $|B_m|$ in Theorem 2.1? Authors claim that $k$ can be infinity, as long as the constraints of $\alpha$ are satisfied. Do authors test the claim numerically?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

See weakness.

Reviewer pdPW2024-08-10

Response

Dear authors, Thank you for addressing my questions, I am happy with the answers and have increased my rating accordingly.

Reviewer oZju2024-08-11

Response to rebuttal

I would like to thank the authors for addressing my questions. I have no further concerns and will maintain my score and support for this work.

Reviewer D4Ej2024-08-12

Official Comment

Thank you for addressing the feedback provided. After reviewing the revisions, I am pleased with the improvements and will be raising my score accordingly.

Reviewer Xj2G2024-08-12

Thanks to the authors for addressing my questions, and points about the presentation. I will maintain my score and positive review.

Program Chairsdecision2024-09-25

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC