On Prediction Feature Assignment in the Heckman Selection Model
Huy Mai, Xintao Wu
TL;DR
The paper addresses MNAR sample selection bias in outcome prediction by identifying the challenge of selecting prediction features from a large set of selection features within the Heckman framework. It introduces Heckman-FA, a data-driven approach that learns a feature-assignment function via Gumbel-Softmax, then extracts a prediction-feature set based on the estimated noise-correlation and goodness-of-fit, before fitting a robust Heckman model augmented with the inverse Mills ratio. Empirical evaluation on CRIME and COMPAS shows that Heckman-FA yields lower testing error than naive baselines and standard RU regression, with Heckman-FA* offering a simpler ranking-based alternative that remains competitive. The method provides a scalable, principled way to mitigate MNAR biases in regression tasks, with practical performance gains and interpretability through the learned feature assignments.
Abstract
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.