Imbalanced Mixed Linear Regression
Authors: Pini Zilber
Boaz Nadler
Venue: NeurIPS 2023 poster (2023)
Keywords: Mixture regression model
Mixture of linear models
Iteratively reweighted least squares
TL;DR: We solve imbalanced mixed linear regression using tools from robust regression

Abstract: We consider the problem of mixed linear regression (MLR), where each observed sample belongs to one of $K$ unknown linear models. 
In practical applications, the mixture of the $K$ models may be imbalanced with a significantly different number of samples from each model. Unfortunately, most MLR methods do not perform well in such settings. Motivated by this practical challenge, in this work we propose Mix-IRLS, a novel, simple and fast algorithm for MLR with excellent performance on both balanced and imbalanced mixtures.
In contrast to popular approaches that recover the $K$ models simultaneously, Mix-IRLS does it sequentially using tools from robust regression. Empirically, beyond imbalanced mixtures, Mix-IRLS succeeds in a broad range of additional settings where other methods fail, including small sample sizes, presence of outliers, and an unknown number of models $K$. Furthermore, Mix-IRLS outperforms competing methods on several real-world datasets, in some cases by a large margin. We complement our empirical results by deriving a recovery guarantee for Mix-IRLS, which highlights its advantage on imbalanced mixtures.

Imbalanced Mixed Linear Regression
Pini Zilber
Boaz Nadler
Faculty of Mathematics and Computer Science Weizmann Institute of Science, Israel {pini. lber, boaz. nadler H@we zmann ac ii
Abstract
We consider the problem of mixed linear regression (MLR), where each observed sample belongs to one of unknown linear models. In practical applications, the mixture of the K models may be imbalanced with significantly different number of samples from each model. Unfortunately, most MLR methods do not perform well in such settings. Motivated by this practical challenge, in this work propose Mix-IRLS, a novel, simple and fast algorithm for MLR with excellent performance on both balanced and imbalanced mixtures. In contrast popular approaches that recover the K models simultancously. Mix IRLS does it sequentially using tools from robust regression. Empirically, beyond imbalanced mixtures, Mix-IRLS succeeds in a broad range of additional settings where other methods fail, including small sample sizes, presence of outliers, and an unknown number of models K. Furthermore, Mix-IRLS outperforms competing methods on several real-world datasets, in some cases by large margin. We complement our empirical results by deriving a recovery guarantee for Mix-IRL IRLS, which highlights its advantage on imbalanced mixtures.
Introduction
In this paper, we consider a simple generalization of the linear regression problem, known as mixed linear regression (MLR) [5, Chapter 14J. In MLR, each sample belongs to one of K unknown linear models, but it is not known to which one. MLR can thus be viewed as a combination of regression and clustering. Despite its simplicity, the presence of multiple linear components makes MLR highly expressive and thus useful data representation model in various applications, including trajectory clustering [20]. health care analysis [16], market segmentation [49], face recognition [7], population clustering [29], drug sensitivity prediction [37] and relating genes to disease phenotypes [8, 47].
Several methods were developed to solve MIR, including expectation maximization [15, 5], alternat- ing minimization [52, 53] and gradient descent [56]. These methods share three common features: they all (i) require as input the number of components K; (ii) estimate the K models simultancously; and (iii) tend perform better on balanced mixtures, where the number of samples from each of the K models are approximately equal. As illustrated in Section 4, given data from an imbalanced mixture, these methods may fail. In addition, most of the theoretical guarantees in the literature assume a balanced mixture. Since imbalanced mixtures are ubiquitous in applications. it is of practical interest develop MLR methods that are able handle such settings, as well as corresponding recovery guarantees.
In this paper, we present Mix -RLS, novel and conceptually different iterative algorithm for MLR, able to handle both balanced and imbalanced mixtures. Mix-IRLS is computationally efficient, simple implement, and scalable to high-dimensional settings. In addition, Mix-IRLS successfully recovers the underlying components with relatively few samples, is robust to noise and outliers, and does require input the number of components K. In Sections 4 and 5 we illustrate the competitive advantage of Mixx x-x-11 IRLS over several other methods, on both synthetic and real data.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

2
16
Figure 1: Illustration of Mix-IRLS. The data is a mixture of K = 3 components. At each round, Mix-IRLS excludes samples with poor and moderate fit (marked by 'X' and "?', respectively), and performs linear regression on the good fit samples. The poor fit samples are passed to the next round.
To motivate our approach, us consider a highly imbalanced mixture, where most samples belong to one model. In this case, the samples that belong to the other models may be regarded as outliers with respect to the dominant one. The problem of finding the dominant model may thus viewed as specific case of robust linear regression, a problem for which there exist several effective algorithms, see e.g. [27, 50]. After finding the dominant model, we remove its associated samples from the observation set and repeat the process to find the next dominant model. This way, the K linear models are found sequentially rather than simultaneously as the aforementioned methods.
For our Mix IRLS sequential approach to succeed, we found it important to allow in its intermediate rounds an "I don't know" assignment some of the samples. Specifically, given an estimate of the dominant model, we partition the samples to three classes, according to their fit to the estimated model: good, moderate and poor fit; see Figure 1. The samples with good are used to re-estimate the model coefficients; those with poor fit are assumed to belong to yet undiscovered model, and hence are passed to the next round; the moderate fit samples, on whose model identity we have only low confidence ("I don't know"), ignored, but used later in refinement phase.
To recover the dominant model each of the K rounds, we perform robust regression using iteratively reweighted least squares (IRLS) [24, 10, 40J. IRLS iteratively solves weighted least squares subproblems, where the weight of each sample depends on its residual with respect the current model estimate. As the iterations proceed, outliers are hopefully assigned smaller and smaller weights, and ultimately ignored; see Figure F.1 for an illustration of this mechanism.
On the theoretical front, in Section 6 we present recovery guarantee for our method. Specifically, show that in population setting with imbalanced mixture of two components, MRLS successfully recovers the linear models. A key novelty in our result is that it holds for sufficiently imbalanced mixture rather than a sufficiently balanced one (or even a perfectly balanced one) as is common in the literature [52, 2, 35]. In addition, unlike most available guarantees, our guarantee allows an unknown K, and it is insensitive the initializatton, allowing it to be arbitrary.
To the best of our knowledge, our work is the first to specifically handle imbalance in the MLR problem, providing both practical algorithm well as theoretical recovery guarantee. Furthermore, we are the first to propose and analyze an adaptation of IRLS to the MLR problem. A related sequential approach solve MLR using robust regression was proposed by [3]. They used random sample consensus (RANSAC) instead of IRLS, and without our "I don't know" concept. To find a component, [3] randomly pick (d (d+2) 2) samples from the data and ordinary least squares (OLS) on them, in hope that they all belong to the same component, As discussed by the authors of [3], their approach feasible only in low dimensional settings, as the probability that all 2 chosen samples belong to the same component decreases exponentially with the dimension d; specifically, [3] considered only d 5. addition,
[3] did not provide theoretical guarantee for their method. In contrast, our Mix- IRLS method is scalable to high dimensions, and enjoys theoretical guarantee.
Notation. For positive integer K, denote [K] = (1...., K). and the set of all permutations over [K] by []] For vector u, denote its Euclidean norm by ||u]] For matrix X, denote its spectral

---

Algorithm 1: Mix-IRLS: main phase
input samples ((1), 1-) number of components K, parameters Win- P. 73 T1
output: estimates 3p ....., gphase-
set S [[n]
2 for k = to K do
initialize 8g randomly
repeat Ti times
compute ri = |x] BR Wol, ViE Sk
compute Wi,k = -11+ n (7) pb)-1 Vi E Sk
compute 3x- -X3.. Wh. X Wa vs. end
set Sk+1 = (i E Sx wi,k S W
set K {p- d samples in Se with largest w)
11
k < K and |Sk+1| < p. then
12
restart Mix- IRLS with Win + With 0.1
13
end
compute Sptase-1 (X Xs,) ) _X ys
14
18 end
norm by |X and its smallest singular value by Omin (X). Given a matrix X E Rxdd and an index set S C [n]. Xg E RRSSkd is the submatrix of X that corresponds to the rows in S. Denote by (iag(w) the diagonal matrix W whose entries W; wy The probability of an event A is P[A]. Denote the expectation and the variance of a random variable x by E[x) and Var(at (ar[x), respectively. The cumulative distribution function of the standard normal distribution /O, 1) is ((-).
2 Problem Setup
Let ((I;, be pairs of explanatory variables v; E Rd and corresponding responses y; E TR. In standard linear regression, one assumes linear relation between the response and the explanatory variables, namely y/- = x f + 6 where 6; IR random noise terms with zero goal is to estimate the coefficient vector 8d Rd.. In mixed linear regression (MLR), in contrast, the assumption is that each response y/ belongs to one of K different linear models (8t) ]==1 Formally,
A common
y: x2 87 + C2. == n,
(1)
where c* = (c)...,,))) E [K]n is the label vector. Importantly, we do not know to which component each pair (zi, 3i) belongs, namely c* is unknown. For simplicity, we assume the number of components K is known, and later on discuss the case where it unknown. Given the n samples {(x,, 84) v:) the goal is estimate 8 ......... , C Rdd, possibly by estimating also c*. See Figure 4(left) for a real-data visualization of MLR in the simplest setting of d = 1 and K = 2.
For the K regression vectors 3 to be identifiable, sufficiently many samples from each component must be observed. The minimal total number of samples depends on the dimension and the mixture proportions. Denote the vector of mixture proportions by p (pi....,PPk PK), with pa Ki [n] : = //nn Then the information limit on the sample size, namely the minimal number of observations required to recover 3 in the absence of noise, is Taato d/ min(p). As formally proven in Appendix B, no method can recover all K models given fewer samples.
The lack of knowledge of the labels c* makes MLR significantly more challenging than standard linear regression. Even the simplified setting of and no noise (e = 0), the problem is in general NP-hard [52]. However, detailed in Section 6, various assumptions render the problem computationally tractable, and enable to derive theoretical guarantees.
with perfect balance (P) = P2 1/2)

---

3 The Mix-IRLS Method
For simplicity, we present our algorithm assuming K is known; the case of unknown K is discussed in Remark 3.3. MiX-IRLS consists of two phases. In its first (main) phase, Mix-IRLS recursively recovers each of the K components Si..... SK by treating the remaining components as outliers. The sequential recovery is the core idea that distinguishes MiX-IRLS from most other methods. In the second phase, we refine the estimates of the first phase by optimizing them simultaneously, similar to existing methods. As discussed below, in many cases accurate estimates are already found in the first phase, in which case the second phase is not necessary. For brevity, we defer the description of the second phase to Appendix C.
Before we dive into details, let us give brief overview of the main phase mechanism. At each round of this phase, MeL- uses techniques from robust regression to estimate the largest component present in the data. Then, partitions the samples to three subsets, according to their fit to the found component: good, moderate and poor. Mix-IRLS refines the component estimate using only the samples with good fit, and proceeds to the next round with the poor fit samples. The moderate fit samples are ignored in the main phase, as we have high uncertainty in their component assignment- they either may or may not belong to the found component. As detailed below Mix- IRLS depends on four parameters: a model fit threshold Whth E (0, 1). an oversampling ration 2 1, tuning parameter ? and the number of IRLS iterations 7i 2 1.
Importantly, the robust regression used by MiX-IRLS succeeds even if the dominant component consists less than 50% of the data. The well-known "breakdown point" of robust regression at 50% assumes nothing on the data samples, and in particular allows for adversarial outliers. In our case, the 'outliers', namely the samples that belong to yet undiscovered component, follow an MLR model. As such, standard high nigh-dimensional robust regression methods can recover the underlying linear model even with 70% and 80% outliers; see Appendix for an empirical illustration.
Next, we present in details the main phase of Mix-IRLS. A pseudocode appears in Algorithm 1. yn) T, we initialize the set of active samples to the entire dataset, S = [n]. Next, we perform the following procedure for K rounds. At round E [K].
Given X = (I1
xx)) and y (y1
starting from random guess 3x for the ke-th vector, we run T₃ iterations of IRLS [24]:
Ti,k= x 3R vil, Vi Sk.
(update residuals)
(2a)
1
Wiik- = I+ F ViE Sk.
(update weights)
(2b)
8x = (X3, ss.) ) X3.
We Us
(update estimate)
(2c)
Se
where Fx = median(ri.e | E S2 and W₁ = diag [I, .... The updated estimate (2c) is the minimizer of the weighted least squares objective ||W (ys, 8k) 1| close to some 37. then samples from other components receive a lower weight and the updated By is
Intuitively, if 3g is
even closer to 33 After Ti iterations of (2a)-(2c), we define the subset St+1 of "poor fit' samples that seem to belong to other components,
S = {i E Sp wi. S Wwh)
This serves as the set of active samples for the next round. In addition, for some p 2 whose choice is discussed below, we define the subset Sy of samples with a 'good fit' to the k-th component,
S = {p- d samples in Sş with largest weights w We then refine the estimate of B& by performing OLS this subset,
golwel
(X xst) -X, Us!
(5)
The subsets Sy and S are in general disjoint, unless the threshold Wi, or the oversampling ratio p are too high. The choice for the value of wn, is discussed in Remark 3.1 and Section 6. A suitable value for p depends on the ratio between the sample size n and the information limit nialo d/ min(p). In the challenging setting of n R Tinto P should be set close to 1, confining the algorithm to use the least possible number of samples at each round. If n. > Ninfa, then P can be set to a higher value, allowing the algorithm to estimate each component more accurately and more robustly to noise by using more samples.

---

-oradio MiXXRS AlMIn 60
oracle MMXIRLS AlMMin
failure probability
0.8
failure probability
0.8
0.5
0.4
0.2
000000000000oo
000000000000000000000
2700
4000
6000
8000
10000
12000
x10
Figure 2: Comparison of various MLR algorithms. Depicted the failure rate, out of 50 random initializations, for which Fiatent >> 2x (see (6)), function of the sample size n. The dimension and noise level are fixed at = 300 and o = 10-2, Mixture: K with p = (0.7, 0.2, 0.1) (left pancl); K 5 with (0.63, 0.2, 0.1, 0.05, 0.02) (right panel).
This concludes the main phase of The second (refinement) phase, described in Appendix C, improves the estimates 3phase-1 using also the moderate fit samples that were ignored in the first phase. Yet. empirically in many cases the second phase is not needed, since the main phase already outputs highly accurate estimates of 3 This is in accordance with the thcoretical result Section 6.
Remark 3.1 (Threshold adaptation). If the threshold Win is too small, the number of poor fit samples [S*+1] of (3) may be insufficient estimate the (k + 1)-th component. To avoid the tuning of wa. we use the following scheme: Mix-IRLS gets as input an initial value for Whn, and if needed, it increases its value and starts over; see Lines 11 13 in Algorithm 1. The value 0.1 is heuristic, and a more theoretically grounded update rule is an open question for future research.
Remark 3.2 (Parameter tuning). Mix-IRLS has four input parameters: P. Ti and an initial value for Wi (see Remark 3.1). As empirically demonstrated in Sections 4 and 5, there is no need to carefully tune these parameters; see Appendix F for an explanation of this point. The values used in our experiments are specified in Appendix G.
Remark 3.3 (Unknown/overrstimaaed K). If K is unknown, or only an upper bound Kmax is given, Mix- IRLS automatically find the true K without cross validation. To find it, we ignore the resetting criterion (Lines 11 to 13 in Algorithm 1), and instead run Mix-IRLS until there are too few samples to estimate the next component, namely |S < pd. We then set K to the number of components with at least pd associated samples, and proceed to the second phase.
Simulation Results
We present simulation results synthetic data in this section, and several real-world datasets in the next one. We compare the performance of Mix-IRLS to the following algorithms: (i) AltMin alternating minimization [52, 53]; (ii) EM expectation maximization [5, Chapter 14], [19]; and (iii) GD gradient descent on a factorized objective [56]. We implemented all methods MATLAB. In some of the simulations, additionally ran version of EM for which the mixture proportions p are given as prior knowledge, but it hardly improved its performance and we did not include it in our results. In addition, we plot the performance of an oracle which is provided with the vector c* of the true labels of all n samples and separately computes the OLS solution for each component.
All methods were given the same random initialization, described shortly. Other initializations did not qualitatively change the results. For EM and GD we added an AltMin refinement step at the end of each algorithm to improve their estimates. In principle, Mix-IRLS has several tuning parameters. However, MiX-IRLS performs well with a fixed set of values (see Remark 3.2). Hence, in the following. we show the results for tuning- free variant of our method. Further technical details, including maximal number of iterations, early stopping criteria and parameter tuning of GD appear in Appendix G.
Similar to [56, 21], each simulation we sampled the entries of the explanatory variables X and of the regression vectors 3* from the standard normal distribution (1, 1). In this section, the dimension is fixed at = 300. The additive noise terms are Gaussian, E; ~ (0 a²). Additional simulations with other values of d and including different noise levels for the K components (i.e.,
MATLAB and Python code implementations of Mix IRLS available www.ttt..oooo...eeeeeeeooeo com//piziber//LLB
5

---

-oracle MixiiRL - IEM GD
-caacl M AlMMn C
median error
median error
100
10
10
10.2
0.02
0.04
0.06
0.08
0.1
0.12
0.14
fraction comppted samples
overparametenization (K- K*)
Figure 3: Robustness of several MLR algorithms to outliers (left panel) and to overparameterization (right panel). Values of dl a, K and p are the same in Figure 2(left), and n = 12000. Qualitatively similar results for a balanced mixture appear in Figures H.15 and 1. X-axis on left panel is fraction of outliers; X-axis on right panel is the difference between the number of components K given to the algorithms and the true K"
different y JK). ). well as simulations with different separation levels between the components, appear in Appendix H. As discussed in the introduction, a central motivation for the development of M IRLS is dealing with imbalanced mixtures. Hence, in this section, the labels of were drawn from a multinomial distribution with highly imbalanced proportions. In Appendix H, we present results for balanced and moderately imbalanced mixtures.
We measure the accuracy of an estimate 3 = {31, ........0 3 by the following quantity:
Flater (3: min max [|9ock) - Btll-
(6)
rE(K]tkEl.....
The minimization above is over all K! permutations, which makes the accuracy measure invariant to the order of the regression vectors in 3. A similar objective was used by [53, 56].
All algorithms were initialized with the same random vectors 31.... 5x. whose entries were sampled from the standard normal distribution (0. 1). For each simulation, we performed 50 independent realizations, each with a different random initialization 33, and report the median errors and the failure probability. The latter quantity is defined as the percentage of runs whose error Fiment (6) is above 2n; see an explanation for this choice in Appendix G. Due to space limits, some figures appear in Appendix H.
In the first simulation, we examine the performance of the algorithms as a function of the sample size n. The results are depicted in Figure 2, and the corresponding runtimes in Figure H.3. Mix- is shown to recover the components with sample size very close to the oracle's 's minimum. All competing methods, in contrast, get stuck in bad local minima unless the sample size is much larger. Importantly, this behavior does not follow from the presence of noise, and as shown in Appendix H, the result does not qualitatively change in a noiseless setting. Moreover, as shown in Figure H.8, the nearly optimal performance of Mix-IRLS seems to be independent of the dimension d. It does depend, however, on the mixture proportions: for a moderately imbalanced mixture, Mix-IRLS does not match the oracle performance. Yet, even in this case, MiX-IRLS still markedly outperforms the other methods; see Figure H.6.
Next, we explore the robustness of the algorithms to additive noise, outliers and overestimation. Figure H. shows that all algorithms are stable to additive noise of various levels, but only Mix-IRLS matches the oracle performance in all runs. To study robustness to outliers, in the following simulation we artificially corrupt a fraction f (0, 1) of the observations. A corrupted response ÿ sampled from a normal distribution with zero mean and variance Σ bJ /n. independently of the original value ys. Figure 3(lfft) shows the error (6) of the algorithms as a function of the corruption fraction f. To let the algorithms cope with outliers while keeping the comparison fair, we made the same modification in all of them: at each iteration, the estimate / is calculated based on the [(1 f)n] samples with smallest residuals. MiX-IRLS, we implemented this modification only in the second phase. As shown in the figure, empirically. MiX-IRLS can deal with a corruption fraction of f 9%, which over 4 times more corrupted samples than the other algorithms. In the balanced setting, Mix IRLS can deal with roughly twice as many corrupted samples (ʃ = 17%). which is almost 6 times more outliers than other methods. These results should not be surprising light of the fact that robust regression is at the heart of TRIS''S mechanism.

---

3.5
tuned
100
percieved tome
0-goo08
median error
2.5
o0000
0oo
00000
ognes o
10
00
14
1.6
1,8
2
22
2.4
2.8
28
number of components K
Figure 4: Left panel: Mix- IRLS estimate for the music perception data [13], where d = 1 and
n 150. Right panel: Comparison of several MLR algorithms on the medical insurance dataset. Note that the y-axis is log scale. Median estimation errors are calculated across 50 random initializations according to (8), and error bars correspond to the median absolute deviation.
The last simulation considered the case of an unknown number of components. Specifically, the various algorithms were given as input number of components K K where K is the true number of mixture components in the data. Here, the error is defined similar to (6), but with K instead of K. Namely, the erтоr is calculated based on the best K vectors in ignoring its other
K K- K" vectors. Figure (rightt) shows that most algorithms have similar performance at the correct parameter value Mix- IRLS is insensitive to overparameterizatioo. and succeeds with unbounded K. This feature is attained thanks the sequential nature of Mix IRLS (Remark 3.3). Similar results hold for balanced mixtures; see Appendix H.
K", and succeeds also at small overparameterizzatoo,,
K 3. Only
5 Real-World Datasets
We begin by analyzing the classical music perception dataset of Cohen [13]. In her thesis, Cohen investigated the human perception of tones by using newly available electronic equipment. The n 150 data points acquired in her experiment are shown in Figure 4(left). Cohen discussed two music perception theories: One theory predicted that in this experimental setting, the perceived tone (y-axis) would be fixed at 2.0, while the other theory predicted an identity function (y = x)). The results, depicted in Figure 4(left), support both theories. As mathematical formulation of this finding, Cohen proposed the MLR model (1) with Bernoulli distribution; see also [15].
where the labels e* are i.i.d. according to
In Figure 4(efft,, the untuned version Mix- IRLS shown capture the two linear trends the data. Notably, untuned Mix- IRLS was not given the number of components, but automatically stopped at K with its default parameters. By increasing the sensitivity Mix-IRLS to new components via the parameter Win. it is possible to find three or even more components; see Figure I.19.
Next, we compare the performance of Mix-IRLS to the algorithms listed in the previous section more challenging problem: the CO2 emission by vehicles in Canada dataset, available on Kaggle (see Appendix G). This dataset approximately follows an MLR model, where the samples can be clustered into components according to their fuel type: regular gasoline, premium gasoline, diesel or ethanol (K = 4). The mixture is imbalanced, with the proportions p (0.49, 0.43, 0.05, 0.03). We ask the following question: given the engine features (engine size, number of cylinders, fuel consumption in city and highway), how close are the methods' clusters to the (ground-true -tuth) fuel type clusters?
We measure the performance of the methods with class-wise balanced accuracy, defined as the average of sensitivity/recall and specificity. We chose this measure it is considered suitable for imbalanced mixtures, but other measures yield qualitatively similar results. As clustering is defined only up to permutation, for each method choose the permutation that maximizes the overall intersection of its estimated clusters with the true ones, defined as
1 i(i e :()) ct
max re <<[K] n
The results, reported Table 1, demonstrate the effectiveness of MiX-IRLS on imbalanced mixtures. While the dominant class (regular gasoline) all methods achieve comparable balanced accuracy,
7

---

Table Results of MLR algorithms on the CO2 emission from vehicles dataset. Reported are median balanced accuracy and median absolute deviation based on 50 random initializations.
Balanced accuracy
Mix-IRLS
AltMin
EM
GD
regular gasoline (49%)
0.58 + 0.00
0.59 + 0.03
0.56 = 0.04
0.57 = 0.07
premium gasoline (43%)
0.59 - 0.00
0.48 = 0.10
0.48 = 0.01
0.40 = 0.03
diesel (5%)
0.89 + 0.00
0.78 + 0.13
0.54 + 0.07
0.70 + 0.12
ethanol (3%)
0.74 + 0.00 0.63 0.07 0.77 0.08
0.61 + 0.05
Mix IRLS significantly outperforms the other methods on the remaining classes. EM is the only method that is slightly better than MiX-IRLS on one of the rare classes, but achieves markedly worse balanced accuracy on the other two rare classes. Moreover, importantly, Mix-IRLS is the only method that is stable with respect to the initialization, demonstrating negligible variance across different random initializations.
Finally, we compare the performance of the algorithms on four of the most popular benchmark datasets for linear regression, all of which are available Kaggle (see Appendix G): medical insurance cost, red wine quality. World Health Organization (WHO) life expectancy, and fish market. The task in each dataset is to predict, respectively: medical insurance cost from demographic details; wine quality from its psychochemical properties; life expectancy from demographic and medical details; and fish weight from its dimensions. For these datasets, MLR is at best an approximate model, and its regression vectors 8 are unknown. Hence, we replace (6) by the following quality measure which represents the fit of an MLR model the data:
1
1 n
Frea (B; (3X, X,y)
min (a. By yx) JELK]
(8)
Varly/ [y]
resembling the K-means objective for clustering [23]. To illustrate the advantage of multi-component model, we also report the results of an ordinary least squares (OLS) solution. In this section, besides vanilla Mix-IRLS we report the results of tuned variant, denoted Mix-IRLS: tuned. The tuning is done using only the observed objective (8), and no oracle is involved in the procedure.
As the number of components K is unknown, we explore the algorithms' performance given different values of K, ranging from 2 to 7. The upper limit 7 was chosen arbitrarily, and the results do not change qualitatively for larger values. All algorithms start from the same random initialization. Additional details appear in Appendix G.
Figure 4(right) shows the performance of the algorithms on the medical insurance dataset in terms of the median error (8) across 50 different random initializations. The results for the other three datasets, as well as the minimal errors achieved across the realizations, are deferred to Appendix I. In general, both MiX-IRLS and Mix- IRIXIRRSS:: : tuned improve upon the other methods, sometimes by 30% or more.
The experiments on the Kaggle datasets demonstrate that empirically Mix-IRLS finds MLR coef- ficients that fit the data significantly better than other methods. Hence, given as input data that (upproximately) does follow an MLR model for some K, our method would often compute accurate estimates of its corresponding mixture coefficients where other methods would fail to do so.
6 Recovery Guarantee for Mix-IRLS
In this section, we theoretically analyze Hix-IRLS in a population setting with an infinite number of samples. For simplicity, we assume the explanatory variables are normally distributed. Without loss of generality, we assume zero mean and identity covariance.
34 (0, Ta).
In the more general case where x; ~ (n, Σ) for some μl and Σ, one may find μ and Σ from the data, and mean center and whiten the samples Tú - E-112 11), before applying our method. The responses yi are assumed to follow model (1) with K 2 components (8i, (1) and labels cj
8

---

generated independently of 2; according to mixture proportions P1 2 P2 > O. This setting was considered in several previous theoretical works that analyzed the method for MLR [2, 14, 33, 36].
We assume the noise terms e are all i.i.d., zero-mean random variables, independent of x; and c. We further assume they are bounded and follow symmetric distribution,
1 S and PlG: < t| = P(q 2 -t] Vt.
(10)
For analysis purposes, we consider a slightly modified variant of
formally described
in Appendix E. In this variant, MiX-IRLS excludes samples with large magnitude, ||:;112 > R where R a fixed parameter. A natural choice for its value is [B[[ = d, e.g. R = 2d. In nigh-dimensional setting with d >> 1, such choice excludes proportion of the samples that is exponentially small in d. For simplicity, we present our result assuming R is large, corresponding to large dimension d; the general result appears in Lemma E.I.
The following theorem states that given a sufficiently imbalanced mixture, namely Pr is large enough, Mix- IRLS successfully recovers the underlying vectors 3i and $1.
Theorem 6.1. Let {(x;: 34) }
100 be i..dd. from a mixture of K = 2 components with proportions
(p1, P2), regression vectors (5i, (2), and noise terms e, that follow (10) with o.. Denote A 3i -32. -22
and y = 5p2/4. Suppose that
9 = g (1+
de |
(11)
/R/
2
Further assume that Niw-IRLS is run with parameters p 0o, and with n, wa that satisfy,
< WPA <
(12)
+ n11 - 9|
1 m A
Then for sufficiently large R, starting from an arbitrary initialization, the first phase of -RLS with at least one iteration (T1 2 1) recovers 8r up to an error that decreases with increasing R,
1
de
max || -821 S 1.2
(13)
VR + *|
Specifically, in the absence of noise (o. = 0), - perfectly recovers the regression vectors.
Theorem 6. considers only the first phase IRLS, as it is sufficient to recover the regression vectors in the described setting. Indeed, empirically, the second phase is often unnecessary. The choice of an oversampling ratio p = oc is suited to the population setting where n = co; see the discussion after Eq. (5). The theorem proof appears in Appendix E.
Remark 6.2 (Different noise levels). Theorem 6. holds also in the case of differrent noise levels T1.02 for each of the two mixture components. In this case, a. is replaced by max(o 12) Indeed, as illustrated empirically in in Figure H.11, Mix-IRLS is able handle such case well. Remark 6.3 (Required imbalance). Due to (11), Theorem 6. holds only for sufficiently imbalanced mixture. In the absence of noise, the factor //4 in the definition of y may be relaxed to any number arbitrarily close to onc, at the expense of increasing R. turn, this implies that our theorem holds for any mixture with the probability of the second component allowed to be arbitrarily close to P2 < 1/2. In simple words, in the noiseless setting our guarantee holds as long as the mixture is not perfectly balanced. More generally, in the presence of noise, there is trade-off between mixture imbalance and noise level: the proportion of the dominant component P1 increases with the noise level d.- This theoretical finding is in agreement with the simulation results in Section 4: Mix- works better as the imbalance increases. We emphasize that empirically, Mix-IRLS works well also on balanced mixtures, P1 = P2 1/2; see Appendix H. Hence is an open problem derive theoretical guarantees for a perfectly balanced mixture.
Remark 6.4 (Allowed range for Wa.). Theorem 6.1 holds for a limited range of values for Wi. see (12). This range depends on q. which in turn depends on the noise level and the mixture imbalance. For example, at noise level a. = 10-2 proportions
(3/4, 1/4), separation of ||| = 1, parameter choice of 1 = and 2 2. the range for which recovery is guaranteed is 0.69 S whh S 0.9. As mentioned previously, empirically we run Mix-IRLS with an initial value of Win 0.1 which is increased by the algorithm (see lines 11-13 in Algorithm 1). Indeed, accordance to this analysis, in many cases the final value is within this range.

---

Remark 6.5 (Overparameterization / unknown K). In practical scenarios, the number of components K is often unknown. Remarkably, Theorem 6.1 can be extended to an overparameterized setting, where the true number of components K* is 2 but Mix-IRLS is given an overestimated number K' > 2, together with a corresponding (arbitrary) initialization (31,...,, This is explicitly discussed in Appendix D (Proposition D.1), and also demonstrated empirically in Figure 3(right).
Theorem 6.1 and Remark 6.5 are in accordance with several empirical findings from previous sections: MiX-IRLS performs better on imbalanced mixtures than on balanced ones; it copes well with an overparameterized K; and it works well starting from a random initialization. Our analysis (Appendix E) sheds light on the inner mechanism of Mix-IRLS that enables these features.
Comparison to prior work. Several works derived MLR recovery guarantees for AltMin [52, 21] and GD [56. 38] in noiseless setting. More related our Theorem 6.1 are works that studied the population EM algorithm in the presence of noise [2, 14, 33, 36, 34). These works assumed a perfectly balanced mixture of two components, P1 = P2 1/2. An exception is [34], who allowed for K' > and an imbalanced mixture. However, their allowed imbalance is limited. In addition, they required a sufficiently accurate initialization. A key novelty in our result is not only that we allow for highly imbalanced mixtures, but that large imbalance actually makes recovery easier for Mix-IRLS, since the quantity q of (11) monotonically decreasing with the mixture imbalance, the allowed range (12) of the parameter wn increases with the imbalance. Furthermore, our result holds for arbitrary initialization. The downside is that Theorem 6. requires sufficient imbalance (see Remark 6.3), and does not provide recovery guarantee for our method on a perfectly balanced mixture, even though empirically, our method works well also in this case. Our result is novel in another aspect as well. In contrast to most existing guarantees, Theorem 6.1 holds also for an arbitrary input number of components K; see Remark 6.5.
recovers the first component by treating the samples from the second component as outliers. In the noiseless setting, our guarantee allows the second component to consist almost 1/2 of the data; see Remark 6.3. For comparison, in the context of robust regression, [40] recently analyzed an IRLS method in noiseless setting, and allowed less than 1/5 corrupted samples. Our higher tolerance is possible thanks to the strong structural assumption of MLR (1).
7 Summary and Discussion
In this work, we presented a novel method to solve MLR, Mix- IRLS, that handles both imbalanced and balanced mixtures. Mix- MiX-IRLS robust to outliers and an overestimated number of components. The latter feature leads to another contribution of our method: under suitable conditions, it can be run with an overestimated K K* and will automatically find the true number of components K*
The basic idea of Mix-IRLS- sequential recovery using tools from robust regression was also employed by [3]. Several important differences between their method and Mix-IRLS were listed in the introduction; first and foremost is the fact that [3] is limited to low dimensions. It is interesting to note that [1] also made connection between MLR and robust regression, but the other way around: as a simplified theoretical model, they assumed the outliers follow a linear model, and applied the EM algorithm to the obtained MLR problem to detect them. Also relevant to our work are recent papers on subspace clustering [48, 18), list-decodable learning [9, 31] and real phase retrieval [42]. We discuss their relation to our work, additional relevant papers and future research directions, in Appendix A.
Our current theoretical analysis is limited to the case of only K 2 components and a population setting with infinitely many samples. While the first assumption is very common in the literature (c.g.. [2, 36), and many others), population analysis is usually followed by a finite-sample one. In our case, a main challenge in such an analysis revolves around the sequential nature of Mix-IRLS. Once Mix- IRLS estimates component, it removes its associated samples. Analyzing modified sample set introduces complex statistical dependencies the data no longer follow normal distribution but rather a conditional distribution. In the population setting. various quantities happen to cancel out nicely; see, for example, the derivation of (E. 17) from (E.16). Under the setting of finite sample size, we need to bound these quantities under a conditional distribution. We hope to be able to overcome these challenges and extend our analysis to finite-sample setting in the future. Finally, another interesting future research direction is to formally prove Mix M 's robustness to outliers, as was empirically demonstrated in Section 4.
10

---

Acknowledgements
The research of P.Z. was partially supported by fellowship for data science from the Isracli Council for Higher Education (CHE). B.N. is the incumbent of the William Petschek Professorial Chair of Mathematics. The research of B.N. was supported part by grant 2362/22 from the Israel Science Foundation. We thank Yuval Kluger and Ofir Lindenbaum for interesting discussions. We thank authors of [56] for sharing their code with us. We also thank the anonymous reviewers for their valuable feedback that improved the quality of our manuscript.
References
[1] Murray Aitkin and Granville Tunnicliffe Wilson. Mixture models, outliers, and the EM algorithm. Technometrics, 22(3):325-3311 1980.
[2] Sivaraman Balakrishnan, Martin J Wainwright, and Bin Yu. Statistical guarantees for the EM algorithm: From population to sample-based analysis. The Annals of Statistics, 45(1):77-120, 2017.
[3] David L Banks, Leanna House, and Kevin Killourhy. Cherry-picking for complex data: robust structure discovery. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 367(1906):4339 4359, 2009.
[4] Babak Barazandeh, Ali Ghafelebashi, Meisam Razaviyayn, and Ram Sriharsha. Efficient algorithms for estimating the parameters of mixed linear regression models. arXiv preprint arXiv:2105.05953. 2021.
[5] Christopher M Bishop. Pattern recognition and machine learning. Springer, 2006.
[6] Arun Tejasvi Chaganty and Percy Liang. Spectral experts for estimating mixtures of linear regressions. In International Conference on Machine Learning, pages 1040-1048. PMLR, 2013.
[7] Xiujuan Chai, Shiguang Shan, Xilin Chen, and Wen Gao. Locally linear regression for pose- invariant face recognition. IEEE Transactions on image processing, 16(7):1716-172,, 2007.
[8] Wennan Chang, Changlin Wan, Yong Zang, Chi Zhang, and Sha Cao. Supervised clustering of high-dimensional data using regularized mixture modeling. Briefings in bioinformatics, 2bbaabbba, 2021.
[9] Moses Charikar, Jacob Steinhardt, and Gregory Valiant. Learning from untrusted data. In Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, pages 47-60, 2017.
[10] Samprit Chatterjee and Martin Mächler. Robust regression: A weighted least squares approach. Communications in Statistics-Theor and Methods, 26(6): 311-1391494 1997.
[11] Sitan Chen, Jerry Li, and Zhao Song. Learning mixtures of linear regressions in subexponential time via fourier moments. In Proceedings of the 52nd Annual ACM SIGACT Symposium Theory of Computing, pages 587-600, 2020.
[12] Yudong Chen, Xinyang Yi, and Constantine Caramanis. A convex formulation for mixed regression with components: Minimax optimal rates. In Conference Learning Theory, pages 560-604, PMLR, 2014.
[13] Elizabeth Ann Cohen. The influence of nonharmonic partials on tone perception. Stanford University, 1980.
[14] Constantinos Daskalakis, Christos Tzamos, and Manolis Zampetakis. Ten steps of EM suffice for mixtures of two gaussians. In Conference Learning Theory, pages 704-710. PMLR, 2017.
[15] Richard De Veaux. Mixtures of linear regressions. Computational Statistics & Data Analysis, 8(3):227-245, 1989.

---

[16] Partha Deb and Ann M Holmes. Estimates of use and costs of behavioural health care: comparison of standard and finite mixture models. Health economics, 9()):455-489 2000,
[17] Theo Diamandis, Yonina Eldar, Alireza Fallah, Farzan Farnia, and Asuman Ozdaglar. A Wasserstein minimax framework for mixed linear regression. In International Conference on Machine Learning, pages 2697-2706. PMLR, 2021.
[18] Tianyu Ding, Zhihui Zhu, Manolis Tsakiris, Rene Vidal, and Daniel Robinson. Dual principal component pursuit for learning a union of hyperplanes: Theory and algorithms. In International Conference on Artificial Intelligence and Statistics, pages 2944-2952 PMIR, 2021.
[19] Susana Faria and Gilda Soromenho. Fitting mixtures of linear regressions. Journal of Statistical Computation and Simulation, 80(2):201-22) 22000-22,,,220 2010.
[20] Scott Gaffney and Padhraic Smyth. Trajectory clustering with mixtures regression models. In Proceedings of the fifth ACM SIGKDD international conference on Knowledge discovery and data mining, pages 63-72, 1999.
[21] Avishek Ghosh and Ramchandran Kannan. Alternating minimization converges super-linearly for mixed linear regression. International Conference on Artificial Intelligence and Statistics, pages 1093-1 103. PMLR, 2020.
[22] Paul Hand and Babhru Joshi. A convex program for mixed linear regression with recovery guarantee for well-separated data. Information and Inference: A Journal of the IMA, 7(3):563- 579, 2018.
[23] Trevor Hastie, Robert Tibshirani, Jerome H Friedman, and Jerome H Friedman. The elements of statistical learning: data mining, inference, and prediction, volume Springer, 2009.
[24] Paul W Holland and Roy Welsch. Robust regression using iteratively reweighted least-squares. Communications in Statistics Theory and Methods, 6(9):813 12-877,,,,, 1977.
[25] Hao Hu, Weixin Yao, and Yichao Wu, The robust EM-type algorithms for log-concave mixtures of regression models. Computational statistics & data analysis, 111:14-26, 2017.
[26] Mian Huang and Weixin Yao. Mixture of regression models with varying mixing proportions: semiparametric approach. Journal of the American Statistical Association, 107(498):711-727
2012.
[27] Peter J. Huber. Robust statistics. John Wiley and Sons, New York, 1981.
[28] David Hunter and Derek S Young. Semiparametric mixtures of regressions. Journal of Nonparametric Statistics, 24(1):19-385 2012.
[29] Salvatore Ingrassia, Simona C Minotti, and Antonio Punzo. Model-based clustering via linear fluster-weighted models. Computational Statistics & Data Analysis, 71:159-182, 2014.
[30] Hansheng Jiang and Adityanand Guntuboyina. A nonparametric maximum likelihood approach to mixture of regression. arXiv preprint 0rXiv:2108.0981 2021.
[31] Sushrut Karmalkar, Adam Klivans, and Pravesh Kothari. List-decodable linear regression. Advances neural information processing systems, 32, 2019.
[32] Abbas Khalili and Jiahua Chen. Variable selection in finite mixture of regression models. Journal of the american Statistical association, 1025 2007.
[33] Jason M Klusowski, Dana Yang. and WD Brinda. Estimating the coefficients of a mixture of two linear regressions by expectation maximization. IEEE Transactions Information Theory, (5(6):3515-3524 2019.
[34] Jeongyeol Kwon and Constantine Caramanis. EM converges for a mixture of many linear regressions. In International Conference on Artificial Intelligence and Statistics, pages 1727- 1736. PMLR, 2020.
12

---

[35] Jeongyeol Kwon, Nhat Ho, and Constantine Caramanis. On the minimax optimality of the EM algorithm for learning two-component mixed linear regression. In International Conference on Artificial Intelligence and Statistics, pages 1405-1413, PMLR, 2021.
[36] Jeongyeol Kwon, Wei Qian, Constantine Caramanis, Yudong Chen, and Damek Davis. Global convergence of the EM algorithm for mixtures of two component linear regression. In Confer- ence on Learning Theory, pages 2055-21 10. PMLR, 2019.
[37] Qianyun Li, Runmin Shi, and Faming Liang. Drug sensitivity prediction with high-dimensional mixture regression. PloS one, 14(2):e0212108, 2019.
[38] Yuanzhi Li and Yingyu Liang. Learning mixtures of linear regressions with nearly optimal complexity. In Conference On Learning Theory, pages 1125-1144. PMLR, 2018.
[39] Geoffrey J McLachlan, Sharon X Lee, and Suren I Rathnayake, Finite mixture models. Annual review of statistics and its application, 6:355-378, 2019.
[40] Bhaskar Mukhoty, Govind Gopakumar, Prateek Jain, and Purushottam Kar. Clobally-convergent iteratively reweighted least squares for robust regression problems. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 313-322. PMLR, 2019.
[41] Soumyabrata Pal, Arya Mazumdar, Rajat Sen, and Avishek Ghosh. On learning mixture of linear regressions in the non-rcalizable setting. In International Conference on Machine Learning. pages 17202-17220. PMLR, 2022.
[42] Liangzu Peng, Christian Kümmerle, and René Vidal. Global linear and local superlinear convergence of IRLS for non-smooth robust regression. Advances in neural information processing systems, 35, 2022.
[43] Daniel Pimentel-Alarcón, Laura Balzano, Roummel Marcia, Robert Nowak, and Rebecca Willett. Mixture regression as subspace clustering. In 2017 International Conference on Sampling Theory and Applications (SampTA), pages 456-459. IEEE, 2017.
[44] Hanie Sedghi, Majid Janzamin, and Anima Anandkumar. Provable tensor methods for learning mixtures of generalized linear models. In Artificial Intelligence and Statistics, pages 1223-1231. PMLR, 2016.
[45] Yanyao Shen and Sujay Sanghavi. Iterative least trimmed squares for mixed linear regression. Advances in Neural Information Processing Systems, 32, 2019.
[46] Weixing Song, Weixin Yao, and Yanru Xing. Robust mixture regression model fitting by laplace distribution. Computational Statistics & Data Analysis, 71:128-137, 2014. 201,
[47] Yifan Sun, Ziye Luo, and Xinyan Fan. Robust structured heterogeneity analysis approach for nigh-dimensional data. Statistics in Medicine, 2022.
[48] Manolis Tsakiris and René Vidal. Hyperplane clustering via dual principal component pursuit. In International conference on machine learning, pages 1482-3481. PMLR, 2017.
[49] Michel Wedel and Wagner A Kamakura. Market segmentation: Conceptual and methodological foundations. Springer Science & Business Media, 2000.
[50] Rand Wilcox. Introduction to Robust Estimation and Hypothesis Testing. Elsevier, 2011.
[51] Weixin Yao, Yan Wei, and Chun Yu. Robust mixture regression using the t-distribution. Computational Statistics & Data Analysis, 71:116-127, 2014.
[52] Xinyang Yi, Constantine Caramanis, and Sujay Sanghavi. Alternating minimization for mixed linear regression. In International Conference on Machine Learning, pages 613-621. PMLR. 2014.
[53] Xinyang Yi, Constantine Caramanis, and Sujay Sanghavi. Solving mixture of many ran- dom linear equations by tensor decomposition and alternating minimization. arXiv preprint 1rX1v:1608.05749, 2016.
13

---

[54] Derek S Young and David R Hunter. Mixtures of regressions with predictor-dependent mixing proportions. Computational Statistics & Data Analysis, 54(10):2253-22 2266, 2010.
[55] Linjun Zhang, Rong Ma, Tony Cai, and Hongzhe Li. Estimation, confidence intervals, and large-scale hypotheses testing for high-dimensional mixed linear regression. arXiv preprint arXiv:2011.03598, 2020.
[56] Kai Zhong, Prateek Jain, and Inderjit S Dhillon. Mixed linear regression with multiple compo- nents. Advances in neural information processing systems, 29, 2016.
14


Reviews:
[{"review_id": "XVl2InLgAM", "reviewer": "Reviewer_LCfo", "summary": "This paper proposes a new method, Mix-IRLS, for mixed linear regression, in which data are drawn from a mixure of linear models whose coefficients we wish to estimate. The method is especially motivated by the failure of existing methods in the case of imbalanced mixtures (i.e., when some components are much rarer than others). After presenting the new method, the paper presents experimental results on simulated and real data, illustrating the improved robustness of Mix-IRLS over existing methods. Finally, the paper proves that, in a particular two-component, infinite-sample example, a slight variant of Mix-IRLS is able to recover both of the true regression coefficients (even without knowing the number of components a priori).", "questions": "1) In addition to each component of c^* being sufficiently large, I think identifiability also requires all of the \\beta's to be distinct (consistent with the 1/||\\Delta|| terms in Eq. (10)). Perhaps this should also be mentioned in the context of identifiability, around Line 94.\n\n2) Lines 97-99, \"the minimal number of observations n required to recover $\\beta^*$ in the absence of noise is $n_{info} = d/min(p)$.\" What is the justification for this sentence, and what assumptions does it require? Is it somehow obvious? Can the authors prove it or provide a reference?\n\n3) Lines 101-102, \"Even in the simplified setting... the problem is NP-hard without further assumptions.\" How is the present work able to get around this? Does it make further assumptions? (If so, which ones?) Does it guarantee only approximate recovery? Or something else?\n\n4) I was confused by the corruption added in the second simulation (corresponding to the left panel of Figure 3). It seems to me that the corrupted samples \\tilde{y}_i are a special case of an MLS component, with \\beta = 0 and a different noise scale. As far as I understand, Mix-IRLS does not know noise level \\sigma, so it cannot distinguish which is the \"real\" data and which is the \"corruption\". So, it seems like Mix-IRLS is arbitrarily selecting one of the two possible MLS models, and is ignoring the other. Is this task well-specified?\n\n5) What is the goal of applying Mix-IRLS to the Kaggle regression benchmark datasets (health insurance, red wine, WHO life expectancy, and fish market)? As far as I could tell, the purpose of Mix-IRLS is parameter (\\beta) estimation, but it's not clear to me what parameters are being estimated in these examples, or how these results could be interpreted.\n\n6) In the Problem Setup (Section 2), the distribution of explanatory variables Xs is not discussed. This is fine, since a precise statistical model is not needed (except in Theorem 6.1), but I feel like I am lacking some intuition for the problem because of this. Sometimes (e.g., Figure 1), their distribution seems to differ for each MLR component, whereas elsewhere (e.g., Theorem 6.1) their distribution is the same. Intuitively, the former (different distributions) setting seems potentially much easier than the latter (same distributions) setting, since the components will be more separated. Is there a particular way the authors suggest thinking about the distribution of Xs?\nI also have a similar question regarding the distributions of \\epsilons -- is it important that the noise distributions be similar across components?", "limitations": "Section 7 discusses some limitations of their theoretical results, but perhaps this could be more extensive (see point 2 under \"Weaknesses\" above).", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The MLR problem is clearly presented and the proposed Mix-IRLS approach is well-motivated. The simulations clearly illustrate the basic strengths of Mix-IRLS (robustness to imbalance, contamination, and unknown K) over existing methods. The application to the music perception dataset of Cohen [13] is compelling and well-explained.", "weaknesses": "1) The introduction doesn't mention any applications motivating the MLR problem. The only real-world example is the music perception dataset of Cohen [13] in Section 5. I would find the paper more compelling if it were clear to the reader why this problem is useful to study. It would be especially compelling if there were real-world examples where class imbalance was a limiting issue for existing methods (the music example appears fairly balanced, although it is hard to tell from the plot).\n\n2) Some of the assumptions in Theorem 6.1 seem arbitrary and unnecessary. Specifically, the assumption X ~ N(0, I_d) is very strong, and the resulting in the theorem doesn't illustrate how the distribution of X impacts the recovery guarantee. Based on existing theory for linear regression, I would expect a much milder assumption, such as a lower bound on the smallest eigenvalue of the covariance matrix of X, to suffice. Also, as far as I can tell, excluding X_i with ||X_i||^2 < R is done only to simplify the analysis and shouldn't really affect the recovery guarantee (indeed, as far as I can tell, the guarantee in Theorem 6.1 only improves as R -> infinity). These assumptions make Theorem 6.1 feel more like an \"example\" than a \"theorem\", and encourage the authors to work towards a more general formulation of Theorem 6.1.\n\nOverall, on initial reading, I lean weakly towards rejection because the real-world significance of the work is not very clear and the theoretical results seem unpolished. But I could change my score if either of these points is improved substantially.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments.\n\n**Weaknesses:**\n1. Thanks for bringing this up. We performed an additional experiment based on a Kaggle CO2 emission from vehicles dataset to address the reviewer's concern. This dataset approximately follows an MLR model, where the samples can be clustered into components according to their fuel type. We compared the performance of the MLR methods given the original mixture, which is imbalanced, and given a subsampled set of observations such that the resulting mixture is balanced. The results imply that the imbalance is, to some extent, a limiting issue for existing methods but not for MixIRLS. Kindly see item 1 in our global response for more details on this experiment. We will include these results and discuss them in the revised manuscript.\n2. We agree with the reviewer that generalizing Theorem 6.1 would greatly improve our work. In fact, since we assume a population setting, we can extend our theorem to the case $X \\sim N(\\mu, \\Sigma)$ without modifying our proof by removing the mean and whitening the samples, $x_i \\to \\Sigma^{-1/2} (x_i - \\mu)$. We thank the reviewer for pointing us in this direction. As a side note, we would like to remark that most works that derived theoretical guarantees for MLR in the finite-sample setting - even for extensively studied methods such as EM - did have the strong assumption $X \\sim N(0, I_d)$. These include [2, 21, 33, 34, 35, 36, 52, 53, 56] and other works.\nThe parameter $R$ is indeed an artifact of our proof. It helps us deal with challenges that are unique to the theoretical analysis of our (sequential) approach; for more details on these challenges, please see item 4 in our global response.\n\n**Questions:**\n1. The reviewer is correct in that separability is required to identify the components. However, separability is unnecessary if our goal is only estimating the $\\beta$ vectors rather than the labels $c_i^*$. Consider, for example, the extreme case of zero separation. Then the mixture is degenerate and the problem is identical to standard linear regression. We will add this discussion to the revised manuscript. Having said that, our current manuscript does lack an analysis of the algorithms' performance with respect to different levels of separability. We ran a suitable set of simulations and described the result in item 2 in our global response.\n2. Kindly see item 3 in our global response.\n3. Yes, we make an additional assumption on the distribution of the explanatory variables: $X_i$ are independent and follow the normal distribution. To the best of our knowledge, this assumption is common to most works on MLR with theoretical guarantees. We will clarify this point in the revised manuscript.\n4. This is an interesting question. Formally, the corrupted samples can be viewed as another MLR component with $\\beta=0$. However, the signal-to-noise ratio in this case is infinitely large. To answer the reviewer's question regarding the behavior of MixIRLS, we need to split the discussion into two settings: known $K$ and unknown $K$. If $K$ is known, treating the corrupted samples as another component would significantly increase the MLR objective $F_\\text{latent}$ (Eq. (6)) due to the large noise of this component. Hence, MixIRLS finds the true $K$ components, and the corrupted samples get low weights and are thus treated as outliers. If $K$ is unknown, the behavior of MixIRLS is the same in the first $K$ components; the key question is whether MixIRLS would incorrectly 'view' the corrupted samples as a $(K+1)$-th component. Since we provide the algorithms with the corruption fraction, the answer is that after finding $K$ components MixIRLS 'knows' it already used all the inliers in data, and does not look for new components. Hence, the corrupted samples will not be treated a $(K+1)$-th component.\n5. Thanks for this question. Kindly see item 1 in our global response.\n6. Thanks for highlighting this unclarity. The reviewer touches on an important question regarding a core MLR setting. In the so-called \"model-based clustering\" variant of the MLR problem, the distribution of the explanatory variables depends on the component (label), $x_i = x_i(c_i^*)$. In our work, we consider the vanilla \"model-free clustering\" setting, where the distribution of the explanatory variables is identical for all components. As noted by the reviewer, this setting is in general more challenging than the former. We apologize for the misleading visualization in Fig. 1, which will be fixed in the revised manuscript.\nAs for the noise distribution, it needs not be similar across the components. On the empirical front, we show in Fig. F.9 that the algorithms (not only ours) are not much affected by variability in the noise levels. On the theoretical front, our theorem also holds in this case; the bound in Eq. (9) would then refer to the maximal noise level. We will clarify this point in the revised manuscript."}, {"review_id": "eUMNmqJggU", "reviewer": "Reviewer_1Ro7", "summary": "In this paper the authors study the problem of mixed linear regression, where instead of a single linear regression model, the response variable comes from a finite mixture of them. The authors propose an algorithm for estimating the regression coefficients, especially in the imbalanced case where the mixing proportions can be very different from each other. The algorithm is efficient and scalable to high dimensions, does not require knowledge of the number of components, and is robust to noise and outliers. The algorithm is tested on both synthetic and real data sets, outperforming several existing methods. Theoretical aspects of the algorithm are also analyzed on the population level in a two-component model. \n", "questions": "1. If I understand correctly, the noise distribution is assumed to be the same across all the components. Does the algorithm apply in the case when the noise distribution is different for each component? Or is there any major modification needed?\n\n2. Let's say we have a two-component model with perfectly balanced weights (p1=p2=1/2). In an intuitive sense, how would the algorithm proceed in this case? I guess the initial step would end up with something like a line right in the middle of the two components, and wonder how the algorithm would correct for this. \n", "limitations": "Theoretical analysis of the algorithm could be strengthened, but the empirical success is already very compelling. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a conceptually simple algorithm. Although the idea of sequentially fitting each component is not new, the incorporation of robust statistics leads to an efficient and scalable algorithm that outperforms several other standard algorithms for mixed linear regression on both synthetic and real data. \n\n2. The authors show that the algorithm does not require the knowledge of the number of components, which is a very favorable property for learning mixture models, along with robustness to noise and outliers. \n\n3. The algorithm handles sufficiently imbalanced mixtures, extending existing literature which is mostly for balanced or sufficiently balanced cases. It also gives comparable and more robust estimation even in the perfectly balanced case. \n\n4. Although the algorithm requires several hyper parameters, the authors show that tuning of the parameters does not lead to major improvement so the algorithm is essentially tuning free.  \n\n5. The theoretical results give a justification of the algorithm on the population level, especially in the sufficiently imbalanced case.  ", "weaknesses": "1. The design of the algorithm seems to rely heavily on the imbalanced mixing proportion assumption to be able to filter out each component, yet it gives competitive or even improved performance in the perfectly balanced case. Some intuition or explanation would be helpful since the balanced case is not covered by the theoretical results.  \n\n2. The theoretical results investigate only the infinite data setting. The paper could be strengthened with a finite sample analysis. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments.\n\n**Weaknesses:**\n1. We thank the reviewer for bringing this unclear point to our attention. It might not be so well-known, but in general, robust regression algorithms succeed even when the number of outliers largely exceeds the number of inliers. The well-known 'breakdown point' at 50% allows for adversarial outliers: consider, for example, a case where all the outliers follow the same linear model. Then the problem is clearly ill-posed. However, the 'outliers' in our case, namely the samples that belong to a yet undiscovered component, are not adversarial, but rather follow an MLR model. In such cases, standard high-dimensional robust regression methods can recover the underlying linear model even with 70% and 80% outliers. This explains the success of MixIRLS in balanced mixtures from an empirical perspective. From a theoretical perspective, we would like to point out that our theorem does include a guarantee for a balanced mixture, provided that the noise is small enough. In the special case of noise-free observations, our theorem holds for all mixture proportions except for the perfectly balanced one. We will elaborate on this point in the revised version of our manuscript. Additionally, we will support our claim regarding the recovery abilities of non-adversarial robust regression with simulations.\n2. Unfortunately, we do not currently have a finite-sample analysis for our algorithm, and we plan to accomplish it in future research. Kindly see item 4 in our global response for more details on the challenges in finite-sample analysis.\n\n**Questions:**\n1. Yes, the algorithm applies also in the case of noise distribution which is different for each component, and no modification is needed. On the empirical front: please see Fig. F.9 in Appendix F. It is shown that the algorithms (not only ours) are not much affected by the variability of the noise levels. On the theoretical front: our theorem also holds in this case; the bound in Eq. (9) would then refer to the maximal noise level. We will clarify this point in the revised manuscript.\n2. This is an interesting question. The algorithm would succeed also in this case. The reason is that even if the mixture is perfectly balanced with $p_1=p_2=1/2$, the different sample magnitudes $||x_i||$ and noise terms $|\\epsilon_i|$ make the problem asymmetric. As such, MixIRLS will gradually tend towards one of the components, and eventually recover it. Only in the population setting perfect balance is a symmetric pathology in which MixIRLS enjoys no guarantee. We will include this discussion in the revised manuscript."}, {"review_id": "7o73EHnjs9", "reviewer": "Reviewer_2aVS", "summary": "In this paper, the authors take into account the problem of recovering a mixture of $K$ unknown linear regressions with imbalanced proportions, i.e. the numbers of samples for components in the mixture are substantially different. Rather than try to recover $K$ components at the same time as in previous work, they propose an algorithm named Mix-IRLS which finds these components sequentially using robust regression. This algorithm is theoretically shown to be able to recover the components with small sample size regardless of information about the true number of components $K$. Additionally, it is also demonstrated to be robust to noise and outliers. Finally, the authors empirically polish the performance of the Mix-IRLS algorithm on imbalanced mixtures by conducting simulation studies on both synthetic and real-word datasets. ", "questions": "1. In line 124, should the number of iterations $T_1$ vary with $k$? In particular, the number of samples of interest decreases after each round, which possibly affects the convergence of parameters.\n\n2. In line 135, I think one of the conditions for S'k and $S_{k+1}$ not being disjoint is that the threshold $\\omega_{th}$ is too high, not too low as claimed in the paper because the higher the $\\omega_{th}$ is the more elements the set $S_{k+1}$ has.\n\n3. Are there any theoretical guarantee that by following the procedures to find the true number of components $K*$, the output is the desired value?\n\n4. What are the main challenges of showing Theorem 6.1 under the settings with the finite number of samples? ", "limitations": "The authors already discussed the limitations of the proposed method in this work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Originality: This is the first work to tackle the imbalance in the mixed linear regression. It provides a novel viewpoint of recovering mixture components sequentially rather than simultaneously as in the literature. \n\n2. Quality: Results in this paper are empirically supported and some of them are theoretically grounded. The proposed method is technically sound and applicable for high-dimensional settings. \n\n3. Clarity: The paper is well-written and easy to follow. Assumptions are discussed clearly. Proofs are presented neatly. \n\n4. Significance: This work is important as it is the pioneer in solving the imbalanced problem of mixed linear regression.\n\n", "weaknesses": "1. The population setting considered in Theorem 6.1 makes this result not practical as we do not have access to an infinite number of samples in practice.\n\n2. Theoretical guarantee for the robustness to outliers and noise of the Mix-IRLS algorithm is missing. \n\n3. There are some claims not associated with any proofs or references. A typical example is a claim in lines 98 & 99: 'No method can recover the models given fewer samples'.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback.\n\n**Weaknesses**:\n1. Unfortunately, we do not currently have a finite-sample analysis or a robustness to outliers guarantee for our algorithm. We agree that these two guarantees would greatly improve the theoretical grounding of our method, and we plan to accomplish them in future research. On the other hand, we do have a robustness to noise guarantee, as described in Section 6. Our guarantee holds for a bounded and symmetric noise.\n2. Addressed in 1 above.\n3. We thank the reviewer for bringing this unclear point to our attention. Kindly see item 3 in our global response. We will further look for and fix any other unjustified claim.\n\n**Questions**:\n1. We agree with the reviewer that a more efficient scheme would lower the value of $T_1$ with $k$. However, it seems nontrivial to devise a suitable scheme provided that the mixture proportions are unknown. Furthermore, the number of iterations done in practice is almost always much less than $T_1$, since we employ an early stopping criterion. This criterion is described in Appendix E.\n2. Thanks for spotting this. We will correct it in the revised manuscript.\n3. Yes. Under the conditions of Theorem 6.1, MixIRLS provably finds the true number of components; please see also Remark 6.4. This claim is formally stated in Appendix C (Proposition C.1), and proved in Appendix D.\n4. Thanks for this question. Kindly see item 4 in our global response."}, {"review_id": "qEJaFV56Vg", "reviewer": "Reviewer_Yapb", "summary": "This paper introduces a novel algorithm for mixed linear regression suitable for balanced and imbalanced mixture settings. It proves the proposed method holds for a sufficiently imbalanced mixture rather than a balanced one through a recovery guarantee on the theoretical front. Empirical results on synthetic data and real-world datasets highlight the superior performance of the proposed algorithm compared to conventional methods. The key contributions of this paper can be summarized as follows:\n+ The paper presents an efficient and straightforward approach to addressing imbalanced mixed linear regression.\n+ It proposes to find the dominant model, remove its associated samples from the observation set, and repeat the process to find the next dominant model.\n+ It demonstrates the effectiveness through empirical results and theoretical analysis.", "questions": "Please refer to the \"Weaknesses\" part.", "limitations": "The authors should clarify the potential negative societal impact of their work.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "+ The paper is well-organized and easy to follow.\n+ The proposed method is efficient and simple. The central idea of simultaneously recovering the models while partitioning the samples into three classes based on their fit to the estimated model is indeed intriguing and adds value to the overall approach. \n+ This approach seems promising for addressing mixed linear regression in both balanced and imbalanced mixture settings, as supported by the empirical results and theoretical analysis presented in the paper.", "weaknesses": "+ My primary concern revolves around the motivation of this paper. While the title focuses on \"imbalanced mixtures,\" the abstract and introduction cover a wide range of settings, such as small sample sizes, the presence of outliers, and an unknown number of models. While it is commendable that the authors aim to demonstrate the robustness of the proposed method, it would be preferable to emphasize one main topic to maintain clarity and focus.\n\n+ Figure 1 visually illustrates the central idea of the proposed method, but further analysis or additional figure tools are required to explain why this method is particularly suitable for imbalanced settings. Providing in-depth insights into the algorithm's performance in such scenarios would enhance the paper's credibility.\n\n+ In Figure 2, a comparison with conventional mixed linear regression methods is presented, but the selected methods mainly date back to before 2016, which might not reflect the state-of-the-art results. Including comparisons with more recent methods would bolster the paper's claims about the effectiveness of the proposed approach. Additionally, line charts, while showing the experiment trends, may not be sufficient to fully convey the experimental findings. Further detailed analysis and discussion of the results are necessary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful comments.\n\n* We thank the reviewer for bringing this issue to our attention. Our goal was to emphasize the advantages of our method, but we understand that the reader might miss our main message. We will revise our abstract and introduction and focus on our method's success in the imbalanced mixture setting.\n* Following the reviewer's comment, we added two new figures to our manuscript. The first one is Fig. 1 in the PDF attached to our global response. This figure illustrates the algorithm's performance given an imbalanced mixture, specifically with the proportions $p = (0.7, 0.2, 0.1)$. Interestingly, at the first iterations, the estimate of MixIRLS is closer to the second component, whose proportion is $p_2 = 0.2$. This must be a consequence of the random initialization, which happened to be closer to the second component. It is further shown that as the iterations proceed, MixIRLS gradually sets its focus to the dominant component ($p_1 = 0.7$). The second figure included in the revised version shows the performance of our IRLS variant as a function of the number of outliers. As expected, IRLS performs better as the number of outliers decreases, corresponding to a higher imbalance in the MLR setting.\n* We compared MixIRLS to popular methods for MLR. Specifically, EM is by far the most popular one, and is considered the state-of-the-art also to these days. More recent works either focused on theory and have no available implementation for their method, or adapted and analyzed EM for other variants of the MLR problem (e.g., sparse MLR), and are thus incomparable to our work.\n* We will further analyze and discuss our experimental results in the revised manuscript. Specifically, we will analyze and interpret our results (i) in the Kaggle datasets and (ii) in the newly added CO2 emission dataset. Kindly see item 1 in our global response for more details."}, {"review_id": "l6Qtqs240Y", "reviewer": "Reviewer_16iH", "summary": "The paper proposes a novel method for estimating parameters in a mixture of linear regression (MLR) model. The authors argue that the maximum likelihood estimator (MLE) does not perform well under extreme imbalanced cases and propose a sequential estimation approach using robust regression methods. The paper presents empirical results to demonstrate the effectiveness of the proposed method.", "questions": "- The mixture model in (1), are the explanatory variables random or fixed? It looks like they are random based on the theory. Should be specified when the model is defined.\n\n- The definition of the mixing proportion is weird in the sense that it depends on the sample size $n$. Should it be $p_k=P(c_i^*=k)$ instead? \n\n- The statement that \"Then the information limit .... is $n_{\\text{info}}\", can you please explain? Similarly for the statment that \"A suitable value for $\\rho$ depends on the ratio ...\". \n\n- $P(\\epsilon_t=t)=P(\\epsilon_t=-t)$ does not mean anything for continuous random variables, should use CDF or pdf instead.", "limitations": "Some limitations and possible future research directions are briefly mentioned in the paper.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The proposed method offers an interesting alternative to the MLE in MLR.\n\n- The method is intuitive, and the paper is well-organized.", "weaknesses": "- Notations and Terminology: There are instances where the notations or terminologies used in the paper are not well-defined. For example:\n    - The term \"significantly imbalanced\" needs clarification.\n    - Line 98: What is the definition of $n_{\\text{info}}$?\n    - Line 190: How is median error defined? Is it the median of $F_{\\text{latent}}$ over 50 repetitions? If so, how is $F_{\\text{latent}}$ defined for the overspecified case?\n\n- Method: While the intuition behind the proposed method is clear, there are several details lacking motivation or explanation. For instance:\n  \n  - What is the intuition behind equation (2b) and the weight in Phase II? How do they differ from directly running a median regression using IRLS?\n  \n  - What happens if no refinement in equation (5) is used? How does the estimator perform?\n\n  - The method seems to rely heavily on hyperparameters such as $\\rho$, $w_{\\text{th}}$, $\\eta$, and $T_1$. However, in Section 4, Line 73, it is claimed that \"Mix-IRLS\" performs well with a fixed set of values for the tuning parameters. It would be helpful to understand why the method is robust to the choice of hyperparameters.\n\n- Experiments: More detailed descriptions of the experiments should be provided to make the results more convincing.\n\n    - What is the definition of \"median error\"? How are the initial values chosen? What does \"highly imbalanced proportions\" mean in Line 83?\n    - It is well-known that fitting mixtures becomes harder when the components are closer to each other. How similar are the different components in the experiments? What happends if the degree of similarity increases? Are these methods stil perform the same?\n    - As mentioned previously, the lack of a clear definition for \"median error\" makes the results difficult to interpret.\n    - Instead of reporting the extreme failure probabilities, it would be better to report the mean squared error (MSE) of the regression estimator. Additionally, the performance of different estimators could be compared using this metric.\n\n- Theory: \n\n    - Providing a definition of \"significantly imbalanced\" would be helpful for understanding the paper. In the abstract, it is mentioned that the method works for imbalanced cases with a \"significantly different\" number of samples. However, in Remark 6.2, it is stated that Theorem 6.1 holds only for a sufficiently imbalanced mixture, and $p_2<1/2$ does not necessarily imply significant imbalance.\n\n    - While theoretical results on the gap between the output of the proposed method and the true value under the population case is interesting, they are not useful in practice where we have a finite number of samples. It would be valuable to demonstrate 1) the consistency of the estimator and 2) the rate of convergence of the proposed method.\n\n    - For cases where the number of components $K$ is unknown, the theoretical results in Appendix C only state that the algorithm stops under the given criteria, but they do not describe any properties of the estimator.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed and constructive feedback. Below we address each point separately. We will carefully revise our manuscript according to the points raised by the reviewer.\n\n**Weaknesses**:\n* Notation and Terminology: \n  - By \"significantly imbalanced\" we mean that the mixture is far from being balanced, e.g. $p=(0.7,0.2,0.1)$. We will reword this term in the revised manuscript to make our point clearer.\n  - The definition of $n_\\text{info}$ is the minimal sample size $n$ that makes the problem identifiable. Kindly see item 3 in our global response for more details.\n  - Yes, the median error is the median of $F_\\text{latent}$ over 50 repetitions. As the reviewer implies, the error is defined differently in the overspecified case, with the true number of components $K^*$ replacing $K$ in Eq. (6); please see Lines 218-219. Namely, in the overspecified case, only the best $K^*$ (out of $K$) components of the estimate are compared to the ground truth.\n* Method:\n  - We thank the reviewer for bringing this unclarity to our attention. Eq. (2b) employs a standard Cauchy weighting, which was found empirically to suit our needs. However, other standard schemes could be used as well. The scheme is similar in phase II but involves two differences tailored to our needs. First, since the weights are calculated for all components, they are subsequently normalized over the components such that $\\sum_{k=1}^K w_{i,k}=1$. Due to this normalization, we do not need the weights to lie between $0$ and $1$ as in standard weighting schemes; so instead of $1/(r^2+1)$ (Cauchy) we use $1/(r^2+\\epsilon_\\text{mp})$ (Eq. (B.1b)). This modification encourages the dominance of one component over the others. We will include this discussion in the revised version of our manuscript.\n  - In practice, the effect on the refinement in Eq. (5) is negligible. However, conceptually, it makes sense to use only the largest-weighted samples - namely, the samples that belong to the estimated component with high confidence - to re-estimate the component. We will clarify this point in the revised manuscript.\n  - Thanks for the question. Let us focus, for example, on the oversampling parameter $\\rho$. To estimate a component, we must use at least $d$ samples, otherwise the system is underdetermined. In our algorithm, we use (the largest-weighted) $\\rho\\cdot d$ samples, $\\rho\\geq 1$; see Eq. (4). If the sample size is large enough, then $\\rho$ has little effect: we can use many samples to estimate a component ($\\rho\\gg 1$), and still have enough samples for the next components. However, in extreme cases where the sample size is close to the information limit, $\\rho$ should be close to 1, so that we use the minimal number of samples for each component. This idea is demonstrated in Fig. 2(right) in the PDF attached to our global response. It is shown that MixIRLS starts to perform worse when $\\rho>n/n_\\text{info}$, as expected. Due to the space limit, we cannot address the other parameters here, but we will do it in the revised manuscript (or during the discussion period, if requested).\n* Experiments:\n  - In the simulations section, the error is defined as $F_\\text{latent}$ (Eq. (6)), and in the experiments section it is $F_\\text{real}$ (Eq. (7)). In both cases, the median error is taken over 50 repetitions, each with a different random initialization (but the same one for all algorithms). The initialization is drawn element-wise from the normal distribution, $\\beta_{k,i}\\sim N(0,1)$. \"Highly imbalanced proportions\" means that the mixture is far from being balanced, e.g. $p=(0.7,0.2,0.1)$, whereas $p=(0.4,0.3,0.3)$ is only a \"moderately imbalanced mixture\".\n  - We thank the reviewer for bringing up this point. In our simulations, the components are well-separated. We agree that examining different separation levels would be very interesting. Kindly see item 2 in our global response.\n  - We hope the definition of \"median error\" is now clear.\n  - In addition to the failure probability, the MSE of all estimators in all simulations is also reported. However, due to space limits, most of these results appear in Appendix F.\n* Theory:\n  - We apologize for this misunderstanding, which will be clarified in the revised version. Empirically, MixIRLS works best for mixtures that are far from being balanced, e.g. $p=(0.7,0.2,0.1)$. However, it works well also for moderately and even perfectly balanced mixtures. Our theory supports these findings: In the general case, our theorem provides a guarantee for MixIRLS only for a sufficiently imbalanced mixture. In the special noise-free case, our theorem holds for an almost perfectly balanced mixture, $p_2<1/2$.\n  - Unfortunately, we do not currently have a finite-sample analysis for our algorithm, but we plan to accomplish it in future research. Kindly see item 4 in our global response for more details on the challenges in finite-sample analysis\n  - Thanks for spotting this. We forgot to explicitly write it in Proposition C.1, but our proof shows that the estimator has the same error in the case of unknown $K$ as in the known $K$ case (in fact, the estimators are identical). We will reword Appendix C accordingly.\n\n**Questions:**\n  - The explanatory variables in the MLR model (1) can be either fixed or random. In our theory, we assume random explanatory variables to simplify the analysis.\n  - Since we do not necessarily assume a random model, the label vector $c^*$ can be fixed. Hence, we cannot define the mixture proportions $p_k$ by $P(c_i^*=k)$. Instead, $p_k$ is defined as the empirical proportion in the observed data, and is thus dependent on the sample size $n$.\n  - For an explanation of $n_\\text{info}$, kindly see item 3 in our global response. For an explanation of the relation between $\\rho$ and $n_\\text{info}$, please see the last item under the \"Method\" title (above in this response).\n  - Thanks, we will replace the probability notation in Eq. (9) with CDF."}]

Author rebuttal:
We thank the reviewers for their time and effort. We appreciate the constructive feedback, which will definitely improve the manuscript. Please find attached a one-page PDF with new figures that help answer some of the reviewers' concerns. Following the reviewers' comments, we describe below few important improvements to our manuscript.

1. **Real-world experiments** [in response to Reviewers LCf0, Yapb, 16iH]. In our Kaggle datasets, it is not known whether the data follow an MLR model. Hence, interpreting the results of MLR algorithms in such cases is a delicate issue. We address this issue in two ways:
    - In the revision, we will more clearly explain our results on these datasets. Our experiments demonstrate the following claim: MixIRLS finds MLR coefficients that fit the data significantly better than other methods. As a consequence, *given that* the data (approximately) follow an MLR model for some $K$, our method will often find the mixture coefficients where other methods fail.
    - We added a new experiment on the CO2 emission by vehicles dataset (also from Kaggle) to complement our empirical study. This dataset approximately follows an MLR model, with the fuel type as the label vector $c^*$ (regular/premium gasoline, diesel or ethanol), such that $K=4$ and $p\approx (0.49,0.43,0.05,0.03)$. As the labels in this dataset are known, we can compare how well each method clusters the data into components. We do so by comparing the maximal intersection over all permutations over 4 items, between the labels found by each method and the true labels: $\max_{\sigma\in [K]!}|\\{i\in [n]:\sigma(c_i)=c_i^*\\}|/n$. Table 1 in the attached PDF demonstrates the favorable performance of MixIRLS on this task: almost 6% higher than the second-best method.
In addition, following Reviewer LCfo's first comment, we further checked whether the success of MixIRLS has to do with the imbalance of the mixture. We randomly subsampled the first three components such that the resulting mixture is approximately balanced (albeit with a much smaller sample size due to the subsampling, resulted in a reduced intersection for all methods compared to the original task). As shown in the Table, in this case MixIRLS performs comparably to other methods. This result advocates the effectiveness of MixIRLS in imbalanced mixtures.

2. **Component separation in simulations** [in response to Reviewer LCfo, 16iH]. As raised by Reviewer 16iH, in our manuscript we did not examine how well the MLR methods perform under different separation levels. The separation is defined as the distance between the components, $\\|\beta^*_i-\beta^*_j\\|$ for $i\neq j$; in case of more than two components ($K>2$) we consider the average separation level. A-priori, we expect to see a small error in two extreme cases: when the separation is very small, because the mixture is close to being degenerate (with the smallest error at $\\|\Delta\\|=0$, which corresponds to a single linear model); and when the separation is large, as the component identification is easier. In between, the error can be larger. We conducted a simulation with several separation levels, and the results are depicted in Fig. 2(left) in the attached PDF. We observed an interesting phenomenon: while both MixIRLS and EM follow the expected behavior, their peak error lies in different values of $\\|\Delta\\|$. At small separation levels, EM has a smaller error than MixIRLS, and at moderate levels it is MixIRLS which has a smaller error. However, as shown in the figure, in all separation levels MixIRLS has a better intersection score (see definition in item 1). The reason is that at smaller separation levels, EM finds an almost perfectly balanced mixture, which is far from the true mixture proportions $p = (0.7,0.2,0.1)$. In contrast, MixIRLS better clusters the samples into components and approximates the mixture proportions well even when the error in $\beta$ is relatively large. We will further elaborate on this matter in the revised manuscript.

3. **Information limit** [in response to Reviewers LCfo, 2aVS, 16iH]. In Lines 97-99 of our manuscript we wrote: "The information limit on the sample size, namely the minimal number of observations n required to recover $\beta^*$ in the absence of noise, is $n_\text{info}=d/\text{min}(p)$. No method can recover the models given fewer samples." We will revise the wording and prove this claim. In linear regression with $d$-dimensional explanatory variables, one must observe at least $d$ samples, otherwise the system is underdetermined. Hence, in MLR, one must observe at least $d$ samples from each component. Since the number of samples associated with the $k$-th component is $n\cdot p_k$, one must have $n\cdot p_k\geq d,\forall k$. In other words, it must hold that $n\cdot p_\text{min}\geq d$. Hence, the information limit, namely the minimal required sample size, is $n_\text{info}=d/p_\text{min}$.

4. **Discussion on finite-sample analysis** [in response to Reviewers LCfo, 1Ro7, 2aVS and 16iH]. Our current work includes only a population setting analysis and lacks a finite-sample one. While performing a finite-sample analysis is currently beyond our reach, we wish to discuss why it is so. A main challenge revolves around the sequential nature of MixIRLS. Once MixIRLS estimates a component, it removes its associated samples. Analyzing a modified sample set introduces challenging dependencies - the data no longer follow a normal distribution but rather a conditional distribution. In the population setting, various quantities happen to cancel out nicely; see, for example, the derivation of Eq. (D.17) from (D.16). Under the setting of finite sample size, we need to bound these quantities under a conditional distribution. This is in contrast to simultaneous approaches, such as EM, which don't have these dependencies as the sample set is fixed throughout the iterations. We hope to find a way to analyze these quantities in the future.

Meta-review: All reviewers are in favour of acceptance, although some limitations regarding the theoretical results were noted during the discussion. For example I recommend the authors clarify what, if any, advantages the proposed method has over the MLE.
Decision: Accept (poster)