Abstract
Under missing-not-at-random (MNAR) sample selection bias, the performance ofa prediction model is often degraded. This paper focuses on one classicinstance of MNAR sample selection bias where a subset of samples havenon-randomly missing outcomes. The Heckman selection model and its variantshave commonly been used to handle this type of sample selection bias. TheHeckman model uses two separate equations to model the prediction and selectionof samples, where the selection features include all prediction features. Whenusing the Heckman model, the prediction features must be properly chosen fromthe set of selection features. However, choosing the proper prediction featuresis a challenging task for the Heckman model. This is especially the case whenthe number of selection features is large. Existing approaches that use theHeckman model often provide a manually chosen set of prediction features. Inthis paper, we propose Heckman-FA as a novel data-driven framework forobtaining prediction features for the Heckman model. Heckman-FA first trains anassignment function that determines whether or not a selection feature isassigned as a prediction feature. Using the parameters of the trained function,the framework extracts a suitable set of prediction features based on thegoodness-of-fit of the prediction model given the chosen prediction featuresand the correlation between noise terms of the prediction and selectionequations. Experimental results on real-world datasets show that Heckman-FAproduces a robust regression model under MNAR sample selection bias.