DRO on Tree-Ensemble Models

This tutorial 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.

Each example follows the same sequence: choose a DRO class and its robustness parameter, provide the tree-booster configuration with update, call fit, and then inspect both predictive accuracy and 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
)

Objective and booster configuration

There are two separate parts of the model setup:

1. Robust objective: the class selects KL-DRO or CVaR-DRO, while eps controls that objective. For fitted sample losses \(\ell_1, \ldots, \ell_n\), the reported KL objective is

   \[\frac{1}{\varepsilon} \log\left(\frac{1}{n}\sum_{i=1}^n \exp(\varepsilon \ell_i)\right),\]

   and the reported CVaR objective is

   \[t + \frac{1}{1-\varepsilon}\frac{1}{n}\sum_{i=1}^n(\ell_i-t)_+,\]

   where \(t\) is the empirical eps quantile of the fitted losses.

2. Booster configuration: the dictionary passed to update controls tree construction. Here, max_depth limits tree depth, learning_rate scales each boosting step, and the required num_boost_round sets the number of boosting iterations. The DRO wrapper supplies the custom training objective during fit.

Before fitting, model.robust_obj is None. After fitting, it contains the robust empirical objective evaluated on the training data. It is a training diagnostic rather than a test-set metric; compare it only when the data, loss, and DRO formulation are held fixed.

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.945
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.930
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.930
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.910
Robust training objective: 0.157335