# On Prediction Feature Assignment in the Heckman Selection Model

arXiv: https://arxiv.org/abs/2309.08043 · HTML: https://mlsanity.com/paper/2309.08043/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 6 of its 24 references that are in the map.
- First posted on arXiv 14 Sep 2023.

## Builds on (most cited first, 6 of 6)

- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [CITRIS: Causal Identifiability from Temporal Intervened Sequences](https://mlsanity.com/paper/2202.03169/) (2022)
- [Near-Optimal Linear Regression under Distribution Shift](https://mlsanity.com/paper/2106.12108/) (2021)
- [Learning from a Biased Sample](https://mlsanity.com/paper/2209.01754/) (2022)
- [On Non-Random Missing Labels in Semi-Supervised Learning](https://mlsanity.com/paper/2206.14923/) (2022)
- [Fair Regression under Sample Selection Bias](https://mlsanity.com/paper/2110.04372/) (2021)

## Abstract (arXiv)

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.