SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Next, we explore the robustness of the algorithms to additive noise, outliers and overestimation. Figure H.10 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 \in (0, 1)$  of the observations. A corrupted response  $\tilde{y}_i$  is sampled from a normal distribution with zero mean and variance  $\sum_{j=1}^{n} y_j^2 / n$ , independently of the original value  $y_i$ . Figure 3(left) 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  $\beta$  is calculated based on the  $\lceil (1 - f)n \rceil$  samples with smallest residuals. In 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 is 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 ( $f = 17\%$ ), which is almost 6 times more outliers than other methods. These results should not be surprising in light of the fact that robust regression is at the heart of Mix-IRLS's mechanism.
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 on a 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 a number of components  $K \geq 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 error is calculated based on the best  $K^{*}$  vectors in  $\beta$ , ignoring its other  $K - K^{*}$  vectors. Figure 3(right) shows that most algorithms have similar performance at the correct parameter value  $K = K^{*}$ , and EM succeeds also at small overparameterization,  $K - K^{*} \leq 3$ . Only Mix-IRLS is insensitive to overparameterization, and succeeds with unbounded  $K$ . This feature is attained thanks to the sequential nature of Mix-IRLS (Remark 3.3). Similar results hold for balanced mixtures; see Appendix H.
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 a mathematical formulation of this finding, Cohen proposed the MLR model (1) with  $K = 2$ , where the labels  $c^*$  are i.i.d. according to a Bernoulli distribution; see also [15].
In Figure 4(left), the untuned version of Mix-IRLS is shown to capture the two linear trends in the data. Notably, untuned Mix-IRLS was not given the number of components, but automatically stopped at  $K = 2$  with its default parameters. By increasing the sensitivity of Mix-IRLS to new components via the parameter  $w_{\mathrm{th}}$ , 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 on a 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-truth) 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 as 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 we choose the permutation that maximizes the overall intersection of its estimated clusters with the true ones, defined as
$$
\max  _ {\sigma \in [ K ]! ]} \frac {1}{n} | \{i \in [ n ]: \sigma \left(c _ {i}\right) = c _ {i} ^ {*} \} |. \tag {7}
$$
The results, reported in Table 1, demonstrate the effectiveness of Mix-IRLS on imbalanced mixtures. While on the dominant class (regular gasoline) all methods achieve comparable balanced accuracy,
Table 1: 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 on 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  $\beta^{*}$  are unknown. Hence, we replace (6) by the following quality measure which represents the fit of an MLR model to the data:
$$
F _ {\text {r e a l}} (\beta ; X, y) = \frac {1}{\operatorname {V a r} [ y ]} \cdot \frac {1}{n} \sum_ {i = 1} ^ {n} \min  _ {j \in [ K ]} \left(x _ {i} ^ {\top} \beta_ {j} - y _ {i}\right) ^ {2}, \tag {8}
$$
resembling the K-means objective for clustering [23]. To illustrate the advantage of a 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 a 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-IRLS: tuned improve upon the other methods, sometimes by  $30\%$  or more.
The experiments on the Kaggle datasets demonstrate that empirically Mix-IRLS finds MLR coefficients that fit the data significantly better than other methods. Hence, given as input data that (approximately) 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 Mix-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,
$$
x _ {i} \sim \mathcal {N} (0, I _ {d}). \tag {9}
$$
In the more general case where  $x_{i} \sim \mathcal{N}(\mu, \Sigma)$  for some  $\mu$  and  $\Sigma$ , one may find  $\mu$  and  $\Sigma$  from the data, and mean center and whiten the samples  $x_{i} \to \Sigma^{-1/2}(x_{i} - \mu)$ , before applying our method. The responses  $y_{i}$  are assumed to follow model (1) with  $K = 2$  components  $(\beta_{1}^{*}, \beta_{2}^{*})$  and labels  $c_{i}^{*}$
generated independently of  $x_{i}$  according to mixture proportions  $p_1 \geq p_2 > 0$ . This setting was considered in several previous theoretical works that analyzed the EM method for MLR [2, 14, 33, 36].
We assume the noise terms  $\epsilon_{i}$  are all i.i.d., zero-mean random variables, independent of  $x_{i}$  and  $c_{i}^{*}$ . We further assume they are bounded and follow a symmetric distribution,
$$
\left| \epsilon_ {i} \right| \leq \sigma_ {\epsilon} \quad \text {a n d} \quad \mathbb {P} [ \epsilon_ {i} \leq t ] = \mathbb {P} [ \epsilon_ {i} \geq - t ] \quad \forall t. \tag {10}
$$
For analysis purposes, we consider a slightly modified variant of Mix-IRLS, formally described in Appendix E. In this variant, Mix-IRLS excludes samples  $x_{i}$  with large magnitude,  $\| x_{i}\|^{2} > R$  where  $R$  is a fixed parameter. A natural choice for its value is  $R\sim \mathbb{E}[\| x_i\|^2 ] = d$ , e.g.  $R = 2d$ . In a high-dimensional setting with  $d\gg 1$ , such a choice excludes a 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.1.
The following theorem states that given a sufficiently imbalanced mixture, namely  $p_1$  is large enough, Mix-IRLS successfully recovers the underlying vectors  $\beta_1^*$  and  $\beta_2^*$ .
Theorem 6.1. Let  $\{(x_i, y_i)\}_{i=1}^{\infty}$  be i.i.d. from a mixture of  $K = 2$  components with proportions  $(p_1, p_2)$ , regression vectors  $(\beta_1^*, \beta_2^*)$ , and noise terms  $\epsilon_i$  that follow (10) with  $\sigma_{\epsilon}$ . Denote  $\Delta = \beta_1^* - \beta_2^*$  and  $\gamma = 5p_2 / 4$ . Suppose that
$$
q \equiv \gamma + \left(1 + \frac {1}{\sqrt {R}}\right) \frac {\sigma_ {\epsilon}}{\| \Delta \|} <   \frac {1}{2}. \tag {11}
$$
Further assume that  $M$  is run with parameters  $\rho = \infty$ , and with  $\eta, w_{th}$  that satisfy,
$$
\frac {1}{1 + \eta (1 - q) ^ {2} \| \Delta \| ^ {2}} <   w _ {t h} <   \frac {1}{1 + \eta q ^ {2} \| \Delta \| ^ {2}}, \tag {12}
$$
Then for sufficiently large  $R$ , starting from an arbitrary initialization, the first phase of  $M\mathfrak{X}$ -IRLS with at least one iteration ( $T_1 \geq 1$ ) recovers  $\beta^*$  up to an error that decreases with increasing  $R$ ,
$$
\max  _ {k = 1, 2} \| \beta_ {k} - \beta_ {k} ^ {*} \| \leq \frac {1}{\sqrt {R}} \frac {\sigma_ {\epsilon}}{\sigma_ {\epsilon} + \gamma \| \Delta \|}. \tag {13}
$$
Specifically, in the absence of noise  $(\sigma_{\epsilon} = 0)$ , Mix-IRLS perfectly recovers the regression vectors.
Theorem 6.1 considers only the first phase of Mix-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  $\rho = \infty$  is suited to the population setting where  $n = \infty$ ; see the discussion after Eq. (5). The theorem proof appears in Appendix E.
Remark 6.2 (Different noise levels). Theorem 6.1 holds also in the case of different noise levels  $\sigma_{1},\sigma_{2}$  for each of the two mixture components. In this case,  $\sigma_{\epsilon}$  is replaced by  $\max \{\sigma_1,\sigma_2\}$ . Indeed, as illustrated empirically in Figure H.11, Mix-IRLS is able to handle such a case as well.
Remark 6.3 (Required imbalance). Due to (11), Theorem 6.1 holds only for a sufficiently imbalanced mixture. In the absence of noise, the factor  $5/4$  in the definition of  $\gamma$  may be relaxed to any number arbitrarily close to one, at the expense of increasing  $R$ . In turn, this implies that our theorem holds for any mixture with the probability of the second component allowed to be arbitrarily close to  $p_2 < 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 a trade-off between mixture imbalance and noise level: the proportion of the dominant component  $p_1$  increases with the noise level  $\sigma_{\epsilon}$ . This theoretical finding is in agreement with the simulation results in Section 4: Mix-IRLS works better as the imbalance increases. We emphasize that empirically, Mix-IRLS works well also on balanced mixtures,  $p_1 = p_2 = 1/2$ ; see Appendix H. Hence it is an open problem to derive theoretical guarantees for a perfectly balanced mixture.
Remark 6.4 (Allowed range for  $w_{\mathrm{th}}$ ). Theorem 6.1 holds for a limited range of values for  $w_{\mathrm{th}}$ , see (12). This range depends on  $q$ , which in turn depends on the noise level and the mixture imbalance. For example, at a noise level  $\sigma_{\epsilon} = 10^{-2}$ , proportions  $p = (3/4, 1/4)$ , separation of  $\| \Delta \| = 1$ , parameter choice of  $\eta = 1$  and  $R \geq 2$ , the range for which recovery is guaranteed is  $0.69 \leq w_{\mathrm{th}} \leq 0.9$ . As mentioned previously, empirically we run Mix-IRLS with an initial value of  $w_{\mathrm{th}} = 0.1$  which is increased by the algorithm (see lines 11-13 in Algorithm 1). Indeed, in 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  $(\beta_{1},\ldots ,\beta_{K})$ . 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, 53, 21] and GD [56, 38] in a noiseless setting. More related to 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,  $p_1 = p_2 = 1/2$ . An exception is [34], who allowed for  $K > 2$  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) is monotonically decreasing with the mixture imbalance, the allowed range (12) of the parameter  $w_{\mathrm{th}}$  increases with the imbalance. Furthermore, our result holds for an arbitrary initialization. The downside is that Theorem 6.1 requires sufficient imbalance (see Remark 6.3), and does not provide a 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.
Mix-IRLS 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-IRLS is robust to outliers and to 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 a 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 (e.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 a component, it removes its associated samples. Analyzing a modified sample set introduces complex statistical 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 (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 a finite-sample setting in the future. Finally, another interesting future research direction is to formally prove Mix-IRLS's robustness to outliers, as was empirically demonstrated in Section 4.
The research of P.Z. was partially supported by a fellowship for data science from the Israeli 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 in part by grant 2362/22 from the Israel Science Foundation. We thank Yuval Kluger and Ofir Lindenbaum for interesting discussions. We thank the 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.
[1] Murray Aitkin and Granville Tunnicliffe Wilson. Mixture models, outliers, and the EM algorithm. Technometrics, 22(3):325-331, 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, Xin Chen, and Wen Gao. Locally linear regression for pose-invariant face recognition. IEEE Transactions on image processing, 16(7):1716-1725, 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, 22(4):bbaa291, 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 Machler. Robust regression: A weighted least squares approach. Communications in Statistics-Theory and Methods, 26(6):1381-1394, 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 on Theory of Computing, pages 587-600, 2020.
[12] Yudong Chen, Xinyang Yi, and Constantine Caramanis. A convex formulation for mixed regression with two components: Minimax optimal rates. In Conference on 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 on Learning Theory, pages 704-710. PMLR, 2017.
[15] Richard D De Vaeux. Mixtures of linear regressions. Computational Statistics & Data Analysis, 8(3):227-245, 1989.
[54] Derek S Young and David R Hunter. Mixtures of regressions with predictor-dependent mixing proportions. Computational Statistics & Data Analysis, 54(10):2253-2266, 2010.
[55] Linjun Zhang, Rong Ma, T 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 components. Advances in neural information processing systems, 29, 2016.