Personalized Losses

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.

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):
    residual = X @ 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',
    epsilon=0.05,
    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)

During training, we call fdro_nn.fit(...) or wdro_nn.fit(...) normally after binding the hooks. For positive-epsilon Wasserstein regression, we use an adversarial generator compatible with continuous targets; the package’s built-in torchattacks path is primarily designed for classification.

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)