XGBoost poor calibration for binary classification on a dataset with high class imbalance

Question

I've read a lot of threads/questions about this issue and I got conflicting answers.

I've trained an XGBoost model on tabular data to predict the risk for a specific event (ie a binary classifier). There are 43169 subjects and only 1690 events. I've taken into account this class imbalance with XGBoost's scale_pos_weight parameter. Hyperparameters were optimized with Bayesian Optimization (note: I have the same problem on a model without optimization).

I'm getting a reasonably well-discriminating model, however calibration looks awful:

Calibration using sklearn's sklearn.calibration.CalibratedClassifierCV doesn't improve the calibration at all (Isotonic and Sigmoid). It looks like XGBoost models cannot be calibrated with these methods.

My questions are:

• Did I do anything obviously wrong (see my code below)?
• Should I try another calibration method?
• Should I try another model (SVM, kNN, ...) and which one could be the most interesting considering that high class imbalance? I always thought that XGBoost was the gold-standard for tabular data. I've created several other models, including on data with class imbalance, and never got such poor calibration.

Thank you!

Load libraries

import pickle
import numpy as np
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
import pandas as pd
import matplotlib as pl
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
%matplotlib inline
import xgboost as xgb
import sklearn
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score
from sklearn.metrics import brier_score_loss, accuracy_score, average_precision_score, precision_score, recall_score, f1_score, roc_auc_score, make_scorer, roc_curve, auc, precision_recall_curve, confusion_matrix, plot_confusion_matrix
from sklearn.calibration import CalibratedClassifierCV, calibration_curve, CalibrationDisplay
from bayes_opt import BayesianOptimization

Create a held-out dataset

X_train, X_test, y_train, y_test = train_test_split(
    X,
    y,
    test_size=0.20,
    random_state=7,
    stratify = y
 )

Create a 5-fold stratified cross-validation

cv = StratifiedKFold(
    n_splits=5,
    shuffle=True,
    random_state=42
    )

Bayesian hyperparameters optimization

Define hyperparameters to explore and their limits

pbounds = {
    'learning_rate': (0.01, 1.0),
    'n_estimators': (10, 1000),
    'min_child_weight':(1, 10),
    'max_depth': (3,12),
    'subsample': (0, 1),  # Change for big datasets
    #'colsample': (0, 1.5),  # Change for datasets with lots of features
    'colsample_bytree': (0.3, 1),
    'gamma': (0, 5),
    'reg_alpha':(1e-5, 0.75),
    'reg_lambda':(1e-5, 0.45)}

Create the function to optimize

def xgboost_hyper_param(learning_rate,
                        n_estimators,
                        min_child_weight,
                        max_depth,
                        subsample,
                        #colsample,
                        colsample_bytree,
                        gamma,
                        reg_alpha,
                        reg_lambda):
    
    max_depth = int(max_depth)
    n_estimators = int(n_estimators)
    
    clf = xgb.XGBClassifier(
        objective='binary:logistic',
        eval_metric = 'auc',
        #tree_method = 'gpu_hist',
        #gpu_id = 0,
        use_label_encoder=False,
        booster = 'gbtree',
        scale_pos_weight = 24,
        learning_rate = learning_rate,
        n_estimators = n_estimators,
        min_child_weight = min_child_weight,
        max_depth = max_depth,
        subsample = subsample,
        #colsample = colsample,
        colsample_bytree = colsample_bytree,
        gamma = gamma,
        reg_alpha = reg_alpha,
        reg_lambda = reg_lambda
        )
    return np.mean(cross_val_score(clf, X_train, y_train, cv=cv, scoring='roc_auc'))

    optimizer = BayesianOptimization(
        f=xgboost_hyper_param,
        pbounds=pbounds,
        random_state=1,
)

Launch nested cross-validation for Bayesian Optimization

optimizer.maximize(init_points=20,
                  n_iter=5)

Get the best parameters

params = optimizer.max['params']
print("Here are the best parameters:")
print(params)

Converting the max_depth and n_estimator values from float to int

params['max_depth']= int(params['max_depth'])
params['n_estimators']= int(params['n_estimators'])

Create model

evals_result ={}
eval_set = [(X_train, y_train), (X_test, y_test)]

LungCancerRisk = xgb.XGBClassifier(

    **params,
    objective='binary:logistic',
    tree_method = 'exact',
    booster = 'gbtree',
    eval_metric=["auc"],
    scale_pos_weight = 24,
    eval_set=eval_set,
    use_label_encoder=False)

LungCancerRisk.fit(
    X_train,
    y_train,
    eval_set=eval_set
)

Assess calibration

clf_list = [
    (LungCancerRisk, "XGBoost"),
]

fig = plt.figure(figsize=(10, 10))
gs = GridSpec(4, 2)
colors = plt.cm.get_cmap("Dark2")

ax_calibration_curve = fig.add_subplot(gs[:2, :2])
calibration_displays = {}
for i, (clf, name) in enumerate(clf_list):
    clf.fit(X_train, y_train)
    display = CalibrationDisplay.from_estimator(
        clf,
        X_test,
        y_test,
        n_bins=10,
        name=name,
        ax=ax_calibration_curve,
        color=colors(i),
    )
    calibration_displays[name] = display

ax_calibration_curve.grid()
ax_calibration_curve.set_title("Calibration plots")

grid_positions = [(2, 0), (2, 1), (3, 0), (3, 1)]
for i, (_, name) in enumerate(clf_list):
    row, col = grid_positions[i]
    ax = fig.add_subplot(gs[row, col])

    ax.hist(
        calibration_displays[name].y_prob,
        range=(0, 1),
        bins=10,
        label=name,
        color=colors(i),
    )
    ax.set(title=name, xlabel="Mean predicted probability", ylabel="Count")

plt.tight_layout()
plt.show()

Answer 1

I think it's a pretty well-known thing that gradient boosted decision trees have a tendency towards poor calibration. However, what you get looks pretty extreme. Part of the answer is right there in the XGBoost documentation, where it says:"If you care about predicting the right probability [...] in such a case, you cannot re-balance the dataset" (i.e. you should not use scale_pos_weight). I suspect if you avoid using it, you could get a predictions that can be calibrated a bit more reasonably.

In terms of other models, I would usually expect some form of logistic regression to be reasonably well calibrated out-of-the-box.