Linear Regression using Heterogeneous Data Batches

Authors: ayush  jain
rajat  sen
weihao  kong
alon  orlitsky
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:54:43

Abstract: 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 [KSS$^+$20] showed that with abundant small-batches, the regression vectors can be learned with only few, $\tilde\Omega( k^{3/2})$, batches of medium-size with $\tilde\Omega(\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.

Main paper: 
 introduction: 
 In numerous applications, including federated learning [WCX + 21], sensor networks [WZ89], crowd-sourcing [SVC16] and recommendation systems [WDVR06], data are collected from multiple sources, each providing a batch of samples. For instance, in movie recommendation systems, users typically rate multiple films. Since all samples in a batch are generated by the same source, they are often assumed to share the same underlying distribution. However, the batches are frequently very small, e.g., many users provide only few ratings. Hence, it may be impossible to learn a different model for each batch.A common approach has therefore assumed [TLW99] that all batches share the same underlying distribution and learn this common model by pooling together the data from all batches. While this may work well for some applications, in others it may fail, or lack personalization. For instance, in recommendation systems, it may not capture the characteristics of individual users.A promising alternative that allows for personalization even with many small batches, assumes that batches can be categorized into k sub-populations with similar underlying distributions. Hence in each sub-population, all underlying distributions are close to and can be represented by a single distribution. Even when k is large, our work allows the recovery of models for sub-populations with a significant fraction of batches. For example, in the recommendation setting, most users can be classified into a few sub-populations such that the distribution of users in the sub-population is close, for instance, those preferring certain genres.In this paper, we focus on the canonical model of linear regression in supervised learning. A distribution D of samples (x, y) follows a linear regression model if, for some regression vector w ∈ R d , the output is y = w • x + η where the input x is a random d dimensional vector and η is a zero-mean noise. While we allow the presence of sub-populations that do not follow the linear regression model or have very few batches, our goal is to recover the regression vectors for all large sub-populations that follow this linear regression model.This setting was first considered in [KSS + 20] for meta-learning applications, where they view and term batches as tasks. [KSS + 20] argue that in meta-learning applications task or batch lengths follow a long tail distribution and in the majority of the batches only a few labeled examples are available. Only a few batches have medium-size labeled samples available, and almost all of them have length ≪ d. Note that similar observations have been made in the recommender system literature where the distribution of a number of ratings per user follows a long-tailed distribution with an overwhelming number of users rating only a few items while rare tail users rating hundreds of items [GKG15]. The same has been observed for the distribution of the number of ratings per item [PT08] . Therefore, it is reasonable to assume that in these applications of interest, a handful of medium-size batches along with a large number of batches of constant size are available. Under this setting our main results allow recovery of all sub-populations that has a significant fraction of batches and follow a linear regression model.Building upon the preceding discussion, we have two sets of batches. Each batch comprises i.i.d. samples drawn from one of the k sub-populations. The distribution of samples for these subpopulations is unknown, as is the identity of the sub-population from which each batch originates. Batches in set B s (denoted as small) consist of at least two samples each, while those in set B m (denoted as medium) are of larger size, containing ≥ n m samples. Let's consider a sub-population indexed by i, which follows a linear regression model with parameter vector w i and output-noise variance ≤ σ 2 . The next theorem presents the guarantees provided by our algorithm for estimating w i .Theorem 1.1 (Informal). Let α > 0 such that both sets B s and B m comprises at least α fraction of batches from a sub-population i. Given |B s | ≥ Ω(d/α 2 ), |B m | ≥ Ω(min( √ k, 1/ √ α)/α), and n m ≥ Ω(min( √ k, 1/ √ α)), our algorithm runs in time poly(d, 1/α, k) and outputs a list L of size Õ(1/α) such that w.h.p., at least one estimate in L is within a distance of o(σ) from w i and has an expected prediction error σ 2 (1 + o(1)) for sub-population i. Moreover, given Ω(log L) samples from sub-population i, we can identify such an estimate from L.In the above theorem, the batches within set B s must contain a minimum of 2 samples, while those within set B m must be of size ≥ n m . Additionally, both sets must include ≥ α fraction of batches originating from the population we want to recover. It follows that our algorithm requires only α(2|B s | + n m |B s |) = Ω(d/α) samples from that specific sub-population. It's worth noting that any algorithm, even when dealing with a single sub-population (k = 1), requires a minimum of Ω(d) samples from said sub-population. To the best of our knowledge, ours is the best sample complexity for recovering the linear regression models in the presence of multiple sub-populations using batch sizes smaller than d.Remark 1.1. When multiple sub-populations meet the criteria outlined for sub-population i, the list L returned by our algorithm includes estimates of the regression vector for each of these subpopulations. Furthermore, since there can be up to min{k, 1/α} distinct sub-populations that satisfy these criteria, any algorithm that aims to provide meaningful guarantees must produce a list of size ≥ min{k, 1/α}, unless the algorithm has access to a sample set known to originate from the specific sub-population i it aims to recover. A similar situation arises in the list-decodable setup [CSV17] , where also the algorithm returns a list of estimates.The only work that provides a polynomial-time algorithm in dimension d, in the same generality as ours is [DJKS22] . They even allow the presence of adversarial batches. However, they require Ω(d/α 2 ) batches from the sub-population, each of them of size Ω(1/α), and therefore, Ω(d/α 3 ) samples in total, which exceeds our sample complexity by a factor of 1/α 2 . Furthermore, their algorithm requires a minimum number of samples per batch, which is quadratically larger than ours. All other works place strong assumptions on the distributions of the sub-population and still require a number of samples much larger than ours, which we discuss next.Most of the previous works [CL13, SJA16, ZJD16, YCS16, LL18, CLS20, DK20, PMSG22] have addressed the widely studied mixed linear regression (MLR) model where all batches are of size 1, and adhere to the following three assumptions: to removing the four assumptions, the algorithm doesn't require all batches in a sub-population to have identical distributions, it only requires them to be close (see Remark 2.1).To recover the regression vector for a sub-population i that follow the linear regression model and has more than α fraction of batches, the algorithm begins by selecting a medium size batch from this sub-population. To obtain such a batch we randomly choose Θ(1/α) batches from the set of medium-size batches B m and repeat our algorithm to obtain one estimate for every batch chosen. With high probability, at least one of the randomly selected medium-sized batches belongs to subpopulation i, and when the algorithm is run for such a batch, w.h.p. it recovers the estimate of w i . Therefore, at least one of the Θ(1/α) estimates accurately recover w i .The regression vector w i minimizes the expected squared loss for the sub-population i. Therefore, we use a gradient-descent-based approach to estimate w i . We start with an initial estimate (all zero) and improve this estimate by performing multiple rounds of gradient descent steps.First, using a large number of smaller batches we estimate a smaller subspace of R d that preserves the norm of the gradient for sub-population i. Next, using the medium-size batch selected at the beginning, we test which of the other medium-size batches have a projection of gradient close to the chosen batch, and use them to estimate the gradient in the smaller subspace. The advantage of subspace reduction is that testing and estimation of the gradient in the smaller subspace is easier, consequently reducing the minimum length of medium-size batches required for testing and the number of medium-size batches required for estimation. Our gradient estimation approach draws inspiration from [KSS + 20]. However, while they employed it to directly estimate regression vectors of all subpopulations simultaneously , we adapt it to sequentially estimate gradients for one sub-population at a time. A crucial new ingredient of our algorithm is the clipping of gradients, which limits the effect of other components and allows the algorithm to work even for heavy-tailed distributions.We describe the algorithm in detail in Section 3 after having presented our main theorems in Section 2. Then in Section 4 we compare our algorithm with the one in [KSS + 20] on simulated datasets, to show that our algorithm performs better in the setting of the latter paper as well as generalizes to settings that are outside the assumptions of [KSS + 20]. 
problem formulation and main results: 
 Consider distributions D 0 , . . . , D k-1 over input-output pairs (x, y) ∈ R d × R each corresponding to one of the k sub-populations. A batch b consists of i.i.d. samples from one of the distributions. Samples in different batches are independent. There are two sets of batches. The batches in B s are small and contain at least two samples each, while batches in B m are of medium size and contain at least n m samples. For any batch the identity of the distribution it is sampled from is unknown. Next, we describe the distributions. To aid it and the remaining paper we first introduce some notation.The L 2 norm of a vector u is denoted by ∥u∥. The norm, or spectral norm, of a matrix M is denoted by ∥M ∥ and is defined as max ∥u∥=1 ∥M u∥. If M is a symmetric matrix, the norm simplifies to ∥M ∥ = max ∥u∥=1 |u ⊺ M u|. The trace of a symmetric matrix M is Tr(M ) := i M ii , the sum of the elements on the main diagonal of M . We will use the symbol S to denote an arbitrary collection of samples. For a batch denoted by b, we will use S b to represent the set of all samples in the batch.Let Σ i := E Di [xx ⊺ ] denote the second-moment matrix of input for distribution D i .To recover the regression vector for a distribution D i , we assume that it satisfies the following assumptions2 :1. (Input distribution) There are constants C and C 1 such that for all i ∈ I,(a) L4-L2 hypercontractivity: For all u ∈ R d , E Di [(x • u) 4 ] ≤ C(E Di [(x • u) 2 ]) 2 .(b) Bounded condition number: For normalization purpose, we assume min ∥u∥=1 u ⊺ Σ i u ≥ 1 and to bound the condition number we assume that ∥Σ i ∥ ≤ C 1 .There is a σ > 0 s.t. for all i ∈ I, y = w i • x + η, where w i ∈ R d is an unknown regression vector, and η is a noise independent of x, with zero mean E Di [η] = 0, and E Di [η 2 ] ≤ σ 2 . Note that by definition, the distribution of η may differ for each i.For a given α s > 0 and α m > 0, our goal is to estimate the regression vector w i for all subpopulations that have at least α s and α m fractions of the batches within B s and B m , respectively, and follow the assumptions described above. Let I ⊆ {0, 1, .., k -1} denote the collection of indices corresponding to all such sub-populations. (Note that in the introduction, we simplified this more general scenario by assuming equal values for α s and α m , both denoted by α.)For other sub-populations i / ∈ I, we only need that the input distribution satisfies ∥Σ i ∥ ≤ C 1 , same as the second half of assumption 1(b). The input-output relation for samples generated by D i for i / ∈ I can be arbitrary, and in particular, does not even need to follow a linear regression model. Moreover, the fraction of batches from these other sub-populations can also be arbitrary.To simplify the presentation, we make two additional assumptions. First, there is a constant C 2 > 0 such that for all components i ∈ {0, 1, .., k -1}, and random sample (x, y) ∼ D i , ∥x∥ ≤ C 2 √ d, a.s. Second, for all i ∈ I and a random sample (x, y) ∼ D i , the noise distribution η = y -w i • x is symmetric around 0. As discussed in Appendix J, these assumptions can be easily removed. Remark 2.1. For simplicity, we assumed that each batch exactly follows one of the k distributions. However, our methods can be extended to more general scenarios. Let D b denote the underlying distribution of batch b. Instead of requiring D b = D i for some i ∈ {0, 1, .., k -1}, our approach can be extended to cases where the expected gradients for D b and D i are close. Additionally, if i ∈ I, the regression vector w i must achieve a mean squared error (MSE) of at most σ 2 for D b . These conditions are satisfied if the following hold:(1) ∥ E D b [xx ⊺ ] -Σ i ∥ is small, (2) for all x ∈ R d , | E D b [y|x] -E Di [y|x]| is small, and (3) if i ∈ I then for all x ∈ R d , E D b [(y -w i • x) 2 |x] ≤ σ 2 .Thus, the strict identity requirement D b = D i can therefore be replaced by these three approximation conditions in our analysis. This extension is feasible as our algorithm and its analysis are based on gradient descent, and can handle stochastic noise. Since the analysis tolerates statistical noise, it can also accommodate small deviations in the gradient distributions across batches in the same subpopulation.We begin by presenting our result for estimating the regression vector of a component D i , for any i ∈ I. This result assumes that in addition to the batch collections B s and B m , we have an extra medium-sized batch denoted as b * which contains samples from D i and w.l.o.g, we assume i = 0. Theorem 2.1. Suppose index 0 is in set I and let b * be a batch of ≥ n m i.i.d. samples fromD 0 . Given δ, ϵ ∈ (0, 1], |B s | = Ω( d αs 2 ϵ 4 ), n m = Ω(min{ √ k, 1 ϵ √ αs } • 1 ϵ 2 ), and |B m | = Ω( 1 αm min{ √ k, 1 ϵ √ αs }), Algorithm 1 runs in polynomial time and returns estimate ŵ, such that with probability ≥ 1 -δ, ∥ ŵ -w 0 ∥ ≤ ϵσ.We provide a proof sketch of Theorem 2.1 and the description of Algortihm 1 in Section 3, and a formal proof in Appendix H. Algorithm 1 can be used to estimate w i for all i ∈ I, and the requirement of a separate batch b * is not crucial. As mentioned in the introduction, it can be obtained by repeatedly sampling a batch from B m and running the algorithm for these sampled b * . Since all the components in I have ≥ α m fraction of batches in B m , then randomly sampling b * from B m , Θ(1/α m ) times would ensure that, with high probability, we have b * corresponding to each component. We can then return a list of size Θ(1/α m ) containing estimates corresponding to each sampled b * . Then, with high probability, the list will have an estimate of the regression vectors for all components. Note that in this case, returning a list is unavoidable as there is no way to assign an appropriate index to the regression vector estimates. The following corollary follows from the above discussion and Theorem 2.1. Corollary 2.2. For given δ, ϵ ∈ (0, 1], if |B s |, n m , and |B m |, meet the requirements outlined in Theorem 2.1, then the above modification of Algorithm 1 runs in polynomial-time and outputs a list L of size Õ(1/α m ) such that with probability ≥ 1 -δ, the list has an accurate estimate for regression vectors w i for each i ∈ I, namely max i∈I min ŵ∈L ∥ ŵ -w i ∥ ≤ ϵσ.In particular, this corollary implies that for any i ∈ I, the algorithm requires only Ω(d/α s ) batches of size two and Ω(min{√ k, 1 √ αs }) medium-size batches of size Ω(min{ √ k, 1 √ αs }) from distribution D i to estimate w i within an accuracy o(σ). Furthermore, it is easy to show that any o(σ) accurate estimate of regression parameter w i achieves an expected prediction error of σ 2 (1 + o(1)) for output y given input x generated from this D i .Note that results work even for infinite k and without any separation assumptions on regression vectors. The min( √ k, 1/ √ α s ) dependence is the best of both words. This dependence is reasonable for recovering components with a significant presence or if the number is few.The total number of samples required by the algorithm from D i in small size batches B s and medium size batches B m are only Õ(d/α s ) and Õ(min{k, 1/α s )). It is well known that any estimator would require Ω(d) samples for such estimation guarantees even in the much simpler setting with just i.i.d. data. Consequently, the total number of samples required from D i by the algorithm is within Õ(1/α s ) factor from that required in a much simpler single component setting.The next theorem shows that for any i ∈ I, given a list L containing estimate of w i and Ω(log(1/α s )) samples from D i , we can identify an estimate of regression vector achieving a small prediction error for D i . The proof of the theorem and the algorithm is in Appendix B. Theorem 2.3. Let index i be in set I and let L be a list of d dimensional vectors. Suppose for a given β > 0 at least one of the vectors in L is within distance β from the regression vector of D i , namely min w∈L ∥w -w i ∥ ≤ β. Given O(max{1, σ 2 β 2 } log L δ ) samples from D i Algorithm 3 identifies an estimate w from the list L, s.t. with probability ≥ 1 -δ, ∥w -w i ∥ = O(β) and it achieves an expected estimation error E Di [( ŵ • x -y) 2 ] ≤ σ 2 + O(β 2 ).Remark 2.2. As a special case the above theorem implies that that for any β ≥ Ω(σ), Algorithm 3 can identify the near best estimate of w i from list L using only O(log L δ ) samples from D i . Combining the above theorem and Corollary 2.2, we get Theorem 2.4. For given δ, ϵ ∈ (0, 1], if |B s |, n m , and |B m |, meet the requirements outlined in Theorem 2.1, then, there exists a polynomial-time algorithm that, with probability ≥ 1 -δ, outputs a list L of size Õ(1/α m ) containing estimates of w i 's for i ∈ I. Furthermore, given |S| ≥ Ω( 1 ϵ 2 log 1 δαm ) samples from D i , for any i ∈ I, Algorithm 3 returns ŵ ∈ L that with probability ≥ 1 -δ satisfies ∥w i -ŵ∥ ≤ O(ϵσ) and achieves an expected estimation errorE Di [( ŵ • x -y) 2 ] ≤ σ 2 + O(ϵ 2 σ 2 ) 
algorithm for recovering regression vectors: 
 This section provides an overview and pseudo-code of Algorithm 1, along with an outline of the proof that achieves the guarantee stated in Theorem 2.1. As per the theorem, we assume that index 0 belongs to I, and we have a batch b * containing n m samples from the distribution D 0 . Note that D 0 satisfies the conditions mentioned in Section 2.1 and that B s and B m each have ≥ |B s |α s and ≥ |B m |α m batches with i.i.d. samples from D 0 . However, the identity of these batches is unknown.Gradient Descent. Note that w 0 minimizes the expected square loss for distribution D 0 . Our algorithm aims to estimate w 0 by taking a gradient descent approach. It performs a total of R gradient descent steps. Let ŵ(r) denote the algorithm's estimate of w 0 at the beginning of step r. Without loss of generality, we assume that the algorithm starts with an initial estimate of ŵ(1) = 0. In step r, the algorithm produces an estimate ∆ (r) of the gradient of the expected square loss for distribution D 0 at its current estimate ŵ(r) . For convenience, we refer to this expected gradient for D 0 at ŵ(r) simply as expected gradient. The algorithm then updates its current estimate for the next round as ŵ(r+1) = ŵ(r) -∆ (r) /C 1 .The main challenge the algorithm faces is the accurate estimation of the expected gradients in each step. Accurately estimating the expected gradients at each step would require Ω(d/ϵ 2 ) i.i.d. samples from D 0 . However, our algorithm only has access to a medium-size batch b * that is guaranteed to have samples from D 0 and this batch contains far fewer samples. And for batches in B s and B m , , ŵ(r) , ϵ1, δ ′ , σ, C, C1) 10:Algorithm: ϵ1 ← Θ(1), ϵ2 ← Θ 1 C 1 √ C+1 (ϵ1 + 1 √ C 1 ) , ℓ ← min{k, 1 2αsϵ 2 2 }, δ ′ ← δ 5R and R ← Θ(C1 log M σ ) 4:EQUATIONs , κ (r) , ŵ(r) , ℓ ) 11:∆ (r) ← GRADEST(B (r) m , S b * ,(r) 2, κ (r) , ŵ(r) , P (r) , ϵ2, δ ′ ) 12:ŵ(r+1) ← ŵ(r) -1 C 1 ∆ (r) 13: end for 14: ŵ ← ŵ(R+1) and Return ŵ the algorithm doesn't know which of the batches has samples from D 0 . Despite these challenges, we demonstrate an efficient method to estimate the expected gradients accurately.The algorithm randomly divides sets B s and B m into R disjoint equal subsets, denoted as {B from D 0 to estimate the gradient at point w (r) . While this division may not be necessary for practical implementation, as shown by our simulation results later, this ensures independence between the stationary point ŵ(r) and the gradient estimate, which facilitates our theoretical analysis and only incurs a logarithmic factor in sample complexity.Next, we describe how the algorithm estimates the gradient and the guarantees of this estimation. Due to space limitations, we provide a brief summary here, and a more detailed description, along with formal proofs, can be found in the appendix. We start by introducing a clipping operation on the gradients, which plays a crucial role in the estimation process.Clipping. Recall that the squared loss of samples (x, y) on point w is (x • w -y) 2 /2 and its gradient is (x • w -y)x. Instead of directly working with the gradient of the squared loss, we work with its clipped version. Given a clipping parameter κ > 0, the clipped gradient for a sample (x, y) evaluated at point w is defined as ∇f (x, y, w, κ) := (x•w-y) |x•w-y|∨κ κx. For a collection of samples S, the clipped gradient ∇f (S, w, κ) is the average of the clipped gradients of all samples in S, i.e., ∇f (S, w, κ) := 1 |S| (x,y)∈S ∇f (x, y, w, κ). The clipping parameter κ controls the level of clipping and for κ = ∞, the clipped and unclipped gradients are the same. The clipping step is necessary to make our gradient estimate more robust, by limiting the influence of the components other than D 0 (in lemma C.1), and as a bonus, we also obtain better tail bounds for the clipped gradients. Theorem C.2 shows that for κ ≥ Ω( E D0 [(y -x • w) 2 ]), the difference between the expected clipped gradient and the ex-pected gradient ∥E D0 [(∇f (x, y, w, κ)] -E D0 [(x • w -y)x]∥ is small. Therefore, the ideal value of κ at point w is Θ( E D0 [(y -x • w) 2 ]).For the estimate ŵ(r) at step r, the choice of the clipping parameter is represented by κ (r) . To estimate a value for κ (r) that is close to its ideal value, the algorithm employs the subroutine CLIPEST (presented as Algorithm 4 in the appendix). The subroutine estimates the expected value of (y -x• ŵ(r) ) 2 by using i.i.d. samples S b * ,(r) 1 from the distribution D 0 . According to Theorem D.1 in Appendix D, the subroutine w.h.p. obtains κ (r) which is close to the ideal value. This ensures that the difference between the expectation of clipped and unclipped gradients is small, and thus, estimating the expectation of clipped gradients can replace estimating the actual gradients.Subspace Estimation. The algorithm proceeds by using subroutine GRADSUBEST (presented as Algorithm 5 in Appendix E) with B = B (r) s , w = ŵ(r) , and κ = κ (r) to estimate a smaller subspace P (r) within R d . The objective of this subroutine is to ensure that, with high probability, the expected clipped gradient for distribution D 0 at point ŵ(r) remains largely unchanged when projected onto the identified subspace. Thus, estimating the expected clipped gradient reduces to estimating its projection onto P (r) , which requires fewer samples due to the lower dimensionality of the subspace. To achieve this, the subroutine constructs a matrix A defined asE[A] = i p i E Di [∇f (x, y, w, κ)] E Di [∇f (x,y, w, κ)] ⊺ , where p i denotes the fraction of batches in B that have samples from D i . Since B (r) s are obtained by randomly partitioning B s , w.h.p. p 0 ≈ α s . It is crucial for the success of the subroutine that the expected contribution of every batch in the above expression is a PSD matrix. The clipping helps in bounding the contribution of other components and statistical noise.The subroutine returns the projection matrix P (r) for the subspace spanned by the top ℓ singular vectors of A, where we choose ℓ = min{k, Θ(1/α s )}. It is worth noting that when 1/α s is much smaller than k (thinking of the extreme case k = ∞), our algorithm still only requires estimating the top ℓ = 1/α s dimensional subspace, since those infinitely many components can create at most (1/α s -1) directions with weight greater than α s . This along with clipping ensures that the direction of D 0 must appear in the top Θ(1/α s ) subspace (see Lemma E.3). Theorem E.1 in Appendix E characterizes the guarantees for this subroutine. Informally, if B ≥ Ω(d/α 2 ), then w.h.p., the expected value of the projection of the clipped gradient on this subspace is nearly the same as the expected value of the clipped gradient, namely ∥ E D0 [P (r) ∇f (x, y, w, κ)] -E D0 [∇f (x, y, w, κ)]∥ is small. We note that our construction of matrix A for the subroutine is inspired by a similar construction Algorithm 2 GRADEST 1: Input: A collection of medium batches B, a collection of samples S * from D0, κ, w, projection matrix P for subspace of R d , parameter ϵ, δ ′ 2: Output: An estimate of clipped gradient at point w. to test whether the expected projection of the clipped gradient for the distribution that generated batch b is close to that of the reference distribution D 0 . With high probability, the algorithm retains batches sampled from D 0 and rejects those from other distributions where the difference in expected projections is significant. This test requires Ω(√ ℓ)samples in each batch, where ℓ is the dimension of the projected clipped gradient, which is same as the dimension of the estimated subspace in the subspace estimation space.After identifying the relevant batches, GRADEST estimates the projection of the clipped gradients using the second half of the samples in these batches. Since the projections of the clipped gradients lie in an ℓ dimensional subspace, Ω(ℓ) samples suffice for the estimation. To obtain high-probability guarantees, the procedure uses the median of means approach for both testing and estimation.The guarantees of the subroutine are described in Theorem F.1, which implies that the estimate ∆ (r) of the gradient satisfies ∥∆ (r) -E D0 [P (r) ∇f (x, y, w, κ)]∥ is small.Estimation guarantees for expected gradient. Using the triangle inequality, we have:∥∆ (r) -E D0 [(x • w -y)x]∥ ≤ ∥ E D0 [∇f (x, y, w, κ)] -E D0 [(x • w -y)x]∥ + ∥ E D0 [∇f (x, y, w, κ)] -E D0 [P (r) ∇f (x, y, w, κ)]∥ + ∥∆ (r) -E D0 [P (r) ∇f (x, y, w, κ)]∥.The first term on the right represents the difference between the expectations of the gradient and the clipped gradient. The second term captures the difference between the expectation of the clipped gradient and its projection onto the estimated subspace. Finally, the last term reflects the difference between the expectation of the projection of the clipped gradient and its estimate.All expectations are taken with respect to the distribution D 0 , and both gradients and clipped gradients are evaluated at the point w. Recall that at the beginning of each gradient descent step, w is updated to the previous round's estimate of w 0 .As previously argued, all three terms on the right side of the inequality are small, which implies that ∆ (r) provides an accurate estimate of the expected value of the gradient for distribution D 0 . Moreover, Lemma G.1 shows that with an accurate estimation of expected gradients, gradient descent reaches an accurate estimation of w 0 after O(log ∥w0∥ σ ) steps. Therefore, setting R = Ω(log ∥w0∥ σ ) suffices. This completes the description of the algorithm and the proof sketch of Theorem 2.1, with a more formal proof available in Appendix H. 
empirical results: 
 Setup. We have sets B s and B m of small and medium size batches and k distributions D i for i ∈ {0, 1, . . . , k -1}. For a subset of indices I ⊆ {0, 1, . . . , k -1}, both B s and B m have a fraction of α batches that contain i.i.d. samples from D i for each i ∈ I. And for each i ∈ {0, 1, . . . , k -1}\I in the remaining set of indices, B s and B m have (1 -|I|α)/(k -|I|) fraction of batches, that have i.i.d samples from D i . In all figures, the output noise is distributed as N (0, 1). All small batches have 2 samples each, while medium-size batches have n m samples each, which we vary from 4 to 32, as shown in the plots. We fix data dimension d = 100, α = 1/16, the number of small batches to |B s | = min{8dk 2 , 8d/α 2 } and the number of medium batches to |B m | = 256. In all the plots, we averaged over 10 runs and report the standard error. Evaluation. Our objective is to recover a small list containing good estimates for the regression vectors of D i for each i ∈ I. We compare our proposed algorithm's performance with that of the algorithm in [KSS + 20]. Given a new batch, we can choose the weight vector from the returned list, L that achieves the best error3 . Then the MSE of the chosen weight is reported on another new batch drawn from the same distribution. The size of the new batch can be either 4 or 8 as marked in the plot. More details about our setup can be found in Appendix L.Setting in [KSS + 20]. We first compare our algorithm with the one in [KSS + 20] in the same setting as the latter paper i.e. with more restrictive assumptions. The results are displayed in Figure 1 , where I = {0, 1, . . . , 15} and all 16 distributions have been used to generate 1/16 fraction of the batches.All the D i 's are equal to N (0, I), and the minimum distance between the regression vectors is comparable to their norm. It can be seen that even in the original setting of [KSS + 20] our algorithm significantly outperforms the other at all the different medium batch sizes plotted on the x-axis.Input distributions. Our algorithm can handle different input distributions for different subgroups.We test this in our next experiment presented in Figure 2 . Specifically, for each i, we randomly generate a covariance matrix Σ i such that its eigenvalues are uniformly distributed in [1, C 1 ], and the input distribution for D i is chosen as N (0, Σ i ). We set C 1 = 4. It can be seen that [KSS + 20] completely fails in this case, while our algorithm retains its good performance.In the interest of space, we provide additional results in Appendix L which include even more general settings: (i) when the minimum distance between regression vectors can be much smaller than their norm (ii) when the number of subgroups k can be very large but the task is to recover the regression weights for the subgroups that appear in a sufficient fraction of the batches. In both these cases, our algorithm performs much better than the baseline. 
conclusion: 
 We study the problem of learning linear regression from batched data in the presence of subpopulations. We removed several restrictive assumptions from prior work and provide better guarantees in terms of overall sample complexity. Moreover, we require relatively fewer medium batches that need to contain less number of samples compared to prior work. Finally, in our empirical results, we show that our algorithm is both practical and more performant compared to a prior baseline.One limitation of the algorithm is that it requires knowledge of parameters C, C 1 , M and σ. However, note that the algorithm works even if their assumed values exceed the actual values. The change in the recovery guarantees is proportional to the disparity between the estimates and the actual parameters. The necessity of knowing reasonable upper bounds on these parameters, even for simpler problems, has been highlighted by previous work, e.g. [CAT + 20].[CL13]Arun Tejasvi Chaganty and Percy Liang, Spectral experts for estimating mixtures of linear regressions, International Conference on Machine Learning (ICML), 2013, pp. 1040-1048.[CLM20] Sitan Chen, Jerry Li, and Ankur Moitra, Learning structured distributions from untrusted batches: Faster and simpler, Advances in Neural Information Processing Systems 33 (2020), 4512-4523.[CLS20] Sitan Chen, Jerry Robust Learning from Batches. [QV18] presented the problem of robust learning of discrete distributions from untrustworthy batches, where a majority of batches share the same distribution and a small fraction are adversarial. They developed an exponential time algorithm for the problem. Subsequent works [CLS20] improved the run-time to quasi-polynomial, while and [JO20b] derived a polynomial time algorithm with an optimal sample complexity. The results were extended to learning one-dimensional structured distributions in [JO21, CLM20], and classification in [JO20a, KFAL20]. [AJK + 22] examined a closely related problem of learning the parameters of an Erdős-Rényi random graph when a portion of nodes may be corrupted and their edges are maybe be chosen by an adversary. 
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: Our experiments are not compute intensive and were run on Macbook pro 16 inch laptop with 16 GB ram and intel processor (2020 model) and each of them took less than 6 hours to finish.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). 
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? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This paper presents work whose goal is to advance the field of Machine Learning There could be many potential societal consequences of our work, none which we feel must be specifically highlighted here.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). 
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: We do not release any new dataset or model.Guidelines:• 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 used an algorithm in a prior work [KSS + 20] and we have cited the original paper.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.• 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. 
 
 Appendix: In this section, we introduce Algorithm 3 and prove that it achieves the guarantees presented in Theorem 2.3.Algorithm 3 SELECTING THE REGRESSION VECTOR 1: Input: Samples S from D i for some i ∈ I, C 1 , a list L of possible estimates of w i , and β ≥ 0s.t. β ≥ min w∈L ∥w -w i ∥. 2: Output: An estimate of w i * from list L 3: Divide S into T 3 = Θ(log(|L|/δ)) equal parts {S j } T3 j=1 4: while max{∥w -w ′ ∥ : w, w ′ ∈ L} ≥ 12C 1 β do 5: pick any w, w ′ ∈ L s.t. ∥w -w ′ ∥ ≥ 12C 1 β. 6: For j ∈ [T 3 ], let a j ← (x,y)∈Sj 1 |Sj | (x • w -y)x • (w -w ′ ) 7: a ← Median{a j : j ∈ [T 3 ]} 8:If a > ∥w -w ′ ∥ 2 /4 remove w, else remove w ′ from L 9: end while 10: Return any of the remaining w ∈ L Without loss of generality, assume i = 0, and let w * = arg min w∈L ∥w -w 0 ∥. From the condition in the theorem, we know that ∥w * -w 0 ∥ ≤ β. The algorithm is given access to |S| = Ω(max 1, σ 2 β 2 log |L| δ ) samples. The algorithm chooses any two vectors w, w ′ ∈ L that are more than 12C 1 β distance apart and tests which of them is more likely to be within β distance from w 0 using samples in S. The algorithm performs T 3 = Θ(log |L| δ ) such tests and takes the majority vote. It retains the vector that is more likely to be closer to w and discards the other from L. The algorithm terminates when all the vectors in L are within a distance of 12C 1 β from each other, by choosing a vector from those remaining in L and returning it as an estimate of w 0 . If w * is retained in L until the end, using the simple triangle inequality for all w that remain in L at the end, we have ∥w -w 0 ∥ ≤ ∥w -w * ∥ + ∥w * -w 0 ∥ ≤ 12C 1 β + β ≤ 13C 1 β = O(β).Therefore, the estimate returned by the algorithm achieves the desired accuracy in estimating w 0 . Hence, it suffices to show that w * is retained at the end with high probability. Suppose w * is not in the final list L. Then it must have been discarded by the test in favor of w ∈ L such that ∥w * -w∥ ≥ 12C 1 β. The following theorem shows that for any w such that ∥w * -w∥ ≥ 12C 1 β, the probability of the testing procedure rejecting w * in favor of w is at most δ/|L|. Theorem B.1. Given β > 0, list L, and samples S from D 0 , if min w∈L ∥w -w 0 ∥ ≤ β and |S| = max{1, σ 2 β 2 } log L δ , then for the parameter a computed in the while loop of Algorithm 3, with probability 1 -δ/|L|, we have a ≤ ∥w -w ′ ∥ 2 /4 if w = w 0 and a > ∥w -w ′ ∥ 2 /4 if w ′ = w 0 .The testing procedure utilized in the algorithm is based on gradients. Specifically, it calculates the average of the gradient computed on samples at point w projected onto the vector (w -w ′ ). The expected value of the gradient at w, and its projection onto (w -w ′ ), are (w -w 0 ) ⊺ Σ 0 and (w -w 0 ) ⊺ Σ 0 (w -w ′ ), respectively. If w ≈ w 0 , then the expected projection will be small. On the other hand. if w ′ ≈ w 0 and w, then expected value of projection is ≈ (w -w ′ ) ⊺ Σ 0 (w -w ′ ) ≳ ∥w -w ′ ∥ 2 . Using these observations, we prove Theorem B.1 in the next subsection.Finally, since the maximum number of comparisons made by the algorithm is |L| -1, a union bound ensures that w * will be retained until the end with probability greater than 1 -δ, completing the proof of Theorem 2.3.Proof. Note that a is the median of the set {a j : j ∈ [T 3 ]}, where each a j is computed using different sets of i.i.d. samples. Consequently, {a j } j∈[T3] are also i.i.d random variables.We begin by calculating the expected value of a j . Using the linearity of expectations, we have:E[a j ] (a) = E (x,y)∈Sj 1 |Sj | (x • w -y)x • (w -w ′ ) (b) = E D0 [(x • w -y)x • (w -w ′ )] = E D0 [(x • (w -w 0 ) -(y -x • w 0 ))x • (w -w ′ )] = E D0 [(x • (w -w 0 )x • (w -w ′ )] -E D0 [(y -x • w 0 )x • (w -w ′ )] = E D0 [x • (w -w 0 )x • (w -w ′ )](1)EQUATIONhere, (a) follows from the definition of a j , (b) follows from the linearity of expectation, and since S j contains i.i.d. samples from D 0 , (c) follows as the noise y -x • w 0 has a zero mean and is independent of x.Next, we compute the variance of a j . Since a j represents the average of (xEQUATION)By applying Chebyshev's inequality, with a probability ≥ 3/4, the following holds for each a j :EQUATIONFirst, we consider the case when w = w * . In this case, we have ∥w 0 -w∥ ≤ β.Using Equation (3), we can express the variance of a j as follows:Var(a j ) ≤ 1 |S j | E D0 ((x • w -y)x • (w -w ′ )) 2 = 1 |S j | E D0 (x • (w -w 0 )x • (w -w ′ ) + (w 0 • x -y)x • (w -w ′ )) 2 ≤ 2 |S j | E D0 (x • (w -w 0 )x • (w -w ′ )) 2 + E D0 ((w 0 • x -y)x • (w -w ′ )) 2 ,where the last step uses the fact that for any u, v ∈ R,(u + v) 2 ≤ 2u 2 + 2v 2 .Next, we bound the two terms on the right. For the first term, we haveEQUATIONFor the second term, we have:EQUATIONwhere the first inequality follows from assumption 1a and the second inequality follows from assumption 1b.Combining the above three equations, we obtain:Var(a j ) ≤ 2 |S j | E D0 (x • (w -w ′ )) 2 C E D0 [(x • (w -w 0 )) 2 ] + σ 2 ≤ 2 |S j | C 1 ∥w -w ′ ∥ 2 CC 1 ∥w -w 0 ∥ 2 + σ 2 ,where the last inequality uses assumption 1b.Using Equation (1), the Cauchy-Schwarz inequality, and assumption 1b, we have:EQUATIONCombining the two equations above, we obtain:E[a j ] + 2 Var(a j ) ≤ ∥w -w ′ ∥ C 1 ∥w -w 0 ∥ + 2 √ 2C 1 |S j | σ + CC 1 ∥w -w 0 ∥ (a) ≤ ∥w -w ′ ∥ C 1 β + √ 8C 1 |S j | σ + √ 8CC 1 |S j | β (b) ≤ 3C 1 ∥w -w ′ ∥β (c) ≤ ∥w -w ′ ∥ 2 4 ,here, in (a), we use w = w * , which implies ∥w -w 0 ∥ ≤ β, in (b), we utilize |S j | ≥ 48C and|S j | ≥ 12σ 2C1β 2 , in (c), we use the fact that for any w and w ′ in the while loop of the algorithm, ∥w -w ′ ∥ ≥ 12C 1 β. Consequently, it follows from Equation (4) that each a j ≤ ∥w-w ′ ∥ 2 4 with probability ≥ 3/4. Hence, with probability ≥ 1 -δ the median of a j is ≤ ∥w-w ′ ∥ 2 4 . Next, we consider the case when w ′ = w * . Firstly, we bound the variance using Equation (3):Var(a j ) ≤ 1 |S j | E D0 ((x • w -y)x • (w -w ′ )) 2 = 1 |S j | E D0 (x • (w -w ′ )) 2 + x • (w ′ -w 0 )x • (w -w ′ ) + (w 0 • x -y)x • (w -w ′ ) 2 (a) ≤ 3 |S j | E D0 ((w -w ′ ) • x) 4 + E D0 (x • (w ′ -w 0 )x • (w -w ′ )) 2 (8) + E D0 ((w 0 • x -y)x • (w -w ′ )) 2 (b) ≤ 3 |S j | C E D0 (x • (w -w ′ )) 2 2 + C E D0 [(x • (w ′ -w 0 )) 2 ] + σ 2 E D0 [(x • (w -w ′ )) 2 ] (c) ≤ 3 |S j | √ C E D0 (x • (w -w ′ )) 2 + C E D0 [(x • (w ′ -w 0 )) 2 ] + σ E D0 [(x • (w -w ′ )) 2 ] 2 .In (a), we use the fact that for any t, u, v ∈ R, (t + u + v) 2 ≤ 3t 2 + 3u 2 + 3v 2 . In (b) the first term is bounded using assumption 1a, the bound on the second term can be obtained similarly to Equation ( 5), and the bound on the last term is from Equation (6). Finally, in (c) we use the fact that for any t, u, v ≥ 0, (t+ u + v) 2 ≤ t 2 + u 2 + v 2 .Using Equation (2) and the equation above, we getE[a j ] -2 Var(a j ) ≥ E D0 [(x • (w -w ′ )) 2 ] + E D0 [x • (w ′ -w 0 )x • (w -w ′ )] - 2 √ 3 |S j | √ C E D0 (x • (w -w ′ )) 2 + C E D0 [(x • (w ′ -w 0 )) 2 ] + σ E D0 [(x • (w -w ′ )) 2 ] (a) ≥ 1 - √ 12C |S j | E D0 [(x • (w -w ′ )) 2 ] -E D0 [(x • (w ′ -w 0 )) 2 ] E D0 [(x • (w -w ′ )) 2 ] - √ 12 |S j | C E D0 [(x • (w ′ -w 0 )) 2 ] + σ E D0 [(x • (w -w ′ )) 2 ] (b) ≥ 1 2 E D0 [(x • (w -w ′ )) 2 ] - 3 2 E D0 [(x • (w ′ -w 0 )) 2 ] - √ 12 |S j | σ E D0 [(x • (w -w ′ )) 2 ] (c) ≥ 1 2 E D0 [(x • (w -w ′ )) 2 ] - 3 2 E D0 [(x • (w ′ -w 0 )) 2 ] -C 1 β E D0 [(x • (w -w ′ )) 2 ],here, in (a) we use the Cauchy-Schwarz inequality, (b) follows from |S j | ≥ 48C, and (c) utilizesEQUATIONhere in (a), we use assumption 1b, (b) relies on w ′ = w * , which implies ∥w ′ -w 0 ∥ ≤ β, and (c) uses the fact that for any w and w ′ in the while loop of the algorithm, ∥w -w ′ ∥ ≥ 12C 1 β and C 1 ≥ 1.Combining the above two equations, we obtainE[a j ] -2 Var(a j ) > 1 4 ∥w -w ′ ∥ 2 .Then from Equation (4) it follows that each a j > ∥w-w ′ ∥ 2 4 with probability ≥ 3/4. Hence, withprobability ≥ 1 -δ/|L| the median of a j is > ∥w-w ′ ∥ 2 4 .The norm of the expected value and covariance of unclipped gradients for components other than D 0 can be significantly larger than D 0 , acting as noise in the recovery process of D 0 . When using unclipped gradients, the algorithm's batch size and the number of batches must increase to limit the effect of these noisy components. And while the norm of the expected value and covariance of the unclipped gradient for D 0 follows desired bounds, the maximum value of the unclipped gradient is unbounded, posing difficulties in applying concentration bounds. The following lemma shows that the clipping operation described in the main paper is able to address these challenges. Lemma C.1. Let S be a collection of random samples drawn from distribution D i for some i ∈ {0, 1, ..., k -1}. For any κ ≥ 0 and w ∈ R d , the clipped gradient satisfies the following properties:1. ∥ E[∇f (S, w, κ)]∥ ≤ κ √ C 1 , 2. ∥Cov(∇f (S, w, κ))∥ ≤ 1 |S| κ 2 C 1 , 3. ∥∇f (S, w, κ)∥ ≤ κC 2 √ d almost surely, 4. E ∥∇f (S, w, κ)∥ 2 ≤ C 1 κ 2 d, 5. for all unit vectors u, E (∇f (S, w, κ) • u) 2 ≤ κ 2 C 1 .This lemma implies that for smaller values of κ, the norm of the expectations and covariance of clipped gradients is bounded by a smaller upper limit. The proof of the lemma is presented in Subsection C.1.The following theorem demonstrates that by appropriately choosing a sufficiently large value of κ, the norm of the expected difference between the clipped and unclipped gradients for distribution D 0 can be small: Theorem C.2. For any ϵ > 0, The theorem shows in order to estimate the expectation of gradients at point w for distribution D 0 , it is sufficient to estimate the expectation of clipped gradients at point w, as long as the clipping parameter κ is chosen to be at least Ωκ 2 ≥ 8CC 1 E D0 [(y -x • w) 2 ]/ϵE D 0 [(y-x•w) 2 ] ϵ .Intuitively, when κ is much larger than E D0 [|y -x • w|], with high probability the clipped and unclipped gradients at point w for a random sample from D 0 will be identical. The proof of the theorem is a bit more nuanced and involves leveraging the symmetry of noise distribution and L4 -L2 hypercontractivity of distribution of x. The proof appears in Subsection C.2.In the algorithm, we set κ to approximately Θ( (E D0 [(y -x • w) 2 ] + σ 2 )/ϵ). This choice ensures that κ is close to the minimum value recommended by Theorem C.2 for preserving the gradient expectation of D 0 . By selecting a small κ, we ensure a tighter upper bound on the expectation and covariance of the clipped gradient for other components, as described in Lemma C.1. The use of the clipping operation also assists in obtaining improved bounds on the tails of the gradient by limiting the maximum possible norm of the gradients after clipping, as stated in item 3 of the lemma. We will now proceed to prove the five claims in the lemma by using these properties.Firstly, we can analyze the expected norm of E Di [∇f (x, y, w, κ)] as follows:∥ E Di [∇f (x, y, w, κ)]∥ = max ∥u∥ ∥ E Di [(∇f (x, y, w, κ) • u)]∥ ≤ max ∥u∥ ∥ E Di [(∇f (x, y, w, κ) • u) 2 ] 1/2 ∥ ≤ max ∥u∥ ∥ E Di [(κx • u) 2 ] 1/2 ∥ ≤ κ C 1 ,here the first inequality follows from the Cauchy-Schwarz inequality and the last inequality follows from assumptions on distributions D i . Combining the above inequality with item a) above proves the first claim in the lemma.Next, to prove the second claim in the lemma, we first establish bounds for the norm of the covariance of the clipped gradient of a random sample:∥Cov Di (∇f (x, y, w, κ))∥ = max ∥u∥ Var Di (∇f (x, y, w, κ) • u) ≤ max ∥u∥ E Di [(∇f (x, y, w, κ) • u) 2 ] ≤ max ∥u∥ E Di [(κx • u) 2 ] ≤ κ 2 C 1 .By using the above bound and combining it with item b), we establish the second claim in the lemma.To prove the third item in the lemma, we first bound the norm of the clipped gradient: We then combine this bound with item c) to prove the third claim in the lemma.Next, we bound the expected value of the square of the norm of the clipped gradient of a random sample,E Di ∥∇f (x, y, w, κ)∥ 2 ≤ E Di κ 2 ∥x∥ 2 = κ 2 Tr(Σ i ) ≤ κ 2 d∥Σ i ∥ ≤ C 1 κ 2 d.This bound, combined with item d), proves the fourth claim in the lemma.Finally, for any unit vector u, we boundE Di (∇f (x, y, w, κ) • u) 2 ≤ κ 2 E Di (x • u) 2 ≤ κ 2 ∥Σ i ∥ ≤ κ 2 C 1 .This bound, combined with item e), shows the fifth claim in the lemma.We will utilize the following auxiliary lemma in the proof of the theorem. This lemma applies to general random variables. Lemma C.3. For any a ∈ R, b > 0 and a symmetric random variable z,E (a + z) - (a + z)b max(|a + z|, b) ≤ 2|a| Pr(z > b -|a|)Proof. We assume a ≥ 0 and prove the lemma for this case. The statement for a < 0 case then follows from symmetry.We rewrite the term inside the expectation in terms of indicator random variables:(a + z) - (a + z)b max(|a + z|, b) = (a + z -b) • 1(z > b -a) + (a + z + b) • 1(z < -b -a) = (a + z -b) • 1(b -a < z ≤ b + a) + (a + z -b) • 1(z > b + a) + (a + z + b) • 1(z < -b -a).Taking the expectation on both sides,E (a + z) - (a + z)b max(|a + z|, b) = E[(a + z -b) • 1(b -a < z ≤ b + a)] + E[(a + z -b) • 1(z > b + a)] + E[(a + z + b) • 1(z < -b -a)] = E[(a + z -b) • 1(b -a < z ≤ b + a)] + 2a Pr(z > b + a),where the last step follows because z is symmetric.Next,E (a + z) - (a + z)b max(|a + z|, b) = E[|a + z -b| • 1(b -a < z ≤ b + a)] + 2|a| Pr(z > b + a) ≤ E[|2a| • 1(b -a < z ≤ b + a)] + 2|a| Pr(z > b + a) ≤ 2|a| Pr(b -a < z ≤ b + a) + 2|a| Pr(z > b + a) = 2|a| Pr(z > b -a).Next, we proceed with the proof of Theorem C.2 using the aforementioned lemma.Proof of Theorem C.2. Let (x, y) be a random sample from distribution D 0 , and let η = y -w 0 • x denote the noise. Recall that η is independent of x.Note that:(w • x -y)x = ((w -w 0 ) • x -η)x.We will now evaluate the expected value of the unclipped gradient.E D0 [(w • x -y)x] = E D0 [((w -w 0 ) • x -η)x] = E D0 [((w -w 0 ) • x)x] = Σ 0 (w -w 0 ). (11)Next, we will bound the norm of the expected value of (w•x-y)x-∇f (x, y, w, κ), which represents the difference between the clipped gradient and the true gradient. We first expand this expression:EQUATIONNext, by applying Lemma C.3, we haveE η ((w -w 0 ) • x -η) - ((w -w 0 ) • x -η) |(w -w 0 ) • x -η| ∨ κ κ ≤ 2|(w -w 0 ) • x| • Pr(η > κ -|(w -w 0 ) • x|).Note that in the above expectation, we fixed x and took the expectation over noise η.Let Z := 1(|x • (w -w 0 )| ≥ κ/2). Observe that Pr(η > κ -|(w -w 0 ) • x|) ≤ Z + Pr(η > κ/2).Combining this observation with the above equation, we have:EQUATIONThen, for any unit vector v ∈ R d , we haveEQUATIONhere the second last inequality follows from Equation (12) and Equation (13). Next, we bound the two terms on the right one by one. We start with the first term:EQUATIONwhere (a) used the Cauchy-Schwarz inequality, (b) used the fact that Z is an indicator random variable, hence, Z 2 = Z, and the Cauchy-Schwarz inequality, (c) uses L4 -L2 hypercontractivity, and (d) follows from the definition of Z and the assumption that ∥Σ 0 ∥ ≤ C 1 .Similarly, we can show thatEQUATION)Applying the Markov inequality to η 2 we get:EQUATION)Similarly, applying the Markov inequality to |x • (w -w 0 )| 4 yields:EQUATION)where the last inequality uses L4 -L2 hypercontractivity.Combining Equations ( 14), (15), ( 16), ( 17) and (18), we have| E D0 [((w • x -y)x -∇f (x, y, w, κ)) • v]| ≤ 8C √ C 1 E D0 [(x • (w -w 0 )) 2 ] 3/2 κ 2 + 8 √ C 1 E D0 [η 2 ] E D0 [(x • (w -w 0 )) 2 ] 1/2 κ 2 ≤ 8C √ C 1 E D0 [(x • (w -w 0 )) 2 ] 1/2 ( E D0 [(x • (w -w 0 )) 2 ] + E D0 [η 2 ]/C) κ 2 (a) ≤ ϵ E D0 [(x • (w -w 0 )) 2 ] 1/2 √ C 1 • ( E D0 [(x • (w -w 0 )) 2 ] + E D0 [η 2 ]/C) E D0 [(y -x • w) 2 ] (b) ≤ ϵ∥w -w 0 ∥ • ( E D0 [(x • (w -w 0 )) 2 ] + E D0 [η 2 ]) E D0 [(y -x • w) 2 ] (c) = ϵ∥w -w 0 ∥,here inequality (a) follows from the lower bound on κ 2 in theorem, inequality (b) follows from Assumption 1b and C ≥ 1, and the last equality follows since y -x • w = x(w -w 0 ) + η, and x and η are independent.Note that the above bound holds for all unit vectors v, therefore,∥ E D0 [((w • x -y)x -∇f (x, y, w, κ))]∥ ≤ max ∥v∥ E D0 [((w • x -y)x -∇f (x, y, w, κ))] • v ≤ ϵ∥w -w 0 ∥.Algorithm 4 CLIPEST 1: Input: A collection of samples S * from D 0 , w, ϵ, δ ′ , σ, C, and C 1 . 2: Output: clipping parameter κ 3: T ← Θ(log 1/δ ′ ) 4: Divide S * into T equal parts randomly, and denote them as {S * j } j∈[T ] 5: θ ← Median 1 |S * j | (x,y)∈S * j (x • w -y) 2 : j ∈ [T ] 6: κ ← 32(C+1)C1(θ+17σ 2 )8(C+1)C1•E D 0 [(y-x•w) 2 ] ϵ ≤ κ ≤ 28 2(C+1)C1(E D 0 [((w-w0)•x) 2 ]+σ 2 ) ϵ .To prove the theorem, we will make use of the following lemma: Lemma D.2. Let S be a collection of m ≥ 64C 2 i.i.d. samples from D 0 and w ∈ R d , then with probability at least 7/8, the following holds:1 4 E D0 [(w • x -y) 2 ] -17σ 2 ≤ 1 m (x,y)∈S (y -w • x) 2 ≤ 3 E D0 [((w -w 0 ) • x) 2 ] + 32σ 2 .Proof. We start by expanding the expression:EQUATION)where the last inequality follows since for any a, b, we havea 2 + 2ab + b 2 ≥ a 2 /2 -b 2 .Similarly, we can show:EQUATION)Since S contains independent samples from D 0 , we have:E   1 m (x,y)∈S (((w -w 0 ) • x) 2   = E D0 [((w -w 0 ) • x) 2 ],andVar 1 m (x,y)∈S ((w -w 0 ) • x) 2 = Var D0 (((w -w 0 ) • x) 2 ) m ≤ E D0 [((w -w 0 ) • x) 4 ] m ≤ C E D0 [((w -w 0 ) • x) 2 ] 2 m ,where the last inequality follows from L4-L2 hypercontractivity.For any a > 0, using Chebyshev's inequality,EQUATION)Using the Markov inequality, for any a > 0, we have:EQUATION)By combining the equations above, we can derive the following inequality:With probability ≥ 1 -2 a 2 , the following holds:1 m (x,y)∈S (w • x -y) 2 ≥ 1 2 E D0 [((w -w 0 ) • x) 2 ](1 -a C √ m ) -a 2 σ 2 , and 1 m (x,y)∈S (w • x -y) 2 ≤ 2 E D0 [((w -w 0 ) • x) 2 ](1 + a C √ m ) + 2a 2 σ 2 .By choosing a = 4 and using m ≥ 64C 2 in the above equation, we can conclude that with probability ≥ 1 -2 a 2 , the following holds:1 4 E D0 [((w -w 0 ) • x) 2 ] -16σ 2 ≤ 1 m (x,y)∈S (w • x -y) 2 ≤ 3 E D0 [((w -w 0 ) • x) 2 ] + 32σ 2 .Next, note thatE D0 [(w • x -y) 2 ] = E D0 [((w -w 0 ) • x -(y -w 0 • x)) 2 ] (a) = E D0 [((w -w 0 ) • x) 2 ] + E D0 [(y -w 0 • x) 2 ] (b) ≤ E D0 [((w -w 0 ) • x) 2 ] + σ 2 ,here (a) follows since x is independent of the output noise y -w 0 •x, and (b) follows since the output noise y -w 0 • x is zero mean and has a variance at most σ 2 . Combining the above two equations completes the proof. Now we prove Theorem D.1 using the above lemma:Proof of Theorem D.1. From the previous lemma and Chernoff bound it follows that with probability ≥ 1 -δ ′ ,1 4 E D0 [(w • x -y) 2 ] -17σ 2 ≤ θ ≤ 3 E D0 [((w -w 0 ) • x) 2 ] + 32σ 2 .Then bound on κ follows from the relation κ = 32(C+1)C1(θ+17σ 2 )ϵ .As a part of gradient estimation in step r, the main algorithm 1 uses subroutine GRADSUBEST for B = B (r)s and w = ŵ(r) . Recall that B (r)s is a random subset of the collection of small batches B s . The purpose of this subroutine is to estimate a smaller subspace of R d such that for distribution D 0 , the expectation of the projection of the clipped gradient onto this subspace closely approximates the 4: A ← 1 2| B| b∈ B ∇f (S b 1 , w, κ)∇f (S b 2 , w, κ) ⊺ + ∇f (S b 2 , w, κ)∇f (S b 1 , w, κ) ⊺ 5: U ← [u 1 , u 2 , ..., u ℓ ],where {u i }'s are top ℓ singular vectors of A 6: Return U U ⊺ true expectation of the clipped gradient, for distribution D 0 . This reduction to a smaller subspace helps reduce the number of medium-sized batches and their required length in the subsequent part of the algorithm.The following theorem characterizes the final guarantee for subroutine GRADSUBEST.Theorem E.1. Let p 0 denote the fraction of batches in B that are sampled from D 0 . For any ϵ, δ ′ > 0, andB = Ω d αsϵ 2 1 αsϵ 2 + C 2 2 C1 log d δ ′ , p 0 ≥ α s /2 and ℓ ≥ min{k, 1 2αsϵ 2 }, with probability ≥ 1 -δ ′ , the projection matrix U U ⊺ returned by subroutine GRADSUBEST satisfy ∥(I -U U ⊺ ) E D0 [∇f (x, y, w, κ)]∥ ≤ 4ϵκ C 1 .The above theorem implies that the difference between the expectation of the clipped gradient and the expectation of projection of the clipped gradient for distribution D 0 is small. Next, we present the description of the subroutine GRADSUBEST and provide a brief outline of the proof for the theorem before formally proving it in the subsequent subsection. .i.d. samples from D i , E[u b ] = E[v b ] = E Di [∇f (x, y, w, κ)]. The subroutine defines A = b∈ B 1 2| B| u b (v b ) ⊺ + v b (u b ) ⊺ .Let p i denote the fraction of batches in B that have samples from D i . Then using the linearity of expectation, we have:E[A] = k-1 i=0 p i E Di [∇f (x, y, w, κ)] E Di [∇f (x, y, w, κ)] ⊺ .It is evident that if the matrix U is formed by selecting the top k singular vectors of E[A], then the projection of E D0 [∇f (x, y, w, κ)] onto U U ⊺ corresponds to itself, and the guarantee stated in the theorem holds. However, we do not have access to E[A], and furthermore, when the number of components k is large, it may be desirable to obtain a subspace of smaller size than k.To address the first challenge, Theorem E.2 in the next subsection shows that ∥A -E[A]∥ is small. This theorem permits the usage of A as a substitute for E[A]. The clipping operation, introduced in the previous subsection, plays a crucial role in the proof of Theorem E.2 by controlling the norm of the expectation and the covariance of the clipped gradient for other components, and the maximum length of clipped gradients across all components. This is crucial for obtaining a good bound on the number of small-size batches required. Additionally, the clipping operation ensures that the subroutine remains robust to arbitrary input-output relationships for other components. Furthermore, the clipping operation assists in addressing the second challenge by ensuring a uniform upper bound on the norm of the expectation of all components, i.e., ∥ E Di [∇f (x, y, w, κ)]∥ ≤ O(κ). Leveraging this property, Lemma E.3 demonstrates that it suffices to estimate the top ≈ 1/p 0dimensional subspace. Intuitively, this is because the infinitely many components can create at most approximately 1/p 0 directions with weights greater than p 0 , indicating that the direction of D 0 must be present in the top Θ(1/p 0 ) subspace.Since B = B (r)s is obtained by randomly partitioning B s into R subsets, and B s contains a fraction of at least α s batches with samples from D 0 , it holds with high probability that p 0 ≳ α s . Consequently, when ℓ ≥ min{k, Ω( 1 αs )}, the subspace corresponding to the top ℓ singular vectors of A satisfies the desired property in the Theorem E.1.We note that the construction of matrix A in subroutine GRADSUBEST is inspired by previous work [KSS + 20]. However, while they employed it to approximate the k-dimensional subspace of the true regression vectors for all components, we focus exclusively on one distribution D 0 at a time and recover a subspace such that, for distribution D 0 , the expectation of the projection of the clipped gradient on this subspace closely matches the true expectation of the clipped gradient.It is worth noting that, in addition to repurposing the subroutine from [KSS + 20], we achieve four significant improvements: 1) A more meticulous statistical analysis and the use of clipping enable our algorithm to handle heavy-tailed distributions for both noise and input distributions. 2) Clipping also facilitates the inclusion of arbitrary input-output relationships for other components. The next two improvements are attributed to an improved linear algebraic analysis. Specifically, our Lemma E.3 enhances the matrix perturbation bounds found in [LL18] and [KSS + 20]. These enhancements enable us to: 3) Provide meaningful guarantees even when the number of components k is very large, 4) reduce the number of batches required when the distance between the regression vectors is small.To prove Theorem E.1, in the following theorem, we will first demonstrate that the term ∥A -E[A]∥ is small when given enough batches.Theorem E.2. For 0 ≤ i ≤ k -1, let z i = E Di [∇f (x,y, w, κ)], and p i denote the fraction of batches in B that have samples from D i . For any ϵ, δ ′ > 0, andB = Ω d αsϵ 2 1 αsϵ 2 + C 2 2 C1 log d δ ′ , with probability at least 1 -δ ′ , A - k-1 i=0 p i z i z ⊺ i ≤ α s ϵ 2 κ 2 C 1 ,where A is the matrix defined in subroutine GRADSUBEST.Proof. Let Z b := ∇f (S b 1 , w, κ)∇f (S b 2 , w, κ) ⊺ . Note that A = 1 2| B| b∈ B (Z b + (Z b ) ⊺ ).Then, from the triangle inequality, we have:A - k-1 i=0 p i z i z ⊺ i ≤ 1 2 1 | B| b∈ B Z b - k-1 i=0 p i z i z ⊺ i + 1 2 1 | B| b∈ B (Z b ) ⊺ - k-1 i=0 p i z i z ⊺ i = 1 | B| b∈ B Z b - k-1 i=0 p i z i z ⊺ i .For a batch b sampled from distribution D i , we have:E[Z b ] = E[∇f (S b 1 , w, κ)∇f (S b 2 , w, κ) ⊺ ] = E[∇f (S b 1 , w, κ)] E[∇f (S b 2 , w, κ) ⊺ ] = E Di [∇f (x, y, w, κ)] E Di [∇f (x, y, w, κ) ⊺ ] = z i z ⊺ i ,where the second inequality follows since samples in S b 1 and S b 2 are independent, and the third equality follows from the linearity of expectation.EQUATIONEQUATIONTo complete the proof, we will prove a high probability bound on the term on the right by applying the Matrix Bernstein inequality. To apply this inequality, we first upper bound |Z b | as follows:∥Z b ∥ = ∥∇f (S b 1 , w, κ)∇f (S b 2 , w, κ) ⊺ ∥ ≤ ∥∇f (S b 1 , w, κ)∥ • ∥∇f (S b 2 , w, κ) ⊺ ∥.From item 3 in Lemma C.1, we have ∥∇f(S b 1 , w, κ)∥ ≤ κC 2 √ d almost surely, and ∥∇f (S b 2 , w, κ)∥ ≤ κC 2 √ d almost surely. It follows that ∥Z b ∥ ≤ κ 2 C 2 2 d. Therefore, ∥Z b - E[Z b ]∥ ≤ 2κ 2 C 2 2 d. Next, we will provide an upper bound for E b∈ B (Z b -E[Z b ]) b∈ B (Z b -E[Z b ]) ⊺ : E b∈ B (Z b -E[Z b ]) b∈ B (Z b -E[Z b ]) ⊺ = E b∈ B (Z b -E[Z b ])(Z b -E[Z b ]) ⊺ ≤ | B| max b∈ B E (Z b -E[Z b ])(Z b -E[Z b ]) ⊺ ≤ | B| max b∈ B E (Z b (Z b ) ⊺ ≤ | B| max b∈ B E[∥∇f (S b 2 , w, κ)∥ 2 ] • E ∇f (S b 1 , w, κ)∇f (S b 1 , w, κ) ⊺ ≤ | B| max b∈ B,u:∥u∥=1 E[∥∇f (S b 2 , w, κ)∥ 2 ] • E (∇f (S b 1 , w, κ) • u) 2 .From item 4 and item 5 in lemma C.1, we have:EQUATIONE (∇f (S b 1 , w, κ) • u) 2 ≤ κ 2 C 1 .Combining these two bounds, wee obtain:E     b∈ B (Z b -E[Z b ])     b∈ B (Z b -E[Z b ])   ⊺   ≤ | B|dκ 4 C 1 2 .Due to symmetry, the same bound holds forE b∈ B (Z b -E[Z b ]) ⊺ b∈ B (Z b -E[Z b ]) . Finally, by applying the Matrix Bernstein inequality, we have:Pr   1 | B| b∈ B (Z b -E[Z b ]) ≥ α s ϵ 2 κ 2 C 1   ≤ 2d exp - | B| 2 θ 2 | B|dκ 4 C 1 2 + | B|θ(2C 2 2 κ 2 d).For B = Ω d αsϵ 2 1 αsϵ 2 + C 2 2 C1 log d δ ′, the quantity on the right-hand side is bounded by δ ′ .Therefore, with probability at least 1 -δ ′ , we have:1 | B| b∈ B (Z b -E[Z b ]) ≤ α s ϵ 2 κ 2 C 1 .Combining the above equation with Equation (23) completes the proof of the Theorem.In the proof of Theorem E.1, we will utilize the following general linear algebraic result:Lemma E.3. For z 0 , z 1 , ..., z k-1 ∈ R d and a probability distribution (p 0 , p 1 , ..., p k-1 ) over k elements, let Z = k-1 i=0 p i z i z ⊺ i . For a symmetric matrix M and ℓ > 0, let u 1 , u 2 , .., u ℓ be top ℓ singular vectors of M and let U = [u 1 , u 2 , ..., u ℓ ] ∈ R d×ℓ , then we have:∥(I -U U ⊺ )z 0 ∥ 2 ≤ 2(ℓ+1)∥M -Z∥+maxj ∥zj ∥ 2 (ℓ+1)p0 ℓ < k 2∥M -Z∥ p0 if ℓ ≥ k.Lemma E.3 provides a bound on the preservation of the component z 0 by the subspace spanned by the top-ℓ singular vectors of a symmetric matrix M . This bound is expressed in terms of the spectral distance between matrices Z and M , the maximum norm of any z i , and the weight of the component corresponding to z 0 in Z. The proof of Lemma E.3 can be found in Section I.Utilizing Lemma E.3 in conjunction with Theorem E.2, we proceed to prove Theorem E.1.Proof of Theorem E.1. From Lemma E.3, we have the following inequality:∥(I -U U ⊺ ) E D0 [∇f (x, y, w, κ)]∥ 2 ≤    2(ℓ+1)∥A-k-1 i=0 piziz ⊺ i ∥+maxj ∥ E D j [∇f (x,y,w,κ)]∥ 2 (ℓ+1)p0 ℓ < k 2∥A-k-1 i=0 piziz ⊺ i ∥ p0 if ℓ ≥ k.By applying Theorem E.2 and utilizing item 1 of Lemma C.1, it follows that with a probability of at least 1 -δ ′ , we have:∥(I -U U ⊺ ) E D0 [∇f (x, y, w, κ)]∥ 2 ≤ 2αsϵ 2 κ 2 C1 p0 + κ 2 C1 (ℓ+1)p0 ℓ < k 2αsϵ 2 κ 2 C1 p0 if ℓ ≥ k.The theorem then follows by using p 0 ≥ α s /2 and ℓ ≥ min{k, 1 2αsϵ 2 }.Recall that in gradient estimation for step r, Algorithm 1 utilizes the subroutine GRADSUBEST to find a projection matrix P (r) for an ℓ-dimensional subspace. In the previous section, we showed that the difference between the expectation of the clipped gradient and the expectation of projection of the clipped gradient on the subspace for distribution D 0 is small. Therefore, it suffices to estimate the expectation of projection of the clipped gradient on the subspace.The main algorithm 1 passes the medium-sized batches B = Bm , the ℓ-dimensional projection matrix P = P (r) , and a collection of i.i.d. samples S * = S b * ,(r) 2 from D 0 to the subroutine GRADEST. Here, Bm is a random subset of the collection of medium-sized batches B m . The purpose of the GRADEST subroutine is to estimate the expected value of the projection of the clipped gradient onto the ℓ-dimensional subspace defined by the projection matrix P . Since the subroutine operates on a smaller ℓ-dimensional subspace, the minimum batch size required for the batches in B m and the number of batches required depend on ℓ rather than d.The following theorem characterizes the final guarantee for the GRADEST subroutine: Theorem F.1. For subroutine GRADEST, let n m denote the length of the smallest batch in B, N denote the number of batches in B that has samples from D 0 and P be a projection matrix for some ℓ dimensional subspace ofR d . If T 1 ≥ Ω(log | B| δ ′ ), T 2 ≥ Ω(log 1 δ ′ ), n m ≥ 4T 1 Ω( √ ℓ ϵ 2 ), |S * | ≥ 2T 1 Ω( √ ℓ ϵ 2 ) samples, and N • n m ≥ T 2 Ω( ℓ ϵ 2 ), then with probability ≥ 1 -2δ ′ the estimate ∆ returned by subroutine GRADEST satisfy ∥∆ -E D0 [P ∇f (x, y, w, κ)]∥ ≤ 9ϵκ C 1 .The above theorem implies that when the length of medium-sized batches is Ω( √ ℓ) and the number of batches in B containing samples from D 0 is Ω(ℓ), the GRADEST subroutine provides a reliable estimate of the projection of the clipped gradient onto the ℓ-dimensional subspace defined by the projection matrix P .Next, we provide a description of the GRADEST subroutine and present a brief outline of the proof for the theorem before formally proving it in the subsequent subsection.In the GRADEST subroutine, the first step is to divide the samples in each batch of B into two equal parts. By utilizing the first half of the samples in a batch b along with the samples S * , it estimates whether the expected values of the projection of the clipped gradient for D 0 and the distribution used for the samples in b are close or not. With high probability, the algorithm retains all the batches from D 0 while rejecting batches from distributions where the difference between the two expectations is large. To achieve this with an ℓ-dimensional subspace, we require Ω( √ ℓ) samples in each batch (see Lemma F.6).Following the rejection process, the GRADEST subroutine proceeds to estimate the projection of the clipped gradients within this ℓ-dimensional subspace using the second half of the samples from the retained batches. To estimate the gradient accurately in the ℓ-dimensional subspace, Ω(ℓ) samples are sufficient (see Lemma F.7). To obtain guarantees with high probability, the procedure employs the median of means approach, both for determining which batches to keep and for estimation using the retained batches.We prove the theorem formally in the next subsection.The following lemma provides an upper bound on the covariance of the projection of the clipped gradients. Lemma F.2. Consider a collection S of m i.i.d. samples from distribution D i . For κ > 0, w ∈ R d and a projection matrix P for an ℓ dimensional subspace of R d , we haveE[P ∇f (S, w, κ)] = E Di [P ∇f (x, y, w, κ)],and ∥Cov(P ∇f (S, w, κ))∥ ≤ 2κ 2 m C 1 and Tr (Cov(P ∇f (S, w, κ))) ≤ ℓ∥Cov(P ∇f (S, w, κ))∥.Proof. Note that, E P ∇f (S, w, κ) = P E[∇f (S, w, κ)]] = P E Di [∇f (x, y, w, κ)] = E Di [P ∇f (x, y, w, κ)], where the second-to-last equality follows from Lemma C.1. This proves the first part of the lemma. To prove the second part, we bound the norm of the covariance matrix:Cov(P ∇f (S, w, κ)) = max ∥u∥≤1 Var(u ⊺ P ∇f (S, w, κ)) = max ∥v∥≤1 Var(v ⊺ ∇f (S, w, κ)) ≤ ∥Cov(∇f (S, w, κ))∥ ≤ κ 2 m C 1 ,where the last inequality follows from Lemma C.1. Similarly,Cov(P ∇f (S ′ , w, κ)) ≤ κ 2 m C 1 .Finally, since random vector P ∇f (S ′ , w, κ) lies in ℓ dimensional subspace of R d , corresponding to projection matrix P , hence its covariance matrix has rank ≤ ℓ. Hence, the relation Tr (Cov(P ∇f (S, w, κ))) ≤ ℓ∥Cov(P ∇f (S, w, κ))∥ follows immediately.The following corollary is a simple consequence of the previous lemma: Corollary F.3. Consider two collections S and S ′ each consisting of m i.i.d. samples from distributions D i and D 0 , respectively. For κ > 0, w ∈ R d and a projection matrix P for an ℓ dimensional subspace of R d , let z = P ∇f (S, w, κ) -∇f (S ′ , w, κ) , we have:E[z] = E Di [P ∇f (x, y, w, κ)] -E D0 [P ∇f (x, y, w, κ)],and ∥Cov(z)∥ ≤ 4κ 2 m C 1 and Tr (Cov(z)) ≤ ℓ∥Cov(z)∥.Proof. The expression for E[z] can be obtained from the previous lemma and the linearity of expectation.To prove the second part, we bound the norm of the covariance matrix of z.Cov(z) = Cov(P ∇f (S, w, κ) -∇f (S ′ , w, κ) ) ≤ 2(Cov(P ∇f (S, w, κ)) + Cov(P ∇f (S ′ , w, κ))).Using the bounds from the previous lemma, we can conclude that ∥Cov(z)∥ ≤ 4κ 2 m C 1 . Finally, since the random vector z lies in the ℓ-dimensional subspace of R d defined by the projection matrix P , its covariance matrix has rank ≤ ℓ. Therefore, we have Tr, (Cov(z)) ≤ ℓ∥Cov(z)∥.The following theorem bounds the variance of the dot product of two independent random vectors. It will be helpful in upper bounding the variance of ζ b j (defined in subroutine GRADEST). Theorem F.4. For any two independent random vectors z 1 and z 2 , we have:Var(z 1 • z 2 ) ≤ 3Tr(Cov(z 1 )) • ∥Cov(z 2 )∥ + 3∥ E[z 1 ]∥ 2 ∥Cov(z 2 )∥ + 3∥ E[z 2 ]∥ 2 ∥Cov(z 1 )∥.Proof. We start by expanding the variance expression:Var(z 1 • z 2 ) = Var(z 1 • z 2 -E[z 1 • z 2 ]) = Var(z 1 • z 2 -E[z 1 ] • E[z 2 ]) = Var((z 1 -E[z 1 ]) • (z 2 -E[z 2 ]) + E[z 1 ] • z 2 + E[z 2 ] • z 1 ) ≤ 3 Var((z 1 -E[z 1 ]) • (z 2 -E[z 2 ])) + 3 Var(E[z 1 ] • z 2 ) + 3 Var(E[z 2 ] • z 1 ) = 3 Var((z 1 -E[z 1 ]) • (z 2 -E[z 2 ])) + 3 E[z 1 ] ⊺ Cov(z 2 ) E[z 1 ] + E[z 2 ] ⊺ Cov(z 1 ) E[z 2 ] ≤ 3 Var((z 1 -E[z 1 ]) • (z 2 -E[z 2 ])) + 3∥ E[z 1 ]∥ 2 ∥Cov(z 2 )∥ + 3∥ E[z 2 ]∥ 2 ∥Cov(z 1 )∥.To complete the proof, we bound the first term in the last expression:Var((z 1 -E[z 1 ]) • (z 2 -E[z 2 ])) = E ((z 1 -E[z 1 ]) • (z 2 -E[z 2 ])) 2 = E[(z 1 -E[z 1 ]) ⊺ Cov(z 2 )(z 1 -E[z 1 ])] ≤ E ∥z 1 -E[z 1 ]∥ 2 • ∥Cov(z 2 )∥ = Tr(Cov(z 1 )) • ∥Cov(z 2 )∥.Using the two previous results, we can establish a bound on the expectation and variance of ζ b j . Lemma F.5. In subroutine GRADEST, let P be a projection matrix of an ℓ dimensional subspace. Suppose S * j has ≥ m i.i.d. samples from D 0 and, S b 1,j and S b 1,j+T1 have ≥ m i.i.d. samples from D i for some i ∈ {0, 1, ..., k -1}. Than we have:E[ζ b j ] = ∥E Di [P ∇f (x, y, w, κ)] -E D0 [P ∇f (x, y, w, κ)]∥ 2 and Var(ζ b j ) ≤ 48 m 2 κ 4 ℓC 1 2 + 24 m κ 2 E[ζ b j ]C 1 .Proof. Let z 1 = P ∇f (S b 1,j , w, κ) -∇f (S * j , w, κ) and z 2 = P ∇f (S b 1,T1+j , w, κ) -∇f (S * T1+j , w, κ) . From Corollary F.3, we know thatE[z 1 ] = E[z 2 ] = E Di [P ∇f (x, y, w, κ)] -E D0 [P ∇f (x, y, w, κ)], Cov(z 1 ) = Cov(z 2 ) = 4κ 2 m C 1 and Tr(Cov(z 1 )) = Tr(Cov(z 2 )) = 4 m ℓκ 2 C 1 . Therefore, ∥ E[∆ i ] -E D0 [P ∇f (x, y, w, κ)]∥ = 1 | B| b∈ B E[z b i ] -E D0 [P ∇f (x, y, w, κ)] ≤ max b∈ B E[z b i ] -E D0 [P ∇f (x, y, w, κ)] ≤ 2ϵκ C 1 . (24)Next, from Lemma F.2,EQUATIONwhere the two equalities follow because for all batches b ∈ B,|S b 2,i | = |S b 2 |/T 2 and |S b 2 | = |S b |/2. Then ∥Cov(∆ i )∥ = 1 | B| max b∈ B ∥Cov(z b i )∥ ≤ 2T 2 κ 2 C 1 | B| • min b∈ B |S b | Since ∆ i lies in an ℓ dimensional subspace of R d , it follows that Tr(Cov(∆ i )) ≤ ℓ∥Cov(∆ i )∥ ≤ 2ℓT 2 κ 2 C 1 | B| • min b∈ B |S b | Note that Var(∥∆ i -E[∆ i ]∥) = Tr(Cov(∆ i )).Then, from Chebyshev's bound:Pr[∥∆ i -E[∆ i ]∥ ≥ ϵκ C 1 ] ≤ Var(∥∆ i -E[∆ i ]∥) ϵ 2 κ 2 C 1 ≤ 2ℓT 2 ϵ 2 | B| • min b∈ B |S b | ≤ 1/8.Combining above with Equation (24),Pr[∥∆ i -E D0 [P ∇f (x, y, w, κ)]∥ ≥ 3ϵκ C 1 ] ≤ 1/4. Let D := {i ∈ [T 2 ] : ∥∆ i -E D0 [P ∇f (x, y, w, κ)]∥ ≤ 3ϵκ √ C 1 ∥}.Then, for T 2 = Ω(log 1 δ ′ ), with probability ≥ 1 -δ ′ , we have|D| ≥ 1 2 T 2 .Recall that in the subroutine, we definedξ i = median{j ∈ [T 2 ] : ∥∆ i -∆ j ∥} and i * = arg min{i ∈ [T 2 ] : ξ i }.From the definition of D, and triangle inequality, for all i, j ∈ D, we have∥∆ i -∆ j ∥ ≤ 6ϵκ √ C 1 . Therefore, if |D| ≥ 1 2 T 2 , then for any i ∈ D, ξ i ≤ 6ϵκ √ C 1 . This would imply ξ i * ≤ 6ϵκ √ C 1 . Furthermore, since |D| ≥ 1 2 T 2 , there exist at least one i ∈ D such that ∥∆ i -∆ i * ∥ ≤ 6ϵκ √ C 1 .Using the definition of D, and the triangle inequality, we can conclude that∥∆i * -E D0 [P ∇f (x, y, w, κ)]∥ ≤ ∥∆ i -E D0 [P ∇f (x, y, w, κ)]∥ + ∥∆ i -∆ i * ∥ ≤ 9ϵκ C 1 .Theorem F.1 then follows by combining lemmas F.6 and F.7.The following lemma shows that with a sufficiently accurate estimation of the expectation of gradients, a logarithmic number of gradient descent steps are sufficient in the main algorithm 1.Lemma G.1. For ϵ > 0, suppose ∥∆ (r) -Σ 0 ( ŵ(r) -w 0 )∥ ≤ 1 2 ∥ ŵ(r) -w 0 ∥ + ϵσ 4 , and R = Ω(C 1 log ∥w0∥ σ ), then ∥ ŵ(r) -w 0 ∥ ≤ ϵσ.Proof. Recall that ŵ(r+1) = ŵ(r) -1 C1 ∆ (r) . Then we have:ŵ(r+1) -w 0 = ŵ(r) -w 0 - 1 C 1 ∆ (r) = ( ŵ(r) -w 0 ) I - 1 C 1 Σ 0 + 1 C 1 (Σ 0 ( ŵ(R+1) -w 0 ) -∆ (r) ).Using triangle inequality, we obtain:∥ ŵ(r+1) -w 0 ∥ ≤ ∥ ŵ(r) -w 0 ∥ I - 1 C 1 Σ 0 + 1 C 1 ∥Σ 0 ( ŵ(r) -w 0 ) -∆ (r) ∥ ≤ ∥ ŵ(r) -w 0 ∥ 1 - 1 C 1 + 1 C 1 ∥ ŵ(r) -w 0 ∥ 2 + ϵσ 4 ≤ ∥ ŵ(r) -w 0 ∥ 1 - 1 2C 1 + ϵσ 4C 1 .Using recursion, we have:∥ ŵ(R+1) -w 0 ∥ ≤ ∥ ŵ(1) -w 0 ∥ 1 - 1 2C 1 R + R-1 i=0 1 - 1 2C 1 i ϵσ 4C 1 ≤ ∥ ŵ(1) -w 0 ∥ exp - R 2C 1 + 2C 1 ϵσ 4C 1 ≤ ϵσ,where the second inequality follows from the upper bound on the sum of infinite geometric series and the last inequality follows from the bound on R and ŵ(1) = 0.Proof of Theorem 2.1. We show that with probability ≥ 1 -δ, for each r ∈ [R], the gradient computed by the algorithm satisfies ∥∆ (r) -Σ 0 ( ŵ(r) -w 0 )∥ ≤ 1 2 ∥ ŵ(r) -w 0 ∥ + ϵσ 4 . Lemma G.1 then implies that for R = Ω(C 1 log M σ ), the output returned by the algorithm ŵ = ŵ(R+1) satisfy ∥ ŵ -w 0 ∥ ≤ ϵσ.To show this, we fix r, and for this value of r, we show that with probability ≥ 1 -δ/R, ∥∆ (r) -Σ 0 ( ŵ(r) -w 0 )∥ ≤ 1 2 ∥ ŵ(r) -w 0 ∥ + ϵσ 4 . Since each round uses an independent set of samples, the theorem then follows by applying the union bound.First, we determine the bound on the clipping parameter. From Theorem D.1, for |S b * |/R = Ω(C 2 log 1/δ ′ ), with probability ≥ 1 -δ ′ , we haveEQUATIONNext, employing Theorem C.2 and utilizing the lower bound on the clipping parameter in the above equation, we obtain the following bound on the norm of the expected difference between clipped and unclipped gradients:EQUATIONRecall that in B s , at least α s fraction of the batches contain samples from D 0 . When B s is divided into R equal random parts, w.h.p. each part B (r) swill have at least α s fraction of the batches containing samples from D 0 .From Theorem E.1, if |B (r)s | = |Bs| R = Ω d αsϵ 2 2 1 αsϵ 2 2 + C 2 2 C1 log d δ ′, then with probability ≥ 1-δ ′ , the projection matrix P (r) satisfies ∥ E D0 [∇f (x, y, w (r) , κ (r) )] -P (r) E D0 [∇f (x, y, w (r) , κ (r) )]∥ ≤ 4ϵ 2 κ (r) C 1 .(28)The above equation shows subroutine GRADSUBEST finds projection matrix P (r) such that the expected value of clipped gradients projection is roughly the same as the expected value of the clipped gradient.Next, we show that subroutine GRADEST provides a good estimate of the expected value of clipped gradients projection. Let N denote the number of batches in B m that have samples from D 0 . If N ≥ Ω(R + log 1/δ ′ ) then with probability ≥ 1 -δ ′ , Bm has Θ(N/R) batches sampled from D 0 . If each batch in B m and batch b * has more than n m samples, nm R = Ω(√ ℓ ϵ 2 2 log( |Bm| δ ′ )),andEQUATIONCombining the above three equations using triangle inequality,EQUATIONwith probability ≥ 1 -5δ ′ .In equation ( 26) using the upper bound,E D0 [(x•(w (r) -w 0 )) 2 ] ≤ ∥w (r) -w 0 ∥ 2 ∥Σ 0 ∥ ≤ C 1 ∥w (r) - w 0 ∥ 2 we get κ (r) ≤ 28 2(C + 1)C 1 2 ∥w (r) -w 0 ∥ 2 + (C + 1)C 1 σ 2 √ ϵ 1 ≤ 28 2(C + 1) √ ϵ 1 (C 1 ∥w (r) -w 0 ∥ + C 1 σ).Combining the two equations,EQUATIONwith probability ≥ 1 -5δ ′ There exist universal constants c 1 , c 2 > 0 such that for ϵ 1 = c 1 andϵ 2 = c2 C1 √ C+1 ϵ + 1 √ C1, the quantity on the right is bounded by ∥w (r) -w 0 ∥/2 + ϵσ/4. We choose these values for ϵ 1 and ϵ 2 and δ ′ = δ 5R .From the above discussion, it follows that ifEQUATIONB m has ≥ 1 nm Ω ℓ ϵ 2 batches sampled from D 0 , then with probability ≥ 1 -δ/R, ∥∆ (r) -Σ 0 ( ŵ(r) -w 0 )∥ ≤ 1 2 ∥ ŵ(r) -w 0 ∥ + ϵσ 4 .Using ℓ = min{k, 1 ϵ 2 αs }, we get the bounds on the number of samples and batches required by the algorithm.To establish the lemma, we first introduce and prove two auxiliary lemmas. Lemma I.1. For k > 0, and a probability distribution (p 0 , p 1 , ..., p k-1 ) over k elements, let Z = k-1 i=0 p 0 z i z ⊺ i , where z i are d-dimensional vectors. Then for all ℓ ≥ 0, ℓ th largest singular value of Z is bounded by max i ∥z i ∥ 2 /ℓ.Proof. Note that Z is a symmetric matrix, so its left and right singular values are the same. Let v 1 , v 2 , ... be the singular vectors in the SVD decomposition of Z, and let a 1 ≤ a 2 ≤ a 3 ≤ ... be the corresponding singular values. Using the properties of SVD, we have:i a i = i v ⊺ i Zv i = i v ⊺ i   k-1 j=0 p j z j z ⊺ j   v i = k-1 j=0 p j i (v i • z j ) 2 ≤ k-1 j=0 p j ∥z j ∥ 2 ≤ max j ∥z j ∥ 2 .Next, we have:i a i ≥ i≤ℓ a i ≥ i≤ℓ a ℓ = ℓ • a ℓ .Combining the last two equations yields the desired result.Lemma I.2. Let u 1 , u 2 , .., u ℓ ∈ R d be ℓ mutually orthogonal unit vectors, and let U = [u 1 , u 2 , ..., u ℓ ] ∈ R d×ℓ . For any set of k vectors z 0 , z 1 , ..., z k-1 ∈ R d , non-negative reals p 0 , p 1 , ..., p k-1 , and reals a 1 , a 2 , ..., a ℓ , we have:∥(I -U U ⊺ )z 0 ∥ 2 ≤ k-1 i=1 p i z i z ⊺ i -j∈[ℓ] a j u j u ⊺ j p 0 . Proof. Let v = (I -U U ⊺ )z 0 .First we show that for all j ∈ [ℓ], the vectors v and u j are orthogonal,u ⊺ j (I -U U ⊺ )z 0 = (u ⊺ j • z 0 ) -(u ⊺ j • z 0 ) = 0.Then,v ⊺ k-1 i=0 p i z i z ⊺ i -j∈[ℓ] a j u j u ⊺ j v = v ⊺ k-1 i=0 p i z i z ⊺ i v = k-1 i=0 p i (z ⊺ i v) 2 ≥ p 0 (z ⊺ 0 v) 2Next, we have:z ⊺ 0 v = z ⊺ 0 (I -U U ⊺ )v + z 0 U U ⊺ v = z ⊺ 0 (I -U U ⊺ )v = v ⊺ v = ∥v∥ 2Combining the last two equations, we obtain:∥v∥ 2 • k-1 i=0 p i z i z ⊺ i -j∈[ℓ] a j u j u ⊺ j ≥ v ⊺ k-1 i=0 z i z ⊺ i -j∈[ℓ] a j u j u ⊺ j v ≥ p 0 ∥v∥ 4 .Dividing both sides by ∥v∥ 2 completes the proof.Next, combining the above two auxiliary lemmas we prove Lemma E.3.Proof of Lemma E.3. Let Λ i (•) denote the i th largest singular value of a matrix. Let M be rank ℓ truncated-SVD of M , then it follows that,∥M -M ∥ = Λ ℓ+1 (M ).First, we consider the case ℓ < k. By applying Weyl's inequality for singular values, we haveΛ ℓ+1 (M ) ≤ Λ ℓ+1 (Z) + Λ 1 (M -Z) ≤ max j ∥z j ∥ 2 ℓ + 1 + ∥M -Z∥,where the last equation follows from Lemma I.1.First applying the triangle inequality, and then using the above two equations, we have∥ M -Z∥ ≤ ∥M -M ∥ + ∥M -Z∥ ≤ max j ∥z j ∥ 2 ℓ + 1 + 2∥M -Z∥.Combining the above equation with Lemma I.2, we have:∥(I -U U ⊺ )z 0 ∥ 2 ≤ 2(ℓ + 1)∥M -Z∥ + max j ∥z j ∥ 2 (ℓ + 1)p 0 .This completes the proof for ℓ < k. To prove for the case ℓ > k, we use Λ ℓ+1 (Z) = 0 in place of the bound Λ ℓ+1 (Z) ≤ maxj ∥zj ∥ 2 ℓ+1 in the above proof for the case ℓ < k.To simplify our analysis, we made two assumptions about the data distributions. We now argue that these assumptions are not limiting.The first additional assumption was that there exists a constant C 2 > 0 such that for all components i ∈ {0, 1, . . . , k -1} and random samples (x, y) ∼ D i , we have ∥x∥ ≤ C 2 √ d almost surely. In the non-batch setting, Cherapanamjeri et al. (2020) [CAT + 20] have shown that this assumption is not limiting. They showed that if other assumptions are satisfied, then there exists a constant C 2 such that with probability ≥ 0.99, we have ∥x∥ ≤ C 2 √ d. Therefore, disregarding the samples for which |x| > C 2 √ d does not significantly reduce the data size. Moreover, it has minimal impact on the covariance matrix and hypercontractivity constants of the distributions. This argument easily extends to the batch setting. In the batch setting, we first exclude samples from batches where ∥x∥ > C 2 √ d. Then we remove small-sized batches with fewer than or equal to 2 samples and medium-sized batches that have been reduced by more than 10% of their original size. It is easy to show that w.h.p. the fraction of medium and small size batches that gets removed for any component is at most 10%. THence, this assumption can be removed with a small increase in the batch size and the number of required samples in our main results.Next, we address the assumption that the noise distribution is symmetric. We can handle this by employing a simple trick. Consider two independent and identically distributed (i.i.d.) samples (x 1 , y 1 ) and (x 2 , y 2 ), wherey i = w * • x i + η i . We define x = (x 1 -x 2 )/ √ 2, y = (y 1 -y 2 )/ √ 2, and η = (η 1 -η 2 )/ √ 2.It is important to note that the distribution of η is symmetric around 0, and the covariance of x is the same as that of x i , while the variance of η is the same as that of η i . Furthermore, we have y = w * • x + η. Therefore, the new sample (x, y) obtained by combining two i.i.d. samples satisfies the same distributional assumptions as before, and in addition, the noise distribution is symmetric. We can combine every two samples in a batch using this approach, which only reduces the batch size of each batch by a constant factor of 1/2. Thus, the assumption of symmetric noise can be eliminated by increasing the required batch sizes in our theorems by a factor of 2.The algorithm requires two distribution parameters: the L4-L2 hypercontractivity parameter C and the condition number of input covariance matrix C 1 . These parameters significantly generalize the input distributions typically considered in mixed linear regression, which have primarily focused on sub-Gaussian distributions with an identity covariance matrix. Specifically, the L4-L2 hypercontractivity parameter generalizes the sub-gaussian distribution assumption and allows for heavier-tailed distributions. The condition-number C 1 of the input covariance matrices generalizes the identity covariance matrix assumption for input distribution, corresponding to C 1 = 1.Additionally, the algorithm requires two population parameters: the number of sub-populations k and the fraction of batches α s present from the relevant sub-population. Note that the algorithm combines these parameters to define a single parameter ℓ.Finally, the algorithm requires two regression parameters: the length M of the regression vector, and the additive noise variance σ.In total the algorithm relies on six parameters, two of which can be combined into one. Although this may seem large, it requires 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 parameter, dividing its possible range logarithmically. This approach increases the size of the resulting list at most logarithmically, ensuring that at least one estimate in the list will be tight.Setup. We have sets B s and B m of small and medium size batches and k distributions D i for i ∈ {0, 1, . . . , k -1}. For a subset of indices I ⊆ {0, 1, . . . , k -1}, both B s and B m have a fraction All small batches have 2 samples each, while medium-size batches have n m samples each, which we vary from 4 to 32, as shown in the plots. We fix data dimension d = 100, α = 1/16, number of small batches to |B s | = min{8dk 2 , 8d/α 2 } and the number of medium batches to |B m | = 256. In all the plots, we average our 10 runs.Evaluation. Our objective is to recover a small list containing good estimates for the regression vectors of D i for each i ∈ I. We compare our proposed algorithm's performance with that of the algorithm in [KSS + 20]. We generate lists of regression vector estimates L Ours and L KSSKO using our algorithm and [KSS + 20], respectively. Then, we create 1600 new batches, each containing n new i.i.d samples randomly drawn from the distribution D i , where for each batch, index i is chosen randomly from I.Each list enables the clustering of the new sample batches. To cluster a batch using a list, we assign it to the regression vector in the list that achieves the lowest mean square error (MSE) for its samples.To evaluate the average MSE for each algorithm, for each clustered batch, we generate additional samples from the distribution that the batch was generated from and calculate the error achieved by the regression vector in the list that the batch was assigned to. We then take the average of this error over all sets. We evaluate both algorithms' performance for new batch sizes n new = 4 and n new = 8, as shown in the plots.Minimum distance between regression vectors. Our theoretical analysis suggests that our algorithm is robust to the case when the minimum distance between the regression vectors are much smaller than their norms. In order to test this, in Figure 3 , we generate half of the regression vectors with elements independently and randomly distributed in U [9, 11], and the other half with elements independently and randomly distributed in U [-11, -9]. Notably, the minimum gap between the vectors, in this case, is much smaller than their norm. It can be seen that the performance gap between our algorithm and the one in [KSS + 20] increases significantly as we deviate from the assumptions required for the latter algorithm to work.Number of different distributions. Our algorithm can notably handle very large k (even infinite) while still being able to recover regression vectors for the subgroups that represent sufficient fraction of the data. In the last plot, we set k = 100 and I = {0, 1, 2, 3} to highlight this ability. In this case, the first four distributions each generate a 1/16 fraction of batches, and the remaining 96 distributions each generate a 1/128 fraction of batches. We provide the algorithm with one additional medium-size batch from D i for each i ∈ I for identification of a list of size I. The results are plotted in Figure 4 , where we can see that the performance gets better with medium batch size as expected. Note that the algorithm in [KSS + 20] cannot be applied to this scenario.• 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.Justification: We have read the code of ethics and our submission abide by that.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).We note that these assumptions are standard in heavy-tailed linear regression where some condition on the anti-concentration of E[xx ⊺ ] is necessary to achieve finite sample complexities [LM16, Oli16, CAT + 20], even in the specific case of a single sub-population.This simple approach showed better empirical performance than Algorithm 3, whose theoretical guarantees we described in Theorem 2.3 


Decision: Five reviewers evaluated the paper, and their overall assessment is positive. I agree with their evaluation and believe the paper makes a strong contribution with compelling results. The reviewers particularly appreciated that the paper addresses an important topic—linear regression with heterogeneous data—and delivers theoretically sound results without relying on unrealistic assumptions on the distribution of the data. Both the reviewers and the AC believe that this is one of the first papers that can handle inputs with non-isotropic and heavy-tailed in heterogeneous settings. Moreover, the paper is exceptionally well-written and easy to read. Overall, the AC holds a very positive opinion of this work.

Reviews: review: paper_summary: paper_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].
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: See weakness.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: What are the main intuitions behind the assumptions made in Remark 2.1? (L. 196)
limitations: limitations: -
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: What are the barriers to implementing/trying this on real data?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See Weaknesses.
limitations: limitations: The limitations have been partially addressed in the paper.

Comments: title: Official Comment
Response to rebuttal
Response
comment: Thanks to the authors for addressing my questions, and points about the presentation. I will maintain my score and positive review.
Thank you for addressing the feedback provided. After reviewing the revisions, I am pleased with the improvements and will be raising my score accordingly.
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.
Dear authors,

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