Linear Regression using Heterogeneous Data Batches
Authors: Ayush Jain
Rajat Sen
Weihao Kong
Abhimanyu Das
Alon Orlitsky
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Machine Learning
Liner Regression
Heterogeneous Data
Mixed linear regression
List decodable learning
Primary area: learning_theory

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.

Linear Regression using Heterogeneous Data Batches
Ayush Jain" Granica Computing Inc. ayush.jaingge jainggraniccan...... ai
Rajat Sen Google Research
Weihao Kong Google Research
janrajat@google.co com
veihaaokongggoogle.com
Abhimanyu Das
Alon Orlitsky UC San Diego alon@ucsd. edu
Google Research
abhidas@google.com com
Abstract
In many learning applications, data collected from multiple sources, each pro- viding a 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 & subgroups, each with its own regression vector. Prior work 20] showed that with abundant with only few, ((k3/22, batches of medium-size with (k) samples each. How- ever, the paper requires that the input distribution for all l subgroups be isotropic Gaussian, and states that removing this assumption is an "interesting and chal- lenging problem". We propose a novel gradient- 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 different, un- known, and heavy-tailed; (2) recovering all subgroups followed by 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.
the regression vectors can be learned
Introduction
In numerous applications, including federated learning [WCX+21], sensor networks [WZ89], crowd-sourcing [SVC16] and recommendation systems [WDVR06], data are collected from mul- tiple sources, each providing batch of samples. For instance, in movie recommendation systems, users typically rate multiple films. Since all samples in batch are generated by the same source, they are often assumed share the same underlying distribution. However, the batches are fre- quently very small, e.g., many users provide only few ratings. Hence, it may be impossible to learn different model for each batch.
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, may not capture the characteristics of individual users.
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
This work was done while the author student UCSD and interning part-time Google Research.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

each popullollt,inn, all underlying distributions are close to and can be represented by a single distribution. Even when & 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 aub-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 distri- bution D of samples (x, y) follows a linear regression model if, for some regression vector w E IRd, the output is y = w 2 + 17 where the imput z a random d dimensional vector and 1) is zero-mean noise. While we allow the presence of sub 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.
1.1 Our Results
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 long tail distribution and in the majority of the batches only 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 number of ratings per user follows a long-tailed distribution with overwhelming number of users rating only 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, 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 aub-populations that has 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 sub- populations is unknown,
is the identity of the sub-population from which each batch originates. Batches in set B, (denoted as small) consist of at least two samples each, while those in set Bm (de- noted as medium) are of larger size, containing 2 Rum samples. Let's consider a sub-population in- dexed by i, which follows a linear regression model with parameter vector w; and output-noise vari- ance 2². The next theorem presents the guarantees provided by our algorithm for estimating wi-
Theorem 1.1 (Informal). Let a > such that both sets B. and Bn comprises at least a fraction of batches from a sub-population i. Given [Ba] 2 (dd/a²), |Bml Z Mminn VT. 1/ La//a)/a), and nm 2 Q(min( / /k, 1/ Va)), our algorithm runs in time poly(d, 1/a,k) and outputs a list L of size ((1/a) such that w.h.p., at least one estimate in L is within a distance of o(a) from w, and has an expected prediction error o² + o(1) for sub-population i. Moreover, given Vlog L) samples from sub-population i, we can identify such an estimate from L.
In the above theorem, the batches within set B, must contain a minimum of 2 samples, while those within set B,, must be of size 2 Rum- Additionally, both sets must include 2 a fraction of batches originating from the population we want to recover. It follows that our algorithm requires only + B=))= = Q(d/a) samples from that specific aub-population. It's worth noting that any algorithm, even when dealing with a single aub-population (k 1). requires a minimum of 2(d) samples from said sub- aub-oopuaation. To the best of our knowledge, ours is the best sample complexity for recovering the linear regression models in the presence of multiple aub-populations using batch sizes smaller than d.
Remark 1.1. When multiple sub- sub-populations meet the criteria outlined for nub-population i, the list L returned by our algorithm includes estimates of the regression vector for each of these sub- populations. Furthermore, since there can be up to min (k, 1/a distinct nub-populations that satisfy these criteria, any algorithm that aims to provide meaningful guarantees must produce list of size 2 miin(k: 1 1/cx). unless the algorithm has access to a sample set known to originate from the specific sub-population it aims to recover. A similar situation arises in the list-decodable setup [CSV17], where also the algorithm returns list of estimates.

---

1.2 Comparison to Prior Work
The only work that provides a polynomial algorithm in dimension d, in the same generality as ours is [DJKS22]. They even allow the presence of adversarial batches. However, they require ((d/a²) batches from the sub-population, each of them of size (((1a)), and therefore, ((d/a²)) samples in total, which exceeds our sample complexity by a factor of 1/a". Furthermore, their algorithm requires 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 number of samples much larger than ours, which we discuss next.
Most of the previous works [CL13, SJA16, ZJDI6, YCS16, LL18, CLS20, DK20, PMSG22] have addressed the widely studied mixed linear regression (MIR) model where all batches are of size 1, and adhere to the following three assumptions:
1. All & sub-populations have 2 a fraction of data. This assumption implies S 1/a.
2. All i distributions follow a linear regression model.
3. All k regression coefficients are well separated, namely || - wyl
A, A i xj.
Even for k = 2, solving MLR, in general, is NP-hard [YCS14]. Hence all these works on mixed linear regression, except [LL18], also made the following assumption:
4. All input distributions (i.e., the distribution over x)) are the same for every sub sub-population, in fact, the same isotropic Gaussian distribution.
With this additional isotropic Gaussian assumption, they provided algorithms that have runtime and sample complexity polynomial in the dimension. However, even with these four strong assumptions, their sample complexity is super-poly poppempplyyomial overall. In particular, the sample complexity in [ZJD16, CLS20, DK20] is polynoill in k and [CL13, SJA16, DK20] require at least quadratic scaling in d. In [CL13, SJAÍ6, YCS16) the sample complexity scales as a large negative power of the minimum singular value of certain moment matrix of regression vectors that can be zero even when the gap between the regression vectors is large. In addition, [ZJD16, YCS16, CLS20] required zero-noise i.e 1) = O. The only work we are aware of that can avoid Assumption 4 and handle different input distributions for different sub- cob-populations under MLR is [LL18]. However, they still require all distributions be Gaussian and 17 = 0, and their sample size, and hence run- run-time is exponential, Ω ((exp(A²)) in k. Recently, [TV23] proposed and analyzed an approximate message- passing algorithm for mixed linear regression. However, their analysis convergence guarantees as the number of samples and dimensions tends to infinity, while assuming a finite number of components. The results also rely on Assumptions 2 and 4.
asymptotie, providing
The work that most closely relates to ours is [KSS+20] which considers batch sizes > 1. While achieves the same dependence and small batches, the sample complexity of the algorithms and the minimum length required for medium-size batches had an additional multiplicative dependence on the inverse separation param- eter 1/A. It also required Assumption 4 mentioned in the section. The follow-up work [KSKO20] which still assumes all four assumptions can handle the presence of small fraction < a²/k² of adversarial batches, but requires O(dk2/a² + k.5 /aa) samples. It also suffers from strong assumptions similar to those in earlier work, and the sum-of-square approach makes it impractical. The sum of the square approach, and stronger isotropic Gaussian assumption, allow it to achieve a better dependence on 1/a on medium-size batch lengths, however, causing a significant increase in the number of medium-size batches required.
on d, k, and 1/a, the length and number of medium
Our improvement over prior work. In contrast, our work avoids all four assumptions, and can recover any sufficiently large sub-populations that follow a linear regression model. In particular: (1) Even when a large number of different sub-populations are present, (e.g., k 2 1/a), we can still recover the regression coefficient of sub-population with sufficient fraction of batches. (2) The distributions do not even need to follow a linear regression model. In particular, our algorithm is robust to the presence of other sub-populations for which the conditional distribution of output given input is arbitrary.
(3) Our work requires no assumption on the separation of regression coefficient A, and our guarantees as well have no dependence on the separation. (4) We allow different input distributions for different sub sub-populations, which can be unknown and heavy tailed, (5) In addition
3

---

to removing the four assumptions, the algorithm doesn't 't require all batches in a sub-population to have identical distributions, it only requires them be close (see Remark 2.1).
1.3 Techniques and Organization
To recover the regression vector for sub-population i that follow the linear regression model and has more than a fraction batches, the algorithm begins by selecting a medium size batch from this sub-population. To obtain such a batch we randomly choose (1/a) batches from the set of medium-size batches Bm 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 sub- population i, and when the algorithm is run for such batch, w.h.p. it recovers the estimate of w;. Therefore, at least one of the (11/a) estimates accurately recover w!s-
The regression vector w, minimizes the expected squared loss for the sub-popun uub-ooulatiion i. Therefore, we use a a seeffinnt-deseen-bbaeeeee approach to estimate wi, 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 I2 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 sub- space reduction is that testing and estimation of the gradient in the smaller subspace is easier, con- sequently 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 + 200. However, while they employed to directly estimate regression vectors of all sub- populations simultaneously we adapt it to sequentially estimate gradients for one sub- 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 Sec- tion 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 well generalizes to settings that are outside the assumptions of [KSS* 20].
2 Problem Formulation and Main Results
Consider distributions Do...., Da-1 over input-output pairs (I, y) E Rd x 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, are small and contain at least two samples each, while batches in B,, are of medium size and contain at least Ntm 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 L2 norm of a vector u is denoted by ||u||. The norm, or spectral norm, of a matrix M is denoted by || and is defined as max Mull If M is symmetric matrix, the norm simplifies to || max]uu]=1 [uTM u/. The trace of symmetric matrix M is Tr(M) := Σ, Mi, the sum of the elements on the main diagonal of M. We will use the symbol to denote an arbitrary collection of samples. For a batch denoted by b, we will use Sh to represent the set of all samples in the batch.
2.1 Data Distributions
Let Σ, : ED, [xx)]] denote the second- moment matrix of input for distribution Dj-
To recover the regression vector for distribution D;, we assume that it satisfies the following assumptions?:
1. (Input distribution) There constants C' and C1 such that for all i I,
We note that these assumptions are standard in teavy-tailed linear regression where some condition on the anti- coneiiennentrtion of B[zr"]] is necessary to achieve finite sample complexities [LM16, Oli16, CAT*20J, even in the specific case of single sub- sub-population.

---

(a) L4-L2 hypercontractivity: For all E R", Ep, u)*] - C(ED, ((x u)
(b) Bounded condition number: For normalization purpose, we assume min ||ll== u E, u 2 and to bound the condition number we assume that || S C₁.
2. (Input-output relation) There is a σ > 0 s.t. for all i E I, w; I + where w; E IR is an unknown regression vector, and v) is noise independent of x, with zero mean E Dp - [n] and E ED, [r] so². Note that by definition, the distribution of may differ for each i.
For given a, O and Om > O, our goal is to estimate the regression vector w; for all sub- populations that have at least 0g and (mm fractions of the batches within B, and Bm respectively, and follow the assumptions described above. Let I S {0, I,..,.., k 1} denote the collection of indices corresponding to all such sub-pope culb-oouuations.. (Note that in the introduction, we simplified this more general scenario by assuming equal values for a, and am, both denoted by a.)
For other sub-populations i E I, we only need that the input distribution satisfies ||E;]| S C1, same as the second half of assumption 1(b). The input-output relation for samples generated by DD; for i & T can be arbitrary, and in particular, does not even need to follow linear regression model. Moreover, the fraction of batches from these other sub-populations can also be arbitrary.
same
To simplify the presentation, we make two additional assumptions. First, there is a constant C2 such that for all components i E (0, k 1}. and random sample (a., y) ~ D:, ||>||| - C2Vd, a.s. Second, for all i E I and a random sample (z, y) ~ Dr, the noise distribution 17 = y - 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 h distributions. However, our methods can be extended to more general scenarios. Let D denote the underlying distribution of batch b. Instead of requiring D = D, for some E {0, 1, 1). our approach can be extended to cases where the expected gradients for Dd and D. are close. Additionally, if i E I, the regression vector must achieve a mean squared error (MSE) of at most o² for D.. These
conditions are satisfied if the following hold: (I) || E [x=]]] Σ is small, (2) for all x Rd
Ep- [y] xz] ED, [y] |=||| is small, and (3) if i E I then for all E R Ep- [(y w; x) ]|1] S o². Thus, the strict identity requirement D D, 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 sub-
population.
2.2 Main Results
We begin by presenting our result for estimating the regression vector of a component D;, for any i E I. This result assumes that in addition to the batch collections B, and Bm, we have an extra medium-sized batch denoted as b" which contains samples from D; and w.l.o.g, we assume = O. Theorem 2.1. Suppose index O is in set I and let b" be batch of 2 Nem i..dd samples from Do. Given d,e E (0, 1]. [Bal = O1( ([,5x)). T'm = 52(min (Vk, (k,,,,,))) 1)) and |Bml = _113 }), Algorithm / runs in polynomial time and returns estimate úú such that 2 2 min( (Vi, c/a.
with probability 2 1 - 5, || - woll
S w eo.
We provide a proof sketch of Theorem 2.1 and the description of Algortihm 1 in Section 3, and formal proof in Appendix H. Algorithm 1 can be used estimate w; for all E I, and the requirement of a separate batch 8² is not crucial. As mentioned in the introduction, it can be obtained by repeatedly sampling a batch from Bm, and running the algorithm for these sampled b*. Since all the components in I have 2 am fraction of batches in Bm. then randomly sampling 6* from Bm. 6(1/am) times would ensure that, with high probability, we have b* corresponding to each component. We can then return a list of size θ ((1/am) 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.
Corollary 2.2. For given d,e E (0, [Bl, Tmms and |Bub meet the requirements outlined in Theorem 2.1, then the above modification of Algorithm runs in polynomial-tin time and outputs
5

---

list L of size (1/amm) such that with probability 2 1 - 6, the list has an accurate estimate for regression vectors w, for each i E I, namely m min I| will S ea.
In particular, this corollary implies that for any i E T, the algorithm requires only (2(//as) batches of size two and Mmmin VE. 3r medium-size batches of size M[(in) Vk, v from distribution D; to estimate w; within an accuracy o(or). Furthermore, it easy show that any o(o) accurate estimate of regression parameter achieves an expected prediction error of o² +o()) for output y given input x generated from this D,.
Note that results work even for infinite and without any separation assumptions on regression vectors. The min( Vk, 1/var, dependence is the best of both words. This dependence is reasonable for recovering components with a significant presence if the number is few.
The total number of samples required by the algorithm from Dj in small size batches B, and medium size batches Bm, are only (d/aa,) and Ó(min 1/. 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 Dj by the algorithm is within (1/aas) factor from that required in a much simpler single component setting.
The next theorem shows that for any E I, given a list L containing estimate of w; and (log(1/a.)) samples from D., we can identify estimate of regression vector achieving small prediction error for D. 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 dimensional vectors. Suppose for a given 0 at least one of the vectors in L is within distance 3 from the regression vector of Di namely min || - will S 3. Given O (max(1, [1, 3 log 4) samples from D; Algorithm 3 identifies an estimate w from the list L, s.t. with probability 2 1-6, - || O(3) and it achieves an expected estimation error ED,[(u- x 1 y) S 0² + (
Remark 2.2. As a special case the above theorem implies that that for any 2 M(r), Algorithm can identify the near best estimate of w; from list L using only Olog ()) samples from D,
Combining the above theorem and Corollary 2.2, get
Theorem 2.4. For given e E (). iB Tom and |Bmb, meet the requirements outlined in The- orem 2.1, then, there exists a polynomial time algorithm that, with probability 2 1-16. 6, outputs a list L size ((1/mm) containing estimates of w; 's for i E 1. Furthermore, given |S| 2 0 =/2 log zdm)) samples from D;, for any E I, Algorithm returns E L that with probubility 2 1 - satisfies lww; x0| v ((ea) and achieves an expected estimation error ED, [w- ==) < a² + 0(2oo))
3 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 Theorem 2.1. As per the theorem, we assume that index 0 belongs to 1, and we have batch 6* containing 7mm samples from the distribution Do. Note that Do satisties the conditions mentioned in Section 2.1 and that B, and Bm each have 2 |Baaa a, and 2 Bm Jam batches with i.i.d. samples from Do. However, the identity of these batches is unknown.
Gradient Descent. Note that Wp minimizes the expected square loss for distribution Do- Our algorithm aims to estimate Wy by taking a gradient descent approach. It performs a total of R gradient descent steps. Let (r) denote the algorithm's estimate of wo at the beginning of step r. Without loss of generality, we assume that the algorithm starts with an initial estimate of x²(1) = 0. In step r, the algorithm produces an estimate A (r) of the gradient of the expected square loss for distribution Do at its current estimate r)(r), For convenience, we refer to this expected gradient for Do at ( simply as expected gradient. The algorithm then updates its current estimate for the next round as ()++1) (r) A11//C,
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/c²) i.i.d. samples from Dn- However, our algorithm only has access to a medium-size batch b* that is guaranteed to have samples from Do and this batch contains far fewer samples. And for batches in B, and B,,.

---

Algorithm MAINALGORITIM
1: Input: Collections of batches B, and Bm, Oss k, medium size hatch b- of i..dd samples from Do. distribution parameters (upper bounds) σ, C, C1, an upper hound M on woll, and 6,
2: Output: Estimate of 1ño
3: 61 <- ((). + O (CrvetT (EI + )
min(k, zarez J. < BR and R <- O(C, log )
4: Partition the collection of batches B, into R disjoint same size random parts {BEm)
5: Similarly partition B' into R disjoint same size random parts {Bm) ]re[R]-
6: Divide samples 5 into 2R disjoint same size random parts {s} <<) )E((R) and (5 (1) )E(m)
7: Initialize (b(1)
8: for r from to
9: n!!") + CCIIEET(SS) (r) ())r) ,,,,"' o,,,,,,CCCO, Ci)
10: p(r) < GRADSUBEST(B" (0) ,00 e))
11: A &- GRADEST(BM.. St (r) ,m 161 pl,),,,00 62,8")
12: <a(++1) + ((()) (())
ci
13: end for
á (R+1)) and Return i
the algorithm doesn't ' know which of the batches has samples from Do. Despite these challenges, we demonstrate an efficient method to estimate the expected gradients accurately.
The algorithm randomly divides sets B, and Bm into R disjoint equal subsets, denoted as {B!") Jr=1 and (D ) respectively, The samples in batch 6* are divided into two collections of equal dis- {sh",(r)
joint subsets, denoted as {5 (r) y
Ir=1 and { (() ( Tr=1. At each iteration r, the algorithm uses the
collections of medium and small batches B and B respectively, along with the two collec- tions of i.i.d. samples s ((r) and se"((r)
(r) from Do to estimate the gradient at point w(r) While this division may not be necessary for practical implementation, as shown by simulation results later, this ensures independence between the stationary point w(()) and the gradient estimate, which facilitates our theoretical analysis and only incurs a logarithmic factor sample complexity.
Next, we describe how the algorithm estimates the gradient and the guarantees of this estimation. Due to space limitations, we provide 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 crucial role in the estimation process.
Clipping. Recall that the squared loss of samples (x, y) on point w is (x- w- y)² /2 and its gradient is (x v).. 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 (æ, y) evaluated at point is defined as V fx, y, w. k) == (a-w-y) [i-w-glVx KZz For a collection of samples S, the clipped gradient VSSS, w. k) is the average of the clipped gradients of all samples in S, i.e. VS(S, w. x) [ (z.w)es V (ar, y. w. x)).
The clipping parameter к controls the level of clipping and for к = oc, the clipped and un- clipped gradients are the same. The clipping step necessary to make our gradient estimate more robust, by limiting the influence of the components other than Do (in lemma C.I), and as a bonus, we also obtain better tail bounds for the clipped gradients. Theorem C.2 shows that for K 2 2( /E Du [(y - x (T). the difference between the expected clipped gradient and the ex- pected gradient |L D. (V f y, w, k]] E D. I( w - y)]|l| is small. Therefore, the ideal value of ς at point is ((v Do (V x ())
For the estimate (r) at step r, the choice of the clipping parameter is represented by k(r) To estimate value for (r))) that is close to its ideal value, the algorithm employs the subroutine CLIPEST (presented as Algorithm in the appendix). The subroutine estimates the expected value of (y- (y-rruuk (r))) by using i.i.d. samples s/: (r)) from the distribution Do- 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 small, and thus, estimating the expectation of clipped gradients can replace estimating the actual gradients.
7

---

Subspace Estimation. The algorithm proceeds by using subroutine GRADSUBEST (presented as Algorithm 5 in Appendix E) with B =B!", w v(() and x x((r) to estimate a smaller subspace p(r) within R.. The objective of this subroutine is to ensure that, with high proba- bility, the expected clipped gradient for distribution Do at point vi(r) remains largely unchanged when projected onto the identified subspace. Thus, estimating the expected clipped gradient re- duces to estimating its projection onto p(r), which requires fewer samples due to the lower di- mensionality of the subspace. To achieve this, the subroutine constructs a matrix A defined as B[A]= = [,,]=,, Pi (VS(,, y,, w, K)] ED, [V y, w; k)]T, where p; denotes the fraction of batches in B that have samples from D. Since B!") are obtained by randomly partitioning B,, w.h.p. Po R ag. It is crucial for the success of the subroutine that the expected contribution of every batch in the above expression is 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 l singular vectors of A, where we choose = min(k, (k, ((1/a,)). It is worth noting that when 1/a,, is much smaller than l (thinking of the extreme case the top & = 1/a, dimensional subspace, since those infinitely many components can create at most -1) directions with weight greater than a., This along with clipping ensures that the direction of Do must appear in the top (1/ox.) subspace (see Lemma E.3). Theorem E. I in Appendix E char- acterizes the guarantees for this subroutine. Informally, if B > ((d/a²), 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 EE,, f y, w, k)] - p. [V of (x. y. W, k)]l is small. We note that our construction of matrix A for the subroutine is inspired by similar construction
00), our algorithm still only requires estimating
Algorithm 2 GRADEST
1: Input: A collection of medium batches B, a collection of samples S from Do. k. w. projection matrix P for subspace of IRd parameter e, 8'
2: Output: An estimate of clipped gradient at point w.
3: Ti 4- (log 1B1) and T- Gloo +)
4: For cach b divide S into two equal random parts Si and S
5: For each further divide St into 2T; equal random parts, and denote them as (St,, (e()))
6: Divide S" into 27, equal random parts, and denote them as (S} bE(TT
7: <} == "J VS()) w. x)) PTP (VI(S),T w. (S(( w,
8: Let B B, median(() 1 [Ti]} S ).
9: For each b divide St into T2 equal parts randomly, and denote them as (S2,, be Tz)
10: For i E [7b], let A; <- [ Σ>eB PV PSS.,, w.).
11: L; <- median [j E
[T>] ||A;- Asll}
12: Let arggainn migminn (i [Ta] }.} and A 2- A..
13: Return A
in [KSS+20], where they used it for directly estimating regression vectors. Our results general- ize the applicability of the procedure to provide meaningful guarantees even when the number of components &o oo. Additionally, Lemma E.3 improves matrix perturbation bounds in Lemma 5.1 of [KSS#201 20]. which is crucial for applying this procedure for heavy-tailed distributions and reducing the number of required batches.
Estimating the expectation of clipped gradient projection. The final subroutine, GRADEST, estimates the expected projection of the clipped gradient. It does this by utilizing two key sets: BKC), one of the R partitions of medium-sized batches, and S .(r), one of the 2R partitions of the batch b", which contains i.i.d. samples from the reference distribution Do- GRADEST begins by splitting each batch in B(()) into two equal parts. For each batch E B(r), the first part of the split is used along with samples from Do in s (r)) 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 Dp- With high probability, the algorithm retains batches sampled from Do and rejects those from other distributions where the difference in expected projections is significant. This test requires π ((V2)
8

---

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 l dimensional subspace, (2( samples suffice for the estimation. To obtain high-probabiliy guarantees, the procedure uses the median of means approach for both testing and estimation.
The guarantees of the subroutine are described in Theorem F.I, which implies that the estimate (r)
of the gradient satisfies || (+) - Ep, (P(r) V fz, w, x)]l is small.
Estimation guarantees for expected gradient. Using the triangle inequality, have:
| (r) - E 1, w w) = Ba, (Vr,,, w. x)]
Ep,1(x w - y)æ=l
+ || EE,,,, D,,,VVVV [V y. w)] - E Dr [ [xx w,)]] xI]] + +1 E [P(r) V [x, Mr w., k)]ll-
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 Do. and both gradients and clipped gra- dients 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 wy-
As previously argued, all three terms on the right side of the inequality small, which implies that (rr) provides an accurate estimate of the expected value of the gradient for distribution Do. More- over, Lemma G.1 shows that with an accurate estimation of expected gradients, gradient descent reaches an accurate estimation of wi after Oloog [ steps. Therefore, setting / = 2ilog |! suffices. This completes the description of the algorithm and the proof sketch Theorem 2.1, with a more formal proof available in Appendix H.
Empirical Results
Setup. We have sets B, and Bm of small and medium size batches and l distributions D, for
1.. For a subset of indices C {0, 1,..., , 1-1}, both B, and Bn have a fraction of a batches that contain i.i.d. samples from D; for each i T. And for each i {0, 1,..,...,4 k 1) \> in the remaining set of indices, B, and Bn, have (1 /]ax)((k= - fraction of batches, that have iid samples from D,, In all figures, the output noise distributed N/ 1). All small batches have 2 samples each, while medium-size batches have 7mm samples each, which we vary from 4 to 32, shown in the plots. We fix data dimension d' 100, a 1/16, the number of small batches
i E {0, l -1-1.
to [B, = min &dk? ad/a²} and the number of medium batches to B,n 256. In all the plots, averased 10 runs and renort the standard error
K5SKD14 samples)
KS5KO(4 samples)
80
KSSKO(8 samples)
KS5KD(8 samples)
Ours14 samples)
Ours(4 samples)
Ours(8 samples)
600
Ours(8 samples)
50
40
a02
20
200
35
botth
medium butch size
Figure 1: Same input dist., k 16, large min. distance between regression vectors.
Figure 2: Different input dist., t = 16, large min. distance between regression vectors.
Evaluation. Our objective is recover a small list containing good estimates for the regression vectors of D; for each i E I. We compare proposed algorithm's performance with that of the

---

algorithm in [KSS 20J. Given new batch, we can choose the weight vector from the returned list, that achieves the best error'. 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 Appendix
Setting in [KSS + 20J. 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 where 1= {0, I...., 15} and all 16 distributions have been used to generate 1/16 fraction of the batches. All the Dj ss are equal (O0, 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 Σj such that its eigenvalues are uniformly distributed in [1, C1]. and the input distribution for D; is chosen as (0, Mx). We set C 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.
5 Conclusion
We study the problem of learning linear regression from batched data in the presence of sub- populations. We removed several restrictive assumptions from prior work and provide hetter guar- antees in terms of overall sample complexity. Moreover, 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
prior baseline.
One limitation of the algorithm is that it requires knowledge of parameters C, C1, M and a. 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, for simpler problems, has been highlighted by previous work, e.g. 2C)T+20).
References
[AJK+22] Jayadev Acharya, Ayush Jain, Gautam Kamath, Ananda Theertha Suresh, and Huanyu Zhang, Robust estimation for random graphs, Conference on Learning Theory, PMLR, 2022, pp. 130-166.
[BDLS17] S. Balakrishnan, S. S. Du, J. Li, and A. Singh, Computationally efficient robust sparse estimation in high dimensions, Proceedings of the 30th Conference on Learning The- ory, COLT 2017, 2017, pp. 169-212.
[BJK15] Kush Bhatia, Prateck Jain, and Purushottam Kar, Robust regression via hard threshold- ing, Advances in neural information processing systems 28 (2015).
[JJKK17]
K. Bhatia, P. Jain, P. Kamalaruban, and Kar, Consistent robust regression, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Informa- tion Processing Systems 2017, 2017, pp. 2107-2116.
[CAT+20] 20]
Yeshwanth Cherapanamjeri, Efe Aras, Nilesh Tripuraneni, Michael I Jordan, Nicolas Flammarion, and Peter L Bartlett, Optimal robust linear regression in nearly linear time, arXiv preprint arXiv:2007.0m 008137 (2020).
'This simple approach showed better empirical performance than Algorithm 3, whose theoretical guarantees described in Theorem 2.3
10

---

[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 un- trusted batches: Faster and simpler, Advances in Neural Information Processing Sys- tems 33 (2020), 4512-4523.
[CLS20] Sitan Chen, Jerry Li, and Zhao Song, Learning mixtures of linear regressions in subexponential time via Fourier moments, STOC, https://arxiv wwreziiv.orr/oeoo pwtg/ppf//// 1tppp/11// /1912. 07629. pdf, 2020.
[CP22]
Yanxi Chen and H. Vincent Poor, Learning mixtures of linear dynamical systems, Proceedings of the 39th International Conference Machine Learning (Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, eds.). Proceedings of Machine Learning Research, vol. 162, PMLR, 17-23 Jul 2022, pp. 3507-3557.
[CSV17] Moses Charikar, Jacob Steinhardt, and Gregory Valiant, Learning from untrusted data, Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, 2017, pp. 47-60.
[DJKS22] Abhimanyu Das, Ayush Jain, Weihao Kong, and Rajat Sen, Efficient list-decodable regression using butches, arXiv preprint arXiv:2211.12744 (2022).
[DK20]
Ilias Diakonikolas and Daniel M Kane, Small covers for near-zero sets of polynomials and learning lutent variable models, 2020 IEEE 61st Annual Symposium on Founda- tions of Computer Science (FOCS), IEEE, 2020, pp. 184-195.
[DKK*19 Ilias Diakonikolas, Gautam Kamath, Daniel M. Kane, Jerry Li, Jacob Steinhardt, and Alistair Stewart, Sever: A robust meta- aegar-algor for stochastic optimization, Proceed- ings of the 36th International Conference on Machine Learning, ICML '19, JMLR, Inc., 2019, pp. 1596-1606.
[DKP*21] Ilias Diakonikolas, Daniel Kane, Ankit Pensia, Thanasis Pittas, and Alistair Stewart, Statistical query lower bounds for list-decodable linear regression, Advances in Neural Information Processing Systems 34 (2021), 3191-3204.
[DKS19] Ilias Diakonikolas, Weihao Kong, and Alistair Stewart, Efficient algorithms and lower bounds for robust linear regression, Proceedings of the Thirtieth Annual ACM-SIAM Symposium on Discrete Algorithms, SIAM, 2019, pp. 2745-2754.
[DT19] Amak Dalalyan and Philip Thompson, Outlier-robust estimation of a sparse linear model using & -penalieed huber's m-estimator, Advances neural information pro- cessing systems 32 (2019).
[FAL17]
Chelsea Finn, Pieter Abbeel, and Sergey Levine, Model-agnostic meta-learning for fast adaptation of deep networks, Proceedings of the 34th International Conference on Machine Learning (ICML), 2017, pp. 1126- 1116-1135.
[Gao20]
Chao Gao, Robust regression via mutivariate regression depth, Bernoulii 26 (2020), no. 2, 1139-1170.
[GKG15] Michael Grottke, Julian Knoll, and Rainer GroB, How the distribution of the number of items rated per user inffuences the quality of recommendations, 2015 15th International Conference on Innovations for Community Services (I4CS), IEEE, 2015, pp. 1-8.
[JLST21] Arun Jambulapati, Jerry Li, Tselil Schramm, and Kevin Tian, Robust regression visited: Acceleration and improved estimation rates, Advances in Neural Information Processing Systems 34 (2021), 4475-4488.
[J020a]
Ayush Jain and Alon Orlitsky. general method for robust learning from batches, arXiv preprint 2iv: 1099 (2020).
[J020b]
Optimal robust learning of discrete distributions from batches, Proceedings of the 37th International Conference Machine Learning, ICML '20, JMLR, Inc., 2020, pp. 4651 16660.
[JO21]
Robust density estimation from batches: The best things in life are (nearly) free, International Conference on Machine Learning, PMLR, 2021, pp. 0698-4708..
11

---

[KFAL20]
Nikola Konstantinov, Elias Frantar, Dan Alistarh, and Christoph Lampert, On the sam- ple complexity of adversarial multi-source pac learning. International Conference on Machine Learning, PMLR, 2020, pp. 5416-5425.
[KKK19] Sushrut Karmalkar, Adam Klivans, and Pravesh Kothari, List-decodable linear regres- sion, Advances in neural information processing systems 32 (2019).
[KKM18] Adam Klivans, Pravesh K Kothari, and Raghu Meka, Efficient algorithms for outlier- robust regression, Conference On Learning Theory, PMLR, 2018, pp. 1420-1430.
[KP19] Sushrut Karmalkar and Eric Price, Compressed sensing with adversarial sparse noise via ll regression, 2nd Symposium on Simplicity in Algorithms, 2019.
[KSAD22] Weihao Kong. Rajat Sen, Pranjal Awasthi, and Abhimanyu Das, Trimmed maximum likelihood estimation for robust learning in generalized linear models, arXiv preprint arXiv:2206.04777 (2022).
[KSKO20] Weihao Kong, Raghav Somani, Sham Kakade, and Sewoong Oh, Robust mela-learning for mixed linear regression with small batches, Advances in Neural Information Pro- cessing Systems 33 (2020).
[KSS+20] Weihao Kong, Raghav Somani, Zhao Song, Sham Kakade, and Sewoong Oh, Meta- learning for mixed linear regression, International Conference on Machine Learning. PMLR, 2020, pp. 5394-5404.
[KZS15] Gregory Koch, Richard Zemel, and Ruslan Salakhutdinov, Siamese neural networks for one-shol image recognition, ICML deep learning workshop, vol. 2, 2015.
[LL18]
Yuanzhi Li and Yingyu Liang. Learning mixtures of linear regressions with nearly optimal complexity, COLT, arXiv preprint 1rXiv:1802.077 2018.
[LM16]
Guillaume Lecué and Shahar Mendelson, Performance of empirical risk minimization in linear aggregation.
[LSLC18] Liu Liu, Yanyao Shen, Tianyang Li, and Constantine Caramanis, High dimensional robust sparse regression, arXiv preprint 1 1664 (2018).
[MGJK19] Bhaskar Mukhoty, Govind Gopakumar, Prateek Jain, and Purushottam Kar, Globally- convergent iteratively reweighted least squares for robust regression problems, The 22nd International Conference on Artificial Intelligence and Statistics, 2019, pp. 313- 322.
[0li16]
Roberto Imbuzeiro Oliveira, The lower tail of random quadratic forms with applica- tions to ordinary least squares, Probability Theory and Related Fields 166 (2016), 1175-1194.
[OLL18]
Boris Oreshkin. Pau Rodríguez López, and Alexandre Lacoste, Tadam: Task depen- dent adaptive metric for improved few-shot learning, Advances in Neural Information Processing Systems, 2018, pp. 721-731.
[PJL20]
Ankit Pensia, Varun Jog, and Po-Ling Loh, Robust regression with covariate filtering: Heavy tails and adversarial contamination, arXiv preprint arXiv:2009.129) 12976 (2020).
[PMSG22]
Soumyabrata Pal, Arya Mazumdar, Rajat Sen, and Avishek Ghosh, On learning mixture of linear regressions in the non-realizable setting, International Conference on Machine Learning, PMLR, 2022, pp. 17202-17220.
[PSBR18]
Adarsh Prasad, Arun Sai Suggala, Sivaraman Balakrishnan, and Pradeep Ravikumar, Robust estimation via robust gradient estimation, arXiv preprint 1rXiv:1802.06444 (2018).
[PT08]
Yoon-Joo Park and Alexander Tuzhilin, The long tail of recommender systems and how to leverage it, Proceedings of the 2008 ACM conference on Recommender systems, 2008, pp. 11-18.
[QV18]
Mingda Qiao and Gregory Valiant, Learning discrete distributions from unirusted batches, Proceedings of the 9th Conference on Innovations in Theoretical Computer Science (Dagstuhl, Germany), ITCS '1, Schloss Dagstuhl-LeiibizzZZennee formatik, 2018, pp. 47:1-47:20.
fuer In-
[RL17] Sachin Ravi and Hugo Larochelle, Optimization as a model for few-shot learning, In- ternational Conference on Representation Learning, 2017.
12

---

[RRS 18] Andrei A Rusu, Dushyant Rao, Jakub Sygnowski, Oriol Vinyals, Razvan Pascanu, Si- mon Osindero, and Raia Hadsell, Meta-learning with latent embedding optimization, arXiv preprint arXiv: 1807 .05960 (2018).
[RY20]
Prasad Raghavendra and Morris Yau, List decodable learning via sum of squares, Pro- ceedings of the Fourteenth Annual ACM-SIAM Symposium Discrete Algorithms, SIAM, 2020, pp. 161-180.
[Sch87]
Jürgen Schmidhuber, Evolutionary principles in self-referential learning, or on learn- ing how learn: the meta-meta... hook, Ph.D. thesis, Technische Universität München, 1987.
[SJA16] Hanie Sedghi, Majid Janzamin, and Anima Anandkumar, Provable tensor methods for learning mixtures of generalized linear models, Artificial Intelligence and Statistics (AISTATS), 2016, pp. 1223-1231.
[SVC16] Jacob Steinhardt, Gregory Valiant, and Moses Charikar, Avoiding imposters and delin- quents: Adversarial crowdsourcing and peer prediction, Advances in Neural Informa- tion Processing Systems 29 (2016).
[TLW99]
Kai Ming Ting, Boon Toh Low, and Ian H Witten, Learning from batched data: Model combination versus data combination, Knowledge and Information Systems (1999), 83-106.
[TP12]
Sebastian Thrun and Lorien Pratt, Learning to learn, Springer Science & Business Media, 2012.
[TV23]
Nelvin Tan and Ramji Venkataramanan, Mixed regression via approximate message passing, Journal of Machine Learning Research 24 (2023), no. 317, 1-44.
[TZD*11) 19] Eleni Triantafillou, Tyler Zhu, Vincent Dumoulin, Pascal Lamblin, Kelvin Xu, Ross Goroshin, Carles Gelada, Kevin Swersky, Pierre-Antoine Manzagol, and Hugo Larochelle, Meta-dataset: A dataset of datasets for learning to learn from few exam- ples, arXiv preprint arXiv: 190:v::193..0000 00096 (2019).
[WCX+21] Jianyu Wang. Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al., A field guide federated optimization, arXiv preprint arXiv:2107.0691 (2021).
[WDVR06] Jun Wang, Arjen P De Vries, and Marcel JT Reinders, Unifying user-based and item- based collaborative filtering approaches by similarity fusion, Proceedings the 29th annual international ACM SIGIR conference on Research and development in infor- mation retrieval, 2006. pp. 501-508.
[WZ89]
Mati Wax and Ilan Ziskind, On unique localization of multiple sourees by passive sen- sor arrays, IEEE Transactions on Acoustics, Speech, and Signal Processing 37 (1989), no. 7, 1009-000
[YCS14]
Xinyang Yi, Constantine Caramanis, and Sujay Sanghavi, Alternating minimization for mixed linear regression, International Conference on Machine Learning, PMLR, 2014, pp. 613-621.
[YCS16]
Solving mixture of many random linear equations by tensor decomposition and alternating minimization, arXiv preprint arXiv:1608.05749 (2016).
[ZJD16]
Kai Zhong, Prateek Jain, and Inderjit S Dhillon, Mixed linear regression with mul- tiple components, Advances in neural information processing systems (NIPS), 2016, pp. 2190-2198.
13

---

A Other related work
Meta Learning. The setting we considered in this paper is closely related to meta learning if we treat each batch task. Meta-learning approaches aim to jointly learn from past experience to quickly adapt to new tasks with little available data [Sch87. TP12]. This is particularly significant in setting when each task is associated with only few training examples. By leveraging structural similarities among those tasks (c.g. sub-popun col--ooulation structure), meta- learning algorithms achieve far better accuracy than what can be achieved for each task in isolation [FAL17, RL17, KZS15, OLL18, TZD+ 19, RRST 18]. Learning mixture of linear dynamical systems has been studied in [CP22].
Robust and I List decodable Linear Regression. Several recent works have focused on obtaining efficient algorithms for robust linear regression and sparse linear regression when a small fraction of data may be advesaarial [BJK15, BJKK17, BDLS17. Gao20, PSBRI& 18, KKMI8, DKK 19, LSLC18, KP19, DT19, MGJK19, DKS19, KKK19, PJL20, CAT JLST21, KSAD22].
In scenarios where over half of the data may be arbitrary or adversarial, it becomes impossible to return single estimate for the underlying model. Consequently, the requirement relaxed to return a small list of estimates such that at least one of them is good estimate for the underlying model. This relaxed framework, called "List decodable learning," " was first introduced in [CSV17]. List- decodable linear regression has been studied by [KKK19, RY20, DKP*21], who have provided exponential runtime algorithms. Additionally, 21] has established statistical query lower bounds, indicating that polynomial-time algorithms may be impossible for this setting. However, as mentioned earlier, the problem can be solved in polynomial time the batch setting as long as the batch size is greater than the inverse of the fraction of genuine data, as demonstrated in [DJKS22]. It's worth noting that an algorithm for list-decodable linear regression can be used to obtain list of regression vector estimates for mixed linear regression.
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 small fraction are adversarial. They developed an exponential time algorithm for the problem. Subsequent works [CLS20] improved the run-time to quasi- polyimpal,nomial, 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, KFAL20J. [AJK 22] examined closely related problem of learning the parameters of an Erdōs- Rényi random graph when portion of nodes may be corrupted and their edges are maybe be chosen by an adversary.
B Selecting regression vector from given list
In this section, we introduce Algorithm 3 and prove that it achieves the guarantees presented in Theorem 2.3.
Algorithm SELECTING THE REGRESSION VECTOR
1: Input: Samples S from D. for some i I, C1, a list L of possible estimates of wi, and 3 2
s.t. 8 2 min z. ||w wyll-
2: Output: An estimate of wi- from list L
3: Divide S into T₃- = (log([
equal parts {S, 1 Tq
4: while [ w'| w. w/ E T.} > 12C, do
5: pick any w. w/ E s.t. || w/ 12C193
6: For j E [Ty), let 02 +-) > (z.))s,
1
(I w y) z (w w²)
[sy]
7: a 4- Median : j E [Ta]]
8: Ifa ||w- w/ 12 remove w, else remove w' from L
9: end while
Return any of the remaining w
Without loss of generality, assume = 0, and let w" arg min || wo |l- From the condition in the theorem, we know that I|w= - woll < B. The algorithm is given access to d² zz log (44) samples. The algorithm chooses any two vectors w, w/ e L that
|S| M(max 1,
14

---

more than B distance apart and tests which of them is more likely to be within 5 distance from wp using samples in S. The algorithm performs T3 (log 14,) 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 distance of 12C from each other, by choosing a vector from those remaining in L and returning it as an estimate of wn- If wt is retained in L until the end, using the simple triangle inequality for all w that remain in L at the end, we have || woll || S ||æc w=I + ||"" woll s 12C1 3 + 3 < 13C = ( Therefore, the estimate returned by the algorithm achieves the desired accuracy in estimating wa. Hence, it suffices 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 E L such that I|w* أا| 2 1.. The following theorem shows that for any i such that I|w" - uill 2 12C18, the probability of the testing procedure rejecting w" in favor of i is at most 5/|L].
woll S 3 and Theorem B.1. Given > list L, and samples S from Do. if min I/
a² 3r log 5. then for the parameter a computed in the while loop of Algorithm 3, with probability 1 8/]], we have a < || w 1//12////// if wo and || w/1///4 wwwl-'113 if if'w wo-
[S] = max (1,
The testing procedure utilized in the algorithm is based on gradients. Specifically, the average of the gradient computed on samples point w projected onto the vector (w w') The expected value of the gradient at w, and its projection onto (w - w'), are (w - TE and (w - wo)T [n(w w'), respectively. If w Wo, then the expected projection will be small. On the other hand. if w' s wp and w, then expected value of projection is R (w - To) - w') Z ||w - w'I... Using these observations, we prove Theorem B.1 in the next subsection.
calculates
Finally, since the maximum number of comparisons made by the algorithm is [L]] -1|-1, 1, a union bound ensures that w² will be retained until the end with probability greater than - 5. completing the proof of Theorem 2.3.
B.1 Proof of Theorem B.1
Note that is the median of the set {az : ; E [Ta]}, where each a, is computed using different sets of i.i.d. samples. Consequently, {ay)se(ra) are also i.i.d random variables.
Proof.
We begin by calculating the expected value of aj. Using the linearity of expectations, we have:
B[ay]
(a)
E Σ (n,))S [9,7 (x - w y) (w w²)
E (x w y) (w w/)]
E Ed, (w wa) (y wo)) (w - [ww']]
E D (w wo)w (w w')l - E Do 1(v w (w w')]
E D. [x (w pow-wo)a
wo)x (w - =w']l
(1)
(c) E n. ((x (w-w')) w') + D. [x (w' - wo)x (w- - w')l,
here, (a) follows from the definition of aj. (b) follows from the linearity of expectation, and since Sj contains i.i.d. samples from Dr, (c) follows the noise y x wo has zero mean and is independent of æ.
Next, compute the variance of aj. Since aj represents the average of (x y)æ (w w') over |Sy] i..dd samples, we have
1
Var(a;) =
Varp, ((x - w y- (w w²))
[Sy]
S
1
E Dy [(x w y) (w - w 1)²
(3)
[s,
By applying Chebyshev's inequality, with a probability 2 3/4, the following holds for each dy:
E B[a,] 2 Var(a,) S aj < E[az] + 2 Var(a;).
(4)
15

---

First, we consider the case when w = ==. In this case, we have l|wo w|l s B.
Using Equation (3), express the variance of aj as follows:
Var(a,) -
|S E D.
((x w- y) (w w') =²))²
Do (xx (w - wo)x (w w') + (wo x -y)z ww::(w - w())²]
|5
2
[E Do (zz- (w- wo)r (w w w'))² E Do ((wo x y)æ (w w²))²)
|5
where the last step uses the fact that for any u, 0 IR, (u + w)² S 2u² 2v. Next, we bound the two terms on the right. For the first term, we have
E D (I (wo wo)x (10 w'))² s v E Do (x (w - wo]] D. - =(]=]
s CE 1. (w- | E 1x (w- - w' ())) (5)
For the second term, we have:
Do [(wo- y)æ (w (w-w')) w'))² E Do [[wo x x- x] 18 Do
(x (w-w")) w'))"
a² Do (xr (w (wwww))² w'))
where the first inequality follows from assumption 1a and the second inequality follows from sumption 1b.
Combining the above three equations, we obtain:
2
Var(a;) 5
Do (xr (w- - w')) (C
Dole : (w - wo)) + +²)
|
2
C
flll w/ ('C w woll oz))
|s,]
where the last inequality uses assumption 1b.
Using Equation (1), the Cauchy-Schwaad inequality, and assumption 1b, we have:
E[a,] & C1 ||w - woll ||w - w/ ||.
Combining the two equations above, we obtain:
2V2C1
a VCCi (CC[|w--oll) moll)
Bja,] + 2/ Var(a,) & ||0 w'll
Cillw- woll+
(a)
v8C2
V8CC1
||l w/ll C
Visşi B
v18
3C,
w w/1
||
w 112
4
here, in (a), we use w", which implics | woll B, in (b), utilize |S 2 48C and 12o² |Sʃ| ž 3 in (c), we use the fact that for any w and w' in the while loop of the algorithm, ||l w'll 120,8. Consequently, it follows from Equation (4) that each ay w--011 with
probability 3/4. Hence, with probability 2 1 1-6 8 the median of aj is < ww-
16

---

Next, we consider the case when w' = w". Firstly, we bound the variance using Equation (3):
Var(a,) -
1
E Do
((x y) (w w ='))²
|s,|
1
E D..
((z (w- - w²)) + x (w' - wwo) (w w') + (wo x - -y)x (w =²))²]
|
(a)
3
E Do ((w - w') x)²| + E Do (1 (w' wp)x (w - (w-w'))²
(8)
+ +ED. Do (uwp x yرæ (w w()
(
3
C Do
(x (o. (w -
w(')²'² (w-
+ (CEE,, ww) + d²) E ED [(xx (w w (d'))1)
[8g
(g)
3
(VCEDs
w') + (VCED,l (VCC- (w0
wo + V E E. (w w ))3]
[8y]
In (a), the fact that for any i, u TR, (1 + u+ u)2 < 362 3u² 3r 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 f, u, V 2 0, (t + u + v)² < t² + u² + v².
Using Equation (2) and the equation above, we get
B(ay] - 2 Var(a;) E 12 (w
w w²1) ED (xr (w' wp)a (w w')]
2/3
(VORR Do (ax (w w'')"] + (VCIDD (w''
(w' wo))?] + o)) v Th Dollo (w - w'))
Ss
(0)
V12C
1-
E ED. (w- w')) v E Do (w' wo w0 DO -oo-w((]²
VISAl
V12
((VCEDo (I. (w' wo)) + to) v E Do (I (w - w'))3] (ww=
3
V12
v E Dol (w - (w--w")
E EDDo Doll (1o' wo))²]
v E [x (w x²))
2
2
V1S₁I"
1
3 v
E DD [1x (w w'))'
EE Do 10 (w'-ww))) vCB
/B (I (w w') ))}
2
here, in (a) we use the Cauchy-Schwarz Schwarz inequality, (b) follows from [S,] > 48C, and (c) utilizes
|s 2 C, Next, we have:
22x²
E Do (a. (w - (w-ww))² w())
2 v EDD (w' [ul-wo))? - V
(9)
2
(4)
1
3
||l w'll
v
C1
1/'w"
woll VCB
2
2
No
(
5
|| - w'll
2 C23
2
@ ||0 w'll
(10)
here in (a), we use assumption 1b, (b) relies on w/ = w", which implies [|w'- -woll < B, and (c) uses the fact that for any u and w' in the while loop of the algorithm, || - w'll 2 12C, and C₁ > 1. Combining the above two equations, we obtain
1
E[a,] 2 Var(a,)
|| w w/ 112.
4
1n-wdl? 12
Then from Equation (4) it follows that each aj probability 2 1 - 8/|L| the median of az is -w- w 1
with probability 2 3/4. Hence, with
17

---

C Properties of Clipped Gradients
The norm of the expected value and covariance of unclipped gradients for components other than Do can be significantly larger than Do. acting as noise in the recovery process of Do. When using unclipped gradients, the algorithm's 's batch size and the number of batches must increase limit the effect of these noisy components. And while the norm of the expected value and covariance of the unclipped gradient for Do 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 collection of random samples drawn from distribution D, for some {0, 1, 1..,. k
1}. For any K 2 O and w E Rd,, the clipped gradient satisfies the following properties:
1. [SV w. x)]l| & RVVL,
2. [Cor(V]((, w, x))1l & 1
3. w,, & #C2vd almost surely,
4. ([,, xt]] s Cridd.
5. for all unit vectors u, [V(( w. n.) u)?]1l S R²C,.
This lemma implies that for smaller values of the norm of the expectations and covariance of clipped gradients bounded by a smaller upper limit. The proof of the lemma is presented in Subsection C.I.
The following theorem demonstrates that by appropriately choosing a sufficiently large value of s. the norm of the expected difference between the clipped and unclipped gradients for distribution Do can be small:
Theorem C.2. For any > 0, x² > SCC, Ev,l(v x -w)3)/e the norm of difference between expected clipped grudient E D. [(V / (xr y, w, x)] and expected unclipped gradient D. (w x - y) a) is at most,
| Do V (x. y/, w. x)] D. [(w x y)=||l S <|lw - woll-
where E D..(w I - y)a] Co (ww-ww wo).
The theorem shows in order to estimate the expectation of gradients point for distribution Dg- it is sufficient to estimate the expectation of clipped gradients at point w, as long as the clipping
E Er, (v-z-w)²
parameter k is chosen to be at least
Intuitively, when is much larger than ED. bly w|]. with high probability the clipped and unclipped gradients at point w for a random sample from Do will be identical. The proof of the theorem is 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 Ol (E D. [(v - x =)]] o2)/e). This choice ensures that is close to the minimum value recommended by Theorem C.2 for preserving the gradient expectation of Do- By selecting small K, we ensure a tighter upper bound on the expectation and covariance of the clipped gradient for other components, as described in Lemma C.I. 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.
C.1 Proof of Lemma C.1
Proof. Since fS,, w, x) average of clipped gradients of [S] independent samples from D., follows that
a) EV (SS w. x)] Ep, VV n:)]
b) Cov V (S, x)) [S| Covp, V W[,
18

---

c) IV w. n:)1 1 ess sup (uppssp)))aa (.,))>,;
| fl (z. w, x)]l ass.
d) E V fS, w, = I E
Vf v. w. x x)] and
e) for all vectors u, E (V w. u) || - D. (Vf(e,),,,)) w
We will now proceed to prove the five claims the lemma by using these properties. Firstly, we can analyze the expected norm of Ep, V / (z, y, w, k)] as follows:
E D. D.D.VV IV y, w, x)]l
max ||]]
E D.
w. K3) u)ll
I max
E p. (V,.
u) 91/2/
||]l
max
E E.,
u) /) S CT.
|||||
here the first inequality follows from the Schwarz inequality and the last inequality follows from assumptions on distributions Dj. Combining the above inequality with item a) above proves the first claim in the lemma.
Next, prove the second claim in the lemma, we first establish bounds for the norm of the covari- ance of the clipped gradient of random sample:
|[Covp, (V (z, y, w. x))1
max Varp, (V (x, y. w, x.)))))))))))) u)
||al|
$ max E D. (V](2, y, w. k) x)
||ui||
max IE D. [Kx u) S C1.
[|a||
By using the above bound and combining 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:
css sup f y, w. R)]l S sup ||>||| S RC2vd.
(a.y)~D.
(w.y)~>D
We then combine this bound with item c) to prove the third claim the lemma.
Next, we bound the expected value of the square of the norm of the clipped gradient of a random sample,
E D. |V y, w. x>][ s ED.
|= D
Tr(,)) dd Cikdd.
This bound, combined with item d), proves the fourth claim in the lemma. Finally, for any unit vector u, we bound
E D. (VS(a,
- - x² B [(x: -
< ~ s C²
This bound, combined with item e), shows the fifth claim in the lemma
C.2 Proof of Theorem C.2
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
R. b O and symmetric random variable 2,
(a + z)b E (a 2) max((a 21,
< 2/a Pr(2 lal)
Proof. We assume 2 and prove the lemma for this case. The statement for < case then follows from symmetry.
19

---

We rewrite the term inside the expectation in terms of indicator random variables:
(a (a+3) 2) (a + 2)b
max([a + z).b)
= () -(=<<bb-a) = 2a - b) 1(b 11bbac < < (a- 10:111=>000) ((2==+))+((++++00000)) 1(2 -b
Taking the expectation on both sides.
(a 2) (a + z)b
max((a + =].b)
-)- =B((a 2 b) [ - a < B b) 1(2 b bta
+ [ b) 1(: < -b a)]
= b) 1(b a < < <b+a] 20 Pr(z > b))
where the last step follows because is symmetric.
Next,
(a 2)b max((a + b].
B (a + 2)
- -000++--000-000000222:00001111-0
< <||[2a| 1 200-00<3<b+t)=22
I 2/a Pr(b ar(b-a<e- (r(=oa<=bbt)> +2 Pr b+a)
2 Pr(=>b-a)
Next, we proceed with the proof of Theorem C.2 using the aforementioned lemma.
Proof of Theorem C.2. Let (ar, y) be a random sample from distribution Do. and let v = -00 x denote the noise. Recall that 17 is independent of æ.
Note that:
(w - x - y)æ = (ww-w)) -
x - n)ar.
We will now evaluate the expected value of the unclipped gradient.
D. [([( (ttw-2 x y) = p, [((w - wo) x n)xr] = D. [((w - wo) x)a] = Eo(w - wn). (11)
Next, we will bound the norm of the expected value of the difference between the clipped gradient and the gradient. We first expand this expression:
VS(r,, w. k), which represents
(w x y)
(w- (w-x-p 1) V FF K)) = (w T 2)
w yl
((w We) I 11) ((w (ww x - 11) (w wo) 17 V x) x)).
(12)
Next, by applying Lemma C.3, we have
((w wo) I 11) [w - w²) I. - 1y K
E n ((w wo) x -n)
< 2 wo) x| Pr(n > (( wo) x1).
Note that in the above expectation, we fixed r and took the expectation over noise 7)-
Let Z 1(1x (w- wo)l 2 x x/2). Observe that Pr(n > k [( wo) x|) < Z Pr(n > x/2). Combining this observation with the above equation, we have:
((w wo) x - m) E, ((w wo) 11) [w wo) 17
2 wo) 21 (Pr(n > x/2) + Z).
(13)
20

---

Then, for any unit vector v E R we have
E Da [(( x - y)ı - (a(,, y, (,)),,)))))))) x)) w]
E DDD BB [(( [ 1)z Vf(r, w,), k)) v]]
E DD. || E Dr [(w x y)æ V/(2, y, w. k)) -ill
E n- Do ( Wa) 2| |z (Z + Pr(n x/2))]
E D. [Z 1( wo) - |a. o|] + 2 Pr(n > B/2) E Dol(w w la ell,
(14)
here 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:
EE Du IZ (xx (w wo o)] WV I/3
(E((2)) E ED (w wo)) v) 1/2
(B[Z)- E E.D. (x (w w.) ( Du (ox v) $11/22 1/2
(E[Z). D,D,D. [(x (w wn))² Do (1r v) w) 1/2
(CC, Pr[[æ (w - ww)l 2 x/2] E D. (w - wo))
(15)
where (a) used the Cauchy-Schwart inequality, (b) used the fact that Z is an indicator random variable, hence, Z Z, and the Cauchy-Schward inequality, (c) uses L4 - L2 hypercontractivity.
and (d) follows from the definition of Z and the assumption that |Eol s C1.
Similarly, we can show that
E E.Dl(uu wo) xl |x u|] s (G1 (,,lx, (w 11) 1/2
(16)
Applying the Markov inequality to nj² we get:
[n.
Pr[in 2 k/2] (x/2)²
(17)
Similarly, applying the Markov inequality to jxt (w - wo)) yields:
E Do || (w - wo)l
CE C.D.ol 1| (w - (oo-wo)l wo)
Pr[[æ (w - mo)1 2 x/2] S
(18)
(n/2))*
(m/2)4
where the last inequality uses L4 L2 hypercontractivity. Combining Equations (14), (15), (16), (17) and (18), we have EE D. ((w y)r V fl v. w. K)) eil
8CVC, [E 1,, (w - w))) 3/2
8/C,E D. in?) D. [x (w w))
x²
8C 8CVCI ( (w (ow-wa)))7 w) 1/2 ( Do (w - wo))) + E ED, (179//C)
<
x2
22) EE D. [I. (w w))))1) 1/2
(B,,,(e (w wo)))l) + ED, In21/C)
VC1
E D - x =
(6) <|lw wpll (LE.,l( (w - wo)) + EE 17 ED, 2
cllw wo ll.
here inequality (a) follows from the lower bound on x² in theorem, inequality (b) follows from Assumption 1b and C > 1, and the last equality follows since y w = s(w wo) + y,, and
and r) are independent.
Note that the above bound holds for all unit vectors therefore,
Da [(w y) f y, w. x)]]|| s max p. K(w y)x Vfle,,,
N. k)]] - <llw -woll
||l||
21

---

D Estimation of clipping parameter
Algorithm 4 CLIPEST
1: Input: A collection of samples S* from Do, w, e, 8', a, C, and C₁.
2: Output: clipping parameter K
3: T < (log 1/8))
4: Divide 5" into T equal parts randomly, and denote them (5) bem
Median {
1951 Lera)es; (æ - - -y)²; j [m]}
5: <
6: (2(C+1))
7: Return
x. In round r, to set K. R3 V V (1- at point w = w((r) the main algorithm runs subrou- tine CLIPEST 4 for S S (r) and w((r) Recall that S -. (r) collection of i.i.d. samples from D-- Using these samples this subroutine estimates E E,, [yy x: w))² median of means and then use it to obtain 6. in the desired range. The following theorem provides the guarantees on the estimation of by this subroutine.
w(ft using the
Theorem D.1. For > 0, T 2 Qllog 1/8') and |Sa| 2 64C and w E Rd With probability 1 8', the clipping parameter returned by subroutine CLIPEST satisfy.
,
2(C+]C (8 o, (1)lw--w))))1--00
< K S 28
To prove the theorem, we will make use of the following lemma: Lemma D.2. Let S be collection of 64C i.i.d. samples from Do and Rd. then with probability at least 7/8, the following holds:
= E D.llww
y) 170²
Σ (2.9)ES
(y - w wowwrr x))² S 33E [[[[[[[[ww wo) x]²] + 320
4
Proof.
We start by expanding the expression:
(w x y) = Σ
1
((w - wo) x + (wo x - y))
m
m
(x.g)ES
(a.y)eS
-
Σ
(((w wwo x)² + 2((w woo x)(wo y) + (wo y)
(m,)ES
((w wo) x)2 - (wo I - y)² w)²)
(19)
2
(F.))S
where the last inequality follows since for any b, we have a² + 2ab + 132 2 a² 12/222 b2.
Similarly, show:
Σ
(w-z
(w-2-v)²² 1
Σ (2((w wo) x)2 2(wo -))
(20)
m
(2,8)SS
(x.y)ES
Since contains independent samples from Do, we have:
Σ [x.v]eS
(((w wo) x)²
EE,,[((w wo) x0)2
and
Var Do - (wo) 2)
Var -
Σ (ww - wo) x)2
(x.8)ES
E Ep,l((w wo) x1
C CE D. DE, - wa) 2)
<
22

---

where the last inequality follows from L4-L2 hypercontractivityy For any a 0, using Chebyshev's inequality,
Pr
1
Σ
((w wo) xx)² - E Do Dol((w - wn) x) 2
C C.D,ll(u - wo) x)
CED,l(w-wo)
s
m (3
vm
a
(21)
Using the Markov inequality, for any >> O, we have:
Pr
(wp x y)² 0²0²
S
E Do (ool(w y)
1
(22)
m
a²0²
a²
(x,8)ES
By combining the equations above, we can derive the following inequality:
With probability 2
2
the following holds:
a²
1
1 2
C
(w x - y)²² 2 E Do [(w wo) . x) - a vm
d'o²,
m
(2,8)ES
and
Σ (*,))ES
x1 (w x y)² 52ED. - +
+ 2a²0²
Vm
By choosing = 4 and using m 2 64C in the above equation, we can conclude that with proba-
2 bility 2 a7, the following holds:
D. [(w - wo) x - 16o S
(w x - y) < 3. [(w wa) xx)2 - 32o
(x:8)ES
Next, note that
E ED, x- --) = Da [(w - wp) x (y - wip x²)
(e) E Da [(( ww)) x) + E E.. (1 wo x)
E E l. wa) x) +0?
here (a) follows since æ is independent of the output noise - Wo xr and (b) follows since the output noise y Wg is zero mean and has a variance most o². Combining the above two equations completes the proof.
Now we prove Theorem D.I using the above lemma:
Proof of Theorem D.I. From the previous lemma and Chernoff bound it follows that with probabil- ity 2 1 8'
E (.. x - y) 1702 s 3E wo) x) + 32o
4
Then bound on follows from the relation 12(C+1)C
E Subspace Estimation
As part of gradient estimation in step r. the main algorithm uses subroutine GRADSUBEST for B B!*)) and (r), Recall that B((r) random subset of the collection of small batches B,-
The purpose of this subroutine estimate a smaller subspace Rdd such that for distribution Dn, the expectation of the projection of the clipped gradient onto this subspace closely approximates the
23

---

Algorithm 5 GRADSUBEST
1: Input: A collection of medium batches B, K, w,
2: Output: A rank projection matrix.
3: For each E B divide its samples 5 into two equal random parts S and S
4: A < 2|
Ebe B (VS(S) w, k) ()V w, x))T VS(S w, k) V [(}], w, x)T)
5: U < [u1, ui..... wel, where {u,} 's are top singular vectors of A
6: Return UUT
true expectation of the clipped gradient, for distribution Do. This reduction to a smaller subspace helps reduce the number of medium-sizze 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 Po denote the fraction of batches in B that sampled Do. For any c, c, O, and B = Q2 a. axe + T. ) log $ po 2 2a/2 and 8 2 min Ierezt J. with probability
1 8', the projection matrix UUT returned by subroutine GRADSUBEST satisfy
II(T - UUT) Da (VI w. w. x)]ll < C1.
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 Do 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.
The subroutine divides samples in each batch b / into two parts, namely S and S.. Then it com- putes the clipped gradients vd == V (S w, k.) and nº == V V(S w, k). From linearity expecta- tion, for any and batch b that contain i.i.d. samples from D., Blu"] = E[c²] = Ep, fxx, y, w, x)]l The subroutine defines A ΣieB 2 (v") + ('")). Let pa denote the fraction of batches in
B that have samples from D,. Then using the lincarity of expectation, we have:
E[A] = Ek=o p: E ff y, w, E. V
f (z, 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 IE P. (V,[VVCC / y, w. #]] onto UUT corresponds to itself, and the guarantee stated in the theorem holds. However, components large, it may be desirable to obtain a subspace of smaller size than
do not have access to E[A], and furthermore, when the number of
To address the first challenge. Theorem E.2 in the next subsection shows that || - [[A]]| is small. This theorem permits the usage of as a substitute for [E[A]. The clipping operation, introduced in the previous subsection, plays 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 good bound on the number of small-size batches required. Additionally, the clipping operation ensures that subroutine remains robust to arbitrary input-output relationships for other components.
Furthermore, the clipping operation assists in addressing the second challenge by ensuring uniform upper bound on the norm of the expectation of all components, i.c., E p. VV f(r, y, x)]l < (K). Leveraging this property. Lemma E.3 demonstrates that it suffices dimensional subspace. Intuitively, this is because the infinitely many components can create at most approximately 1/Po directions with weights greater than Po. indicating that the direction of Do must be present in the top (11/po) subspace.
estimate the top 1/Po-
Since B P²(") is obtained by randomly partitioning B, into R subsets, and B, contains fraction of at least a, batches with samples from Dp, it holds with high probability that Pa 2 ag,- Conse- quently, when satisfies the desired property in the Theorem E.
min{k; 2( (((_.), }}. the subspace corresponding the top l singular vectors of A
We note that the construction matrix A subroutine GRADSUBEST is inspired by previous work -KSS-200 20]. However, while they employed it to approximate the k-dimensional subspace of
24

---

the true regression vectors for all components, we focus exclusively on one distribution Do at a time and recover a subspace such that, for distribution Do. the expectation 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 repurposing the subroutine from [KSS* SSS$$220 20], 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 20J. These enhancements enable 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.
E.1 Proof of Theorem E.1
To prove Theorem 1, in the following theorem, will first demonstrate that the term - R[A]]] is small when given enough batches.
Theorem E.2. For S i S t - 1, let 2, = E V f y, w, k)], and denote the fraction of batches in B that have samples from Di. For any e, 8' 0, and B = 2 0.5" 0,6² with probability at least 8',
1
+
C2
C2
log 4
A
- t=0
S Ca
where the matrix defined subroutine GRADSUBEST.
Proof. Let th Vf(S), w. x K)).
Note that
A = 2
Σ (z (Z9)T)
BE
Then, from the triangle inequality, we have:
k
k-
Σ
2
S
2
Σ
+ē
Σ
(2
Σ
2:2,2
13
B
be
be
k-1
Pizzi
T
B
bEB
For batch sampled from distribution D,, we have:
E E = E f(st R V w. K
E f(St R)] E (S) w. K]
E D. Vf K]] V w. =] =
where the second inequality follows since samples in o and S are independent, and the third equality follows from the linearity of expectation.
follows that
k-1 BI
Σ
B
bE
1=0
25

---

and
A
s
Z BE
B(z%) bE
Σ
(23)
To complete the proof, will prove high probability bound the term the right by applying the Matrix Bernstein inequality. To apply this inequality, we first upper bound |Z as follows:
|za|| V fS². we V)V w, x fVS w, K)ll |1 f( I,
From item 3 Lemma C.I, we have [S w, x)]l - xC2 vd almost surely, and
| f fSS, w, x)ll & xC2Vd almost surely. follows that || S n C2d. Therefore, I || & Cdd.
| Next, we will provide an upper bound for [ (Enen (Z+ - (En - 1|
'1l
| E Zd -[-"")("--EE""11
[ max
E 77 E 122 []]
Be
B max
(Z"
beB
B max
[E[]] f ur [| |V[ w, VS) w
w, II) [ | [[ w.) I))
BEB
B
max
be Ba:lu:
From item and item 5 in lemma C.I, we have:
f w, 1)] S Cnrdd,
and
E [[ x C²
Combining these two bounds, obtain:
]
Σ 2 ( 6
Σ (3 -(2))1121))
E
be i3
Due to symmetry, the same bound holds for Finally, by applying the Matrix Bernstein inequality, have:
T E - (Z'
12
Pr
(Z - B[Z")
2 2
< cdc
B
BBldkA + 100(223222)
For B 22 (+)
(1
log d / the quantity on the right-hand side bounded by 8'.
Therefore, with probability at least 1 8', we have:
Σ h B
S 2a,en?2
Combining the above equation with Equation (23) completes the proof of the Theorem.
In the proof of Theorem E.I. we will utilize the following general linear algebraic result:
26

---

Lemma E.3. For 20, 21, ..... 2x-1 E R and probability distribution (Po, Pi.....P-- over k el- ements, let Z Ex pizzzz. For symmetric matrix M and i > O, let u1, u2... us be top & singular vectors of M and let U [un, U2, ue] Rdxx, then we have:
2 1) - maxx 1|2:1
t< k
[[II -UT)2ll
(1+1)p
M Po
ift k.
Lemma E.3 provides a bound on the preservation the component 2g by the subspace spanned by the top-l singular vectors of symmetric matrix M. This bound is expressed in terms of the spectral distance between matrices Z and M, the maximum norm of any 2, and the weight the component corresponding to 2q in Z. The proof of Lemma E.3 can be found in Section I.
Utilizing Lemma E.3 conjunction with Theorem E.2, we proceed to prove Theorem E.1.
Proof of Theorem E.I. From Lemma E.3, we have the following inequality:
2(++1))A- pia. + p, IV ((=,,,,>>)
8<
- UUT) D (,, x)]]1 S
(2+1)po
21 AA- P,
if 1 k.
Po
By applying Theorem E.2 and utilizing item 1 of Lemma C.I. it follows that with probability of at least 8', we have:
C²Cr (E+1)po
l << k
(UTTUU)EDD U,
if 2 k.
The theorem then follows by using Po 2 02/2 and e min(k, 20re
L
F
Grad Estimation
Recall that in gradient estimation step r, Algorithm utilizes the subroutine GRADSUBEST to find a projection matrix p(r) for an é-dimensional subspace. 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 Do is small. Therefore, suffices estimate the expectation of projection of the clipped gradient the subspace.
The main algorithm passes the medium-sized batches B (r) the Ê-dimensional projection matrix p(r), and collection of i.i.d. samples S = s .() from Do to the subroutine GRADEST. Here, B!r is random subset of the collection of medium-sizer batches B..-
The purpose of the GRADEST subroutine is 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 smaller (-dimensional subspace, the minimum batch size required for the batches in B,, and the number of batches required depend on l rather than d.
The following theorem characterizes the final guarantee for the GRADEST subroutine:
Theorem F.I. For subroutine GRADEST, let nmm denote the length of the smallest batch in B, N denote the number of batches in B that has samples from Do and be a projection matrix for B some & dimensional subspace of Rdd Ti 2 2(log 8 T2 2 22oog $), nn 2 47,00 L |S*| 2 22T29( ( samples, and N Thm
(722(4) then with probability > 28' the estimate A
returned by subroutine GRADEST satisfy
A - E D.
w.). Cr
The above theorem implies that when the length of medium-sized batches is π(VL) and the number of batches in B containing samples from Do is (2(),, the GRADEST subroutine provides a reliable estimate of the projection of the clipped gradient onto the E-dimensional subspace defined by the projection matrix P.
27

---

Next, we provide description of the GRADEST subroutine and present a brief outline of the proof for the theorem before formally proving 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 Da and the distribution used for the samples in b are close or not. With high probability, the algorithm retains all the batches from Do while rejecting batches from distributions where the difference between the two expectations is large. To achieve this with an E-dimensional subspace, we require ((v)) samples in each batch (see Lemma F.6).
Following the rejection process, the GRADEST subroutine proceeds 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 f-dimensional subspace, (2(2) 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 keep and for estimation using the retained batches.
We prove the theorem formally in the next subsection.
F1 Proof of Theorem F.1
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 Di. For > 0, w Rd and a projection matrix P foor an E dimensional subspace of RRd, we have
E[PV [(S, w. x)] I.V,IP/I w, k)],
and (V x=))]l S 2xd C1 and 17 (Co w. x))) S (O
Proof. Note that,
[I, E[P(VJ(S, w, x))] = PE[VS(SS w, x)]] = PE D. x)] = E [PV (2, y, k)],
where the second-to-lass equality follows from Lemma C.I.
This proves the first part of the lemma. To prove the second part, we bound the norm of the covari- ance matrix:
Cov(PV w. x)) max Var(uTPV w. k))
[[
max Vf w. k))
||
[[Co((V((S,, k))l
R2
C1,
m
where the last inequality follows from Lemma C.I. Similarly,
Cov(PVf(S', k)) S R?
C1.
m
Finally, since random vector PVf((' k) lies in l dimensional subspace of Rd, correspond- ing to projection matrix P, hence its covariance matrix has rank < E. Hence, the relation Tr (Cov(PVJ(S) w, x))) S él[Cov((P) f(S, w, x))ll follows immediately.
The following corollary is simple consequence of the previous lemma: Corollary F.3. Consider two collections and S' each consisting of m i.i.d. samples from distri- butions DD and Do. respectively. For σ. > 0, IRd and projection matrix P for an dimensional subspace of Rd, let 2 = PVVf((( w. k) - ( w. k)). we have:
E[2] E D. PV m E D. IPV fx, W. k)],
and [|Cor())]] < 2x2 C1 and Tr (Cov(z)) < < Cov()((
28

---

Proof. The expression for B|2] can be obtained from the previous lemma and the linearity of expec- tation.
To prove the second part, we bound the norm of the covariance matrix of 2.
Cov(z) Cov(P(Vf(S, w, v w, x)) = 2(Cov(PV/(S,, us, R)) Cov(PV w, k))). Using the bounds from the previous lemma,
conclude that ||Cov(=)]| An2 C₁.
Finally, since the random vector z lies in the 2-dimensional subspace of IR defined by the projection matrix P, its covariance matrix has rank v Therefore, we have Tr, (Cov(z)) - (=vlll
The following theorem bounds the variance of the dot product of two independent random vectors. will helpful in upper bounding the variance of ç) (defined subroutine GRADEST).
Theorem F.4. For any two independent random vectors 21 and ø2, have:
Var(z1 22) & 37r(Cov(zz)) - ||Cov[z))]l + 3|| Cor(()))) + 3|
Proof. We start by expanding the variance expression:
Var(21 22) Var(21 - E[21 - 22))
Var(21 22 E[21] [2(22)
Var((21 - E[z1]) (22 B[22]) + [[2]] 22 [[22] z1)
S 3 Var((21 E[2]]) (22 E[22]) + Var(B[al] 22) Var(B 21)
= 3 Var((=1 - [E[z1] (22
+ [[23]))))] B[2]] Cov(zi) B[z]]
S 3 Var((=1 - [E[=]]) (22 -B[[=]]) + 3 B[x]]1]]Cooo + 3| B
To complete the proof, we bound the first term the last expression:
Var((z) - [[a]]) (z2 E(22l))
= E ((21 B[ail) (22 - -[
= IE[(21
- B(ai))
S EE [[ [[<il] -[=.ll"-- |[Cov(22))
= Tr(Cov(=1)) ||Cov(z2)]]-
Using the two previous results, we can establish bound on the expectation and variance of CJ- Lemma F.5. In subroutine GRADEST, let be a projection matrix of an dimensional subspace. Suppose Sj has 2 m i.i.d. samples from Do and, Si and St. Si,it,, have 2 i..d. samples from for some 10,1...., 1}. Than have:
E(CH |E [PV (....) K)] E ED, y. w))l
(.,
and
48
24 EC
Var(C;) - m²
E|c7/C1.
- m
Proof. Let = P(VS(S) (, WK)) - VS(S,, w. x)) and 22 = P(V ff(S) W. k)).
From Corollary E3, we know that
E[a]] E[ay] = Ep, (PP
E
Cov(z1) Cov(z2) = 4x2 Ci
and
4 Ln 2n=C0
Tr(Cov(21)) Tr(Cov(22))
m
29

---

Note that 5 21 22. Then bound on the variance of ç) follows by combining the above bounds with Theorem F.4. Finally, the expected value of c} is:
E B[zi] [ |[=]]]]
The following lemma provides characterization of the minimum batch length in B and the size of the collection S" required for successful testing in subroutine GRADEST.
Lemma F.6. In subroutine GRADEST, let be projection matrix of an dimensional subspace, |Bl and each batch b B has at least |S 4T,2 samples, and |S*|
T 2 20log 8'
27,2( ( Then with probability 2 - 8', the subset B in subroutine GRADEST satisfy the following:
1. B retains all the batches in B that had samples from Do-
2. B does not contain any batch that had samples from Dj if is such that
[E IPV (,, k)] - LPV.PP( x)]ll 2e
Proof. The lower bound on |SS| in the lemma ensures that for each batch and all E [2Ti]. have St 2 2( ), and the lower bound on |S ensures that for all E [2Ti]. S, = 0 ()
First, consider the batches that have samples from the distribution Do-
For any such batch and j E [Ti]. from Lemma F.5, we have [[C]] = and Var O(d
o(d
| Therefore, for Ti 2 S2olo ). it follows that with probability 2 1 62/2 for every batch b B that has samples from Do the median of {c} (C)))en Je(m) will be less than x²x and it will be retained
in B. This completes the proof of the first part.
Next, consider the batches that have samples from any distribution D; for which
| [PV y. w. x]] E p. (PV f(z, w. x.]ll > 2
For any such batch b and j [Ti), according to Lemma F.5, have E[Cy) 2 4c2 and Var(c)) = (EE Hence, for Ti 2 22oog ااا ), follows that with probability at least 1 - 8' /2, the median of { [E( for every batch will be greater than e&x²C,, and those batches will not be included B. This completes the proof of the second part.
The following theorem characterizes the number samples required in for an accurate estimation of A.
Lemma F.7. Suppose the conclusions in Lemma F.6 hold for ß defined in subroutine GRADEST,
T2 2 2(log 1.) ), each batch b E B has size 2 nm, and [B] Rom 2 2T (2), then with probability
> 1 the estimate A returned by subroutine GRADEST satisfy
|A D. (PV.PPP k)ll s 9enV C₁.
Proof. Recall that in subroutine GRADEST, defined
A 1
Σ PV PVS W.
bEB
Let 2t PV w, R). From Lemma F.6, for all E B, we have
E D. f2, w. k)] < 2Ex CT
30

---

Therefore,
1
E A E
k)]1
E[a
ED,
(VV yr
PV
B
BE
I max bEB
||
E
(PVS
x)]l
20// Ci. (24)
Next, from Lemma F.2,
R² S
T2n²
2I2h2
2T2R² C1
[[Cov<>]]| -
C1
C1
C1 -
(25)
2,
|
min
||51
where the two equalities follow because for all batches E B, |S and |S
/
Then
2T
[Cov(A.))l
max b
[|Cov
B min S
Since A lies in an l dimensional subspace of Rd, it follows that
20T2n² C1
Tr(Cov(A.)) < E|]Cov(A))) <
B min BEB |
Note that Var(I - B[A,]]]) Tr(Cov(A)) (4:).. Then, from Chebyshev's bound:
Var(1 -F(A,]])
2fT₂
Pr|l|: E 2 as C1 <
I
S 1/88.
exxC1
22) 2 min |
Combining above with Equation (24),
Pr[lA: - E D. [P. y, K)]] 3Ex C 1/4
Let D = E E [Tb]
A - E E IPV fx, u. x)]l 3 Then, for T2 22log }), with
probability 8', we have
|D| 2 2 T2.
Recall that in the subroutine, defined & median (j [T2] || and i = arg minfi [25] : }.
From the definition of D, and triangle inequality, for all i, E D, we have |A; A,ll S Gexv CT. Therefore, if [D] 2 72, then for any E D, & S 6ex/Ci. This would imply & I 6exv CT. Furthermore, since [D] 2 1T2, there exist at least one E D such that | A A. I bek /eK/C1
Using the definition of D, and the triangle inequality,
conclude that
|| ED.
x)ll ||A EE (VV wlll + ||A: A / CT.
Theorem F.l then follows by combining lemmas F.6 and F.7.
G Number of steps required
The following lemma shows that with a sufficiently accurate estimation of the expectation of gradi- ents, a logarithmic number of gradient descent steps are sufficient in the main algorithm
Lemma G.I. For e 0, suppose (r) Co(((r) wo)ll < 1 (1(0(()) woll + 2f and R
(CC) log [wlll] ) then (ra woll ea.
31

---

Proof.
Recall that w0++1) = w(r) - 1 Ar). Then we have:
((++1) wD w Wo
1 A (()(r) wo)
(1
1
1
Eo
(Eo(w -wo) - A
C1
Ci
C1
Using triangle inequality, we obtain:
1
1
|| woll ||(((r) woll
1
Eg
+
- -wb) - 1
C1
C1
|10(r)
wo
- (1(0()) - woll
1 C₁
+
Ci
2
S (110(() - woll
1 201
60
4C1
Using recursion, we have:
[( woll S (1)00)) woll (
R
R-
+
20,
20
40,
i=0
I (110(1)) woll exp
R
+ 2C1
60
2C1
4C1
So,
where the second inequality follows from the upper bound on the sum of infinite geometric series and the last incquality follows from the bound on R and w(()) O.
H Final Estimation Guarantees
Proof of Theorem 2.1. We show that with probability
8, for each [R]. the gradient
computed by the algorithm satisfies || | - (w I 1 (1 wo | 4. Lemma G.1
then implies that for R OCC₁ log M ). the output returned by the algorithm e (( satisfy || woll S Eo.
To show this, we fix r, and for this value of r, we show that with probability 2 1 - 6/ (1
sa Eo wo)] < || - woll + 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.I, for |86 I/R 2CC log 1/8",, with probability
8¹, we have
8(C+ 1)C, (E B..I( (1 x w()))p))
2(C + 1 D. [(x (w(r) wn)) + 02)
S (r) < 28
E1
(26)
Next, 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:
E ED w ((())) En(w(() - wo) s <1llw/
woll-
(27)
Recall that in B,, at least a, fraction of the batches contain samples from Da- When B, is divided into R equal random parts, w.h.p. each part B((r) will have at least ag fraction of the batches containing samples from Do-
From Theorem 1. if (B! =
|
R
22 d
0,e3
a,e
+
Ci
C² log 8.)
then with probability 2 1-8',
the projection matrix p()) satisfies
E P. f y, w(rr, (())),)())))) p(r) E Do
Vf w) ,((6))11 s 4r
(28)
32

---

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 good estimate of the expected value of clipped gradients projection. Let N denote the number of batches in Bm that have samples from Do- If N 2 SR + log 1/8') then with probability 1 1 8'. Bhh has (NN/R) batches sampled from Do- If each batch in Bm and batch b" has more than 7m samples, 7.m = 22 Vi log( c 8' [Bml)), and
R
Nonm 2 2 2 log 1/8"), then from Theorem F.I, with probability 2 1 - 8'
a (r) E D, f w() ,e | S 9e2n!! VG.
(29)
Combining the above three equations using triangle inequality,
|| (r) - Eo(((rr) - wo)ll s 13c2k(7) (r) VCr + f - woll:
(30)
with probability 2 1 - 58'.
In equation (26) using the upper bound, D. [(x (< (w( wo)) s wol we get
woll - Cillw(
2(C 1)C1 (|(((])))))))))))))
woll? 2 + (C 28 /8C +1)
(())) s 28
(C woll + (1o).
VEI
VEI
Combining the two equations,
364e, v 2(C + 1)C
|| (r)(() Eo(((() wo)ll S
(Cillw woll (Cao t (1ll woll,
Vi
(31)
with probability 2 1 58' There exist universal constants C1 0 such that for 61 C1 and
CIVC=T c v() the quantity on the right is bounded by ||w((r) - woll/2 + ca/4. We
choose these values for 6 and 62 and 8' 5
5R
From the above discussion, it follows that if [. = =6 =i((.4.t) 7mm 2 50 and Bm has 2
δ batches sampled from Do, then with probability 2 5/R.
1 [10(()) -woll + ca
A (r) En(w0)) - wo -
2
4
Using l = min(k, Foa. }. we get the bounds on the number of samples and batches required by the algorithm.
I Proof of Lemma E.3
To establish the lemma, we first introduce and prove two auxiliary lemmas.
Lemma L.1. For 0, and a probability distribution (Po.PI, P1, Pk-1) over elements, let Ex=0 pozzz,, where zı are d-dimensional vectors. Then for all 2 0, 8 largest singular value of is bounded by max; 2//E
Proof. Note that Z is symmetric matrix, so its left and right singular values are the same. Let U1 02 be the singular vectors in the SVD decomposition of 2, and let a S a2 S 0g S be the corresponding singular values. Using the properties of SVD, we have:
Σ
P32zz
Σ Pi 1=0
Σ i
Uss
2,)²
1 max 1|2/|1
1=0
33

---

Next, we have:
Σn Σ isE
24 2
2
Σ SSE
og l.
at-
Combining the last two equations yields the desired result.
Lemma 1.2. Let u1, U2.... [u1, U2; Ug] E Rdxx For any set of vectors 20. 21, ... 2k-1 E IR non-negative reals Po, PL, ..., Pk-1, and reals a1, 02 ... dg, we have:
IR. be f. mutually orthogonal wnit vectors, and let U
E Pizzz Esgggauyuu
- -UTT)
po
Proof. catv=(1-U0)
First show that for all j E [$]. the vectors and uj are orthogonal,
u UUT)=0 (u] 20) (u] za) O.
Then,
|OT EEio Pizizi - " auu,))) = =||0T (2 Piziz o01| - (2T1))| > ||po(=
Next, we have:
o/ UUUT) + ~0U
2] UUTY vTv ||aa|
Combining the last two equations, we obtain:
||a||2 /E p Eser u 2 | (Et= 2, - Es pollu/14
Dividing both sides by ||'l| completes the proof.
Next, combining the above two auxiliary lemmas we prove Lemma E.3.
Proof of Lemma E.3. Let (.).) denote the 5th largest singular value of a matrix. Let M be rank truncatedSSDD of M, then follows that,
|| Ml] Ae4a
First, we consider the case l < k. By applying Weyl's inequality for singular values, we have
max, ||
Aes (M) < 1 (Z) + (M Z) <
+ +IIM - Z11.
t+1
where the last equation follows from Lemma 1.1.
First applying the triangle inequality, and then using the above two equations, we have
max, ||, + 2|IM Z||-
A - 2 S M M + |M 2||
t+1
Combining the above equation with Lemma 1.2, we have:
- (UT)0 2(E + 1)I[M - Z + max, ||,
(e 1)po
This completes the proof for & < k. To prove for the case é > k. we use Ac+1(Z) 0 in place of the bound Ae+1(Z) < max,
||, in the above proof for the case l < k.
2+1
34

---

Removing the Additional Assumptions
To simplify our analysis, we made two assumptions about the data distributions. We now argue that these assumptions not limiting.
The first additional assumption was that there exists constant C2 > such that for all components
i E {0, ........ 1} and random samples (x. y) Di we have ||||| S C2vd 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 C2 such that with probability 2 0.99, we have ||=||| Σ C2vd. Therefore, disregarding the samples for which |x| > C2vd does not significantly reduce the data size. Moreover, has minimal impact on the covariance matrix and hypercontracctivtyy constants of the distributions. This argument easily extends to the batch setting. In the batch setting, we first exclude samples from batches where ||=||| > C2vd. Then we remove small-sized batches with fewer than or equal to 2 samples and medium-sieee sized batches that have been reduced by more than 10% of their original size. is easy to show that w.h.p. the fraction of medium and small size batches that gets removed for any component most 10%. THence, this assumption can be removed with a small increase the batch size and the number of required samples in our main results.
Next, address the assumption that the noise distribution is symmetric. We can handle this by employing a simple trick. Consider two independent and identically distributed (ii.d.) samples (I1, 91) and (x2, 82). where y/4 = w" I; + 7%. We define x = (01 - /22/v2, y = (#1 - 922/ and n = (m1- /22/ 22. It is important to note that the distribution of v7 is symmetric around O, and the covariance of is the same as that of 24, while the variance of r7 is the same that of mn- Furthermore, we have y w² - + 17. Therefore, the new sample (xr y) obtained by combining two i.i.d. samples satisfies the same distributional assumptions 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 constant factor of 1/2. Thus, the assumption of symmetric noise can be eliminated by increasing the required batch sizes in our theorems by factor of 2.
K
Algorithm's Parameter Requirements
The algorithm requires two distribution parameters: the L4-L2 hypercontractivity parameter C and the condition number of input covariance matrix C1. These parameters significantly generalize the input distributions typically considered in mixed linear regression, which have primarily focused on Gabb-Gassiann distributions with an identity covariance matrix. Specifically. the L4-L2 hypercontrac- tivity parameter generalizes the sub- gaussian distribution assumption and allows for heavier-tailed distributions. The condition-number Ci of the input covariance matrices generalizes the identity covariance matrix assumption for input distribution, corresponding to C1 1.
Additionally, the algorithm requires two population parameters: the number of sub-populations and the fraction of batches present from the relevant sub-population. Note that the algorithm combines these parameters to define single parameter L.
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 good estimate. If for any parameter we not have 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.
L More Simulation Details
Setup. We have sets B2 and Bm of small and medium size batches and distributions Dj for i E (0, 1,......
..... ,,--
For a subset of indices 7 C {0, 1,...,A k 1}, both B2 and Bm, have a fraction
35

---

190
Ours/4 samples)
KS9KO14 semples)
Ours/8 samples)
KSSKD/8 samples)
630
Ourst4 samples)
CurstB samples)
330
MSE
230
200
20
Eiediam batzk 1220
Figure 4: Different input dist, A 100, large minimum distance between regression vectors, recovering 4 components that have 1/16 frac- tion of batches each
Figure 3: Same input dist. (standard normal), = 16, small minimum distance between re- gression vectors, recovering all
of a batches that contain i.i.d. samples from D; for each i E I. And for each i E {0. 1........ ............ h- 1 in the remaining set of indices, B, and Bm, have (1 - / (k: ( [/]) fraction of batches, that have i.i.d samples from D. In all figures the output noise is distributed as (N 1).
All small batches have 2 samples each, while medium-size batches have Rem samples each, which we vary from to 32, as shown in the plots. We fix data dimension d = 100, a = 1/16, number of small batches to [B. min min(8dk², 8d/a") and the number of medium batches Bm 256. In all the plots, we average our 10 runs.
Evaluation. Our objective is recover small list containing good estimates for the regression vectors of D; for each é E I. We compare our proposed algorithm's performance with that of the algorithm in [KSS $ 20]. We generate lists of regression vector estimates Lours and LKSSKO using our algorithm and [KSS* 20), respectively. Then, we create 1600 new batches, each containing Nnew i.i.d samples randomly drawn from the distribution D,, where for each batch, index is chosen randomly from I.
Each list enables the clustering of the new sample batches. To cluster batch using list, we assign it to the regression vector 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 Nnew 4 and Nzew 8, as shown in the plots.
Minimum distance between regression vectors. Our theoretical analysis suggests that algo- rithm 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[, -).. Notably, the minimum gap between the vec- tors, 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 sigmificantly as we deviate from the assumptions required for the latter algorithm to work.
Number of different distributions. 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, to highlight this ability. In this case, the first four distributions each generate 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 Dj for each E for identification of 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.
Our algorithm notably handle very large (even infinite)
36

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
These claims are justified by our theorems and the experimental results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We present such a discussion in the last section of the main paper.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the ap- proach. For example, a facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Ys]]
37

---

Justification:
All assumptions and proofs are clearly mentioned in the paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theo- rems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We include the details to reproduce our experiments in the paper. We also include the scripts to reproduce all our results in supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models. it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
38

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our experiments are on synthetic data and we include the scripts to reproduce all our results in supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (nttps://nips.cm cc/
wtt.miiisssnoiciiiiiioici for more details. wttpii//guiidee/oeeeeeeeeee
While we encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips. cic.ips.ccppubbii c/gui deuidde//oe Cooe//ubeeeeieeeeeeeee ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g.. data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We provide these details in the experiment sections in appendix and the main paper. We also provide the code used to generate the figures in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification:
We report standard error over 10 runs in all our plots.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
39

---

The method for calculating the error bars should be explained (closed form formula, call 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.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative enror 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.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
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 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 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 failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https http///eurrrp..peee wttpp//cc///pcecieeeeeee httpbb/iieettiiccccceeeeeeeeeeee
Answer: [Yes]
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 NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal 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.
40

---

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), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is 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, 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 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 unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion 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).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release any new dataset or model.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere usage guidelines restrictions access the model or implementing safery 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 challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We used algorithm in prior work [KSS 20] and we have cited the original
paper.
Guidelines:
The answer NA means that the paper does not 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 used and, 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.
41

---

If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswithcode com/aatasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and sereenshots, applicable, well as details about compensation (if any)?
Question:
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper deseribe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
42

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere 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.
43


Reviews:
[{"review_id": "jpxDvAt9dW", "reviewer": "Reviewer_ntyP", "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].", "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? \n\nAuthors claim that $k$ can be infinity, as long as the constraints of $\\alpha$ are satisfied. Do authors test the claim numerically?", "limitations": "See weakness.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The introduction of the algorithm is mostly clear, despite its complexity. The intuitions are well-explained. The gradient clipping technique is novel compared with [KSS+20]. Also, the use of gradient descent in the algorithm presents an interesting contribution.", "weaknesses": "In line 123, the authors claim that the model does not rely on the previous assumptions. Authors can also briefly discuss how their algorithm innovations circumvents these assumptions. \n\n\n\nSome minor issues about writing:\n\nThe sentence in line 332 is very long and hard to follow.\n\nIn line 836, there is an error.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the careful review and constructive feedback, and for appreciating the contribution, novelty, presentation clarity, and intuition. \n\nRegarding your comments and questions:\n\nLine 123: Section 1.3 summarizes the algorithm’s innovations, and we will elaborate on how they help overcome assumptions in prior works.\n\nLine 332: We will break it into several sentences so it is easier to follow.\n\nLine 836: We will add the missing reference.\n\nNext, we answer the two questions raised in the review.\n\nFor an example motivating the number of small batches, observe that in many recommendation systems, most users rate only a few items. The ratings by any such user may serve as a small batch. Since these systems have many users, the number of small batches will also be large. Observe also that typically relatively few users provide a fair number of ratings, and they can serve as medium-sized batches.\n\nFor a test justifying that $k$ can be large, Figure 4 in the appendix shows an experiment where the number $k$ of subpopulations is large (100), while the number of subpopulations with sufficiently many batches for enabling recovery is small (4). Our theoretical analysis shows that this experiment can be repeated with even larger $k$, without impacting the algorithm’s performance.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ezy2sfocek", "reviewer": "Reviewer_Xj2G", "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.", "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 . \n\nDuchi 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": "Yes.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "There's a lot to like about the overall approach to the problem. I would compliment it as being \"applied theory,\" in contrast to a lot of prior work which takes a more purely theoretical assumption-laden mindset (as commented upon also by the authors, in describing the many assumptions that they are able to remove). Algorithmic techniques like clipping and median-finding are coupled with appropriate theoretical choices, to achieve a sample complexity that's nearly optimal for these regimes. \n\nThe applicability of the problem itself is more than meets the eye; theory can motivate usage of certain algorithmic techniques. In this case, working in the dual gradient domain to do clipping has a very helpful effect, as the authors prove -- and this is currently not standard for the problem.", "weaknesses": "Sec. 1.2 and 1.1 cover overlapping material, and in addition they could be switched with each other. I found 1.2 much clearer about the top-level contributions (e.g. removal of assumptions), and 1.1 more of a detailed theoretical comparison. Overall, some rearrangements of the initial section 1 would improve presentation (e.g. \"Our work\" line 34 unclear reference at that point for the reader). Would suggest moving forward \"improvement over prior work\" subsection even within Sec. 1.2. It takes 3.5 pages to formally define the setting with batches, continuing the reader's ambiguity somewhat.\n\nFor 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive review, the appreciation of our contribution, and the two suggestions that we discuss next.\n\nThe first suggestion concerns rearranging Sections 1.1 and 1.2. Section 1.1 discusses the paper’s main results, and section 1.2 addresses the results’ improvement over prior work. We agree with the reviewer that the comparison in Section 1.2 results in some duplication of results outlined in Section 1.1, and will rearrange some of the material to minimize the repetition. \n\nThe second comment concerns R-way partitioning of the data. More concretely, our algorithm runs in R rounds, and our theoretical analysis requires the use of independent samples in each round. \n\nWe thank the reviewer for mentioning that Duchi’s work might help eliminate the need for this R-way partitioning in our analysis either entirely or under appropriate constraints. \n\nIn searching papers on Ergodic Gradient Descent by Duchi et. al. we found the paper “Ergodic Mirror Descent” that we believe the reviewer referred to. While this paper performs gradient descent using ergodic samples, and hence generalizes the independent samples we use, it still does not eliminate the need of new samples in each round of gradient descent. In the limited rebuttal period, we could not find a way to extend the analysis in the paper to allow using the same samples in each algorithm round, and we leave that for future work.\n\nNote however that to achieve estimation accuracy $\\epsilon \\sigma$ with noise variance $\\le \\sigma^2$, and regression vectors of length at most $M$, the number of rounds $R$ required by the algorithm is just $\\log(M/\\sigma)$. Hence the factor $R$ increase in the sample complexity implied by our theoretical analysis is at most a logarithmic factor. \nFurthermore, in our numerical experiments, the algorithm recovers the regression vectors even when the same samples are used in each round. However, we agree with the reviewer that without a theoretical proof it is unclear if it will not be necessary for some extreme datasets.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1uTYYpYkmI", "reviewer": "Reviewer_oZju", "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.\n\nThe 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.\n\nI 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.", "questions": "What are the main intuitions behind the assumptions made in Remark 2.1? (L. 196)", "limitations": "-", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written and well-motivated. The problem addressed is natural and aligns with numerous modern applications in machine learning and data science. The review of prior works is thorough, and the authors have effectively positioned their work within the existing literature.\n\nThe authors claim significant advancements in their bounds concerning the relationship to certain parameters and dimensions compared to previous works. For example, their bounds require at least $\\Omega\\left(k^{3/2}\\right)$ batches of data, where each batch contains at least $\\Omega\\left(k^{1/2}/\\epsilon^2\\right)$ samples. As can be seen, the number of batches is independent of $\\epsilon$ (the estimation error). In comparison, some of the most notable prior works require the number of batches to grow significantly with $\\epsilon$.\n\nThe authors have mitigated several significant constraints present in prior works, including eliminating the assumption of a Gaussian distribution for feature vectors. Additionally, this work removes the assumption that all subgroups must follow a linear model. In fact, the authors extend the concept to accommodate scenarios where the number of subgroups can even approach infinity, ensuring the recovery of the \"majority\" of subgroups with linear models.", "weaknesses": "There are too many hyper-parameters in Algorithm 1, which are provided to the algorithm as input. However, I found the discussion on how to tune these hyper-parameters in practice to be lacking. It would be beneficial to include more detailed guidelines or heuristics for selecting appropriate values for these hyper-parameters, potentially based on empirical studies or theoretical insights. Additionally, some of the assumptions in Section 2.1, such as the distributional properties of the input data and the specific conditions under which the theoretical guarantees hold, may need more clarification. Providing concrete examples or scenarios illustrating these assumptions would help in understanding their practical implications and ensuring they are not overly restrictive.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for many positive comments. The reviewer’s main concerns involve the number of parameters and distributional assumptions, both addressed below.\n\nWhile the number of parameters may seem large, each has a natural and important contribution.\n\nTwo accuracy parameters $\\epsilon$ and $\\delta$ that are part and parcel of any PAC learning algorithm. They specify the accuracy we are looking for and determine the number of samples needed to achieve it. \n\nTwo regression parameters: the length $M$ of the regression vector, and the additive noise variance $\\sigma$, reflect the complexity of the input-output relation.\n\nTwo population parameters: the number $k$ of sub-populations and the fraction of batches $\\alpha$ present from the relevant sub-population. \nNote that the algorithm combines these parameters to set one parameter $\\ell$.\n\nFinally, two distribution parameters: the L4-L2 hypercontractivity parameter $C$ and the condition number of input covariance matrix $C_1$. These parameters reflect the complexity/richness of the distribution class and substantially generalize input distributions considered in the past for mixed linear regression, as they only considered sub-gaussian distributions with the identity covariance matrix. \n\n\nThe L4-L2 hypercontractivity parameter significantly generalizes the sub-gaussian distribution assumption and allows for heavier-tailed distributions. \nThe condition-number $C_1$ of the input covariance matrices significantly generalizes the identity covariance matrix assumption for input distribution that corresponds to $C_1=1$. \n\nThough the number of parameters seems large, the algorithm needs only an upper bound on their values. If the upper bounds are tight, they will result in a good estimate.\nIf for any parameter we do not have a tight upper bound, we can run the algorithm for different values of the parameters, dividing their possible range logarithmically, and \nchoose the parameters that achieve a small error on a holdout set.\nWe can add this approach in more detail in the final version of the paper and formally define error computation and selection for our setting.\n\nFurthermore, it is easy to see that these upper bounds are needed. For example, even for the simplest case when all samples are from the same sub-population, some condition on the anti-concentration of  $XX^T$ as above is needed to obtain finite sample complexities [Lecué and Mendelson, 2016, Oliveira, 2016] in heavy-tailed linear regression, and this corresponds to a low condition number. \n\nThe reviewer also asks about the intuition behind Remark 2.1. Our algorithm and its analysis are based on gradient descent, and can handle stochastic noise. Therefore we require only the distribution of gradients to be similar for batches in the same sub-population and since the analysis can tolerate statistical noise it can also tolerate deviation in the distribution of gradients across batches as long as it is small.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jyOxGqZuSG", "reviewer": "Reviewer_pdPW", "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.", "questions": "What are the barriers to implementing/trying this on real data?", "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.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "Originality and significance: The paper considers a novel way to tackle the problem of linear regression with heterogeneous batches. The work is significant as their algorithm removes a number of assumptions in prior literature, and seems more performant than previous algorithms. The theory provided is solid and sound. The idea is good, and the description of the proof method and the grad clipping idea was good. \n\nQuality and clarity: The problem is well motivated, and the contextualization with respect to prior works is made clear.", "weaknesses": "Clarity: \n\n- 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.\n- Theorem 2.3: The first sentence does not make sense to me. What does it mean?\n\nMore comments: \n\n- 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.\n- 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.\n- 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? \n- The input distributions can be different, but they must be known? Whether yes or no, I would mention this clearly in the contribution.\n- 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?\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful review and positive comments. The reviewer’s comments are addressed herein.\n\n1. We thank the reviewer for pointing out the typos, and will correct them along with some others we found. \n\n2. The condition part of Theorem 2.3 can be more clearly stated as follows: \nLet index $i$ be in set $I$ and let list $L$ be a list of $d$ dimensional vectors. Suppose for a given $\\beta>0$ at least one of the vectors in $L$ is within distance $\\beta$ from the regression vector of distribution $\\mathcal D_i$, namely $\\min_{w\\in L} ||w-w_i||\\le \\beta$.\n\n3. We thank the reviewer for the reference to the recent work on mixed linear regression. It is clearly relevant, and we will include it in our paper.\n\n 4. (see next point)\n5. Addressing 4 and 5 together, note that, as stated in corollary 2, the $L$ identified by our algorithm is at most $\\tilde O(1/\\alpha_m)$, even for arbitrary large $k$. \nFurthermore, if all $k$ sub-populations follow the assumption in Theorem 2.1 the argument preceding corollary 2.2 can be modified as follows to show that the list $L$ identified by our algorithm can be trimmed down to have at most $k$ elements. In such a case, for any choice of $b^*$ from $B_m$ to run algorithm 1, the regression vector estimates $\\hat w$ will satisfy $||\\hat w-w_i||\\le \\epsilon \\sigma$ for some $i\\in [k]$. \nA simple argument using the triangle inequality shows that removing the estimates from $L$ until no two estimates are within distance $2\\epsilon\\sigma$ from the list $L$ will result in a list $L$ of size $\\le k$ such that for all $i$ we have  $||w_i-\\hat w||\\le 3\\epsilon \\sigma$. \nAlso, we note that we already included the literature on list decoding. We referred to a recent work of list decodable linear regression using batches [DJKS22] in section 1.2 and other references in the appendix. We can discuss these connections in more detail as per the reviewer’s suggestion.\n\n6. The input distributions need not be known. We will mention it clearly in the contributions following your suggestion.\n\n7. If for two distributions the regression vectors have no separation then when the algorithm is run using a medium-sized batch as $b^*$ from either of two distributions we will be estimating the same regression vector. The upper bound on the size of the output list is independent of the separation.\n\n8. As shown by the synthetic experiments in Figures 1-4, the algorithm is practical even for reasonable-size datasets and dimension $d$. As such there are no barriers to trying it out on real data. We leave it to future work to find interesting and relevant datasets for this setting.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CdKQ2ELhTJ", "reviewer": "Reviewer_D4Ej", "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.", "questions": "See Weaknesses.", "limitations": "The limitations have been partially addressed in the paper.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The general quality of the paper is commendable, and it addresses all pertinent issues effectively. It adequately covers related works and provides a comprehensive perspective on the topic.\n\nThe main topic is an important research area with applications in a wide range of real-world projects. Existing works in this area often rely on several assumptions regarding distribution and batch sizes, which can limit their efficiency and accuracy. This paper offers a more general solution that does not depend on restrictive assumptions, thereby enhancing its applicability and robustness.", "weaknesses": "The main concerns about the paper include:\n\nIt 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.\nIn 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).\n\nWhile 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.\n\nAdditionally, 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.\n\nIn 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for many positive comments about our paper. Below we address the remaining comments of the reviewer one by one.\n\n1. Note that even when $k$ is arbitrarily large, the number of batches and batch sizes required in Theorem 2.1 and Corollary 2.2 are reasonable as in that case they depend on $\\alpha$ instead of $k$, implying that batches can be recovered even for arbitrary large $k$. We also provide numerical experiments to further confirm this fact (see Figure 4 in the appendix.) Our main observation that helps in such a recovery for large $k$ is that even when $k=\\infty$, a small subspace can be estimated that contains gradients for all components with a significant fraction present.\n\n2. As discussed in line 112 of section 1.2, our sample complexity is lower than Kong et al. (2020). Our time complexity $\\tilde O(|B_s| d^2/\\alpha_m)$ is at most $\\tilde O(1/\\alpha_m)$ times the complexity of Kong et al. (2020). But as is clear from the value and our numerical experiments show, this complexity is still reasonable in practice.  \n\n\n3. Note that of all previous works, only Kong et al. (2020)  take advantage of the batch structure. All other works mentioned in the paper do not, and hence they will naturally fare worse than the baseline we considered. For that reason we did not include them in our results.\n\n4. We note that in Figures 1 and 2, $k$ is kept constant at 16. The small gap between the two algorithms for small values of medium batch size is because if it is too small then it falls below the requirement of Theorem 2.1. In that case, both our algorithm and baselines fail to recover the regression vector.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (spotlight)