SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pini Zilber Boaz Nadler
Weizmann Institute of Science, Israel
{pini.zilber, boaz.nadler}@weizmann.ac.il
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.
In this paper, we consider a simple generalization of the linear regression problem, known as mixed linear regression (MLR) [5, Chapter 14]. 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 a 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 MLR, including expectation maximization [15, 5], alternating 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 simultaneously; and (iii) tend to 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 to develop MLR methods that are able to handle such settings, as well as corresponding recovery guarantees.
In this paper, we present Mix-IRLS, a novel and conceptually different iterative algorithm for MLR, able to handle both balanced and imbalanced mixtures. Mix-IRLS is computationally efficient, simple to 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 not require as input the number of components  $K$ . In Sections 4 and 5 we illustrate the competitive advantage of Mix-IRLS over several other methods, on both synthetic and real data.
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, let 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 be viewed as a 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 in 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 to 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 fit are used to re-estimate the model coefficients; those with poor fit are assumed to belong to a 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"), are ignored, but used later in a refinement phase.
To recover the dominant model at each of the  $K$  rounds, we perform robust regression using iteratively reweighted least squares (IRLS) [24, 10, 40]. IRLS iteratively solves weighted least squares subproblems, where the weight of each sample depends on its residual with respect to 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 a recovery guarantee for our method. Specifically, we show that in a population setting with an imbalanced mixture of two components,  $\mathrm{Mix}$ -IRLS successfully recovers the linear models. A key novelty in our result is that it holds for a 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 to the initialization, 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 a practical algorithm as well as a theoretical recovery guarantee. Furthermore, we are the first to propose and analyze an adaptation of IRLS to the MLR problem. A related sequential approach to 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 + 2)$  samples from the data and run 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 is feasible only in low dimensional settings, as the probability that all  $d + 2$  chosen samples belong to the same component decreases exponentially with the dimension  $d$ ; specifically, [3] considered only  $d \leq 5$ . In addition, [3] did not provide a theoretical guarantee for their method. In contrast, our Mix-IRLS method is scalable to high dimensions, and enjoys a theoretical guarantee.
Notation. For a positive integer  $K$ , denote  $[K] = \{1, \dots, K\}$ , and the set of all permutations over  $[K]$  by  $[K]!$ . For a vector  $u$ , denote its Euclidean norm by  $\| u\|$ . For a matrix  $X$ , denote its spectral
Algorithm 1: Mix-IRLS: main phase
input: samples  $\{(x_i,y_i)\}_{i = 1}^n$  , number of components  $K$  parameters  $w_{\mathrm{th}},\rho ,\eta ,T_1$  output: estimates  $\beta_{1}^{\mathrm{phase - I}},\ldots ,\beta_{K}^{\mathrm{phase - I}}$    
1 set  $S_{1} = [n]$    
2 for  $k = 1$  to  $K$  do   
3 initialize  $\beta_{k}$  randomly   
4 repeat  $T_{1}$  times   
5 compute  $r_{i,k} = |x_i^\top \beta_k - y_i|$ $\forall i\in S_k$    
6 compute  $w_{i,k} = (1 + \eta r_{i,k}^{2} / \bar{r}_{k}^{2})^{-1},\quad \forall i\in S_{k}$    
7 compute  $\beta_{k} = (X_{S_{k}}^{\top}W_{k}X_{S_{k}})^{-1}X_{S_{k}}^{\top}W_{k}y_{S_{k}}$    
8 end   
9 set  $S_{k + 1} = \{i\in S_k:w_{i,k}\leq w_{\mathrm{th}}\}$    
10 set  $S_k^\prime = \{\rho \cdot d$  samples in  $S_{k}$  with largest  $w_{i,k}\}$    
11 if  $k <   K$  and  $|S_{k + 1}| <   \rho \cdot d$  then   
12 restart Mix-IRLS with  $w_{\mathrm{th}}\gets w_{\mathrm{th}} + 0.1$    
13 end   
14 compute  $\beta_k^{\mathrm{phase - I}} = (X_{S_k'}^\top X_{S_k'})^{-1}X_{S_k'}^\top y_{S_k'}'$    
15 end
norm by  $\| X\|$  and its smallest singular value by  $\sigma_{\mathrm{min}}(X)$ . Given a matrix  $X\in \mathbb{R}^{n\times d}$  and an index set  $S\subseteq [n]$ ,  $X_{S}\in \mathbb{R}^{|S|\times d}$  is the submatrix of  $X$  that corresponds to the rows in  $S$ . Denote by  $\operatorname {diag}(w)$  the diagonal matrix  $W$  whose entries are  $W_{ii} = w_{i}$ . The probability of an event  $A$  is  $\mathbb{P}[A]$ . Denote the expectation and the variance of a random variable  $x$  by  $\mathbb{E}[x]$  and  $\operatorname {Var}[x]$ , respectively. The cumulative distribution function of the standard normal distribution  $\mathcal{N}(0,1)$  is  $\Phi (\cdot)$ .
Let  $\{(x_i, y_i)\}_{i=1}^n$  be  $n$  pairs of explanatory variables  $x_i \in \mathbb{R}^d$  and corresponding responses  $y_i \in \mathbb{R}$ . In standard linear regression, one assumes a linear relation between the response and the explanatory variables, namely  $y_i = x_i^\top \beta^* + \epsilon_i$  where  $\epsilon_i \in \mathbb{R}$  are random noise terms with zero mean. A common goal is to estimate the coefficient vector  $\beta^* \in \mathbb{R}^d$ . In mixed linear regression (MLR), in contrast, the assumption is that each response  $y_i$  belongs to one of  $K$  different linear models  $\{\beta_k^*\}_{k=1}^K$ . Formally,
$$
y _ {i} = x _ {i} ^ {\top} \beta_ {c _ {i} ^ {*}} ^ {*} + \epsilon_ {i}, \quad i = 1, \dots , n, \tag {1}
$$
where  $c^* = (c_1^*, \ldots, c_n^*)^\top \in [K]^n$  is the label vector. Importantly, we do not know to which component each pair  $(x_i, y_i)$  belongs, namely  $c^*$  is unknown. For simplicity, we assume the number of components  $K$  is known, and later on discuss the case where it is unknown. Given the  $n$  samples  $\{(x_i, y_i)\}_{i=1}^n$ , the goal is to estimate  $\beta^* \equiv \{\beta_1^*, \ldots, \beta_K^*\} \subset \mathbb{R}^d$ , 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  $\beta^{*}$  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 = (p_{1},\ldots ,p_{K})$ , with  $p_k = |\{i\in [n]:c_i^* = k\} | / n$ . Then 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_{\mathrm{info}} = 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 in the simplified setting of  $K = 2$  with perfect balance  $(p_1 = p_2 = 1/2)$  and no noise  $(\epsilon = 0)$ , the problem is in general NP-hard [52]. However, as detailed in Section 6, various assumptions render the problem computationally tractable, and enable to derive theoretical guarantees.
3 The Mix-IRLS Method
For simplicity, we present our algorithm assuming  $K$  is known; the case of an 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  $\beta_1^*,\ldots ,\beta_K^*$  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 a brief overview of the main phase mechanism. At each round of this phase, Mix-IRLS uses techniques from robust regression to estimate the largest component present in the data. Then, it 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  $w_{\mathrm{th}} \in (0,1)$ , an oversampling ration  $\rho \geq 1$ , a tuning parameter  $\eta$  and the number of IRLS iterations  $T_{1} \geq 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 a yet undiscovered component, follow an MLR model. As such, standard high-dimensional robust regression methods can recover the underlying linear model even with  $70\%$  and  $80\%$  outliers; see Appendix H for an empirical illustration.
Next, we present in details the main phase of Mix-IRLS. A pseudocode appears in Algorithm 1. Given  $X = \left(x_{1} \cdots x_{n}\right)^{\top}$  and  $y = \left(y_{1} \cdots y_{n}\right)^{\top}$ , we initialize the set of active samples to the entire dataset,  $S_{1} = [n]$ . Next, we perform the following procedure for  $K$  rounds. At round  $k \in [K]$ , starting from a random guess  $\beta_{k}$  for the  $k$ -th vector, we run  $T_{1}$  iterations of IRLS [24]:
$$
r _ {i, k} = \left| x _ {i} ^ {\top} \beta_ {k} - y _ {i} \right|, \forall i \in S _ {k}, \quad (\text {u p d a t e r e s i d u a l s}) \tag {2a}
$$
$$
w _ {i, k} = \frac {1}{1 + \eta \cdot r _ {i , k} ^ {2} / \bar {r} _ {k} ^ {2}}, \forall i \in S _ {k}, \quad (\text {u p d a t e w e i g h t s}) \tag {2b}
$$
$$
\beta_ {k} = \left(X _ {S _ {k}} ^ {\top} W _ {k} X _ {S _ {k}}\right) ^ {- 1} X _ {S _ {k}} ^ {\top} W _ {k} y _ {S _ {k}}, \quad (\text {u p d a t e e s t i m a t e}) \tag {2c}
$$
where  $\bar{r}_k = \mathrm{median}\{r_{i,k} \mid i \in S_k\}$ , and  $W_k = \mathrm{diag}(w_{1,k}, w_{2,k}, \ldots)$ . The updated estimate (2c) is the minimizer of the weighted least squares objective  $\| W^{\frac{1}{2}}(y_{S_k} - X_{S_k} \beta_k) \|^2$ . Intuitively, if  $\beta_k$  is close to some  $\beta_i^*$ , then samples from other components receive a lower weight and the updated  $\beta_k$  is even closer to  $\beta_i^*$ . After  $T_1$  iterations of (2a)-(2c), we define the subset  $S_{k+1}$  of 'poor fit' samples that seem to belong to other components,
$$
S _ {k + 1} = \left\{i \in S _ {k}: w _ {i, k} \leq w _ {\mathrm {t h}} \right\}. \tag {3}
$$
This serves as the set of active samples for the next round. In addition, for some  $\rho \geq 1$  whose choice is discussed below, we define the subset  $S_{k}^{\prime}$  of samples with a 'good fit' to the  $k$ -th component,
$$
S _ {k} ^ {\prime} = \left\{\rho \cdot d \text {s a m p l e s i n} S _ {k} \text {w i t h l a r g e s w e i g h t s} w _ {i, k} \right\}. \tag {4}
$$
We then refine the estimate of  $\beta_{k}$  by performing OLS on this subset,
$$
\beta_ {k} ^ {\text {p h a s e - I}} = \left(X _ {S _ {k} ^ {\prime}} ^ {\top} X _ {S _ {k} ^ {\prime}}\right) ^ {- 1} X _ {S _ {k} ^ {\prime}} ^ {\top} y _ {S _ {k} ^ {\prime}}. \tag {5}
$$
The subsets  $S_{k}^{\prime}$  and  $S_{k + 1}$  are in general disjoint, unless the threshold  $w_{\mathrm{th}}$  or the oversampling ratio  $\rho$  are too high. The choice for the value of  $w_{\mathrm{th}}$  is discussed in Remark 3.1 and Section 6. A suitable value for  $\rho$  depends on the ratio between the sample size  $n$  and the information limit  $n_{\mathrm{info}} = d / \min(p)$ . In the challenging setting of  $n \approx n_{\mathrm{info}}$ ,  $\rho$  should be set close to 1, confining the algorithm to use the least possible number of samples at each round. If  $n \gg n_{\mathrm{info}}$ , then  $\rho$  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.
Figure 2: Comparison of various MLR algorithms. Depicted is the failure rate, out of 50 random initializations, for which  $F_{\mathrm{latent}} > 2\sigma$  (see (6)), as a function of the sample size  $n$ . The dimension and noise level are fixed at  $d = 300$  and  $\sigma = 10^{-2}$ . Mixture:  $K = 3$  with  $p = (0.7, 0.2, 0.1)$  (left panel);  $K = 5$  with  $p = (0.63, 0.2, 0.1, 0.05, 0.02)$  (right panel).
This concludes the main phase of Mix-IRLS. The second (refinement) phase, described in Appendix C, improves the estimates  $\beta_{k}^{\mathrm{phase-I}}$  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  $\beta^{*}$ . This is in accordance with the theoretical result in Section 6.
Remark 3.1 (Threshold adaptation). If the threshold  $w_{\mathrm{th}}$  is too small, the number of poor fit samples  $|S_{k + 1}|$  of (3) may be insufficient to estimate the  $(k + 1)$ -th component. To avoid the tuning of  $w_{\mathrm{th}}$ , we use the following scheme: Mix-IRLS gets as input an initial value for  $w_{\mathrm{th}}$ , and if needed, it increases its value and starts over; see Lines 11 to 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:  $\eta, \rho, T_1$  and an initial value for  $w_{\mathrm{th}}$  (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/overestimated  $K$ ). If  $K$  is unknown, or only an upper bound  $K_{\mathrm{max}}$  is given, Mix-IRLS can 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_{k + 1}| < \rho d$ . We then set  $K$  to the number of components with at least  $\rho d$  associated samples, and proceed to the second phase.
4 Simulation Results
We present simulation results on synthetic data in this section, and on several real-world datasets in the next one. We compare the performance of Mix-IRLS to the following algorithms: (i) A1tMin - 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 in MATLAB. In some of the simulations, we additionally ran a 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, as 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 a 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], in each simulation we sampled the entries of the explanatory variables  $X$  and of the regression vectors  $\beta^{*}$  from the standard normal distribution  $\mathcal{N}(0,1)$ . In this section, the dimension is fixed at  $d = 300$ . The additive noise terms are Gaussian,  $\epsilon_{i} \sim \mathcal{N}(0,\sigma^{2})$ . Additional simulations with other values of  $d$  and  $\sigma$ , including different noise levels for the  $K$  components (i.e.,
$^{1}$ MATLAB and Python code implementations of Mix-IRLS are available at github.com/pizilber/MLR.
Figure 3: Robustness of several MLR algorithms to outliers (left panel) and to overparameterization (right panel). Values of  $d, \sigma, K$  and  $p$  are the same as in Figure 2(left), and  $n = 12000$ . Qualitatively similar results for a balanced mixture appear in Figures H.15 and H.16. 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  $\sigma_{1},\ldots ,\sigma_{K})$  , as 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 Mix-IRLS is dealing with imbalanced mixtures. Hence, in this section, the labels  $c_{i}^{*}$  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  $\beta \equiv \{\beta_{1},\dots ,\beta_{K}\}$  by the following quantity:
$$
F _ {\text {l a t e n t}} (\beta ; \beta^ {*}) = \min  _ {\sigma \in [ K ]!} \max  _ {k = 1, \dots , K} \| \beta_ {\sigma (k)} - \beta_ {k} ^ {*} \|. \tag {6}
$$
The minimization above is over all  $K!$  permutations, which makes the accuracy measure invariant to the order of the regression vectors in  $\beta$ . A similar objective was used by [53, 56].
All algorithms were initialized with the same random vectors  $\beta_{1},\ldots ,\beta_{K}$ , whose entries were sampled from the standard normal distribution  $\mathcal{N}(0,1)$ . For each simulation, we performed 50 independent realizations, each with a different random initialization  $\beta$ , and report the median errors and the failure probability. The latter quantity is defined as the percentage of runs whose error  $F_{\mathrm{latent}}$  (6) is above  $2\sigma$ ; 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-IRLS is shown to recover the components with sample size very close to the oracle'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.