Personalization

In this notebook, we personalize the loss as a convex piecewise-affine loss. We use the same loss with exact \(f\)-DRO and Wasserstein DRO models, then carry its asymmetric decision costs into the neural approximation classes.

The tree-ensemble section illustrates the same 3:1 cost ratio with a cost-sensitive binary cross-entropy loss in XGBoost.

For post-training worst-case-loss bounds, see the Neural DRO robustness certificates guide.

A general convex piecewise-affine loss

Let \(t=\theta^\top x+b\) be the prediction and \(z=t-y\) its residual. We parameterize the loss as

\[L(z)=\max_{j=1,\ldots,J}\{a_jz+c_j\}.\]

A maximum of affine pieces is convex. Changing the arrays of slopes \(a_j\) and offsets \(c_j\) creates many losses without rewriting the DRO formulation. LAD, pinball loss, and epsilon-insensitive loss are special cases.

Original newsvendor loss

This loss covers a lot of operational examples. For example, the original asymmetric newsvendor loss

\[L(z)=3(-z)^+ + z^+=\max\{-3z,z\},\]

where under-prediction costs three times as much as over-prediction. Thus its piecewise-affine slopes are [-3, 1] and both offsets are zero. The helper functions below remain general: substitute any finite slope and offset arrays to define another convex piecewise-affine loss.

from types import MethodType

import cvxpy as cp
import numpy as np

from dro.linear_model.chi2_dro import Chi2DRO
from dro.linear_model.wasserstein_dro import WassersteinDRO

UNDERAGE_COST = 3.0
OVERAGE_COST = 1.0
NEWSVENDOR_SLOPES = np.array([-UNDERAGE_COST, OVERAGE_COST])
NEWSVENDOR_OFFSETS = np.zeros(2)

X = np.array([
    [1.0, 1.0],
    [2.0, 1.0],
    [3.0, 1.0],
    [4.0, 1.0],
])
y = np.array([1.0, 1.0, 0.0, 0.0])

def piecewise_affine_numpy(residual, slopes, offsets):
    pieces = slopes[:, None] * residual[None, :] + offsets[:, None]
    return np.max(pieces, axis=0)

def piecewise_affine_cvxpy(residual, slopes, offsets):
    pieces = cp.vstack([
        slope * residual + offset
        for slope, offset in zip(slopes, offsets)
    ])
    return cp.max(pieces, axis=0)

def newsvendor_loss(self, X, y):
    prediction = X @ self.theta + self.b
    residual = prediction - y
    return piecewise_affine_numpy(
        residual, NEWSVENDOR_SLOPES, NEWSVENDOR_OFFSETS
    )

def newsvendor_cvx_loss(self, X, y, theta, b, design_matrix=None):
    features = X if design_matrix is None else design_matrix
    residual = features @ theta + b - y
    return piecewise_affine_cvxpy(
        residual, NEWSVENDOR_SLOPES, NEWSVENDOR_OFFSETS
    )

Exact \(f\)-DRO

For an \(f\)-divergence ambiguity set, the inner problem changes the probabilities assigned to the observed samples. It only needs the numerical per-sample newsvendor loss (_loss) and its CVXPY representation (_cvx_loss). The same two hooks can be attached to KLDRO, Chi2DRO, TVDRO, or CVaRDRO; Chi-squared DRO is used here as a representative.

fdro_model = Chi2DRO(
    input_dim=X.shape[1], model_type='lad', solver='CLARABEL'
)
fdro_model._loss = MethodType(newsvendor_loss, fdro_model)
fdro_model._cvx_loss = MethodType(newsvendor_cvx_loss, fdro_model)
fdro_model.update({'eps': 0.10})

fdro_parameters = fdro_model.fit(X, y)
fdro_parameters

{'theta': [-0.47101600668852506, 0.9420320649788056], 'b': array(0.94203196)}

Exact Wasserstein DRO

Wasserstein DRO also needs the loss’s Lipschitz penalty. For ground cost

\[d((x,y),(x',y'))=\lVert A(x-x')\rVert_p+\kappa|y-y'|,\]

let \(q\) be dual to \(p\). On an unrestricted input-output space, the residual loss above has penalty

\[\max_j|a_j|\,\max\left\{\lVert A^{-1}\theta\rVert_q,\frac{1}{\kappa}\right\}.\]

When kappa='inf', labels cannot move and the \(1/\kappa\) term is omitted. This is the Lipschitz reformulation behind Theorem 4 of Regularization via Mass Transportation. The helper below targets the linear-kernel case; a custom kernel representation must supply its corresponding feature norm.

def newsvendor_penalization(self, theta):
    if self.kernel != 'linear':
        raise NotImplementedError(
            'This tutorial penalty is written for a linear feature map.'
        )

    if self.p == 1:
        dual_norm = np.inf
    elif self.p == 'inf':
        dual_norm = 1
    else:
        dual_norm = 1.0 / (1.0 - 1.0 / self.p)

    residual_lipschitz = float(np.max(np.abs(NEWSVENDOR_SLOPES)))
    feature_term = cp.norm(
        self.cost_inv_transform @ theta,
        dual_norm,
    )
    if self.kappa == 'inf':
        return residual_lipschitz * feature_term
    return residual_lipschitz * cp.maximum(
        feature_term,
        1.0 / self.kappa,
    )

wdro_model = WassersteinDRO(
    input_dim=X.shape[1], model_type='lad', solver='CLARABEL'
)
wdro_model._loss = MethodType(newsvendor_loss, wdro_model)
wdro_model._cvx_loss = MethodType(newsvendor_cvx_loss, wdro_model)
wdro_model._penalization = MethodType(newsvendor_penalization, wdro_model)
wdro_model.update({'eps': 0.05, 'p': 2, 'kappa': 'inf'})

wdro_parameters = wdro_model.fit(X, y)
wdro_parameters

{'theta': [-0.4999999974041001, 3.802290229295084e-14], 'b': array(2.)}

Neural approximation with a decision-aware piecewise loss

The same newsvendor decision problem can be approximated with a neural predictor. More generally, suppose under-prediction and over-prediction lead to different downstream costs. The corresponding unbalanced L1 loss is

\[\ell(\hat y,y)=c_{\mathrm{under}}(y-\hat y)^+ + c_{\mathrm{over}}(\hat y-y)^+.\]

For positive costs, only one term can be active, so this is equivalently the maximum of two affine pieces. It is therefore a personalized, convex, piecewise-affine decision loss. The loss function below returns one value per sample: neural \(f\)-DRO robustly aggregates that vector, while neural Wasserstein DRO evaluates it after generating perturbed inputs.

import torch

from dro.neural_model.fdro_nn import Chi2NNDRO
from dro.neural_model.fdro_utils import RobustLoss
from dro.neural_model.wdro_nn import WNNDRO

def unbalanced_l1(
    predictions,
    targets,
    underage_cost=UNDERAGE_COST,
    overage_cost=OVERAGE_COST,
):
    predictions = predictions.squeeze(-1)
    targets = targets.reshape_as(predictions)
    pieces = torch.stack([
        underage_cost * (targets - predictions),
        overage_cost * (predictions - targets),
    ])
    return torch.maximum(pieces[0], pieces[1])

Neural \(f\)-DRO

Subclass RobustLoss to replace its individual loss while retaining its chi-squared or CVaR probability optimization. Bind the resulting scalar robust criterion to the neural model.

class DecisionAwareRobustLoss(RobustLoss):
    def __init__(self, underage_cost, overage_cost, **kwargs):
        super().__init__(is_regression=True, **kwargs)
        self.underage_cost = underage_cost
        self.overage_cost = overage_cost

    def _compute_individual_loss(self, outputs, targets):
        return unbalanced_l1(
            outputs,
            targets,
            underage_cost=self.underage_cost,
            overage_cost=self.overage_cost,
        )

decision_robust_loss = DecisionAwareRobustLoss(
    underage_cost=UNDERAGE_COST,
    overage_cost=OVERAGE_COST,
    geometry='chi-square',
    size=0.2,
    reg=0.1,
)

fdro_nn = Chi2NNDRO(
    input_dim=2,
    num_classes=1,
    task_type='regression',
    model_type='linear',
    size=0.2,
    reg=0.1,
    device=torch.device('cpu'),
)

def decision_aware_fdro_criterion(self, outputs, labels):
    return decision_robust_loss(outputs, labels)

fdro_nn._criterion = MethodType(decision_aware_fdro_criterion, fdro_nn)

Neural Wasserstein DRO

WNNDRO._loss must return the per-sample loss vector. Its existing criterion handles adversarial input generation and then averages these personalized losses. The rest of the training API is unchanged.

wdro_nn = WNNDRO(
    input_dim=2,
    num_classes=1,
    task_type='regression',
    model_type='linear',
    robust_objective='transport_penalty',
    gamma=2.0,
    adversarial_steps=2,
    adversarial_step_size=0.02,
    device=torch.device('cpu'),
)

def decision_aware_wdro_loss(self, outputs, labels):
    return unbalanced_l1(
        outputs,
        labels,
        underage_cost=UNDERAGE_COST,
        overage_cost=OVERAGE_COST,
    )

wdro_nn._loss = MethodType(decision_aware_wdro_loss, wdro_nn)

Call fdro_nn.fit(...) or wdro_nn.fit(...) after binding the hooks. The Wasserstein regression model above uses robust_objective='transport_penalty', so both its inner perturbation search and outer objective use the custom loss with continuous targets. gamma penalizes squared feature transport; it is not a norm-ball radius.

batch_predictions = torch.tensor([[0.8], [1.6], [2.5]], requires_grad=True)
batch_targets = torch.tensor([1.0, 1.4, 3.0])

fdro_objective = fdro_nn._criterion(batch_predictions, batch_targets)
wdro_individual_losses = wdro_nn._loss(batch_predictions, batch_targets)
combined_check = fdro_objective + wdro_individual_losses.mean()
combined_check.backward()

print('f-DRO robust objective:', float(fdro_objective.detach()))
print('WDRO individual losses:', wdro_individual_losses.detach().numpy())
print('gradient shape:', tuple(batch_predictions.grad.shape))

f-DRO robust objective: 1.0904502868652344
WDRO individual losses: [0.59999996 0.20000005 1.5       ]
gradient shape: (3, 1)

Matching the Wasserstein inner search to a custom loss

The preceding wdro_nn._loss check evaluates the custom loss without running the inner search. The model above uses transport-penalty mode. By comparison, with the default robust_objective='norm_ball' and positive epsilon, Torchattacks generates perturbations using classification cross-entropy; overriding _loss changes the outer loss only. For continuous targets, robust_objective='transport_penalty' uses _loss in both the inner search and outer objective. The example below differentiates that complete objective on a small batch.

wdro_transport = WNNDRO(
    input_dim=2,
    num_classes=1,
    task_type='regression',
    model_type='linear',
    robust_objective='transport_penalty',
    gamma=2.0,
    adversarial_steps=2,
    adversarial_step_size=0.02,
    device=torch.device('cpu'),
)
wdro_transport._loss = MethodType(decision_aware_wdro_loss, wdro_transport)

transport_inputs = torch.tensor([[0.2, 0.8], [0.6, 0.4], [0.9, 0.1]])
transport_targets = torch.tensor([0.3, 0.5, 0.7])
wdro_transport.current_inputs = transport_inputs
transport_outputs = wdro_transport.model(transport_inputs)
transport_objective = wdro_transport._criterion(transport_outputs, transport_targets)
transport_objective.backward()
print('Transport-penalty objective:', float(transport_objective.detach()))

Transport-penalty objective: 3.6209213733673096

Custom data loaders, optimizers, and schedulers

The neural classes now accept train_loader, val_loader, optimizer, and scheduler in fit(...). Both loaders are required on this path, and X and y are omitted. Each batch must yield (inputs, labels). Build the optimizer after calling update_user_mode(...) so it follows the replacement network’s parameters. A scheduler requires an explicit optimizer; scheduler_step='epoch' advances it once per epoch, while 'batch' advances it after every training batch.

The next toy classification data stay in [0, 1], which also meets the input range expected by Torchattacks in the HR-DRO example. The class-1 examples have twice the cost in the personalized \(f\)-DRO loss below.

from torch.utils.data import DataLoader, TensorDataset
import torch.nn.functional as F

from dro.neural_model.fdro_nn import CVaRNNDRO
from dro.neural_model.hrdro_nn import HRNNDRO

torch.manual_seed(7)
train_positions = torch.linspace(0.05, 0.95, steps=16)
val_positions = torch.tensor([0.15, 0.35, 0.65, 0.85])

def make_classification_data(positions):
    inputs = torch.stack([positions, 1.0 - positions], dim=1)
    labels = (positions >= 0.5).long()
    return inputs, labels

train_inputs, train_labels = make_classification_data(train_positions)
val_inputs, val_labels = make_classification_data(val_positions)
train_loader = DataLoader(
    TensorDataset(train_inputs, train_labels), batch_size=4, shuffle=True
)
val_loader = DataLoader(TensorDataset(val_inputs, val_labels), batch_size=4)

Chi-squared and CVaR neural \(f\)-DRO

Both models use the same per-example cost-sensitive classification loss. Their RobustLoss geometry determines how those losses are aggregated. The ResidualClassifier below is a user-defined torch.nn.Module: its skip connection makes it different from the package’s default MLP. Pass an instance to update_user_mode(...) before fitting. Its forward method maps each two-feature input to two class logits. Each run also supplies an SGD optimizer and a StepLR scheduler to the shared fitting loop. The learning-rate value printed at the end confirms that the scheduler was stepped.

class CostSensitiveRobustLoss(RobustLoss):
    def __init__(self, **kwargs):
        super().__init__(is_regression=False, **kwargs)

    def _compute_individual_loss(self, outputs, targets):
        per_sample_ce = F.cross_entropy(outputs, targets, reduction='none')
        return (1.0 + targets.float()) * per_sample_ce

class ResidualClassifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.skip = torch.nn.Linear(2, 2)
        self.hidden = torch.nn.Sequential(
            torch.nn.Linear(2, 8), torch.nn.Tanh(), torch.nn.Linear(8, 2)
        )

    def forward(self, inputs):
        return self.skip(inputs) + self.hidden(inputs)

f_dro_results = {}
for name, model_class, geometry in [
    ('chi-squared', Chi2NNDRO, 'chi-square'),
    ('CVaR', CVaRNNDRO, 'cvar'),
]:
    neural_model = model_class(
        input_dim=2, num_classes=2, task_type='classification',
        model_type='mlp', size=0.5, reg=0.1, device=torch.device('cpu'),
    )
    custom_network = ResidualClassifier()
    neural_model.update_user_mode(
        input_dim=2, num_classes=2, model=custom_network,
        task_type='classification', device=torch.device('cpu'),
    )
    robust_loss = CostSensitiveRobustLoss(
        geometry=geometry, size=0.5, reg=0.1
    )

    def cost_sensitive_criterion(self, outputs, labels, loss_module=robust_loss):
        return loss_module(outputs, labels)

    neural_model._criterion = MethodType(cost_sensitive_criterion, neural_model)
    optimizer = torch.optim.SGD(neural_model.model.parameters(), lr=0.05)
    scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer, step_size=1, gamma=0.8
    )
    metrics = neural_model.fit(
        train_loader=train_loader, val_loader=val_loader,
        optimizer=optimizer, scheduler=scheduler, scheduler_step='epoch',
        epochs=2, verbose=False,
    )
    f_dro_results[name] = {
        'metrics': metrics,
        'final_learning_rate': optimizer.param_groups[0]['lr'],
    }

f_dro_results

{'chi-squared': {'metrics': {'acc': 0.5, 'f1': 0.3333333333333333},
  'final_learning_rate': 0.03200000000000001},
 'CVaR': {'metrics': {'acc': 0.5, 'f1': 0.3333333333333333},
  'final_learning_rate': 0.03200000000000001}}

HR-DRO with a replacement network

HRNNDRO.update_user_mode(...) also rebuilds its attack against the replacement network. The optimizer and scheduler are created afterward. HR-DRO still uses its built-in cross-entropy loss here; the cost-sensitive \(f\)-DRO hook above does not alter HR-DRO’s CVXPY weighting or Torchattacks objective. Set train_batch_size to the loader batch size to avoid needless CVXPY problem rebuilds. A shorter last batch is handled dynamically.

hr_dro = HRNNDRO(
    input_dim=2, num_classes=2, task_type='classification',
    model_type='linear', alpha=0.1, r=0.01, epsilon=0.02,
    learning_approach='HD', train_batch_size=4,
    adversarial_params={
        'method': 'FGSM', 'norm': 'l-inf', 'steps': 1, 'step_size': 0.01,
    },
    device=torch.device('cpu'),
)
hr_network = torch.nn.Sequential(
    torch.nn.Linear(2, 8), torch.nn.ReLU(), torch.nn.Linear(8, 2)
)
hr_dro.update_user_mode(
    input_dim=2, num_classes=2, model=hr_network,
    task_type='classification', device=torch.device('cpu'),
)
assert hr_dro.attack.model is hr_dro.model
hr_optimizer = torch.optim.AdamW(hr_dro.model.parameters(), lr=0.01)
hr_scheduler = torch.optim.lr_scheduler.StepLR(
    hr_optimizer, step_size=1, gamma=0.8
)
hr_metrics = hr_dro.fit(
    train_loader=train_loader, val_loader=val_loader,
    optimizer=hr_optimizer, scheduler=hr_scheduler, scheduler_step='epoch',
    epochs=2, verbose=False,
)
print('HR-DRO validation:', hr_metrics)
print('Final learning rate:', hr_optimizer.param_groups[0]['lr'])

HR-DRO validation: {'acc': 0.75, 'f1': 0.7333333333333334}
Final learning rate: 0.0064

Tree-based ensemble with a personalized loss

For binary XGBoost KL-DRO, loss(preds, labels) supplies one loss per training example. The wrapper uses those values for DRO sample weights and its reported robust objective. Replacing loss() alone does not change the ordinary cross-entropy gradients and Hessians used by boosting. To train with a prediction-dependent cost-sensitive loss, we also override _kl_dro_loss().

Here positive-class log loss costs three times as much as negative-class log loss, matching the 3:1 ratio above. The Hessian uses the same weighted diagonal approximation as the built-in KL-DRO objective. This example covers binary classification; other losses require matching derivatives and a compatible booster objective. The small toy run limits numerical-library threads during fitting and prediction.

from scipy.special import expit
from threadpoolctl import threadpool_limits
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split

from dro.tree_model.xgb import KLDRO_XGB

class CostSensitiveKLDRO_XGB(KLDRO_XGB):
    def __init__(self, positive_cost, negative_cost, eps=0.1):
        super().__init__(eps=eps, kind='classification')
        self.positive_cost = positive_cost
        self.negative_cost = negative_cost

    def loss(self, preds, labels):
        probabilities = np.clip(np.asarray(preds).reshape(-1), 1e-8, 1 - 1e-8)
        labels = np.asarray(labels).reshape(-1)
        costs = np.where(labels == 1, self.positive_cost, self.negative_cost)
        log_loss = -labels * np.log(probabilities) - (1 - labels) * np.log1p(-probabilities)
        return costs * log_loss

    def _kl_dro_loss(self, raw_scores, dtrain, epsilon=0.1):
        labels = dtrain.get_label()
        probabilities = expit(raw_scores)
        losses = self.loss(probabilities, labels)
        scaled_losses = epsilon * losses
        weights = np.exp(scaled_losses - np.max(scaled_losses))
        weights *= len(labels) / weights.sum()
        costs = np.where(labels == 1, self.positive_cost, self.negative_cost)
        gradients = weights * costs * (probabilities - labels)
        hessians = weights * costs * probabilities * (1 - probabilities)
        return gradients, hessians

tree_X, tree_y = make_classification(
    n_samples=120, n_features=2, n_informative=2, n_redundant=0,
    random_state=7,
)
tree_X_train, tree_X_test, tree_y_train, tree_y_test = train_test_split(
    tree_X, tree_y, test_size=0.25, stratify=tree_y, random_state=7,
)
tree_model = CostSensitiveKLDRO_XGB(
    positive_cost=UNDERAGE_COST, negative_cost=OVERAGE_COST, eps=0.1,
)
tree_model.update({
    'objective': 'binary:logistic', 'max_depth': 2,
    'learning_rate': 0.2, 'num_boost_round': 6,
    'nthread': 1, 'seed': 7,
})
with threadpool_limits(limits=1):
    tree_model.fit(tree_X_train, tree_y_train)
    tree_accuracy = (tree_model.predict(tree_X_test) == tree_y_test).mean()
print('Losses at probability 0.5:', tree_model.loss(
    np.array([0.5, 0.5]), np.array([0, 1]),
))
print(f'Test accuracy: {tree_accuracy:.3f}')
print(f'Robust training objective: {tree_model.robust_obj:.3f}')

Losses at probability 0.5: [0.69314718 2.07944154]
Test accuracy: 0.933
Robust training objective: 0.445