A transformer-based deep learning framework to predict employee attrition

Overviews

Figure 1 summarizes the major steps in our study. First, the collected dataset was divided into two parts: training data (85%) and test data (15%). While the training data are then used for model optimization and development, test data are kept separately to avoid any data leakage. To train deep learning models, 15% of total training data are used to create a validation set, and the rest of the data are used for model development. The validation set plays a role in finding the optimal models. To train the machine learning model, the training data are used without creating an independent validation set. A 5-fold cross-validation is applied to the training data to compute the average performance corresponding to specific sets of parameters to find the best hyperparameters. The machine learning models are then retrained with training data and hyperparameters. Finally, the deep learning models and machine learning models are benchmarked using the independent test set.

Data preparation

IBM Analytics provides the IBM HR Analytics Employee Attrition dataset (Aizemberg, 2019) which was used in this study. The dataset contains information about various factors that can contribute to employee attrition, including demographic information, job satisfaction, job involvement, performance ratings, and other factors. The dataset had 1,470 samples with 16.1% (237 samples) who left their companies while 83.9% (1,233 samples) kept their job position. The dataset has 35 variables with 26 continuous and nine categorical variables. Stratified random sampling was used to generate the development and test sets, with percentages of 85% and 15%, respectively. In addition, the sampling method also was used to produce the validation set from 10% of the development data. The validation set was assigned to supervise the training process to find the optimal model. independent test set and the development set were checked to assure that no data duplicates were found in both sets. After utilizing one-hot encoding on categorical variables, normalization was employed to ensure that all variables are on a common scale and eliminates the influence of large values on the model’s predicted outcomes. The data can be used to build predictive models to identify employees who are at risk of leaving the company, as well as to develop strategies for improving employee satisfaction and engagement. Table 1 gives information on the distribution of samples of development, validation, and test sets.

Model architecture

Figure 2 describes the architecture of our proposed method. Overall, the model was constructed with one column embedding layer, one N-stacked Transformer layer, and one multilayer perceptron (Huang et al., 2020). A Transformer block is specified by two components, including a multi-head self-attention layer and a position-wise feed-forward layer. For a feature-target pair (x, y) where x ≡ {x_categorical, x_continuous}, x_categorical and x_continuous represent for all categorical features and continuous features. Hence, x_categorical = {x₁, x₂, x₃, …, x_m} contain a set of categorical features x_i which are independently embedded using the column embedding of dimension d. The embedding of feature x_i is denoted as e_ϕi(x_i) ∈ℝ^d. For all the existing categorical features, the set of embeddings is defined as E_ϕ(x_categorical) = {e_ϕ1(x₁), e_ϕ2(x₂), e_ϕ3(x₃), …, e_ϕm(x_m)} which are passed through the first Transformer layer. The first Transformer layer’s output is then fetched to the second Transformer layer as input. The process is successively repeated in the same manner until passing the Nth Transformer layer. Each embedding is finally converted into contextual embedding after exiting the last Transformer layer. Briefly, the series of Transformer layers can be defined as a function f_ω which map embeddings E_ϕ(x_categorical) = {e_ϕ1(x₁), e_ϕ2(x₂), e_ϕ3(x₃), …, e_ϕm(x_m)} to create the corresponding contextual embeddings H_ϕ(x_categorical) = {h₁, h₂, h₃, …, h_m} ∈ ℝ^d. The contextual embeddings H_ϕ(x_categorical) are then concatenated with the continuous features x_continuous to generate a new vector of dimension (d + m + c). This vector is operated by the multilayer perceptron layer (denoted as g_ϵ) afterward to return the predicted probability for target y. The loss function (denoted as K) for model training is expressed as: (1)${\displaystyle L\left(x,y\right)\equiv K\left(g_{ϵ}\left(f_{\omega }\left(E_{\varphi }\left(x_{\text{categorical}}\right)\right),x_{\text{continuous}}\right),y\right).}$

Column embedding

For a column of categorical feature i, a lookup table for embedding e_ϕi(.) is defined. If a categorical feature x_i has d_i classes, the lookup table e_ϕi(.) is specified by (d_i + 1) embeddings in which the additional embedding corresponds to a missing value. The feature x_i = j ∈ [0, 1, 2, 3, …, d_i] is encoded with a new embedding e_ϕi(j) = [c_ϕj, w_ϕi,j] where c_ϕj ∈ ℝ^l and w_ϕi,j ∈ ℝ^d−l with the dimension c_ϕj and l are hyperparameters for tuning. The distinct identifier c_ϕj ∈ ℝ^l differentiates the classes from one column to other columns. Positional encodings are not employed for tabular data since there is no feature order available, in contrast to language models where the embeddings are inserted element by element with the positional encoding of the word in the sentence.

Modeling experiments

The Adam optimizer (Kingma & Ba, 2014) was used to repeatedly adjust network parameters at a learning rate of 5 × 10⁻⁵. During the training process, the network was monitored to find its optimal state when its validation loss bottomed. After 80 training epochs, we found that our model’s validation loss converged at epoch 27th. Hence, we selected the model at epoch 27th as our optimal model. Since the problem is binary classification, we chose binary cross-entropy as our loss function: (2)${\displaystyle \text{Loss}=\sum _{i=1}^n{y}_i\times log{\widehat{y}}_i+\left(1-y_i\right)\times log\left(1-{\widehat{y}}_i\right),}$

where y is the actual label and $\widehat{y}$ is the predicted probability. In our modeling experiments, we designed and tested all deep learning models under the PyTorch 1.3.1 platform. These models were trained on an Intel i7-12700 CPU with 64GB RAM and an NVIDIA GeForce RTX 3090 Ti GPU. One epoch required around 1.2 s for training and 0.2 s for testing.

Assessment metrics

In this study, a number of metrics were used to evaluate the performance of the models, including the area under the receiver operating characteristic curve (AUCROC), the area under the precision–recall curve (AUCPR), and accuracy (ACC). These metrics were calculated based on the True Positive, False Positive, True Negative, and False Negative values. The AUCROC and AUCPR are significant metrics in machine learning for evaluating classification models. The AUCROC measures the model’s ability to accurately classify positive and negative instances, while the AUCPR focuses on precision and recall trade-offs, particularly in imbalanced datasets. Both metrics provide comprehensive assessments of the model’s performance, are robust to class imbalance, and simplify the evaluation process. They help researchers and practitioners make informed decisions about model selection and optimization, providing valuable insights into the classifier’s discriminative ability.

Model assessment

To train the model, stratified random sampling was used to select 10% of the original development set to construct a validation set to monitor the training, while the remaining data were used for training. The models’ validation loss converged around epoch 27th. Both training and validation loss continues to slightly decrease after 20 epochs. The best model was chosen at which epochs resulted in the lowest validation loss. To enhance the performance of the model, the optimal model was reloaded and then trained with one, two, and three additional epochs using the whole original development set. The learning rate was set at 1 × 10⁻⁵. For comparison, the following three models were obtained: (a) the model trained with one additional epoch, (b) the model trained with two additional epochs, and (c) the model trained with three additional epochs. The model of setup (a) shows better performance than the models of setups (b), and (c). The variations in AUCROC values, however, are not significantly different. Models of setup (a), (b), and (c) have AUCPR values of around 0.38. The AUCROC value of 0.75 is consequently incredibly meaningful for addressing the problem of employee attrition prediction (Table 2).

Comparative analysis

We also trained several state-of-the-art models using conventional machine learning algorithms to predict employee attrition. The computational frameworks were constructed with Gradient Boosting (Friedman, 2001; Friedman, 2002), Extreme Gradient Boosting (Chen & Guestrin, 2016), Random Forest (Breiman, 2001), and Extremely Randomized Trees (Geurts, Ernst & Wehenkel, 2006) which are abbreviated as GB, XGB, RF, and ERT, respectively. All models were tuned with selected parameters via the grid search method to get the best model performance. We performed grid searches with 5-fold cross-validation on the whole original development set to tune the GB, XGB, RF, and ERT models. For the deep learning model, the validation set (split from the original development set) was used separately to determine the optimal epoch in which the model exhibits the smallest validation loss. The comparative analysis provides evidence that our proposed method can work better than other conventional machine learning models which have been commonly used to address these types of problems (Table 2). In addition, our research suggests constructing more effective prediction tools by using advances in deep learning. Since there is no significant difference in performances found among these setups (a), (b), and (c), we selected the model of setup (a) as our official model to conduct a comparative analysis with the other state-of-the-art approaches. Achieving a test AUCROC value of about 0.75, our model demonstrates a substantial improvement in predictive efficiency over other models. For an imbalanced-class dataset, the AUCPR value is more valuable than the accuracy when evaluating a binary classification model. Our models yielded an AUCPR value of 0.38, whereas other methods obtain AUCPR values of at most 0.37. Figure 3 visualizes the areas under the curves of all the models.

Statistical analysis

To investigate the robustness of the models, we replicated the experiments with different random seeds to avoid sampling bias. Table 3 gives information on the performance of all models over ten trials. The results indicated that our model outperforms the other machine learning models with AUCROC and AUCPR values of 0.7632 and 0.4792, respectively. In terms of AUCROC, the ERT and GB models show fairly equivalent performance, followed by the RF and XGB models. While the ERT model is still the best conventional approach compared to the others, the RF model achieves higher AUCPR values than the GB and XGB models. Also, to assess the statistical significance of the results, we used two-tailed independent t-tests with a confidence interval of 0.95 to compare the performance of our model to that of each machine learning model (Table 4). The p-values of these pairwise comparisons between our model and the other models confirm the statistical significance of these results.