# On Prediction Feature Assignment in the Heckman Selection Model

Huy Mai, Xintao Wu  Department of Electrical Engineering and Computer Science
University of Arkansas
Fayetteville, Arkansas, USA
{huymai<EMAIL_ADDRESS>
###### 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.

###### Index Terms:

sample selection bias, missing-not-at-random, Heckman selection model, robust
linear regression

##  I Introduction

Regression is sensitive to dataset shift [16], where the training and testing
sets come from different distributions. Dataset shift can arise due to sample
selection bias, where a sample is non-uniformly chosen from a population for
training a model. This type of bias can cause a subset of training samples to
be partially observed, where any of the covariates or outcome of a sample is
missing, or completely unobserved. Consequently, the performance of the model
after training on this biased set will be degraded when the model is deployed.
Most approaches such as [3], [13], and [19] handle the missing-at-random (MAR)
setting, where the selection of training samples is assumed to be independent
from the outcome given the observed variables in the training set. However,
these approaches cannot properly account for the missing-not-at-random (MNAR)
setting, where the selection of training samples is assumed to not be
independent from the outcome given the observed variables in the training set.

Figure 1: Illustration of the effect of MNAR sample selection bias on the
predictions of a linear model. Solid (dashed) line represents the regression
line fitted on the biased (unbiased and fully observed) set.

In this work, we focus on the problem of MNAR sample selection bias on the
outcome. We provide an example of this bias scenario in Figure 1 by
considering the relationship between SAT score (feature) and the amount of
scholarship offered by a certain university (outcome), where some students
have missing values of scholarship. There could be some hidden mechanism
behind the missing outcomes. For instance, amounts of scholarship offered to
students who have not declared their majors are not collected. When the
undeclared students are omitted from the training, a biased model is produced
and could be very different from the ground truth model that would have been
trained had scholarship amounts of all students been collected. However, their
feature information are available during training. In order for the trained
model to be close to the ground truth model, we leverage the observed feature
information of those records with missing outcomes in the training process.

The Heckman selection model [7] is a Nobel-Prize winning approach that
addresses MNAR sample selection bias on the outcome. The method involves two
equations: the prediction equation and the selection equation. The prediction
equation describes the relationship between the covariates and the outcome of
interest. The selection equation specifies the probability that a sample is
selected in the training data. Estimation of the selection equation requires
an exclusion restriction, where the selection equation includes one or more
variables that are not included in the outcome equation. To handle MNAR sample
selection bias on the outcome, the Heckman model considers the respective
noise terms of the prediction and selection equations, which follow a
bivariate normal distribution.

Although the presence of an exclusion restriction avoids multicollinearity for
the prediction model [14], the process to identify a valid exclusion
restriction is often difficult in practice. This is first due to the lack of
clear theoretical knowledge on which selection features should be included in
the prediction model [6]. Moreover, using the Heckman selection model with an
invalid exclusion restriction can lead to non-robust regression on the biased
training set [23]. Choosing from features that affect the selection, one way
to address these challenges is to search through all combinations of selection
features to find a suitable set of prediction features for the Heckman
selection model. However, this search process becomes computationally
expensive as we deal with a large number of selection features in real-world
data.

TABLE I: Description of main symbols used in paper. Notation | Description
---|---
$m(n)$ |  size of $\mathcal{D}_{s}(\mathcal{D}_{tr})$
$\mathcal{D}_{s}$ ($\mathcal{D}_{u}$) | training samples with observed (missing) outcome
$\bm{x}_{i}$, $y_{i}$, $s_{i}$ |  features, outcome, selection for $i$th training sample
$\bm{x}_{i}^{(s)}$ ($\bm{x}_{i}^{(p)}$) | selection (prediction) features of a sample
$K$ ($J$) | number of selection (prediction) features
$\rho$ | correlation coefficient between the noise terms
$\psi$ | assignment function
$\bm{Z}$ | categorical sample matrix
$\pi_{2k}(\pi_{1k})$ |  probability of (not) assigning $k$th selection feature
$T$ | number of training epochs
$B$ | number of Gumbel-Softmax samples (extraction)

###  I-A Problem Formulation

We provide a list of important symbols used throughout the paper in Table I.
In our work, we generally use unbolded letters to denote scalars, bold lower-
case letters to denote vectors, and bold upper-case letters to denote
matrices. For accented letters, we use a hat to denote estimations, tilde to
denote approximations, and underline to denote augmented vectors or matrices.
As exceptions to these conventions, we use $\bm{X}_{\cdot k}$ to denote the
$k$th column of any matrix and $\bm{\pi}$ to denote the probability matrix.

We let $\mathcal{X}$ be the feature space and $\mathcal{Y}$ be the continuous
target attribute. We also denote $\mathcal{D}_{tr}=\\{\bm{t}_{i}\\}_{i=1}^{n}$
as the training set of $n$ samples that are originally sampled from the
population to be modeled yet biased under MNAR sample selection bias. We
define each sample $\bm{t}_{i}$ as

$\bm{t}_{i}=\begin{cases}(\bm{x}_{i},y_{i},s_{i}=1)&1\leq i\leq m\\\
(\bm{x}_{i},s_{i}=0)&m+1\leq i\leq n\end{cases}$ (1)

where the binary variable $s_{i}$ indicates whether or not $y_{i}$ is
observed. We define $\mathcal{D}_{s}$ as the set containing the first $m$
training samples where each sample is fully observed and $\mathcal{D}_{u}$ as
the set that contains the remaining $n-m$ training samples with unobserved
outcome.

######  Definition I.1 (MNAR Sample Selection)

Missing-not-at-random (MNAR) sample selection occurs for a sample $\bm{t}_{i}$
if $s_{i}$ is not independent of $y_{i}$ given $\bm{x}_{i}$, i.e.
$P(s_{i}|\bm{x}_{i},y_{i})\neq P(s_{i}|\bm{x}_{i})$.

MNAR means that both the probability of data missingness and the value of
missing observations are affected by unobserved variables. For the Heckman
model, the selection mechanism models the missingness of $y_{i}$ in terms of a
set of selection features. Thus the following assumptions are additionally
made in this work:

  * (i)

For all $\bm{t}_{i}$, $\bm{x}^{(s)}_{i}$ consists of every prediction feature
and additional features that do not affect the outcome. The selection features
can be specified by domain users or simply learned via goodness-of-fit as all
feature information are available in the training data.

  * (ii)

Given selection features $\bm{x}_{i}^{(s)}\subseteq\bm{x}_{i}$ of the $i$th
training sample, $P(s_{i}|\bm{x}_{i},y_{i})$ is approximated by computing
$P(s_{i}|\bm{x}_{i}^{(s)})$.

Problem Statement. To perform regression, a linear model
$h(\bm{x}_{i}^{(p)};\bm{\beta})$ with prediction features $\bm{x}^{(p)}_{i}$
and parameters $\bm{\beta}$ is fitted by learning to minimize a loss function
over $\mathcal{D}_{tr}$. Given that $\mathcal{D}_{tr}$ is biased due to MNAR
sample selection, we seek to choose prediction features from a set of
selection features using an assignment function $\psi$. Based the extracted
assignment of prediction features, a model
$h(\underline{\bm{x}}^{(p)}_{i};\underline{\bm{\beta}})$ is fitted after
running the Heckman model, where $\underline{\bm{x}}^{(p)}_{i}$ denotes the
extracted prediction features augmented with the inverse Mills ratio (IMR) and
$\underline{\bm{\beta}}$ denotes the Heckman coefficients.

###  I-B Contributions

In this work, we present the Heckman selection model with Feature Assignment
(Heckman-FA) as a framework that finds a suitable assignment of prediction
features for the Heckman model to robustly handle MNAR sample selection bias
on the outcome. The core contributions of our work are summarized as follows.
First, Heckman-FA trains an assignment function that determines whether a
selection feature is assigned as a prediction feature or not. The assignment
function is defined in terms of samples from the Gumbel-Softmax distribution
[9]. Second, using the parameters of the trained assignment function, Heckman-
FA extracts the set of prediction features for the Heckman model based on
goodness-of-fit and the correlation between the noise terms of the prediction
and selection equations. Third, we apply our method to real-world datasets and
compare the performance of Heckman-FA against other regression baselines. We
empirically show that Heckman-FA produces robust prediction models against
MNAR sample selection bias and outperforms these regression baselines.

##  II Related Work

###  II-A Incorporating the Heckman Selection Model

The Heckman selection model has been widely utilized to handle MNAR sample
selection bias in different applications. For example, [2] evaluated the
utility of the Heckman selection model in criminology by testing the method on
sentencing data. In epidemiology, [1] proposed a Heckman-type selection model
to account for HIV survey nonparticipation and HIV status. Variants of the
Heckman selection model have also been proposed (see a comprehensive survey
[21]). In the area of fair machine learning, [4] applied the Heckman model to
correct MNAR sample selection bias while achieving fairness for linear
regression models. Very recently, [10] extended the Heckman model to model
multiple domain-specific sample selection mechanisms and proposed to jointly
learn the prediction model and the domain selection model to achieve
generalization on the true population. For approaches that incorporate the
Heckman selection model and its variants, the prediction features are manually
chosen beforehand. Our work makes a data-driven choice for the prediction
features based on a trained feature assignment function.

Empirical analysis has often been used to examine the effect of exclusion
restrictions on the performance of the Heckman selection model. [14] conducted
Monte Carlo experiments with and without exclusion restrictions and found that
the Heckman model is susceptible to collinearity issues in the absence of
exclusion restrictions. Recently, [23] conducted a simulation showing that
naive ordinary least squares (OLS) on the biased training set outperforms the
Heckman model based approaches that do not have a valid exclusion restriction.

For the Heckman selection model, the correlation between the two noise terms
carries information about the sample selection process. Because the noise
terms are unobserved, the true value of the correlation is unknown. [6] used a
range of values for the correlation to derive uncertainty intervals for the
parameters of the prediction model. In our work, we also consider a range of
values for the correlation. However, rather than defining the range of
correlation values for a fixed set of prediction features, we specify the
range of correlation values as we dynamically assign prediction features for
the Heckman selection model.

The idea of variable assignment based on reparametrized random samples has
also been explored in previous work. For instance, [15] proposed to train an
assignment function that maps each latent variable to some causal factor as
part of learning a causal representation. Similar to our work, the assignment
function is trained by sampling an assignment from the Gumbel-Softmax
distribution. In our work, however, we do not map variables to causal factors.
Instead, we map each variable to a value that indicates whether or not the
variable is assigned as a feature for the predictive model.

###  II-B Learning under Biased Training Data

Machine learning on missing training data has been well-studied. There are a
number of techniques that handle MAR sample selection bias in the training
set. Importance weighting [20] is commonly used to handle the MAR setting to
reweigh the training samples. However, it can result in inaccurate estimates
due to the influence of data instances with large importance weights on the
reweighted loss. To address this drawback, recent MAR approaches have been
constructed based on distributionally robust optimization of the reweighted
loss. [3] proposed a robust bias-aware regression approach, which considers a
worst-case shift between the training and testing sets to adversarially
minimize the reweighted loss. [19] introduced Rockafellar-Uryasev (RU)
regression to produce a model that is robust against bounded MAR sample
selection bias, where the level of distribution shift between the training and
testing sets is restricted. Based on the assumption for MNAR sample selection
bias, methods that account for MAR bias would not properly model the MNAR data
mechanism. Thus we expect methods that handle MAR bias to not be robust
against MNAR sample selection bias. On the other hand, our approach uses the
Heckman selection model to model the MNAR data mechanism, where the selection
of samples is expressed as a linear equation. Moreover, unlike the above MAR
approaches, we consider observed features of training samples with missing
outcomes when modeling the missing data mechanism.

In particular, there are recent approaches that address the problem of MNAR
labels in training data. In recommender learning, [22] proposed the joint
learning of an imputation model and a prediction model to estimate the
performance of rating prediction given MNAR ratings. [12] adopted two
propensity scoring methods into its loss function to handle bias of MNAR
feedback, where user feedback for unclicked data can be negative or positive-
yet-unobserved. While the approaches in [12] and [22] also use separate
propensity estimation models to predict the observation of a label, they
consider matrix factorization as the prediction model, which is not for linear
regression on tabular data.

In semi-supervised learning, [8] employed class-aware propensity score and
imputation strategies on the biased training set toward developing a semi-
supervised learning model that is doubly robust against MNAR data. We
emphasize that our problem setting is different than semi-supervised learning.
In semi-supervised learning, unlabeled samples are separated into clusters
based on similarities. However, in our problem setting, we do not perform
clustering on the samples with missing labels.

##  III Heckman Selection Model Revisited

Formally, the Heckman selection model [7] models the selection and prediction
equations as follows. For any
$(\bm{x}_{i},y_{i})\in\mathcal{X}\times\mathcal{Y}$, the selection equation of
the $i$th sample is

$d_{i}=\bm{x}_{i}^{(s)}\bm{\gamma}+u_{i}^{(s)}$ (2)

where $\bm{\gamma}$ is the set of regression coefficients for selection,
$\bm{x}_{i}^{(s)}$ denotes the features for sample selection, and
$u_{i}^{(s)}\sim\mathcal{N}(0,1)$ is the noise term for the selection
equation. The selection value of the $i$th sample $s_{i}$ is defined as:

$s_{i}=\left\\{\begin{aligned} &1\hskip 14.22636ptd_{i}>0\\\ &0\hskip
14.22636ptd_{i}\leq 0\\\ \end{aligned}\right.$ (3)

The model learns the selection based on Eq. (3) and the prediction of the
$i$th sample based on linear regression, with

$y_{i}=\bm{x}_{i}^{(p)}\bm{\beta}+u_{i}^{(p)}$ (4)

Assuming $u_{i}^{(p)}\sim\mathcal{N}(0,\sigma^{2})$, we define
$u_{i}^{(p)}=\sigma\epsilon_{i}$ where $\epsilon_{i}\sim\mathcal{N}(0,1)$.
Moreover, the noise terms $u_{i}^{(p)}$ and $u_{i}^{(s)}$ are correlated with
a correlation coefficient of $\rho$. If $\rho$ differs from zero, there is an
indication that the missing observations are MNAR.

To correct the bias in $\mathcal{D}_{tr}$, the conditional expectation of the
predicted outcome

$\begin{split}\mathbb{E}[y_{i}|s_{i}=1]&=\mathbb{E}[y_{i}|d_{i}\geq 0]\\\
&=\mathbb{E}[\bm{x}_{i}^{(p)}\bm{\beta}+u_{i}^{(p)}|\bm{x}_{i}^{(s)}\bm{\gamma}+u_{i}^{(s)}>0]\\\
&=\bm{x}_{i}^{(p)}\bm{\beta}+\mathbb{E}[u_{i}^{(p)}|\bm{x}_{i}^{(s)}\bm{\gamma}+u_{i}^{(s)}>0]\\\
&=\bm{x}_{i}^{(p)}\bm{\beta}+\mathbb{E}[u_{i}^{(p)}|u_{i}^{(s)}>-\bm{x}_{i}^{(s)}\bm{\gamma}]\end{split}$
(5)

is computed over all samples in $\mathcal{D}_{s}$. Because
$u_{i}^{(p)}\sim\mathcal{N}(0,\sigma^{2})$ and
$u_{i}^{(s)}\sim\mathcal{N}(0,1)$ are correlated,

$\mathbb{E}[u_{i}^{(p)}|u_{i}^{(s)}>-\bm{x}_{i}^{(s)}\bm{\gamma}]=\lambda_{i}\rho\sigma$
(6)

where
$\displaystyle\lambda_{i}=\frac{\phi(-\bm{x}_{i}^{(s)}\bm{\gamma})}{1-\Phi(-\bm{x}_{i}^{(s)}\bm{\gamma})}=\frac{\phi(-\bm{x}_{i}^{(s)}\bm{\gamma})}{\Phi(\bm{x}_{i}^{(s)}\bm{\gamma})}$
is the inverse Mills ratio (IMR). We denote $\phi(\cdot)$ as the standard
normal density function and $\Phi(\cdot)$ as the standard normal cumulative
distribution function. Thus Eq. (5) is rewritten as

$\mathbb{E}[y_{i}|s_{i}=1]=\bm{x}_{i}^{(p)}\bm{\beta}+\lambda_{i}\rho\sigma=\underline{\bm{x}}^{(p)}_{i}\underline{\bm{\beta}}$
(7)

where $\underline{\bm{x}}^{(p)}_{i}=[\bm{x}^{(p)}_{i},\lambda_{i}]$,
$\underline{\bm{\beta}}=[\bm{\beta},\beta_{H}]$, and $\beta_{H}=\rho\sigma$.

We provide pseudocode of executing the Heckman model in Algorithm 1. The
Heckman model follows two steps:

Step 1. Since there is no prior knowledge of the true value of $\lambda_{i}$
for each sample in $\mathcal{D}_{s}$, $\lambda_{i}$ is estimated as
$\hat{\lambda}_{i}$ by first computing $\hat{\bm{\gamma}}$ using probit
regression over $\mathcal{D}_{tr}$. As indicated in line 2 of Algorithm 1, we
estimate $\hat{\bm{\gamma}}$ after maximizing

$\begin{split}\mathcal{L}(\bm{\gamma})&=\prod_{i=1}^{n}P(s_{i}=1)^{s_{i}}\cdot
P(s_{i}=0)^{1-s_{i}}\\\
&=\prod_{i=1}^{n}\Phi(\bm{x}^{(s)}_{i}\bm{\gamma})^{s_{i}}(1-\Phi(\bm{x}^{(s)}_{i}\bm{\gamma}))^{1-s_{i}}\end{split}$
(8)

over $\mathcal{D}_{tr}$, where $\mathcal{L}(\cdot)$ is the likelihood. As
shown in line 4, $\hat{\bm{\gamma}}$ is then used to compute
$\hat{\lambda}_{i}$ for each $\bm{t}_{i}$ in $\mathcal{D}_{s}$.

Step 2. Using $\hat{\lambda}_{i}$, the prediction model is

$\begin{split}\hat{y}_{i}&=\bm{x}_{i}^{(p)}\bm{\beta}+\hat{\lambda}_{i}\beta_{H}\end{split}$
(9)

The estimated set of coefficients $\underline{\hat{\bm{\beta}}}$ is computed
by minimizing $\displaystyle\sum_{i=1}^{m}(y_{i}-\hat{y}_{i})^{2}$ over
$\mathcal{D}_{s}$. As a result, $\underline{\hat{\bm{\beta}}}$ is computed
using the closed-form solution

$\underline{\hat{\bm{\beta}}}=(\underline{\bm{X}}^{(p)T}\underline{\bm{X}}^{(p)})^{-1}\underline{\bm{X}}^{(p)T}\bm{y}$
(10)

as indicated in line 7.

Algorithm 1 Heckman

Input:  Training set
$\mathcal{D}_{tr}=\\{(\bm{x}_{i},y_{i},s_{i}=1)\\}_{i=1}^{m}\cup\\{(\bm{x}_{i},s_{i}=0)\\}_{i=m+1}^{n}$,
selection feature matrix $\bm{X}^{(s)}$, prediction feature matrix
$\bm{X}^{(p)}$

Output: estimated Heckman coefficients $\underline{\hat{\bm{\beta}}}$, IMR
$\\{\lambda_{i}\\}_{i=1}^{m}$

1: $\mathcal{D}_{s}\leftarrow\\{(\bm{x}_{i},y_{i},s_{i}=1)\\}_{i=1}^{m}$

2: Estimate $\hat{\bm{\gamma}}$ after maximizing Eq. (8) on $\mathcal{D}_{tr}$
using $\bm{X}^{(s)}$

3: for $\bm{t}_{i}\in\mathcal{D}_{s}$ do

4:
$\hat{\lambda}_{i}\leftarrow\frac{\phi(-\bm{x}_{i}^{(s)}\hat{\bm{\gamma}})}{\Phi(\bm{x}_{i}^{(s)}\hat{\bm{\gamma}})}$

5:
$\underline{\bm{x}}_{i}^{(p)}\leftarrow[\bm{x}_{i}^{(p)},\hat{\lambda}_{i}]$

6: end for

7: Obtain $\underline{\hat{\bm{\beta}}}$ by evaluating Eq. (10) on
$\mathcal{D}_{s}$ using $\underline{\bm{X}}^{(p)}$

8: return $\underline{\hat{\bm{\beta}}}$ and
$\\{\hat{\lambda}_{i}\\}_{i=1}^{m}$

Exclusion Restriction. The Heckman selection model assumes that the selection
features consist of every prediction feature and additional features that do
not affect the outcome. The selection and prediction features are generally
not identical as it can introduce multicollinearity to the prediction model.
Specifically, if the selection and prediction features are the same, the IMR
$\lambda_{i}$ would be expressed as

$\begin{split}\lambda_{i}&=\frac{\phi(-\bm{x}_{i}^{(s)}\bm{\gamma})}{\Phi(\bm{x}_{i}^{(s)}\bm{\gamma})}=\frac{\phi(-\bm{x}_{i}^{(p)}\bm{\gamma})}{\Phi(\bm{x}_{i}^{(p)}\bm{\gamma})}\end{split}$
(11)

This would mean that Eq. (7) is properly identified through the nonlinearity
of $\lambda_{i}$. However, $\lambda_{i}$ is roughly linear over a wide range
of values for $\bm{x}^{(s)}_{i}\bm{\gamma}$ [18]. Hence, Step 2 of the Heckman
model would yield unrobust estimates for $\underline{\hat{\bm{\beta}}}$ due to
the multicollinearity between $\bm{x}^{(p)}_{i}$ and $\lambda_{i}$.

##  IV Methodology

When choosing prediction features from a set of $K$ selection features, we
encounter the following challenges. First, there are $2^{K}-1$ possible
choices to make for the set of prediction features. For any dataset that has a
large number of selection features, searching for a suitable set of prediction
features becomes computationally expensive. Second, the Heckman selection
model does not perform well for exclusion restrictions that are not valid. In
other words, some choices for the set of prediction features are not helpful
when using the Heckman model to handle MNAR sample selection bias on the
outcome.

We introduce Heckman-FA, a framework for using the Heckman selection model via
a learned feature assignment. Heckman-FA first learns the weights of an
assignment function $\psi$ that draws samples from the Gumbel-Softmax
distribution [9]. This function outputs an assignment of prediction features
given a matrix $\bm{\pi}$ of probabilities of including a selection feature in
the prediction model. The framework then uses this assignment of prediction
features to run the Heckman model and compute the mean absolute error (MAE) of
predictions on $\mathcal{D}_{s}$ when using the Heckman selection model. To
optimize $\psi$, we minimize the MAE with respect to $\bm{\pi}$. This results
in an estimated probability matrix $\hat{\bm{\pi}}$. To extract the prediction
features, Heckman-FA looks through different prediction feature matrices by
sampling from the Gumbel-Softmax distribution using $\hat{\bm{\pi}}$. When
determining the extracted prediction feature matrix, we first consider whether
or not the estimated correlation between the noise terms is within a range
that is user-defined based on prior domain knowledge. We further consider
goodness-of-fit to ensure that the prediction model is of quality. Using the
selection and extracted prediction features, we run the Heckman model to fit a
robust prediction model under MNAR sample selection bias on the outcome.

###  IV-A Assignment Function

We formally introduce an assignment function $\psi(k)$, which is defined as

$\begin{split}\psi(k)&=\begin{cases}1&k\text{th selection feature is
assigned}\\\ 0&k\text{th selection feature is not
assigned}\end{cases}\end{split}$ (12)

In general, an assignment function determines which of the $K$ selection
features are also features in the prediction equation.

Assignment Computation. In Algorithm 2, we provide the pseudocode for
computing a matrix $\bm{X}^{(p)}$ of assigned prediction features from the
selection feature matrix $\bm{X}^{(s)}$. We assume that both $\bm{X}^{(p)}$
and $\bm{X}^{(s)}$ have $K$ columns. However, we define the $J$ features
assigned for prediction as the columns in $\bm{X}^{(p)}$ that are not equal to
the zero vector $\bm{0}$.

To obtain $\bm{X}^{(p)}$ based on selection features, we compute the
assignment $\psi(k)$ for the $k$th selection feature by drawing samples from a
categorical distribution. Let $\bm{Z}\in\mathbb{R}^{2\times K}$ be a
categorical sample matrix such that each element is either 0 or 1. The Gumbel-
Max trick is used to efficiently draw categorical samples. Following the steps
of [9], a sample $\bm{Z}_{\cdot k}$, the $k$th column of $\bm{Z}$, is drawn
from a categorical distribution with class probabilities $\pi_{1k}$ and
$\pi_{2k}$, where $\pi_{2k}$ ($\pi_{1k}$) is the probability that the $k$th
selection feature is (not) assigned to the prediction model. Formally,
$\bm{Z}_{\cdot k}$ is expressed as

$\bm{Z}_{\cdot
k}=\text{one\\_hot}(\underset{q}{\text{argmax}}[g_{qk}+\log\pi_{qk}])$ (13)

where $g_{qk}\sim\text{ Gumbel}(0,1)$ and $q\in\\{1,2\\}$. Thus, in terms of
$\bm{Z}_{\cdot k}$, we express the assignment $\psi$ for the $k$th selection
feature as

$\begin{split}\psi(k;\pi)&=\begin{bmatrix}0\\\ 1\\\
\end{bmatrix}^{T}\bm{Z}_{\cdot k}\end{split}$ (14)

and compute

$\begin{split}x_{ik}^{(p)}&=x^{(s)}_{ik}\odot\psi(k;\bm{\pi})\end{split}$ (15)

for each selection feature as indicated in lines 5 and 6 of Algorithm 2, using
$\odot$ to denote elementwise multiplication.

Algorithm 2 FeatureAssign

Input:  Training set
$\mathcal{D}_{tr}=\\{(\bm{x}_{i},y_{i},s_{i}=1)\\}_{i=1}^{m}\cup\\{(\bm{x}_{i},s_{i}=0)\\}_{i=m+1}^{n}$,
number of selection features $K$, probability matrix $\bm{\pi}$

Output: prediction feature matrix $\bm{X}^{(p)}$

1: $\mathcal{D}_{s}\leftarrow\\{(\bm{x}_{i},y_{i},s_{i}=1)\\}_{i=1}^{m}$

2: Compute $\bm{Z}$ using Eq. (13) based on $\bm{\pi}$

3: for all $\bm{t}_{i}\in\mathcal{D}_{s}$ do

4: for $k\in\\{1,\ldots,K\\}$ do

5: $\psi(k;\bm{\pi})\leftarrow\begin{bmatrix}0\\\ 1\\\
\end{bmatrix}^{T}\bm{Z}_{\cdot k}$

6: $x_{ik}^{(p)}\leftarrow x^{(s)}_{ik}\odot\psi(k;\bm{\pi})$

7: end for

8: end for

9: return $\bm{X}^{(p)}$

Backpropagation. To train the assignment function $\psi$, we learn
$\hat{\pi}_{1k}$ and $\hat{\pi}_{2k}$ for each selection feature. However,
$\bm{Z}_{\cdot k}$ is expressed in terms of argmax, which is non-
differentiable. Thus we cannot derive $\nabla_{\bm{\pi}}\bm{Z}$ in order to
learn $\hat{\bm{\pi}}$. On the other hand, we can use the Gumbel-Softmax
distribution [9] to approximate the categorical sample $\bm{Z}_{\cdot k}$. By
the Straight-Through Gumbel Estimator, we compute
$\tilde{\bm{Z}}\in\mathbb{R}^{2\times K}$ as a continuous, differentiable
approximation of argmax, where

$\tilde{z}_{qk}=\frac{\text{exp}((\log\pi_{qk}+g_{qk})/\tau)}{\sum_{r=1}^{2}\text{exp}((\log\pi_{rk}+g_{rk})/\tau)}$
(16)

with $\tau$ as the softmax temperature. Notice that as $\tau\rightarrow 0$,
Eq. (16) will approximate the argmax function in Eq. (13), and the Gumbel-
Softmax sample vector will approach a one-hot vector. Hence
$\nabla_{\bm{\pi}}\bm{Z}\approx\nabla_{\bm{\pi}}\tilde{\bm{Z}}$, where

$\begin{split}\frac{\partial\tilde{z}_{qk}}{\partial\pi_{qk}}&=\frac{\prod_{r=1}^{2}\text{exp}((\log\pi_{rk}+g_{rk})/\tau)}{(\sum_{r=1}^{2}\text{exp}((\log\pi_{rk}+g_{rk})/\tau))^{2}}\cdot\frac{1}{\tau\pi_{qk}}\end{split}$
(17)

Therefore, although we use Eq. (14) to obtain assignments, we use
$\tilde{\bm{Z}}$ instead of $\bm{Z}$ when performing backpropagation to update
$\bm{\pi}$.

Based on Eq. (17), $\nabla_{\bm{\pi}}\tilde{\bm{Z}}$ is well-defined, so we
are able to learn the assignment function and estimate parameters
$\hat{\bm{\pi}}$. Formally, we compute a probability matrix $\hat{\bm{\pi}}$
such that

$\small\begin{split}\hat{\bm{\pi}}&=\underset{\bm{\pi}}{\text{argmin}}\,\mathcal{L}_{MAE}\\\
&=\underset{\bm{\pi}}{\text{argmin}}\left(\frac{1}{m}\sum_{i=1}^{m}\bigg{|}y_{i}-(\bm{x}_{i}^{(p)}\bm{\beta}+\hat{\lambda}_{i}\beta_{H})\bigg{|}\right)\end{split}$
(18)

where $x_{ik}^{(p)}=x^{(s)}_{ik}\odot\psi(k;\bm{\pi})$ for the $k$th selection
feature. To ensure that the Heckman model is a quality prediction model, we
consider the predictive performance of the Heckman model on $\mathcal{D}_{s}$
to learn the assignment function. In this work, we choose to minimize the MAE
over $\mathcal{D}_{s}$ in order to obtain $\hat{\bm{\pi}}$. For
$\nabla_{\bm{\pi}}\mathcal{L}_{MAE}$, using on Eq. (17), we see that

$\small\begin{split}\frac{\partial\mathcal{L}_{MAE}}{\partial\pi_{qk}}&=\frac{\partial\mathcal{L}_{MAE}}{\partial\hat{y}_{i}}\cdot\frac{\partial\hat{y}_{i}}{\partial
x^{(p)}_{ik}}\cdot\frac{\partial
x^{(p)}_{ik}}{\partial\psi(k)}\cdot\frac{\partial\psi(k)}{\partial\tilde{z}_{qk}}\cdot\frac{\partial\tilde{z}_{qk}}{\partial\pi_{qk}}\\\
&=-\frac{1}{m}\sum_{i=1}^{m}\frac{y_{i}-\hat{y}_{i}}{|y_{i}-\hat{y}_{i}|}\cdot\beta_{k}\cdot
x^{(s)}_{ik}\cdot\frac{\partial\psi(k)}{\partial\tilde{z}_{qk}}\\\
&\,\,\cdot\frac{\prod_{r=1}^{2}\text{exp}((\log\pi_{rk}+g_{rk})/\tau)}{(\sum_{r=1}^{2}\text{exp}((\log\pi_{rk}+g_{rk})/\tau))^{2}}\cdot\frac{1}{\tau\pi_{qk}}\end{split}$
(19)

where $\displaystyle\frac{\partial\psi(k)}{\partial\tilde{z}_{2k}}=1$. Thus
$\nabla_{\bm{\pi}}\mathcal{L}_{MAE}$ is well-defined. Moreover,
$\nabla_{\bm{\pi}}\mathcal{L}_{MAE}$ is not a zero matrix, meaning that
$\bm{\pi}$ is updated during backpropagation.

Discussion. Although other metrics such as mean squared error (MSE) and root
mean squared error (RMSE) are typically used to evaluate the performance of
linear regression, we cannot minimize these metrics in order to obtain
$\hat{\bm{\pi}}$. We show that if our objective is to minimize the MSE or
RMSE, then $\bm{\pi}$ does not change at all during backpropagation. As an
example, consider $\mathcal{L}_{MSE}$, which denotes the MSE loss function.
For $\nabla_{\bm{\pi}}\mathcal{L}_{MSE}$, we see that

$\small\begin{split}\frac{\partial\mathcal{L}_{MSE}}{\partial\pi_{qk}}&=\frac{\partial\mathcal{L}_{MSE}}{\partial\hat{y}_{i}}\cdot\frac{\partial\hat{y}_{i}}{\partial
x^{(p)}_{ik}}\cdot\frac{\partial
x^{(p)}_{ik}}{\partial\psi(k)}\cdot\frac{\partial\psi(k)}{\partial\tilde{z}_{qk}}\cdot\frac{\partial\tilde{z}_{qk}}{\partial\pi_{qk}}\\\
&=\frac{\partial\mathcal{L}_{MSE}}{\partial\psi(k)}\cdot\frac{\partial\psi(k)}{\partial\tilde{z}_{qk}}\cdot\frac{\partial\tilde{z}_{qk}}{\partial\pi_{qk}}\\\
\end{split}$ (20)

where

$\small\begin{split}\frac{\partial\mathcal{L}_{MSE}}{\partial\psi(k)}&=-\frac{2}{m}\sum_{i=1}^{m}(y_{i}-\hat{y}_{i})\cdot\beta_{k}\cdot
x^{(s)}_{ik}\\\ &=-\frac{2}{m}\beta_{k}\sum_{i=1}^{m}(y_{i}-\hat{y}_{i})\cdot
x^{(s)}_{ik}\\\ &=-\frac{2}{m}\beta_{k}\sum_{i=1}^{m}v_{i}\cdot
x^{(s)}_{ik}\end{split}$ (21)

where $v_{i}$ is the error of predicting the outcome of
$\bm{t}_{i}\in\mathcal{D}_{s}$ using the Heckman selection model. Given that
$\bm{\beta}$ has length $K$, consider two cases on the $k$th selection
feature. First, for any $k$th selection feature not assigned for prediction,
$\bm{X}^{(p)}_{\cdot k}=\bm{0}$, where $\bm{X}^{(p)}_{\cdot k}$ is the $k$th
column of $\bm{X}^{(p)}$. Thus $\beta_{k}=0$ after running the Heckman
selection model and hence
$\displaystyle\frac{\partial\mathcal{L}_{MSE}}{\partial\psi(k)}=0$. Second,
for any $k$th selection feature assigned for prediction,
$x^{(s)}_{ik}=x^{(p)}_{ik}$ for all $\bm{t}_{i}\in\mathcal{D}_{s}$. Thus

$\small\begin{split}\frac{\partial\mathcal{L}_{MSE}}{\partial\psi(k)}=-\frac{2}{m}\beta_{k}\sum_{i=1}^{m}v_{i}\cdot
x^{(s)}_{ik}&=-\frac{2}{m}\beta_{k}\sum_{i=1}^{m}v_{i}\cdot x^{(p)}_{ik}\\\
&=-\frac{2}{m}\beta_{k}\left(\bm{v}^{T}\bm{X}^{(p)}_{\cdot k}\right)\\\
\end{split}$ (22)

where $\bm{v}$ is the error vector. Now $\bm{X}^{(p)}_{\cdot k}$ is considered
part of the input space that is used to compute $\hat{\bm{y}}$ and in turn
$\bm{v}$. Because the values in $\bm{\beta}$ are estimated by minimizing
$\bm{v}^{T}\bm{v}$, then $\bm{v}$ is orthogonal to the columns of the input
space used to compute $\hat{\bm{y}}$, which includes $\bm{X}^{(p)}_{\cdot k}$.
Thus $\bm{v}^{T}\bm{X}^{(p)}_{\cdot k}=0$, and
$\displaystyle\frac{\partial\mathcal{L}_{MSE}}{\partial\psi(k)}=0$ for all
selection features assigned for prediction. Based on these two cases,
$\displaystyle\frac{\partial\mathcal{L}_{MSE}}{\partial\pi_{qk}}=0$ for all
$K$ selection features. Hence $\nabla_{\bm{\pi}}\mathcal{L}_{MSE}$ is a zero
matrix.

###  IV-B Extraction of Suitable Assignment

After we train the assignment function $\psi$, we utilize the estimated
parameters $\hat{\bm{\pi}}$ to extract a suitable set of features for the
prediction model. We propose a sampling-based strategy to extract
$\bm{X}^{(p)}$ using $\hat{\bm{\pi}}$. We base this strategy on the estimated
correlation $\hat{\rho}$ between the noise terms $u^{(s)}_{i}$ and
$u^{(p)}_{i}$ and the adjusted $R^{2}$ value, denoted as $R^{2}_{a}$.
Specifically, a suitable assignment of prediction features is extracted based
on a user-defined range of values for the correlation. The range of values is
provided based on prior domain knowledge on a given dataset. Moreover, we
observe that multiple prediction feature assignments can correspond to a
correlation that is within the user-defined range. To further decide on which
prediction feature matrix to extract, we also consider goodness-of-fit to
ensure that the prediction model is of quality.

In general, $\rho$ carries information about the nature of the sample
selection process. However, because $u^{(s)}_{i}$ and $u^{(p)}_{i}$ are
unobserved for all $\bm{t}_{i}$, the true value of $\rho$ is unknown for a
dataset. We must instead consider the estimated correlation $\hat{\rho}$. To
compute $\hat{\rho}$, we note that a consistent estimator of $\sigma^{2}$ can
be derived based on the Heckman selection model. First, define
$v_{i}=y_{i}-\hat{y}_{i}$ as the error of predicting the outcome of the $i$th
sample using the Heckman selection model. The true conditional variance of
$v_{i}$ is

$\small\begin{split}\mathbb{E}[v^{2}_{i}|s_{i}=1]&=\sigma^{2}(1+\rho^{2}(\lambda_{i}(-\bm{x}^{(s)}_{i}\bm{\gamma})-\lambda^{2}_{i}))\\\
&=\sigma^{2}+(\rho\sigma)^{2}(\lambda_{i}(-\bm{x}^{(s)}_{i}\bm{\gamma})-\lambda^{2}_{i})\end{split}$
(23)

Consider the average conditional variance over all samples in
$\mathcal{D}_{s}$, where

$\small\begin{split}\text{plim
}\frac{1}{m}\sum_{i=1}^{m}\mathbb{E}[v^{2}_{i}|s_{i}=1]&=\sigma^{2}\left(1+\frac{\rho^{2}}{m}\sum_{i=1}^{m}\lambda_{i}(-\bm{x}^{(s)}_{i}\gamma)-\lambda^{2}_{i}\right)\\\
&=\sigma^{2}+\frac{(\rho\sigma)^{2}}{m}\sum_{i=1}^{m}\lambda_{i}(-\bm{x}^{(s)}_{i}\gamma)-\lambda^{2}_{i}\end{split}$
(24)

with plim denoting convergence in probability. Now the average conditional
variance over all samples in $\mathcal{D}_{s}$ is estimated using the mean
squared error

$\begin{split}\frac{1}{m}\sum_{i=1}^{m}v^{2}_{i}&=\frac{1}{m}\sum_{i=1}^{m}(y_{i}-\hat{y}_{i})^{2}\end{split}$
(25)

of predicting the outcome using the Heckman model. Thus, using Eq. (24) and
(25), $\sigma^{2}$ can be estimated by computing

$\small\begin{split}\hat{\sigma}^{2}&=\frac{1}{m}\sum_{i=1}^{m}(y_{i}-\hat{y}_{i})^{2}-\frac{\hat{\beta}_{H}^{2}}{m}\sum_{i=1}^{m}\hat{\lambda}_{i}(-\bm{x}^{(s)}_{i}\hat{\bm{\gamma}})-\hat{\lambda}^{2}_{i}\end{split}$
(26)

where $\hat{\beta}_{H}$ is an estimate of $\rho\sigma$. Therefore, using Eq.
(26), we obtain $\displaystyle\hat{\rho}=\hat{\beta}_{H}/\hat{\sigma}$.

We consider a set of prediction features to be suitable for the Heckman
selection model if $\hat{\rho}$ is within the range $[\rho_{min},\rho_{max}]$,
where the values of $\rho_{min}$ and $\rho_{max}$ are user-specified. We work
with this user-specified range because the true value of $\rho$ is unknown. We
expect that the values of $\rho_{min}$ and $\rho_{max}$ are appropriately
chosen to indicate that the Heckman model properly handles MNAR sample
selection bias. On one hand, the range $[\rho_{min},\rho_{max}]$ should not
contain 0 since $\rho=0$ indicates MAR sample selection bias. On the other
hand, $\rho_{min}$ and $\rho_{max}$ should not be too negative or too positive
since strong correlation renders the Heckman model unstable [17].

Because there may be multiple values of $\hat{\rho}$ that are in
$[\rho_{min},\rho_{max}]$, we also obtain $R^{2}_{a}$ by computing

$\begin{split}R^{2}_{a}&=1-\frac{(1-R^{2})(m-1)}{m-J-1}\end{split}$ (27)

where $R^{2}$ is the coefficient of determination and $J$ is the number of
prediction features. Having the largest number of prediction features such
that $\hat{\rho}$ is in $[\rho_{min},\rho_{max}]$ does not imply that the
linear model is the best fitted. Thus $R^{2}_{a}$ is helpful since it also
factors in the number of prediction features when measuring the goodness-of-
fit for a linear model.

Algorithm 3 Extraction

Input:  Training set
$\mathcal{D}_{tr}=\\{(\bm{x}_{i},y_{i},s_{i}=1)\\}_{i=1}^{m}\cup\\{(\bm{x}_{i},s_{i}=0)\\}_{i=m+1}^{n}$,
number of selection features $K$, estimated parameters $\hat{\bm{\pi}}$, lower
correlation threshold $\rho_{min}$, upper correlation threshold $\rho_{max}$,
number of Gumbel-Softmax samples $B$

Output:  Prediction feature matrix $\bm{X}^{(p)}$

1: $R^{2*}_{a}\leftarrow-\infty$

2: for $B$ iterations do

3: $\bm{X}^{(p)}_{temp}\leftarrow$ FeatureAssign($\mathcal{D}_{tr}$, $K$,
$\hat{\bm{\pi}}$)

4: $\underline{\hat{\bm{\beta}}},\\{\hat{\lambda}_{i}\\}_{i=1}^{m}\leftarrow$
Heckman($\mathcal{D}_{tr}$, $\bm{X}^{(s)}$, $\bm{X}^{(p)}_{temp}$)

5: Compute $\hat{\sigma}^{2}$ using Eq. (26)

6: $\hat{\rho}\leftarrow\hat{\beta}_{H}/\hat{\sigma}$

7: Compute $R^{2}_{a}$ using Eq. (27)

8: if $\hat{\rho}\in[\rho_{min},\rho_{max}]$ and $R^{2}_{a}>R^{2*}_{a}$ then

9: $\bm{X}^{(p)}\leftarrow\bm{X}^{(p)}_{temp}$

10: $R^{2*}_{a}\leftarrow R^{2}_{a}$

11: end if

12: end for

13: return $\bm{X}^{(p)}$

As our strategy to extract $\bm{X}^{(p)}$ for the Heckman selection model, we
propose to collect $B$ Gumbel-Softmax samples based on $\hat{\bm{\pi}}$. Out
of the $B$ samples, we choose the prediction feature matrix such that the
prediction model has the highest value of $R^{2}_{a}$ given that $\hat{\rho}$
is in $[\rho_{min},\rho_{max}]$. The pseudocode for this process is listed in
Algorithm 3. In lines 2-12, we iteratively look for $\bm{X}^{(p)}$. In line 3,
we sample $\bm{X}^{(p)}_{temp}$ by executing Algorithm 2 based on
$\hat{\bm{\pi}}$. In line 4, we execute the Heckman model based on
$\underline{\bm{X}}^{(p)}_{temp}$ and get $\underline{\hat{\bm{\beta}}}$ and
$\\{\hat{\lambda}_{i}\\}_{i=1}^{m}$. In line 5, we compute
$\hat{\sigma}^{(p)2}$ using Eq. (26). In line 6, we calculate $\hat{\rho}$ by
dividing $\hat{\beta}_{H}$ by $\hat{\sigma}$. In line 7, we compute
$R^{2}_{a}$. Throughout each iteration, we check whether or not $\hat{\rho}$
is in $[\rho_{min},\rho_{max}]$ and $R^{2}_{a}$ is larger than the current
largest $R^{2*}_{a}$ value. If this condition is satisfied, then we update
$\bm{X}^{(p)}$, as indicated in line 9. In line 13, $\bm{X}^{(p)}$ is
returned.

Algorithm 4 Heckman-FA

Input:  Training set
$\mathcal{D}_{tr}=\\{(\bm{x}_{i},y_{i},s_{i}=1)\\}_{i=1}^{m}\cup\\{(\bm{x}_{i},s_{i}=0)\\}_{i=m+1}^{n}$,
number of selection features $K$, initial fixed value $c$, number of training
epochs $T$, learning rate $\alpha$, lower correlation threshold $\rho_{min}$,
upper correlation threshold $\rho_{max}$, number of Gumbel-Softmax samples $B$

Output:  Augmented prediction feature matrix $\underline{\bm{X}}^{(p)}$,
estimated Heckman coefficients $\underline{\hat{\bm{\beta}}}$

1: for $K$ selection features do

2: Initialize $\pi_{2k}=c$ and $\pi_{1k}=1-c$

3: end for

4: for $T$ epochs do

5: $\bm{X}^{(p)}\leftarrow$ FeatureAssign($\mathcal{D}_{tr}$, $K$, $\bm{\pi}$)

6: $\underline{\hat{\bm{\beta}}},\\{\hat{\lambda}_{i}\\}_{i=1}^{m}\leftarrow$
Heckman($\mathcal{D}_{tr}$, $\bm{X}^{(s)}$, $\bm{X}^{(p)}$)

7: Compute $\mathcal{L}_{MAE}$ using $\underline{\hat{\bm{\beta}}}$ and
$\\{\hat{\lambda}_{i}\\}_{i=1}^{m}$

8:
$\bm{\pi}\leftarrow\bm{\pi}-\displaystyle\alpha\nabla_{\bm{\pi}}\mathcal{L}_{MAE}$

9: $\hat{\bm{\pi}}\leftarrow\bm{\pi}$

10: end for

11: $\bm{X}^{(p)}\leftarrow$ Extraction($\mathcal{D}_{tr}$, $K$,
$\hat{\bm{\pi}}$, $\rho_{min}$, $\rho_{max}$, $B$)

12: $\underline{\hat{\bm{\beta}}},\\{\hat{\lambda}_{i}\\}_{i=1}^{m}\leftarrow$
Heckman($\mathcal{D}_{tr}$, $\bm{X}^{(s)}$, $\bm{X}^{(p)}$)

13:
$\underline{\bm{X}}^{(p)}\leftarrow[\bm{X}^{(p)},\\{\hat{\lambda}_{i}\\}_{i=1}^{m}]$

14: return $\underline{\bm{X}}^{(p)}$ and $\underline{\hat{\bm{\beta}}}$

TABLE II: Features used for selection. Dataset | Selection Features
---|---
CRIME | householdsize, racepctblack, racePctWhite,
| racePctAsian, racePctHisp, agePct12t21, agePct12t29,
| agePct16t24, agePct65up, numbUrban, pctUrban,
| medIncome, pctWWage, pctWFarmSelf, pctWInvInc,
| pctWSocSec, pctWPubAsst, pctWRetire, medFamInc,
| perCapInc, whitePerCap, HispPerCap, blackPerCap,
| indianPerCap, AsianPerCap, OtherPerCap
COMPAS | sex, age, juv_fel_count, juv_misd_count,
| priors_count, two_year_recid, age_cat_25 - 45,
| age_cat_Greater than 45, race_Caucasian,
| c_charge_degree_M
TABLE III: Comparison of Heckman-FA against baselines using extracted prediction features. Extracted prediction features | $\hat{\rho}$ | $R^{2}_{a}$ | Naive | RU[19] | Heckman-FA
---|---|---|---|---|---
Train MSE | Test MSE | Train MSE | Test MSE | Train MSE | Test MSE
CRIME
1, 2, 3, 4, 5, 7, 8, |  |  |  |  |  |  |  |
11, 15, 16, 21, 22, 23, 24 | 0.0423 | 0.4780 | 0.0159 | 0.0218 | 0.0109 | 0.0212 | 0.0159 | 0.0203
COMPAS
1, 2, 5, 7, 8, 9, 10 | 0.2364 | 0.3453 | 0.0757 | 0.2560 | 0.0730 | 0.2589 | 0.0757 | 0.2506

###  IV-C Heckman Selection Model with Feature Assignment

Algorithm 4 gives the pseudocode of Heckman-FA. In lines 1-3, we initialize
$\bm{\pi}$. In this work, we use a fixed value $c\in(0,1)$ to initialize
$\pi_{2k}$ and $\pi_{1k}=1-\pi_{2k}$ for the $k$th selection feature. This
means that each selection feature has an equal probability of being assigned
as a prediction feature when we start training $\psi$. For $c$, we give users
some flexibility to choose which value to use. However, we suggest for users
to not use values that are extremely close to 0 or 1 for $c$. This ensures
that $\psi$ is trained on a variety of sets of prediction features based on
random Gumbel-Softmax samples. In lines 5-9, $\psi$ is trained over $T$
epochs. In line 5, we obtain the prediction feature matrix $\bm{X}^{(p)}$ by
executing Algorithm 2. In line 6, we execute the steps of the Heckman model to
get $\underline{\hat{\bm{\beta}}}$ and $\\{\hat{\lambda}_{i}\\}_{i=1}^{m}$. In
line 7, we use $\underline{\hat{\bm{\beta}}}$ and
$\\{\hat{\lambda}_{i}\\}_{i=1}^{m}$ to compute $\mathcal{L}_{MAE}$. In line 8,
we update $\bm{\pi}$ by computing $\nabla_{\pi}\mathcal{L}_{MAE}$ using Eq.
(19). In line 11, using $\hat{\bm{\pi}}$, we extract the assigned prediction
feature matrix $\bm{X}^{(p)}$ by calling Algorithm 3. In line 12, we run the
Heckman model using the extracted prediction features and obtain
$\underline{\hat{\bm{\beta}}}$ and $\\{\hat{\lambda}_{i}\\}_{i=1}^{m}$. After
having $\underline{\bm{X}}^{(p)}$ as the concatenation of $\bm{X}^{(p)}$ and
$\\{\hat{\lambda}_{i}\\}_{i=1}^{m}$ in line 13, Heckman-FA returns
$\underline{\bm{X}}^{(p)}$ and $\underline{\hat{\bm{\beta}}}$.

Heckman-FA*. We also present Heckman-FA* as alternative option to Heckman-FA,
where users extract $\bm{X}^{(p)}$ by simply ranking the selection features
based on the largest value of $\hat{\pi}_{2k}$ instead of using Algorithm 3.
In other words, we rank the likeliest selection features to be assigned as
prediction features based on the objective in Eq. (18) that computes
$\hat{\bm{\pi}}$. We then examine the first $J$ selection features in the
ranking for all $J\in\\{1,\ldots,K-1\\}$. Letting the first $J$ selection
features in the ranking be prediction features, we run the Heckman selection
model on the training set and obtain $\underline{\hat{\bm{\beta}}}$. We then
compute $\hat{\sigma}^{2}$, $\hat{\rho}$, and $R^{2}_{a}$ using the same steps
indicated in lines 5-7 in Algorithm 3. Finally, out of all
$J\in\\{1,\ldots,K-1\\}$, we choose the set of top $J$ selection features in
the ranking as prediction features such that
$\hat{\rho}\in[\rho_{min},\rho_{max}]$ and $R^{2}_{a}$ is at a maximum given
$\hat{\rho}$. As a result, the columns of $\bm{X}^{(p)}$ corresponding to
these features do not equal $\bm{0}$.

Computational Complexity. To derive the computational complexity of Algorithm
4 (Heckman-FA), we first have to consider the complexity of Algorithm 2. The
assignment computation takes $O(mK)$ time. During the training of $\psi$ in
lines 4-10, the assignment computation is repeated for $T$ epochs. Thus the
complexity of lines 4-10 is $O(nKT)$. Moreover, the complexity of Algorithm 3
is $O(nKB)$ as the assignment computation is repeated for $B$ Gumbel-Softmax
samples. Therefore, Heckman-FA has a computational complexity of $O(nK(T+B))$.

We also consider the complexity of Heckman-FA*. Similar to Heckman-FA, we
first see that $\psi$ is trained in $O(nKT)$ time when running Heckman-FA*.
However, the complexity of extraction is different for Heckman-FA* than for
Heckman-FA. Because the Heckman selection model is called for $K-1$ sets of
selection features, the extraction process has a complexity of $O(m(K-1))$.
Thus the complexity of Heckman-FA* is $O(nKT+m(K-1))$.

##  V Experiments

###  V-A Setup

We evaluate the performance of our framework on CRIME [5] and COMPAS [11]
datasets. The CRIME dataset contains socio-economic, law enforcement, and
crime information for $1,994$ communities in the United States. For each
community, we predict the total number of violent crimes committed (per
100,000 population). The COMPAS dataset consists of $5,278$ records collected
from defendants in Florida between 2013 and 2014. Given attributes such as
race, age, and priors count, we predict the defendant’s decile score.

For each dataset, we split the dataset to include 70% of samples in
$\mathcal{D}_{tr}$. We then construct $\mathcal{D}_{s}$ based on
$\mathcal{D}_{tr}$. For the CRIME dataset, we create sample selection bias in
$\mathcal{D}_{tr}$ by selecting communities such that the proportion of people
under poverty is less than 0.05. As a result, 976 out of 1,395 communities in
$\mathcal{D}_{tr}$ are in $\mathcal{D}_{s}$. For the COMPAS dataset, we select
defendants in $\mathcal{D}_{tr}$ who have a violent decile score of less than
10. As a result, there are 2,585 samples in $\mathcal{D}_{s}$. We provide the
set of selection features used for each dataset in Table II, where $K=26$ for
the CRIME dataset and $K=10$ for the COMPAS dataset.

Baselines and Hyperparameters. We compare our approach to the following
baselines: (1) naive linear regression (Naive) on $\mathcal{D}_{s}$ and (2)
Rockafellar-Uryasev (RU) regression [19], which involves fitting two neural
networks with the RU loss to train a robust model under bounded MAR sample
selection bias. In our experiments, we examine the RU regression baseline
where the distribution shift between training and testing sets are restricted
given $\Gamma=3$. Unlike Heckman-FA, the baselines do not have access to
features that model the selection of training samples. As a result, we should
expect the baselines to not be as effective as Heckman-FA in handling
regression under MNAR sample selection bias.

As we compute $\hat{\bm{\pi}}$ using Heckman-FA, we work with the following
hyperparameters. When initializing $\bm{\pi}$, we set $c=0.75$. We then train
$\psi$ for $T=4,000$ epochs with the softmax temperature $\tau=1$ for both
datasets. We set the learning rate $\alpha$ equal to 0.75 for the CRIME
dataset and 0.05 for the COMPAS dataset. For the extraction of prediction
features, we draw $B=1,000$ Gumbel-Softmax samples. In all experiments, the
range $[\rho_{min},\rho_{max}]$ is set to be $[0.01,0.1]$ for the CRIME
dataset and [0.1, 0.3] for the COMPAS dataset. All models are implemented
using Pytorch and executed on the Dell XPS 8950 9020 with a Nvidia GeForce RTX
3080 Ti.

TABLE IV: Results of paired $t$-test on the difference between the testing MSE of Heckman-FA and other baselines. Comparison | Test MSE Difference | P-value
---|---|---
CRIME
Heckman-FA - Naive | -0.0008 $\pm$ 0.0004 | 0.0001
Heckman-FA - RU | -0.0007 $\pm$ 0.0007 | 0.0053
COMPAS
Heckman-FA - Naive | -0.0060 $\pm$ 0.0008 | 7.546e-10
Heckman-FA - RU | -0.0093 $\pm$ 0.0012 | 1.039e-9

###  V-B Results on Heckman-FA

In Table III, we report the training and testing MSE to compare Heckman-FA to
the other baselines when using the extracted prediction features. When
evaluating the methods using the extracted prediction features, we observe
that the training MSE is equal for Heckman-FA and Naive. However, we see that
the testing MSE is lower for the Heckman-FA than Naive. For instance, the
testing MSE for Heckman-FA on the CRIME and COMPAS datasets is 0.0203 and
0.2506, respectively, which is lower than the testing MSE for Naive. For both
datasets, we see that Heckman-FA produces a model that outperforms Naive given
the extracted prediction features. When comparing Heckman-FA and RU, the
training MSE of RU is 0.0050 and 0.0027 lower than Heckman-FA on the CRIME and
COMPAS datasets, respectively. This is expected as RU fits a non-linear model
for regression. On the other hand, we see that the testing MSE for Heckman-FA
is 0.0009 and 0.0083 lower than RU for the CRIME and COMPAS datasets,
respectively.

We also run a paired $t$-test on 10 different assignments of prediction
features to analyze the significance of the comparison between Heckman-FA and
the other baselines. Table IV shows results of the test. We see that the
p-value is very small after running the hypothesis test on both datasets. For
instance, on the CRIME dataset, the p-value is 0.0001 and 0.0053 when
comparing Heckman-FA with Naive and RU, respectively. Given that Heckman-FA
significantly outperforms Naive and RU, the results in Tables III and IV show
that Heckman-FA outputs a robust regression model under MNAR sample selection
bias.

TABLE V: Testing MSEs of Heckman-FA using different values of $T$ and $c$. $c$ |  $T$ training epochs
---|---
100 | 500 | 1,000 | 2,000
CRIME
0.25 | 0.0205 | 0.0210 | 0.0209 | 0.0211
0.5 | 0.0211 | 0.0214 | 0.0211 | 0.0207
0.75 | 0.0217 | 0.0216 | 0.0212 | 0.0208
COMPAS
0.25 | 0.2514 | 0.2506 | 0.2530 | 0.2507
0.5 | 0.2502 | 0.2505 | 0.2502 | 0.2502
0.75 | 0.2502 | 0.2502 | 0.2506 | 0.2502
TABLE VI: Execution times (in seconds) and testing MSEs of Heckman-FA using different values of $T$ and $B$. $B$ |  $T$ training epochs
---|---
Execution time | Testing MSE
100 | 500 | 1,000 | 100 | 500 | 1,000
CRIME
100 | 5.32 | 17.67 | 29.33 | 0.0209 | 0.0206 | 0.0210
500 | 16.02 | 26.06 | 44.23 | 0.0213 | 0.0210 | 0.0209
1,000 | 37.84 | 56.87 | 75.28 | 0.0212 | 0.0206 | 0.0209
COMPAS
100 | 7.70 | 21.30 | 39.32 | 0.2503 | 0.2502 | 0.2505
500 | 22.16 | 37.39 | 54.05 | 0.2506 | 0.2506 | 0.2503
1,000 | 49.72 | 70.46 | 97.64 | 0.2502 | 0.2506 | 0.2502
TABLE VII: Performance of baselines against Heckman-FA* (Top $J$ Features). Extracted prediction features | $\hat{\rho}$ | $R^{2}_{a}$ | Naive | RU[19] | Heckman-FA*
---|---|---|---|---|---
Train MSE | Test MSE | Train MSE | Test MSE | Train MSE | Test MSE
CRIME
2, 3, 8, 12, 16, 18, 21, 24 | 0.0609 | 0.4375 | 0.0172 | 0.0236 | 0.0142 | 0.0232 | 0.0172 | 0.0216
COMPAS
2, 5, 6, 8, 10 | 0.2000 | 0.3432 | 0.0762 | 0.2575 | 0.0751 | 0.2617 | 0.0762 | 0.2527

Sensitivity Analysis. We perform sensitivity analysis on Heckman-FA by testing
the approach over different values for the number of epochs $T$, fixed initial
value $c$, and number of Gumbel-Softmax samples $B$ drawn during assignment
extraction. Table V gives the testing MSE of Heckman-FA across different
values of $T$ and $c$ while fixing $B=1,000$. For each combination of $T$ and
$c$ listed in Table V, the testing MSE of Heckman-FA is almost equal to the
testing MSE provided in Table III for both datasets. We have a similar
observation for each combination of $T$ and $B$ as shown in the right three
columns of Table VI, which give the testing MSE of Heckman-FA across different
values of $T$ and $B$ while fixing $c=0.75$. This shows that the performance
of Heckman-FA is not sensitive to changes in how $\bm{\pi}$ is initialized and
the number of Gumbel-Softmax samples examined during extraction.

TABLE VIII: Testing MSEs of Heckman-C and Heckman-FA* on the CRIME dataset across different values of $J$ for the number of prediction features based on top $J$ selection features. Underlined testing MSE corresponds to extracted set. $J$ | Heckman-C | Heckman-FA*
---|---|---
6 | 0.0490 | 0.0290
7 | 0.0563 | 0.0290
8 | 0.0479 | 0.0216
9 | 0.0500 | 0.0215
10 | 0.0498 | 0.0216
11 | 0.0555 | 0.0216
12 | 0.0524 | 0.0216

We also look at $\hat{\bm{\pi}}$, which is used to draw Gumbel-Softmax samples
for assignment computation, to consider the condition
$\hat{\rho}\in[\rho_{min},\rho_{max}]$. We compare using the learned
$\hat{\bm{\pi}}$ to fixing all elements of $\hat{\bm{\pi}}$ as 0.5 when
drawing Gumbel-Softmax samples during extraction. To fairly compare the two
settings, we obtain $\hat{\bm{\pi}}$ after having $c=0.5$. For each setting
for $\hat{\bm{\pi}}$, we compute the proportion of Gumbel-Softmax samples that
correspond to $\hat{\rho}\in[\rho_{min},\rho_{max}]$ out of 1,000 total
samples. We repeat the experiment 40 times using the COMPAS dataset. We find
that the average proportion of having $\hat{\rho}\in[\rho_{min},\rho_{max}]$
is 0.0601 after using a fixed value of 0.5 for all elements of
$\hat{\bm{\pi}}$. However, when using the learned $\hat{\bm{\pi}}$ from
Heckman-FA, the average proportion of having
$\hat{\rho}\in[\rho_{min},\rho_{max}]$ increases by 0.0211. This result
indicates that compared to letting each element of $\hat{\bm{\pi}}$ be 0.5,
Heckman-FA is more likely to extract assignments of prediction features such
that $\hat{\rho}\in[\rho_{min},\rho_{max}]$ when the Gumbel-Softmax samples
are drawn using the learned $\hat{\bm{\pi}}$.

Execution Time. We report the execution time after running Heckman-FA across
different values of $T$ and the number of Gumbel-Softmax samples $B$ drawn
during the assignment extraction in left three columns of Table VI. We observe
that for both datasets, Heckman-FA runs fast for each combination of $T$ and
$B$. For instance, when $T=100$ and $B=100$, Heckman-FA is completed after
5.32 and 7.70 seconds for the CRIME and COMPAS datasets, respectively. While
the execution time increases as $T$ and $B$ increase, the testing MSE of
Heckman-FA remains close to the testing MSE reported in Table III. This result
shows that Heckman-FA is fast while maintaining quality performance on the
testing set.

###  V-C Results on Heckman-FA*

In Table VII, we compare the performance of baselines to Heckman-FA*. When
using the extracted prediction features, we see that while the training MSEs
of Heckman-FA* and Naive are equal for both datasets, the testing MSE for
Heckman-FA* is 0.0020 and 0.0048 lower than Naive on the CRIME and COMPAS
datasets, respectively. This shows that by simply choosing $J$ prediction
features after ranking selection features based on $\hat{\pi}_{2k}$, Heckman-
FA* is robust against MNAR sample selection bias. Moreover, in terms of the
testing MSE, Heckman-FA* outperforms RU by 0.0016 and 0.0090 on the CRIME and
COMPAS datasets, respectively.

Comparison with Correlation-Based Ranking. To further show the effectiveness
of Heckman-FA*, which requires training $\psi$, we also compare the ranking of
selection features based on $\hat{\pi}_{2k}$ to the ranking of selection
features based on their strength of correlation with the outcome. We use the
name Heckman-C to describe the correlation-based ranking of selection
features. Unlike the ranking based on $\hat{\pi}_{2k}$, Heckman-C does not
rely on training the assignment function beforehand. For Heckman-FA* and
Heckman-C, we run each approach on the CRIME dataset using the top $J$
selection features across different values of $J$. Table VIII provides the MSE
of the models on the testing set. We consider the values of $J=6$ through
$J=12$. The underlined testing MSE corresponds to the size of the final set of
prediction features. For Heckman-FA*, the final set of prediction features
consists of $8$ prediction features. For Heckman-C, there is no testing MSE
underlined as there is no set of top $J$ selection features extracted for the
final set of prediction features. In other words, when ranking the selection
features based on their correlation with the outcome using the CRIME dataset,
there is no set of $J$ features from the ranking that ensures the robustness
of Heckman-C under MNAR sample selection bias. In our experiment, we also find
that the condition $\hat{\rho}\in[\rho_{min},\rho_{max}]$ is satisfied for
Heckman-FA* for the range $J=8$ through $J=12$. However, for Heckman-C, no set
of $J$ prediction features satisfies this condition for any value of $J$. This
indicates that when assigning prediction features for the Heckman model,
ranking selection features based on $\hat{\pi}_{2k}$ after training $\psi$ is
more effective than ranking based on the correlation of features with the
outcome.

##  VI Conclusion

In this paper, we introduced Heckman-FA, a novel data-driven approach that
obtains an assignment of prediction features for the Heckman selection model
to robustly handle MNAR sample selection bias. Given a set of features that
are used to fit the selection of samples, our approach first trains an
assignment function by minimizing the MAE on the set of fully observed
training samples. Heckman-FA finds a set of prediction features for the
Heckman model by drawing a number of Gumbel-Softmax samples using the learned
probability of assignment for each selection feature. This set is extracted
based on the prediction model’s goodness-of-fit and the estimated correlation
between noise terms. We observed that the Heckman-FA produces a robust
regression model under MNAR sample selection bias on the outcome after
training the model on real-world datasets. In the future, we plan to extend
our approach to the problem of learning MNAR outcomes/labels in non-tabular
data.

Reproducibility. The source code can be downloaded using the link
https://tinyurl.com/25s786z6.

## Acknowledgements

This work was supported in part by NSF 1946391 and 2137335.

## References

  * [1] Till Bärnighausen, Jacob Bor, Speciosa Wandira-Kazibwe, and David Canning.  Correcting hiv prevalence estimates for survey nonparticipation using heckman-type selection models.  Epidemiology, pages 27–35, 2011.
  * [2] Shawn Bushway, Brian D Johnson, and Lee Ann Slocum.  Is the magic still there? the use of the heckman two-step correction for selection bias in criminology.  Journal of quantitative criminology, 23:151–178, 2007.
  * [3] Xiangli Chen, Mathew Monfort, Anqi Liu, and Brian D Ziebart.  Robust covariate shift regression.  In AISTATS, pages 1270–1279. PMLR, 2016.
  * [4] Wei Du, Xintao Wu, and Hanghang Tong.  Fair regression under sample selection bias.  In Big Data, pages 1435–1444. IEEE, 2022.
  * [5] Dheeru Dua and Casey Graff.  UCI machine learning repository.  http://archive.ics.uci.edu/ml, 2017.
  * [6] Minna Genbäck, Elena Stanghellini, and Xavier de Luna.  Uncertainty intervals for regression parameters with non-ignorable missingness in the outcome.  Statistical Papers, 56:829–847, 2015.
  * [7] James J Heckman.  Sample selection bias as a specification error.  Econometrica: Journal of the econometric society, pages 153–161, 1979.
  * [8] Xinting Hu, Yulei Niu, Chunyan Miao, Xian-Sheng Hua, and Hanwang Zhang.  On non-random missing labels in semi-supervised learning.  arXiv:2206.14923, 2022.
  * [9] Eric Jang, Shixiang Gu, and Ben Poole.  Categorical reparameterization with gumbel-softmax.  arXiv:1611.01144, 2016.
  * [10] Hyungu Kahng, Hyungrok Do, and Judy Zhong.  Domain generalization via heckman-type selection models.  In ICLR, 2023.
  * [11] J. Larson, S. Mattu, L. Kirchner, and J. Angwin.  Compas dataset.  https://github.com/propublica/compas-analysis/, 2017.
  * [12] Jae-woong Lee, Seongmin Park, and Jongwuk Lee.  Dual unbiased recommender learning for implicit feedback.  In SIGIR, pages 1647–1651, 2021.
  * [13] Qi Lei, Wei Hu, and Jason Lee.  Near-optimal linear regression under distribution shift.  In ICML, pages 6164–6174, 2021.
  * [14] Siu Fai Leung and Shihti Yu.  On the choice between sample selection and two-part models.  Journal of econometrics, 72(1-2):197–229, 1996.
  * [15] Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Stratis Gavves.  Citris: Causal identifiability from temporal intervened sequences.  In ICML, pages 13557–13603. PMLR, 2022.
  * [16] Jose G. Moreno-Torres, Troy Raeder, Rocío Alaiz-Rodríguez, Nitesh V. Chawla, and Francisco Herrera.  A unifying view on dataset shift in classification.  Pattern recognition, 45(1):521–530, 2012.
  * [17] Kazumitsu Nawata.  A note on the estimation of models with sample-selection biases.  Economics Letters, 42(1):15–24, 1993.
  * [18] Patrick Puhani.  The heckman correction for sample selection and its critique.  Journal of economic surveys, 14(1):53–68, 2000.
  * [19] Roshni Sahoo, Lihua Lei, and Stefan Wager.  Learning from a biased sample.  arXiv:2209.01754, 2022.
  * [20] Hidetoshi Shimodaira.  Improving predictive inference under covariate shift by weighting the log-likelihood function.  Journal of statistical planning and inference, 90(2):227–244, 2000.
  * [21] Francis Vella.  Estimating models with sample selection bias: a survey.  Journal of Human Resources, pages 127–169, 1998.
  * [22] Xiaojie Wang, Rui Zhang, Yu Sun, and Jianzhong Qi.  Doubly robust joint learning for recommendation on data missing not at random.  In ICML, pages 6638–6647, 2019.
  * [23] Sarah E Wolfolds and Jordan Siegel.  Misaccounting for endogeneity: The peril of relying on the heckman two-step method without a valid instrument.  Strategic Management Journal, 40(3):432–462, 2019.


# Deep Learning Model Explainability for Inspection Accuracy Improvement in
the Automotive Industry
††thanks: Supported by Faurecia

Anass EL HOUD  Faurecia Clean Mobility
Ecole Centrale de Lyon
Bavans, France
<EMAIL_ADDRESS>Charbel EL HACHEM  FEMTO-ST Institute, CNRS
Univ. Bourgogne Franche-Comte (UBFC)
Belfort, France
<EMAIL_ADDRESS>Loïc PAINVIN  R&D Center
Faurecia Clean Mobility
Bavans, France
<EMAIL_ADDRESS>
###### Abstract

The welding seams visual inspection is still manually operated by humans in
different companies, so the result of the test is still highly subjective and
expensive. At present, the integration of deep learning methods for welds
classification is a research focus in engineering applications. This work
intends to apprehend and emphasize the contribution of deep learning model
explainability to the improvement of welding seams classification accuracy and
reliability, two of the various metrics affecting the production lines and
cost in the automotive industry. For this purpose, we implement a novel hybrid
method that relies on combining the model prediction scores and visual
explanation heatmap of the model in order to make a more accurate
classification of welding seam defects and improve both its performance and
its reliability. The results show that the hybrid model performance is
relatively above our target performance and helps to increase the accuracy by
at least 18%, which presents new perspectives to the developments of deep
Learning explainability and interpretability.

###### Index Terms:

Visual Inspection, Industry 4.0, Deep Learning, Model Explainability Heatmap,
Hybrid Classification

##  I Introduction

Welding is a manufacturing process consisting in joining two or more elements
in a permanent way while ensuring continuity between these elements. The
assembly is done either by heating, by pressure, or by the 2 combined. This
process has an important role in productions lines, especially in the
automotive industry, and is surely being influenced by the advent of digital
technologies such as artificial intelligence, IoT, and intelligent
manufacturing [1, 2, 3].
Welding defect is an issue that affects Faurecia plants worldwide. Currently,
there is a lot of research work being done on the automation of quality
control [4] and visual inspection of welding seams. Faurecia has developed
automatic systems based on artificial intelligence to check the quality of
welding seams in plants. Until now, these systems have been evaluated solely
by their accuracy, the number of true positive and true negative. However,
they have some shortcomings and errors in the classification of welding seam
quality even when their accuracy is high [4]. Since these deep learning models
work like black boxes, research is progressing on the explainability and
interpretability of the results delivered by AIs [5]. Nowadays, model
interpretation is becoming a primary evaluation metric as well as its
performance, and a good compromise between explainability and accuracy is more
and more necessary in industrial applications.
In the study, the followed approach consists on:

  1. 1.

Taking advantage of a pre-trained model (MobileNet) for welding seams
classification at Faurecia, tested by El Hachem et al. [6].

  2. 2.

Implementing another deep learning model (ResNet-50) for comparison.

  3. 3.

Using visual explanation heatmap to analyse and interpret the results of these
models and to seek improvement opportunities.

  4. 4.

Developing a hybrid approach taking into account model explainability to
improve the overall accuracy.

The developed inspection system should be able to classify the defects with at
least 97% accuracy target, as requested by the client.
This paper is organized as follows. Section II introduces the problem
statement and motivation to set up this study. Section III presents an
overview of the existing work done in the area of welding seam inspection and
model explainability. In Section IV, the novel hybrid approach is introduced.
In Section V, the implementation details are mentioned. In Section VI, the
analysis of the experimental results is presented. The last section gives the
concluding remarks.

##  II Problem Statement

###  II-A Automated Visual Inspection

The goal of automation is to use digital technologies in order to reduce the
workload of the production operator, to improve the productivity and quality
of the work performed, more specifically in the automotive industry. These
systems take advantage of the vast evolution in the fields of artificial
intelligence, computer vision, and robotics. To remain competitive, the
companies use automated systems to make their plants more autonomous and
reduce the human impact on the reliability of quality control. Automation also
allows overcoming several other problems such as fatigue, security, data
integrity, reliability, and availability. Hence the recent interest of
companies to invest in this technological transition [7].

Faurecia produces exhaust systems that contain many welding zones. A leak test
is done to check the quality of the welds after the welding operation is
completed. Four welds are not checked by this test. That is why their
automatic visual inspection is necessary to ensure the quality of the whole
manufactured product.

Figure 1: Reference weld (OK) and five possible defects (NOK)

Figure 1 shows the different possible scenarios of the quality of the welds:
only the first case is considered as OK.

###  II-B Actual Model Accuracy

During all the tests done on different deep learning trained models, only a
few were able to achieve the accuracy needed. The major source of low accuracy
is that the two classes are not equally represented in the sample. Since weld
defects are not potentially always present in production, photos from both the
OK and NOK classes will not behave the same occurrence, resulting in the
problem of unbalanced data which can be solved using resampling techniques
[8]. For this application, the problem was initially solved by using the data
augmentation technique with various filters and transformations. Although, the
accuracy can be further corrected by applying other modifications to the
overall approach.

###  II-C Model Reliability

As many deep learning applications are life or cost crucial, their accuracy is
not the only evaluation metric that matters [9], especially in an industrial
context. The model reliability is also considered as an important criterion to
choose the best model. In fact, even if the accuracy target is achieved, the
model might use a biased part of the image to classify it. This behavior
represents a problem that is hard to detect and can generate a lot of damage
since the model appears to perform well when in reality it is completely
biased and its accuracy drops when faced with new data.

Figure 2: Model’s activation heatmap of Misclassified NOK Weld

Figure 2 shows the model’s activation heatmap of a misclassified weld. In
reality, the weld is defective (NOK) but the model classifies it as of good
quality (OK). It can clearly be seen that this is due to the false region on
which the model relies to make its decision. Yet the model has good overall
accuracy on the dataset. It is therefore necessary to study the reliability of
the model in addition to its accuracy and to exploit the correlation between
the activation heatmap results and the other scores of the model to make a
more accurate decision.

##  III Related Works

###  III-A Welding Seam Classification

El Hachem et al. [6] applied data augmentation to improve the accuracy of 4
welding seams classification. A deep learning model based on Mobilenet
architecture was trained on a set of pictures of augmented weld defects. The
image is initially fed to a convolutional neural network (feature extractor).
After adjusting the weights, the network returns two scores for each image to
classify it. The predicted class is the one with the highest prediction score.
In general, the two prediction scores can either be class probabilities or
class scores, which accept values larger than 1 (Fig.3).

Figure 3: Old decision making approach

The best results of this old approach are shown in table 1. Although the
targeted accuracy (97%) was achieved in some welds, it remains a challenge on
other welds. Hence the need to look for other methods to overcome this
problem.

TABLE I: Initial Accuracy based on standard approach | Accuracy
---|---
Weld 1 | 97,5%
Weld 2 | 57,5%
Weld 3 | 80,5%
Weld 4 | 96,5%

Yang et al. [10] have constructed a novel deep learning framework to classify
weld defects. With this new approach, the classification accuracy was improved
by 3.18% and 4.33%. Even with all the work done, it is believed that there is
still a large room for improving the classification performance over the
generic DNN models by taking into account other parameters.

###  III-B Visual Explanation Heatmap

Knowing how to explain the reasons that led to a decision gives it more
credibility and shows a strong capacity for reasoning and intelligence. The
explainability of the AI model reinforces the credibility of the obtained
results and also allows us to evaluate the reliability of the model and its
behavior if a partial change of the data occurred. For example, in health
care, the model verification and interpretation by medical experts is a
necessity [11]. There are different methods that make the AI model
explainable: Selvaraju et al. [12] have first proposed the Grad-CAM method, a
visualization technique able to tell which parts of a given image led the
trained convolutional network to its final classification decision. This
method makes it easy to debug the decision process of a CNN, especially in the
case of misclassification. The result of this method is an activation heatmap
indicating the parts of the image that have contributed the most to the final
decision of the network.

Figure 4: Example of Grad-CAM heatmap on Dog & Cat Classification

Figure 4 shows the application of Grad-CAM heatmap on Dog & Cat
Classification. The model is accurate because it relies on the correct region
of the image to classify it.
In the same context, Zhang et al. [13] have applied model explainability to
interpret the deep learning models trained to classify multiple sclerosis
types in the brain using the Grad-CAM method. The experimental results showed
that Grad-CAM gives the best heatmap localizing ability, and CNNs with a
global average pooling layer and pre-trained weights had the best
classification performance.

###  III-C Hybrid Methods

Ahlawat et al. [14] have implemented a hybrid method using CNN and SVM
classifier for handwritten digit recognition: CNN works as an automatic
feature extractor and SVM works as a binary classifier. Their results showed
that the hybrid approach achieved an accuracy of 99.28% for the MNIST dataset.
Soumaya et al. [15] have tested a hybrid classification model using a genetic
algorithm and SVM to detect Parkinson’s disease. Their method attempts to give
an accuracy of 80% and 72.50% using two kernels of SVM. The hybrid method
seems to ensure the optimization of the classification system by minimizing
the dimension of the features vector and maximizing the accuracy.
In the same context, Ahammad et al. [16] have suggested a new CNN-deep
segmentation-based boosting classifier for spinal cord injury prediction. This
method gives 10% improvement on the classification rate.
Another hybrid method is developed by Liu et al. [17] for CO2 welding defects
detection by using CNN for primary features extraction and LSTM for feature
fusion. The algorithm reaches 94% of prediction accuracy.

###  III-D Frozen Graph Representation

The automated visual systems shall be deployed on the plants. For this, it is
preferable to have a trained model file as light as possible. The Tensorflow
framework proposes the concept of frozen graph which is used to save what is
necessary for the model while performing inference: the model is said to be
frozen when keeping only the parts of it that matter to the deployment phase.
However, access to the frozen model weights becomes more difficult. Hence the
need to develop a solution to save the model as a frozen graph without making
it harder to extract or restore the data from its layers when performing a
model explainability method. A frozen graph in TensorFlow is saved as a file
with the extension ”.pb”. It stands for protocol buffers, a language-neutral,
platform-neutral extensible mechanism for serializing structured data as
defined by its developer Google [18].

##  IV Methods

###  IV-A Convolutional Neural Networks

Convolutional Neural Networks (CNN) have undergone a remarkable evolution with
the development of new computational tools. These neural networks perform
themselves all the tedious work of feature extraction and description: during
the training phase, the classification error is minimized in order to optimize
the parameters and the feature extraction. Moreover, the specific architecture
of the network allows the extraction of features of different complexity, from
the simplest to the most sophisticated in the image, hence the name ”feature
extractor”.
CNN achieved very interesting results on different challenging datasets such
as ImageNet [19]. Which shows its high performance in solving such vision
problems.
For classification problems, the neural network calculates from the input
image a score (or probability) for each class. The class attributed to the
input object corresponds to the one with the highest score. The values of the
layer weights are learned by backpropagation of the gradient [20]: the
parameters that minimize the regularized loss function are computed
progressively, for each layer, starting from the end of the network. The
optimization is done with stochastic gradient descent.

###  IV-B Model Explainability

The main approach of this work is based on the Grad-CAM method, introduced by
Selvaraju et al. [12]. The method assigns importance to each position in the
last convolutional layer by computing the linear combination of activations,
weighted by the corresponding output weights for the predicted class. The
resulting class activation mapping is then resampled to the size of the input
image.
This technique will be used as a reliability test of the deep learning model
by comparing the area with high color intensity to the real area of interest
in the image. The model is more reliable if the two areas have closer values
and positions.
The implemented algorithm generates an activation heatmap showing the region
with the highest impact on the final decision (fig. 5).

Figure 5: Model’s Activation heatmap of a welding seam

The second step is the calculation of the Red Color Ratio (RCR) in the
heatmap. A K-nearest neighbor algorithm (KNN) is used to classify and identify
precisely the pixels’ closest color in the heatmap with K = 3 representing the
RGB color system. In the end, the Red Color Ratio is calculated by dividing
the number of pixels considered as red to the :

$RCR=\frac{\text{Number of Red Cluster pixels}}{\text{Total number of
pixels}}\times 100$

By plotting the Red Color Ratio (RCR) in the heatmap as a function of the
class scores, a high correlation is identified. This reinforces the idea that
weld defects are often labeled by a limited feature zone, compared to non-
defective welds. RCR is considered as a reliability score of the model
prediction.

Figure 6: Distribution of images per class scores and RCR

The figure 6 shows the distribution of the weld 1 images per class scores and
RCR. Each blue dot is a non-defective weld in reality (OK Class), and the
orange triangle is a defective weld (NOK Class). This identified correlation
was the major motivation to investigate this direction (fig. 6). The dataset
is now more separable than when represented by only two scores.

The use of 3 scoring parameters (Class 1 Score, Class 2 Score, Red Color
Ratio) may help to get more accurate thresholding for the classification of
these points. This is the main reason behind the development of this new model
explainability-based approach (fig. 7).

Figure 7: Novel decision making approach

##  V Implementation Details

In this study, two deep learning models were used: the MobileNet architecture
that has been trained on the same welding seams dataset by El Hachem et al.
and a ResNet-50 model that is trained during this study for comparison.

###  V-A Experimental Environement

The experimental environment is powered by Intel i5 CPU, 2.30 GHz with 64-bit,
Windows 10 system and 8 GB memory. The software programming environment is
Python. It uses both Keras and Tensorflow as backend.

####  V-A1 Feature Extractors

For the ResNet-50 model, Adam is selected as optimizer of the CNN with a
learning rate set to $10^{-5}$. A pre-trained version of the model weights
trained on the ImageNet database has been loaded. The input of the CNN is
represented by the vector (batch size, height, width, depth). The batch size
(equal to 32) defines the number of samples that will be propagated through
the whole network. The other parameters represent the height, width, and depth
of the image. The input has the following values (32, 128, 128, 3). For the
Mobilenet model, Rmsprop is selected as the optimizer of the CNN. The chosen
learning rate decay type is exponential starting with a value of 0.01 and
ending with a value equal to 0.0001. The model has been trained with 9000
epochs.

####  V-A2 Machine Learning Classifiers

The XGboost Classifier is a gradient boosting algorithm that offers a very
large panel of hyperparameters, it is thus possible, thanks to this diversity
of parameters, to have total control over the implementation of Gradient
Boosting [21]. The chosen booster is ”gbtree” which uses a tree-based
boosting. The step size shrinkage is set to 0,5 and the used sampling method
is the uniform selection.

The decision tree uses the Gini function to measure the quality of the split.
The max depth is set to the default value so that the tree nodes are expanded
until all leaves are pure or contain less than 2 elements.

On the other hand, Support Vectors Machines Classifier (SVM) are a family of
machine learning algorithms that solve classification and regression problems.
They are known for their strong theoretical guarantees and their great
flexibility. The idea of the Support Vector Machine is to project the data
into a higher-dimensional space and to make them separable. SVMs become a
universal approximator [22]: with enough data, the algorithm can always find
the best possible boundary to separate two classes.
There are several formulations of SVM, each formulation makes a specific
transformation of the input data and has a specific shape of the decision
boundaries, these formulations are called kernels. In this study, two SVM
kernels are used:

  * •

Linear Kernel: This is the case of a linear classifier, without space change.
The data are assumed to be linearly separable, i.e. there is a hyperplane that
separates the data without error. This is not quite the case observed on the
correlation curve above. However, it is still worthwhile to see the effect of
the separation by this kernel on the final accuracy.

  * •

5-Polynomial Kernel: This is the case where we change the space by applying a
polynomial transformation (in this case a polynomial of degree 5). The choice
of the decision boundaries will be based on the new space.

In total, four different machine learning classifiers are tested for the
improvemeent of the decision making part on the two CNN feature extractors.

##  VI Experimental Results

Figure 8 shows the results obtained on the weld 1, weld 2, weld 3, and weld 4
using MobileNet and ResNet-50 models based on the old approach. The target is
only reached with the weld 1 using MobileNet architecture.

$0.8$$1$$1.2$$1.4$$1.6$$1.8$$2$$2.2$$2.4$$2.6$$2.8$$3$$3.2$$3.4$$3.6$$3.8$$4$$4.2$$50$$55$$60$$65$$70$$75$$80$$85$$90$$95$$100$Weld IndexAccuracy (%)MobileNetResNet-50 Figure 8: Comparison between MobileNet & ResNet - Old Approach TABLE II: Mobilenet Model Accuracy - New Approach Welding | Accuracy per Classifier
---|---
Seams | XGboost | Decision Tree | SVM Linear | SVM Poly5
Weld 1 | 98,4% | 98,4% | 99,3% | 98,9%
Weld 2 | 99,3% | 98,1% | 98,8% | 98,8%
Weld 3 | 98,7% | 97,2% | 98,1% | 98,7%
Weld 4 | 98,1% | 97,9% | 98,1% | 98,1%

Table II proves that adding a statistical machine learning classifier after
the feature extractor and class activation heatmap seems to increase the
overall accuracy of the model for all of the four welds. The target accuracy
is acheived by the four machine learning classifiers, with a tiny advantage to
XGboost classifier.

On the other hand, Table III explains how hybrid ResNet-50 behaves when the
network is followed by a machine learning classifier. The accuracy is better
than the standard approach but still smaller than Mobilenet.

TABLE III: ResNet-50 Model Accuracy - New Approach Welding | Accuracy per Classifier
---|---
Seams | XGboost | Decision Tree | SVM Linear | SVM Poly5
Weld 1 |  97,1% | 95,1% | 96,1% | 95,3%
Weld 2 | 98,4% | 95,1% | 96,8% | 98,8%
Weld 3 | 97,3% | 92,8% | 96,7% | 98,1%
Weld 4 | 97,9% | 95,9% | 97,2% | 97,9%

##  VII Conclusion & Future directions

In this paper, a hybrid approach of CNN-ML Classifier is proposed for welds
defects classification. This approach adds a new reliability score calculated
using the model explainability methods, more specifically Grad-CAM heatmap.
This new score makes the dataset more separable. The hybrid approach proposed
in this paper seems to perform very well on weld defects classification. The
highest accuracy improvement was by +71% for weld 2 using MobileNet-XGboost
classifier and by +18% for the weld 3 using MobileNet-SVM Poly5 Kernel.
MobileNet architecture performs better than ResNet-50 in extracting the
feature from the dataset images. For the classifier part, XGBoost achieves
higher accuracy and gives promising results comparing to SVM Linear kernel,
Poly5 kernel, and Decision Tree at last.
This work presents new perspectives on the developments of new model-driven
optimization methods to improve the accuracy of vision systems. In future
work, it can be extended to other classification problems with highly
unbalanced datasets.

## Acknowledgment

This work was done as a part of a CIFRE (N 2018/1029) project with Faurecia,
funded by the Ministry of Higher Education and Research of France, managed by
the Association Nationale de la Recherche et de la Technologie (ANRT) and was
partially supported by the EIPHI Graduate School (contract ”ANR-17-EURE-
0002”).

## References

  * [1] Baicun Wang, S Jack Hu, Lei Sun, and Theodor Freiheit.  Intelligent welding system technologies: State-of-the-art review and perspectives.  Journal of Manufacturing Systems, 56:373–391, 2020.
  * [2] Johannes Egger and Tariq Masood.  Augmented reality in support of intelligent manufacturing–a systematic literature review.  Computers & Industrial Engineering, 140:106195, 2020.
  * [3] Ernie Mazuin Mohd Yusof, Mohd Ismail Yusof, Rafidah Ali, Izzat Hilmi Harjimi, and Qasidah Kamarul Bahrin.  Welding station monitoring system using internet of thing (iot).  Indonesian Journal of Electrical Engineering and Computer Science, 18(3):1319–1330, 2020.
  * [4] Charbel El Hachem, Gilles Perrot, Loïc Painvin, and Raphaël Couturier.  Automation of quality control in the automotive industry using deep learning algorithms.  In 2021 International Conference on Computer, Control and Robotics (ICCCR), pages 123–127. IEEE, 2021.
  * [5] Alex John London.  Artificial intelligence and black-box medical decisions: accuracy versus explainability.  Hastings Center Report, 49(1):15–21, 2019.
  * [6] Charbel El Hachem, Gilles Perrot, Loïc Painvin, and Raphaël Couturier.  Welding seam classification in the automotive industry using deep learning algorithms.  IEEE, 2021.
  * [7] Rehan Syed, Suriadi Suriadi, Michael Adams, Wasana Bandara, Sander JJ Leemans, Chun Ouyang, Arthur HM ter Hofstede, Inge van de Weerd, Moe Thandar Wynn, and Hajo A Reijers.  Robotic process automation: contemporary themes and challenges.  Computers in Industry, 115:103162, 2020.
  * [8] Petr Mrozek, John Panneerselvam, and Ovidiu Bagdasar.  Efficient resampling for fraud detection during anonymised credit card transactions with unbalanced datasets.  In 2020 IEEE/ACM 13th International Conference on Utility and Cloud Computing (UCC), pages 426–433. IEEE, 2020.
  * [9] Basemah Alshemali and Jugal Kalita.  Improving the reliability of deep neural networks in nlp: A review.  Knowledge-Based Systems, 191:105210, 2020.
  * [10] Lu Yang and Hongquan Jiang.  Weld defect classification in radiographic images using unified deep neural network with multi-level features.  2020\.
  * [11] Wojciech Samek, Thomas Wiegand, and Klaus-Robert Müller.  Explainable artificial intelligence: Understanding, visualizing and interpreting deep learning models.  arXiv preprint arXiv:1708.08296, 2017.
  * [12] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra.  Grad-cam: Visual explanations from deep networks via gradient-based localization.  In Proceedings of the IEEE international conference on computer vision, pages 618–626, 2017.
  * [13] Yunyan Zhang, Daphne Hong, Daniel McClement, Olayinka Oladosu, Glen Pridham, and Garth Slaney.  Grad-cam helps interpret the deep learning models trained to classify multiple sclerosis types using clinical brain magnetic resonance imaging.  Journal of Neuroscience Methods, 353:109098, 2021.
  * [14] Savita Ahlawat and Amit Choudhary.  Hybrid cnn-svm classifier for handwritten digit recognition.  Procedia Computer Science, 167:2554–2560, 2020.
  * [15] Zayrit Soumaya, Belhoussine Drissi Taoufiq, Nsiri Benayad, Korkmaz Yunus, and Ammoumou Abdelkrim.  The detection of parkinson disease using the genetic algorithm and svm classifier.  Applied Acoustics, 171:107528, 2021.
  * [16] Sk Hasane Ahammad, V Rajesh, Md Zia Ur Rahman, and Aimé Lay-Ekuakille.  A hybrid cnn-based segmentation and boosting classifier for real time sensor spinal cord injury data.  IEEE Sensors Journal, 20(17):10092–10101, 2020.
  * [17] Tianyuan Liu, Jinsong Bao, Junliang Wang, and Yiming Zhang.  A hybrid cnn–lstm algorithm for online defect recognition of co2 welding.  Sensors, 18(12):4369, 2018.
  * [18] Jianhua Feng and Jinhong Li.  Google protocol buffers research and application in online game.  In IEEE conference anthology, pages 1–4. IEEE, 2013.
  * [19] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  Advances in neural information processing systems, 25:1097–1105, 2012.
  * [20] Robert Hecht-Nielsen.  Theory of the backpropagation neural network.  In Neural networks for perception, pages 65–93. Elsevier, 1992\.
  * [21] Tianqi Chen and Carlos Guestrin.  Xgboost: A scalable tree boosting system.  CoRR, abs/1603.02754, 2016.
  * [22] Junping Wang, Quanshi Chen, and Yong Chen.  Rbf kernel based support vector machine with universal approximation and its application.  In International symposium on neural networks, pages 512–517. Springer, 2004.