Under missing-not-at-random (MNAR) sample selection bias, the performance of a prediction model is often degraded. This paper focuses on one classic instance of MNAR sample selection bias where a subset of samples have non-randomly missing outcomes. The Heckman selection model and its variants have commonly been used to handle this type of sample selection bias. The Heckman model uses two separate equations to model the prediction and selection of samples, where the selection features include all prediction features. When using the Heckman model, the prediction features must be properly chosen from the set of selection features. However, choosing the proper prediction features is a challenging task for the Heckman model. This is especially the case when the number of selection features is large. Existing approaches that use the Heckman model often provide a manually chosen set of prediction features. In this paper, we propose Heckman-FA as a novel data-driven framework for obtaining prediction features for the Heckman model. Heckman-FA first trains an assignment function that determines whether or not a selection feature is assigned as a prediction feature. Using the parameters of the trained function, the framework extracts a suitable set of prediction features based on the goodness-of-fit of the prediction model given the chosen prediction features and the correlation between noise terms of the prediction and selection equations. Experimental results on real-world datasets show that Heckman-FA produces a robust regression model under MNAR sample selection bias.

翻译：在非随机缺失（MNAR）样本选择偏差下，预测模型的性能通常会下降。本文聚焦于MNAR样本选择偏差的一个经典实例，即部分样本的结果存在非随机缺失。赫克曼选择模型及其变体常被用于处理此类样本选择偏差。该模型通过两个独立方程分别对样本预测和选择进行建模，其中选择特征包含所有预测特征。使用赫克曼模型时，必须从选择特征集中合理选取预测特征。然而，为赫克曼模型选择合适的预测特征极具挑战性，尤其当选择特征数量庞大时。现有采用赫克曼模型的方法通常依赖人工预选预测特征集。本文提出Heckman-FA这一新型数据驱动框架，用于为赫克曼模型自动获取预测特征。Heckman-FA首先训练一个分配函数，以判定某个选择特征是否应被分配为预测特征；随后利用训练函数的参数，基于给定预测特征下预测模型的拟合优度以及预测方程与选择方程噪声项间的相关性，提取出合适的预测特征集。在真实数据集上的实验结果表明，Heckman-FA能在MNAR样本选择偏差下生成稳健的回归模型。