Tree DRO

This notebook fits distributionally robust tree ensembles with XGBoost and LightGBM. The examples use the same binary classification data and booster settings so that the KL-DRO and CVaR-DRO setup is easy to compare.

See the tree-model tutorial for the objective formulations, booster configuration, and interpretation of the fitted robust objective.

Setup

Import the two robust objectives for each supported boosting backend.

from sklearn.model_selection import train_test_split
from sklearn.datasets import make_classification
from dro.tree_model.xgb import KLDRO_XGB, CVaRDRO_XGB
from dro.tree_model.lgbm import KLDRO_LGBM, CVaRDRO_LGBM

Dataset

X, y = make_classification(
    n_samples=1000,
    n_features=5,
    n_informative=3,
    n_redundant=1,
    n_classes=2,
    random_state=42,
)

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

XGBoost

KL-DRO

Instantiating KLDRO_XGB selects the KL-robust objective. The eps argument belongs to the DRO formulation, whereas xgb_config contains the ordinary boosting parameters.

kl_xgb = KLDRO_XGB(eps=0.001)
xgb_config = {
    "max_depth": 2,
    "learning_rate": 1,
    "num_boost_round": 4,
}
kl_xgb.update(xgb_config)
kl_xgb.fit(X_train, y_train)

predictions = kl_xgb.predict(X_test)
print(f"Test accuracy: {(predictions == y_test).mean():.3f}")
print(f"Robust training objective: {kl_xgb.robust_obj:.6f}")

Test accuracy: 0.950
Robust training objective: 0.130190

CVaR-DRO

Switching to CVaRDRO_XGB changes the robust objective without changing the fitting workflow. Here, eps=0.1 sets the CVaR quantile level used to emphasize observations with larger fitted losses.

cvar_xgb = CVaRDRO_XGB(eps=0.1)
xgb_config = {
    "max_depth": 2,
    "learning_rate": 1,
    "num_boost_round": 4,
}
cvar_xgb.update(xgb_config)
cvar_xgb.fit(X_train, y_train)

predictions = cvar_xgb.predict(X_test)
print(f"Test accuracy: {(predictions == y_test).mean():.3f}")
print(f"Robust training objective: {cvar_xgb.robust_obj:.6f}")

Test accuracy: 0.945
Robust training objective: 0.137117

LightGBM

The LightGBM wrappers use the same separation between the DRO constructor and the booster configuration. The following examples keep the configuration aligned with the XGBoost examples.

KL-DRO

kl_lgbm = KLDRO_LGBM(eps=0.1)
lgbm_config = {
    "max_depth": 2,
    "learning_rate": 1,
    "num_boost_round": 4,
}
kl_lgbm.update(lgbm_config)
kl_lgbm.fit(X_train, y_train)

predictions = kl_lgbm.predict(X_test)
print(f"Test accuracy: {(predictions == y_test).mean():.3f}")
print(f"Robust training objective: {kl_lgbm.robust_obj:.6f}")

Test accuracy: 0.940
Robust training objective: 0.139228

CVaR-DRO

As above, the class and eps define the robust objective; the configuration continues to describe only how the ensemble is built.

cvar_lgbm = CVaRDRO_LGBM(eps=0.1)
lgbm_config = {
    "max_depth": 2,
    "learning_rate": 1,
    "num_boost_round": 4,
}
cvar_lgbm.update(lgbm_config)
cvar_lgbm.fit(X_train, y_train)

predictions = cvar_lgbm.predict(X_test)
print(f"Test accuracy: {(predictions == y_test).mean():.3f}")
print(f"Robust training objective: {cvar_lgbm.robust_obj:.6f}")

Test accuracy: 0.925
Robust training objective: 0.157335