Stochastic EM for Shuffled Linear Regression
This work addresses a practical issue in data analysis where samples are permuted, offering an improved method for researchers and practitioners dealing with shuffled datasets, though it is incremental as it builds on existing EM frameworks.
The paper tackles the problem of linear regression with shuffled data by treating the unknown permutation as a latent variable and using a stochastic EM approach. It shows that this method reduces parameter error, is less sensitive to initialization, and performs better on partially shuffled datasets compared to hard EM, with modest error recovery in real experiments.
We consider the problem of inference in a linear regression model in which the relative ordering of the input features and output labels is not known. Such datasets naturally arise from experiments in which the samples are shuffled or permuted during the protocol. In this work, we propose a framework that treats the unknown permutation as a latent variable. We maximize the likelihood of observations using a stochastic expectation-maximization (EM) approach. We compare this to the dominant approach in the literature, which corresponds to hard EM in our framework. We show on synthetic data that the stochastic EM algorithm we develop has several advantages, including lower parameter error, less sensitivity to the choice of initialization, and significantly better performance on datasets that are only partially shuffled. We conclude by performing two experiments on real datasets that have been partially shuffled, in which we show that the stochastic EM algorithm can recover the weights with modest error.