Home Artificial Intelligence Improving Generalization in Survival Models

Improving Generalization in Survival Models

0
Improving Generalization in Survival Models

Traditional Approach

Many existing implementations on survival evaluation start off with a dataset containing one commentary per individual (patients in a health study, employees within the attrition case, clients within the client churn case, and so forth). For these individuals we typically have two key variables: one signaling the event of interest (an worker quitting) and one other measuring time (how long they’ve been with the corporate, as much as either today or their departure). Along with these two variables, we then have explanatory variables with which we aim to predict the danger of every individual. These features can include the job role, age or compensation of the worker, for instance.

Moving on, most implementations on the market take a survival model (from simpler estimators resembling Kaplan Meier to more complex ones like ensemble models and even neural networks), fit them over a train set after which evaluate over a test set. This train-test split is frequently performed over the person observations, generally making a stratified split.

In my case, I began with a dataset that followed several employees in an organization monthly until December 2023 (in case the worker was still at the corporate), or until the month they left the corporate — the event date:

Taking the last record of every worker — Image by creator

With the intention to adapt my data to the survival case, I took the last commentary of every worker as shown in the image above (the blue dots for energetic employees, and the red crosses for workers who left). At that time for every worker, I recorded whether the event had occurred at that date or not (in the event that they were energetic or in the event that they had left), their tenure in months at the moment, and all their explanatory variables. I then performed a stratified train-test split over this data, like this:

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

# We load our dataset with several observations (record_date) per worker (employee_id)
# The event column indicates if the worker left on that given month (1) or if the worker was still energetic (0)
df = pd.read_csv(f'{FILE_NAME}.csv')

# Making a label where positive events have tenure and negative events have negative tenure - required by Random Survival Forest
df_model['label'] = np.where(df_model['event'], df_model['tenure_in_months'], - df_model['tenure_in_months'])

df_train, df_test = train_test_split(df_model, test_size=0.2, stratify=df_model['event'], random_state=42)

After performing the split, I proceeded to suit a model. On this case, I selected to experiment with a Random Survival Forest using the scikit-survival library.

from sklearn.preprocessing import OrdinalEncoder
from sksurv.datasets import get_x_y
from sksurv.ensemble import RandomSurvivalForest

cat_features = [] # list of all the specific features
features = [] # list of all of the features (each categorical and numeric)

# Categorical Encoding
encoder = OrdinalEncoder(handle_unknown='use_encoded_value', unknown_value=-1)
encoder.fit(df_train[cat_features])

df_train[cat_features] = encoder.transform(df_train[cat_features])
df_test[cat_features] = encoder.transform(df_test[cat_features])

# X & y
X_train, y_train = get_x_y(df_train, attr_labels=['event','tenure_in_months'], pos_label=1)
X_test, y_test = get_x_y(df_test, attr_labels=['event','tenure_in_months'], pos_label=1)

# Fit the model
estimator = RandomSurvivalForest(random_state=RANDOM_STATE)
estimator.fit(X_train[features], y_train)

# Store predictions
y_pred = estimator.predict(X_test[features])

After a fast run using the default settings of the model, I used to be thrilled with the test metrics I saw. To start with, I used to be getting a concordance index above 0.90 within the test set. The concordance index is a measure of how well the model predicts the order of events: it reflects whether employees predicted to be at high risk were indeed those leaving the corporate first. An index of 1 corresponds to perfect prediction accuracy, while an index of 0.5 indicates a prediction no higher than random likelihood.

I used to be particularly desirous about seeing if the workers who left within the test set matched with essentially the most dangerous employees in line with the model. Within the case of the Random Survival Forest, the model returns the danger scores of every commentary. I took the share of employees who left the corporate within the test set, and used it to filter essentially the most dangerous employees in line with the model. The outcomes were very solid, with the workers flagged with essentially the most risk matching almost perfectly with the actual leavers, with an F1 rating above 0.90 within the minority class.

from lifelines.utils import concordance_index
from sklearn.metrics import classification_report

# Concordance Index
ci_test = concordance_index(df_test['tenure_in_months'], -y_pred, df_test['event'])
print(f'Concordance index:{ci_test:0.5f}n')

# Match essentially the most dangerous employees (in line with the model) with the workers who left
q_test = 1 - df_test['event'].mean()

thr = np.quantile(y_pred, q_test)
risky_employees = (y_pred >= thr) * 1

print(classification_report(df_test['event'], risky_employees))

Getting +0.9 metrics on the primary run should set off an alarm: was the model really in a position to predict whether an worker was going to remain or leave with such confidence? Imagine this: we submit our predictions saying which employees are almost certainly to go away. Nevertheless, a pair months go by, and HR then reaches us fearful, saying that the individuals who left through the last period, didn’t exactly match with our predictions, no less than at the speed it was expected from our test metrics.

We’ve got two fundamental problems here: the primary one is that our model isn’t extrapolating quite in addition to we thought. The second, and even worse, is that we weren’t in a position to measure this lack of performance. First, I’ll show a straightforward way we are able to estimate how well our model is really extrapolating, after which I’ll speak about one potential reason it might be failing to accomplish that, and how one can mitigate it.

Estimating Generalization Capabilities

The important thing here is gaining access to panel data, that’s, several records of our individuals over time, up until the time of event or the time the study ended (the date of our snapshot, within the case of worker attrition). As a substitute of discarding all this information and keeping only the last record of every worker, we could use it to create a test set that can higher reflect how the model performs in the longer term. The thought is sort of easy: suppose we’ve got monthly records of our employees up until December 2023. We could move back, say, 6 months, and pretend we took the snapshot in June as a substitute of December. Then, we’d take the last commentary for workers who left the corporate before June 2023 as positive events, and the June 2023 record of employees who survived beyond that date as negative events, even when we already know a few of them eventually left afterwards. We’re pretending we don’t know this yet.

We take a snapshot in June 2023 and use the next period as our test set — Image by creator

As the image above shows, I take a snapshot in June, and all employees who were energetic at the moment are taken as energetic. The test dataset takes all those energetic employees at June with their explanatory variables as they were on that date, and takes the newest tenure they achieved by December:

test_date = '2023-07-01'

# Choosing training data from records before the test date and taking the last commentary per worker
df_train = df[df.record_date < test_date].reset_index(drop=True).copy()
df_train = df_train.groupby('employee_id').tail(1).reset_index(drop=True)
df_train['label'] = np.where(df_train['event'], df_train['tenure_in_months'], - df_train['tenure_in_months'])

# Preparing test data with records of energetic employees on the test date
df_test = df[(df.record_date == test_date) & (df['event']==0)].reset_index(drop=True).copy()
df_test = df_test.groupby('employee_id').tail(1).reset_index(drop=True)
df_test = df_test.drop(columns = ['tenure_in_months','event'])

# Fetching the last tenure and event status for workers within the test dataset
df_last_tenure = df[df.employee_id.isin(df_test.employee_id.unique())].reset_index(drop=True).copy()
df_last_tenure = df_last_tenure.groupby('employee_id').tail(1).reset_index(drop=True)

df_test = df_test.merge(df_last_tenure[['employee_id','tenure_in_months','event']], how='left')
df_test['label'] = np.where(df_test['event'], df_test['tenure_in_months'], - df_test['tenure_in_months'])

We fit our model again on this latest train data, and once we finish we make our predictions for all employees who were energetic on June. We then compare these predictions to the actual consequence of July — December 2023 — that is our test set. If those employees we marked as having essentially the most risk left through the semester, and people we marked as having the bottom risk didn’t leave, or left slightly late within the period, then our model is extrapolating well. By shifting our evaluation back in time and leaving the last period for evaluation, we are able to have a greater understanding of how well our model is generalizing. After all, we could take this one step further and perform some variety of time-series cross validation. For instance, we could iterate this process again and again, every time moving 6 months back in time, and evaluating the model’s accuracy over several time frames.

After training our model once more, we now see a drastic decrease in performance. To start with, the concordance index is now around 0.5 — corresponding to that of a random predictor. Also, if we attempt to match the ‘n’ most dangerous employees in line with the model with the ‘n’ employees who left within the test set, we see a really poor classification with a 0.15 F1 for the minority class:

So clearly there’s something mistaken, but no less than we at the moment are in a position to detect it as a substitute of being misled. The fundamental takeaway here is that our model performs well with a conventional split, but doesn’t extrapolate when doing a time-based split. It is a clear sign that a while bias could also be present. Briefly, time-dependent information is being leaked and our model is overfitting over it. That is common in cases like our worker attrition problem, when the dataset comes from a snapshot taken at some date.

Time Bias

The issue cuts all the way down to this: all our positive observations (employees who left) belong to past dates, and all our negative observations (currently energetic employees) are all measured on the identical date — today. If there may be a single feature that reveals this to the model, then as a substitute of predicting risk we can be predicting if an worker was recorded in December 2023 or before. This could possibly be very subtle. For instance, one feature we could possibly be using is the engagement rating of the workers. This feature could well show some seasonal patterns, and measuring it at the identical time for energetic employees will certainly introduce some bias within the model. Possibly in December, through the holiday season, this engagement rating tends to diminish. The model will see a low rating related to all energetic employees, so it might learn to predict that every time the engagement runs low, the churn risk also goes down, when in truth it needs to be the alternative!

By now, a straightforward yet quite effective solution for this problem needs to be clear: as a substitute of taking the last commentary for every energetic worker, we could just pick a random month from all their history throughout the company. This may strongly reduce the possibilities of the model picking on any temporal patterns that we don’t want it to overfit on:

For the energetic employees, we take random records slightly than their last one — Image by creator

In the image above we are able to see that we at the moment are spanning a broader set of dates for the energetic employees. As a substitute of using their blue dots at June 2023, we take the random orange dots as a substitute, and record their variables on the time, and the tenure they’d to this point in the corporate:

np.random.seed(0)

# Select training data before the test date
df_train = df[df.record_date < test_date].reset_index(drop=True).copy()

# Create an indicator for whether an worker eventually churns throughout the train set
df_train['indicator'] = df_train.groupby('employee_id').event.transform(max)

# Isolate records of employees who left, and store their last commentary
churn = df_train[df_train.indicator==1].reset_index(drop=True).copy()
churn = churn.groupby('employee_id').tail(1).reset_index(drop=True)

# For workers who stayed, randomly pick one commentary from their historic records
stay = df_train[df_train.indicator==0].reset_index(drop=True).copy()
stay = stay.groupby('employee_id').apply(lambda x: x.sample(1)).reset_index(drop=True)

# Mix churn and stay samples into the brand new training dataset
df_train = pd.concat([churn,stay], ignore_index=True).copy()
df_train['label'] = np.where(df_train['event'], df_train['tenure_in_months'], - df_train['tenure_in_months'])
del df_train['indicator']

# Prepare the test dataset similarly, using only the snapshot from the test date
df_test = df[(df.record_date == test_date) & (df.event==0)].reset_index(drop=True).copy()
df_test = df_test.groupby('employee_id').tail(1).reset_index(drop=True)
df_test = df_test.drop(columns = ['tenure_in_months','event'])

# Get the last known tenure and event status for workers within the test set
df_last_tenure = df[df.employee_id.isin(df_test.employee_id.unique())].reset_index(drop=True).copy()
df_last_tenure = df_last_tenure.groupby('employee_id').tail(1).reset_index(drop=True)

df_test = df_test.merge(df_last_tenure[['employee_id','tenure_in_months','event']], how='left')
df_test['label'] = np.where(df_test['event'], df_test['tenure_in_months'], - df_test['tenure_in_months'])

We then train our model once more, and evaluate it over the identical test set we had before. We now see a concordance index of around 0.80. This isn’t the +0.90 we had earlier, but it surely definitely is a step up from the random-chance level of 0.5. Regarding our interest in classifying employees, we’re still very far off the +0.9 F1 we had before, but we do see a slight increase in comparison with the previous approach, especially for the minority class.

LEAVE A REPLY

Please enter your comment!
Please enter your name here