Team Datanaut Capstone Project: Predicting Employee Attrition Using Machine Learning
As we know, employee attrition is one of the most significant challenges that any organization faces nowadays. The primary reason is that it becomes expensive to fill vacant positions with new candidates, which leads to a decrease in overall performance.
In this article, Team Datanaut will perform EDA and find out which factors are related to attrition and construct a classification model that can predict if an employee is going to leave the organization (binary classification: attrition/no attrition) based on demography, position, salary, and other relevant variables.
Goals:
1. To build a classification model that predicts whether an employee will leave the company or stay.
2. To determine the factors with high importance in predicting employee attrition, so the organization can act on them.
The basic procedures used to achieve these goals are:
1. Gathering the Data
The dataset used for this project is the IBM HR Analytics Employee Attrition dataset, obtained from Kaggle.
First, the required libraries were imported:
Then the data was loaded into the notebook;
Next was to check the shape of the dataset;
The dataset has 1470 rows and 35 columns.
We then check for the data information:
A quick look at the data with HR.info() shows the columns span both numerical (26 int64 columns) and categorical (9 object columns) data . Things like Age, DailyRate, Department, EducationField, JobRole, MaritalStatus, OverTime, and the target column, Attrition.
2. Data Cleaning and Validation
Before any analysis can be done, the data must be void of inconsistencies or errors unless you’re aiming at a flawed result.
The first step is to check for null values:
There were no missing values anywhere in the dataset. Every one of the 35 columns returned a count of 0.
Next, duplicate rows were checked for:
This also returned 0 . The data is void of duplicate rows.
While exploring the data further, three columns turned out to be constants across every single row, and one turned out to be a pure identifier:
- EmployeeCount : Always 1
- Over18 : Always "Y"
- StandardHours : Always 80
- EmployeeNumber : 1470 unique values (a pure ID, one per row)
None of these carry any predictive signal, so they were flagged to be dropped before modeling.
The target variable, Attrition, was also checked:
3. Exploratory Data Analysis (EDA)
This is the core of any data science project. EDA helps to discover hidden patterns in data and uncovers relationships between variables.
Univariate Analysis
Attrition is checked as a proportion of the whole dataset:
Roughly 84% of employees stayed (No) and about 16% left (Yes). This shows that the data is imbalanced, which means accuracy alone won’t be a reliable metric for evaluating the models later . Recall and ROC-AUC matter more here.
Bivariate Analysis : Factors associated with employee attrition
Interpretation: The boxplot indicates that employees who remained with the company tend to be older than those who left. The median age is higher for employees with no attrition, while employees who left are generally younger. Although there are a few older employees who also left the company, the overall pattern suggests that younger employees have a higher likelihood of attrition.
Interpretation: The KDE plot shows that most employees live within a short distance of the workplace, with the highest concentration between 2 and 5 units from home. Employees who remained with the company are more concentrated at shorter commuting distances, while employees who left tend to have relatively longer commuting distances. This suggests that longer commutes may be associated with higher attrition.
Interpretation: Education level 3 has the highest number of employees. It also has the highest number of attrition cases, but this is likely because there are more employees in this category overall.
Interpretation: Research & Development has the highest number of employees, and consequently the highest number of attrition cases . Likely because there are more employees in this category overall, not necessarily because the department is worse for retention.
4. Summary of EDA findings
- Age: Older employees tend to stay; younger employees are more likely to leave.
- DistanceFromHome: Employees who remained are concentrated at shorter commuting distances, while those who left tendk to commute farther.
- Gender: Males are more likely to quit than females.
- Department: Research & Development has the most attrition cases, but this tracks with it also being the largest department.
- Education: Employees at education level 3 leave more than any other level.
- EducationField: Employees from Life Sciences are more likely to quit than those from other fields.
- EnvironmentSatisfaction: Employees with the lowest satisfaction level (1) leave the most.
- JobInvolvement: Employees with a job involvement level of 3 leave the most.
- JobSatisfaction: Employees with a job satisfaction level of 3 leave the most.
- MaritalStatus: Single employees are more likely to quit than married or divorced employees.
- JobRole: Laboratory Technicians, Sales Representatives, and Human Resources see disproportionately more attrition than other roles.
- OverTime: Employees who work overtime are far more likely to leave than those who don’t.
- MonthlyIncome: Employees who stayed generally earn more than those who left.
- YearsInCurrentRole: Employees who stay a shorter time in their current role are more likely to leave.
Now that there’s a good idea of what’s going on in the data, the next step is to prepare it for modeling.
4. Feature Preprocessing
After the whole EDA, the data needs to be prepared for modeling. This is an essential step because models only accept data processed in a certain format.
i. Dropping unhelpful columns. As noted during data cleaning, EmployeeCount, Over18, and StandardHours are constants across the whole dataset, and EmployeeNumber is just a row identifier. None of these carry predictive signal, so they were dropped.
ii. Encoding categorical features. ML models only accept data in numerical format, so the categorical columns (Attrition, BusinessTravel, Department, EducationField, Gender, JobRole, MaritalStatus, Over18, OverTime) were encoded using Scikit-learn’s LabelEncoder:
Feature correlation. A heatmap was used to check the correlation between features, since ideally the features should be reasonably independent of each other:
Train-test split. The dataset was split into explanatory variables ( X), and target variable ( y), then further split into train and test sets in a ratio of 80:20 respectively. Splitting the data helps to assess the model’s performance on unseen data after being trained on the training data:
This produced X_train: (1176, 30), X_test: (294, 30), y_train: (1176,), y_test: (294,).
Get Gideon Silas’s stories in your inbox
Join Medium for free to get updates from this writer.
Feature scaling. Before modeling, the data needs to be scaled to handle skewed features. Scikit-learn’s StandardScaler ensures that for each feature the mean is 0 and the variance is 1, bringing all the features to the same magnitude. The scaler is fit only on the training data, then reused on the test data, to avoid leaking test-set information into the transform:
5. Modelling
The employee attrition problem is a classification problem, therefore about five classification models were trained and compared.
Training Model 1 : Logistic Regression
Interpretation:
Overall accuracy came out to 0.874 . That is about 87.4% of employees in the test set were correctly classified. At first glance this looks like an excellent model. However, accuracy alone can be misleading, especially when the classes are imbalanced (84% stayed vs. 16% left).
- True Negative : 239 (non-attrition rightly predicted as non-attrition)
- False Positive : 8 (non-attrition predicted as attrition)
- True Positive : 18 (attrition rightly predicted as attrition)
- False Negative : 29 (attrition predicted as non-attrition)
ROC-AUC: 0.675 : The model performs fairly against the dataset, but there’s clear room to improve how well it catches the minority (attrition) class.
Training Model 2 : XGBClassifier
Accuracy: 0.861 · ROC-AUC: 0.607
Training Model 3 : Random Forest Classifier
Accuracy: 0.837 · ROC-AUC: 0.550
Training Model 4 : Decision Tree Classifier
Accuracy: 0.782 · ROC-AUC: 0.595
The training-set report for this model was a perfect 1.00 across the board — a classic sign of overfitting, since performance drops sharply on the test set.
Training Model 5 : Gradient Boosting Classifier
Accuracy: 0.840 · ROC-AUC: 0.578
6. Model Evaluation
Since the classes are highly imbalanced (84% stayed vs. 16% left), the models were compared using ROC-AUC rather than accuracy alone .
So far, Logistic Regression has the highest ROC-AUC score and was chosen as the model to tune further.
Tuning the Logistic Regression model
1. Using class_weight = "balanced"
Accuracy: 0.752 · ROC-AUC: 0.757
Balancing the class weights pushed recall on the attrition class up sharply (from 0.38 to 0.77) at the cost of precision and overall accuracy. A reasonable trade-off, since for this problem missing a leaver is generally more costly than flagging a few extra stayers.
2. Using SMOTE (Synthetic Minority Over sampling)
SMOTE creates synthetic examples of the minority class (Attrition = 1) so the model sees a more balanced training set. It’s applied only to the training data, resampling before the split, or resampling the test set, would leak synthetic information into evaluation and give an unrealistically optimistic score.
Before SMOTE: {0: 986, 1: 190} → After SMOTE: {0: 986, 1: 986}
Accuracy: 0.8741 · ROC-AUC: 0.8061
SMOTE and class_weight = ‘balanced’ both address the same problem , the model barely seeing attrition cases during training , in different ways, and they gave similar recall gains here.
3. Using GridSearchCV
Best Parameters: {’C’: 1, 'class_weight’: None, 'penalty’: 'l2’, 'solver’: 'liblinear’}
Best ROC-AUC (cross-validated): 0.8069
Accuracy: 0.8741 · ROC-AUC: 0.8061
GridSearchCV essentially confirmed the model’s default hyperparameters were already close to optimal for this data. The ROC-AUC only nudged up slightly (0.8057 - 0.8061), and accuracy stayed flat.
Choosing the classification threshold systematically
A default 0.5 threshold isn’t necessarily the right cutoff, especially for an imbalanced problem where missing a leaver (false negative) is usually more costly than flagging a stayer (false positive). Instead of picking a threshold by hand, the precision-recall curve was used to find the threshold that maximizes the F2 score , which weights recall twice as heavily as precision:
Best threshold (max F2): 0.150. Precision at this threshold: 0.352. Recall at this threshold: 0.787
Applying a threshold of 0.35 (a slightly more conservative cutoff chosen from the same precision/recall trade-off curve) to the predicted probabilities:
This threshold gives a much better balance for the business goal: recall on the attrition class rose from 0.38 (at the default 0.5 cutoff) to 0.55, at only a small cost to overall accuracy.
7. Feature Importance
The logistic regression coefficients were used to identify which variables most strongly influence employee attrition, sorted by absolute value so the largest effects .
Positive or negative come first:
Interpretation of the features: Overtime emerged as the strongest positive predictor, indicating that employees who work overtime are more likely to leave the company. Other factors associated with increased attrition include years since last promotion, number of companies previously worked at, department, marital status, and distance from home. Conversely, employees with more years in their current role, longer tenure with their current manager, higher job satisfaction, higher job level, greater work experience, better environment satisfaction, higher job involvement, and older age were less likely to leave. Overall, the findings suggest that workload, career progression, and employee engagement play significant roles in employee retention.
Conclusion
As a team, we have worked through the full data science journey , from understanding the problem, through EDA, to building and tuning a model that predicts employee attrition.
Employees who work overtime, haven’t been promoted recently, have switched employers frequently, and are early in their tenure with the company or their current manager are the ones most at risk of leaving. Logistic Regression, tuned via GridSearchCV with a custom decision threshold, gave the best balance of ROC-AUC (0.806) and recall (0.55) on the minority attrition class among the five models tested.
Recommendations
- Monitor and manage overtime load, especially for employees showing early signs of disengagement.
- Revisit promotion cadence : Employees who haven’t been promoted in a while are a flight risk.
- Pay closer attention to newer employees and those newly assigned to a manager or role, since shorter tenure correlates with higher attrition.
- Track environment and job satisfaction scores as a leading indicator, not just an HR survey afterthought.
All codes for this project can be found on:
