DOCSEARCHbyAlgolia
Submit Docs

xgbse (XGBoost Survival Embeddings)

repository·main·Indexed 43 days ago

https://github.com/loft-br/xgboost-survival-embeddings

A library for survival analysis that combines XGBoost with embeddings and debiased estimators to provide calibrated survival curves and confidence intervals. It follows the scikit-learn API and includes models like XGBSEKaplanTree, XGBSEKaplanNeighbors, and XGBSEDebiasedBCE, along with a BootstrapEstimator for non-parametric CIs and survival-specific metrics such as the concordance index and Brier score.

Tokens
24K
Snippets
81
Records
100
Agent score
50%

What's inside xgbse

  1. Overview of XGBoost Survival Embeddings (xgbse)

    main

    XGBoost Survival Embeddings (xgbse)

    xgbse is a library designed to bridge the gap between statistical rigor and computational efficiency in Survival Analysis. It builds upon the expressive and efficient xgboost implementation by adding a layer of statistical properties.

    Key Features:

    • Calibrated Survival Curves: Provides unbiased survival curves with confidence intervals rather than just point predictions.
    • High Predictive Power: Maintains performance competitive with vanilla xgboost.
    • Efficiency: Offers an easy-to-use and computationally efficient implementation.
    • Explainability: Supports model explainability through the use of prototypes.

    This package is suitable for applications such as predictive maintenance, customer churn, credit risk, and asset liquidity risk.

  2. Overview of xgbse

    main

    What is xgbse?

    xgbse (XGBoost Survival Embeddings) is a library designed to bridge the gap between traditional statistical survival models and highly efficient algorithmic models like XGBoost.

    It provides:

    • Calibrated survival curves: Generates unbiased survival curves with confidence intervals rather than just point predictions.
    • High predictive power: Maintains performance competitive with vanilla xgboost.
    • Efficiency: An easy-to-use implementation for industrial applications like predictive maintenance, customer churn, and credit risk.
    • Explainability: Uses prototypes to provide insights into model decisions.

  3. XGBSEKaplanTree: Efficient single-tree Kaplan-Meier

    main

    The XGBSEKaplanTree is a high-efficiency simplification of the nearest neighbors approach.

    How it works:

    • Instead of an ensemble, it fits a single tree via XGBoost.
    • It calculates Kaplan-Meier curves for every leaf in that single tree.
    • At scoring time, a query simply identifies which leaf the sample falls into to retrieve the pre-calculated KM curve.

    Trade-offs:

    • Pros: Extremely efficient and scales easily to millions of examples.
    • Cons: Lower predictive power than ensemble-based methods. To mitigate this, you can use XGBSEBootstrapEstimator to fit a forest of XGBSEKaplanTree models to reduce variance.

  4. XGBSEDebiasedBCE: Logistic regression with time windows

    main

    The XGBSEDebiasedBCE approach uses XGBoost embeddings as input for a set of logistic regressions.

    How it works:

    • It trains logistic regressions to predict survival at discrete, user-defined time windows.
    • To avoid bias caused by censoring, the logistic regressions estimate the point probabilities ($di/ni$) used in the Kaplan-Meier formula.
    • It produces nearly unbiased survival curves and can provide confidence intervals via bootstrapping.
    • Training and scoring are parallelized using joblib, making it scalable to millions of samples.

  5. XGBSEStackedWeibull: Stacking XGBoost with Weibull AFT

    main

    The XGBSEStackedWeibull approach combines the discrimination power of XGBoost with the statistical rigor of a parametric model.

    How it works:

    • An XGBoost survival model is trained to act as a risk estimator.
    • The predicted risk metric is then used as the sole independent variable in a Weibull Accelerated Failure Time (AFT) regression.
    • This results in calibrated survival curves and benefits from the strong extrapolation capabilities of the Weibull distribution.
    • This method is generally faster to train than XGBSEDebiasedBCE because it involves stacking with a single-variable model.

  6. Estimate survival curves and confidence intervals

    main

    Survival curves and confidence intervals (CIs) can be estimated using the Exponential Greenwood formula.

    For XBGSEKaplanTree and XBGSEKaplanNeighbors

    These models support CI estimation out-of-the-box using the return_ci=True argument in .predict(). Note that for XBGSEKaplanTree, the ci_width parameter must be passed during .fit(), whereas for XBGSEKaplanNeighbors, it is passed during .predict().

    For XGBSEDebiasedBCE

    This model does not support CIs out-of-the-box. Instead, use the XGBSEBootstrapEstimator to obtain non-parametric confidence intervals via bootstrapping.

    Using XGBSEBootstrapEstimator

    You can wrap any base model (including XBGSEKaplanTree or XBGSEKaplanNeighbors) with XGBSEBootstrapEstimator to get more stable, non-parametric CIs. This estimator also allows building CIs for interval probabilities, which the Exponential Greenwood formula does not support.

    # Example for KaplanNeighbors with CI
# ci_width is passed at .predict()
model = XGBSEKaplanNeighbors(PARAMS_XGB_AFT, N_NEIGHBORS)
model.fit(X_train, y_train, time_bins=TIME_BINS)
mean, upper_ci, lower_ci = model.predict(X_valid, return_ci=True, ci_width=0.99)

# Example for KaplanTree with CI
# ci_width is passed at .fit()
xgbse_model = XGBSEKaplanTree(PARAMS_TREE)
xgbse_model.fit(X_train, y_train, time_bins=TIME_BINS, ci_width=0.99)
mean, upper_ci, lower_ci = xgbse_model.predict(X_valid, return_ci=True)

# Example using Bootstrap Estimator
base_model = XGBSEDebiasedBCE(PARAMS_XGB_AFT, PARAMS_LR)
bootstrap_estimator = XGBSEBootstrapEstimator(base_model, n_estimators=20)
bootstrap_estimator.fit(X_train, y_train, validation_data=(X_valid, y_valid), early_stopping_rounds=10, time_bins=TIME_BINS)
mean, upper_ci, lower_ci = bootstrap_estimator.predict(X_valid, return_ci=True)

  7. XGBSEKaplanNeighbors: Kaplan-Meier on nearest neighbors

    main

    The XGBSEKaplanNeighbors approach treats the XGBoost leaf embeddings as a way to define similarity between samples.

    How it works:

    • It uses Hamming distance on the sparse, high-dimensional leaf embeddings to find similar elements.
    • For each sample, it identifies a set of $N$ neighbors and applies the Kaplan-Meier estimator to that neighbor-set to derive survival estimates.
    • It provides vectorized survival curves including upper and lower confidence intervals based on the Exponential Greenwood formula.

    Note: This method can be computationally expensive at very large scales due to the nearest neighbor search requirements during both training and scoring. It is recommended to use dart as the booster to maintain the integrity of the leaf co-occurrence logic.

  8. How xgbse works: Solving XGBoost survival analysis limitations

    main

    Standard XGBoost survival implementations (Cox and AFT) are efficient at ordering individuals by risk (C-index), but they suffer from three main statistical shortcomings:

    1. Lack of survival curves: They typically output point estimates (expected survival time) rather than full survival curves.
    2. No confidence intervals: They do not provide estimates of uncertainty.
    3. Poor calibration: AFT model outputs are highly sensitive to hyperparameters like aft_loss_distribution_scale. Small changes in these parameters can drastically shift predicted survival times even if the C-index remains stable.

    xgbse addresses these issues by using XGBoost as a feature transformer. It extracts embeddings from the terminal nodes (leaves) of the tree ensemble. These embeddings are then used as inputs for more statistically rigorous models (Logistic Regression, Weibull AFT, or Kaplan-Meier estimators) to provide calibrated survival curves and confidence intervals.

  9. Get confidence intervals for predicted survival curves

    main

    You can obtain confidence intervals for predicted survival curves using two different approaches in xgbse:

    1. Numerical Form (Exponential Greenwood formula): Use XGBSEKaplanTree or XGBSEKaplanNeighbors. These models support the return_ci=True parameter in the .predict() method to return the mean survival curve along with upper and lower bounds.

    2. Non-parametric Form (Bagging): Use XGBSEBootstrapEstimator to wrap any XGBSE model. This meta-estimator uses bagging to estimate confidence intervals, which can improve performance at the cost of increased computation time.

    To use these, ensure your data is converted to the structured format using convert_to_structured(T, E) and that you provide time_bins during the .fit() call.

    # Numerical Form example
from xgbse import XGBSEKaplanTree

model = XGBSEKaplanTree(params)
model.fit(X_train, y_train, time_bins=TIME_BINS)
mean, upper_ci, lower_ci = model.predict(X_test, return_ci=True)

# Non-parametric Form example
from xgbse import XGBSEBootstrapEstimator

base_model = XGBSEKaplanTree(params)
bootstrap_estimator = XGBSEBootstrapEstimator(base_model, n_estimators=100)
bootstrap_estimator.fit(X_train, y_train, time_bins=TIME_BINS)
mean, upper_ci, lower_ci = bootstrap_estimator.predict(X_test, return_ci=True)

  11. Basic usage of XGBoost Survival Embeddings

    main

    The xgbse package follows the scikit-learn API. To use it, you must provide X (features) and y (survival data). Because survival data requires both time and event information, y must be a numpy structured array. You can use xgbse.converters.convert_to_structured to transform your duration and event columns into this required format.

    .predict() returns a pandas DataFrame where each column represents a time window and the values represent the probability of survival before or exactly at that time window.

    # importing dataset from pycox package
from pycox.datasets import metabric

# importing model and utils from xgbse
from xgbse import XGBSEKaplanNeighbors
from xgbse.converters import convert_to_structured

# getting data
df = metabric.read_df()

# splitting to X, y format
X = df.drop(['duration', 'event'], axis=1)
y = convert_to_structured(df['duration'], df['event'])

# fitting xgbse model
xgbse_model = XGBSEKaplanNeighbors(n_neighbors=50)
xgbse_model.fit(X, y)

# predicting
event_probs = xgbse_model.predict(X)
event_probs.head()

  12. Use XGBSEStackedWeibull for survival prediction

    main

    The XGBSEStackedWeibull class implements a stacking approach that combines the discrimination power of XGBoost with the statistical rigor of a Weibull Accelerated Failure Time (AFT) model.

    How it works:

    1. An XGBoost model is trained to predict a risk metric from the input features.
    2. This risk metric is then used as the sole independent variable in a Weibull AFT regression.

    Key Characteristics:

    • Speed: Faster training compared to XGBSEDebiasedBCE because it uses a single-variable parametric model.
    • Extrapolation: Better extrapolation capabilities than XGBSEKaplanNeighbors or XGBSEKaplanTree because it avoids the cure fraction problem.
    • Assumptions: It assumes a specific parametric shape for the survival curve (Weibull distribution).

    Usage Workflow:

    1. Initialize with xgb_params and weibull_params.
    2. Call .fit(X, y, ...) where y is a structured array containing event indicators and times.
    3. Call .predict(X) to get survival probabilities.
    from xgbse import XGBSEStackedWeibull

# Initialize
model = XGBSEStackedWeibull(
    xgb_params={'max_depth': 3},
    weibull_params={'disp': False}
)

# Fit (y must be a structured array with event and time)
model.fit(X, y, time_bins=[1, 5, 10])

# Predict
survival_probs = model.predict(X_test)