Other Neural DRO

This notebook demonstrates how to call the neural-network implementations of chi-square DRO, CVaR-DRO, Wasserstein DRO, and Holistic Robust DRO. All four examples use the same binary classification data, MLP architecture, and training settings so that the API differences are easy to compare. The focus here is code usage rather than mathematical formulations.

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

Shared setup

The Wasserstein and Holistic Robust estimators create bounded adversarial perturbations, so the synthetic features are scaled to the interval [0, 1]. Set verbose=False to keep progress-bar output out of the generated documentation.

import numpy as np
import torch
from sklearn.datasets import make_classification
from sklearn.preprocessing import MinMaxScaler

from dro.neural_model import Chi2NNDRO, CVaRNNDRO, HRNNDRO, WNNDRO

np.random.seed(42)
torch.manual_seed(42)

X, y = make_classification(
    n_samples=200,
    n_features=10,
    n_informative=6,
    n_redundant=2,
    weights=[0.7, 0.3],
    random_state=42,
)
X = MinMaxScaler().fit_transform(X).astype(np.float32)
y = y.astype(np.int64)

fit_options = {
    "train_ratio": 0.8,
    "lr": 1e-3,
    "batch_size": 32,
    "epochs": 3,
    "verbose": False,
}

print(f"X shape: {X.shape}; class counts: {np.bincount(y)}")

X shape: (200, 10); class counts: [139  61]

Chi-square neural DRO

size controls the uncertainty-set radius and reg controls regularization of the sample weights. Training uses the standard neural fit() interface.

torch.manual_seed(42)
chi2_model = Chi2NNDRO(
    input_dim=X.shape[1],
    num_classes=2,
    task_type="classification",
    model_type="mlp",
    size=0.1,
    reg=0.1,
    max_iter=100,
)
chi2_metrics = chi2_model.fit(X, y, **fit_options)
print("Chi-square DRO:", chi2_metrics)

Chi-square DRO: {'acc': np.float64(0.175), 'f1': 0.14893617021276595}

CVaR neural DRO

For CVaRNNDRO, size is the CVaR fraction and must be between zero and one. The remaining training call is identical to the chi-square example.

torch.manual_seed(42)
cvar_model = CVaRNNDRO(
    input_dim=X.shape[1],
    num_classes=2,
    task_type="classification",
    model_type="mlp",
    size=0.2,
    reg=0.1,
    max_iter=100,
)
cvar_metrics = cvar_model.fit(X, y, **fit_options)
print("CVaR-DRO:", cvar_metrics)

CVaR-DRO: {'acc': np.float64(0.175), 'f1': 0.14893617021276595}

Wasserstein neural DRO

WNNDRO exposes the perturbation radius, attack step count, step size, norm, and attack method in its constructor. This compact example uses three projected-gradient steps on tabular features in [0, 1]. Torchattacks clips adversarial inputs to that range.

torch.manual_seed(42)
wasserstein_model = WNNDRO(
    input_dim=X.shape[1],
    num_classes=2,
    task_type="classification",
    model_type="mlp",
    epsilon=0.1,
    adversarial_steps=3,
    adversarial_step_size=0.02,
    adversarial_norm="l2",
    adversarial_method="PGD",
)
wasserstein_metrics = wasserstein_model.fit(X, y, **fit_options)
print("Wasserstein DRO:", wasserstein_metrics)

Wasserstein DRO: {'acc': np.float64(0.175), 'f1': 0.14893617021276595}

Custom WNNDRO architecture and training controls

The constructor’s model_type selects an initial network. update_user_mode(...) replaces it with a user-defined PyTorch module; in norm_ball mode, WNNDRO also binds its attack to the new network. This classifier accepts the same 10-feature batches as above and returns two class logits. Construct the optimizer after replacing the network so that it updates the correct parameters. The external loaders and learning-rate scheduler are optional training controls.

from torch.utils.data import DataLoader, TensorDataset

class ResidualTabularClassifier(torch.nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.features = torch.nn.Sequential(
            torch.nn.Linear(input_dim, 16),
            torch.nn.ReLU(),
            torch.nn.Linear(16, 2),
        )
        self.skip = torch.nn.Linear(input_dim, 2)

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

torch.manual_seed(42)
custom_wasserstein_model = WNNDRO(
    input_dim=X.shape[1], num_classes=2, model_type="linear",
    epsilon=0.05, adversarial_steps=2, adversarial_step_size=0.02,
    adversarial_norm="l2", device=torch.device("cpu"),
)
custom_network = ResidualTabularClassifier(X.shape[1])
custom_wasserstein_model.update_user_mode(
    input_dim=X.shape[1], num_classes=2, model=custom_network,
    device=torch.device("cpu"),
)

features = torch.from_numpy(X)
labels = torch.from_numpy(y)
train_loader = DataLoader(
    TensorDataset(features[:160], labels[:160]), batch_size=32, shuffle=True,
)
val_loader = DataLoader(
    TensorDataset(features[160:], labels[160:]), batch_size=32,
)
optimizer = torch.optim.Adam(custom_network.parameters(), lr=0.05)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)
custom_metrics = custom_wasserstein_model.fit(
    train_loader=train_loader, val_loader=val_loader, optimizer=optimizer,
    scheduler=scheduler, epochs=10, verbose=False,
)
print("Custom WNNDRO:", custom_metrics)
print("Learning rate after training:", scheduler.get_last_lr()[0])

Custom WNNDRO: {'acc': 0.85, 'f1': 0.765625}
Learning rate after training: 0.01743392200500001

Holistic Robust neural DRO

HRNNDRO combines its robustness settings with an adversarial-training configuration. train_batch_size is kept consistent with the batch size passed to fit().

torch.manual_seed(42)
holistic_model = HRNNDRO(
    input_dim=X.shape[1],
    num_classes=2,
    task_type="classification",
    model_type="mlp",
    alpha=0.1,
    r=0.01,
    epsilon=0.05,
    learning_approach="HD",
    adversarial_params={
        "steps": 3,
        "step_size": 0.02,
        "norm": "l2",
        "method": "PGD",
    },
    train_batch_size=32,
)
holistic_metrics = holistic_model.fit(X, y, **fit_options)
print("Holistic Robust DRO:", holistic_metrics)

Holistic Robust DRO: {'acc': np.float64(0.175), 'f1': 0.14893617021276595}

Compare the returned metrics

Every estimator returns the same validation-metric dictionary for this classification example, which makes downstream comparison straightforward.

results = {
    "Chi-square DRO": chi2_metrics,
    "CVaR-DRO": cvar_metrics,
    "Wasserstein DRO": wasserstein_metrics,
    "Holistic Robust DRO": holistic_metrics,
}

for method, metrics in results.items():
    print(f"{method:24s} accuracy={metrics['acc']:.3f}, f1={metrics['f1']:.3f}")

Chi-square DRO           accuracy=0.175, f1=0.149
CVaR-DRO                 accuracy=0.175, f1=0.149
Wasserstein DRO          accuracy=0.175, f1=0.149
Holistic Robust DRO      accuracy=0.175, f1=0.149

We also consider another toy example: ### Two-feature attack examples This small classification problem exercises the norm_ball attacks on a two-dimensional feature matrix: \(\ell_2\) PGD (Torchattacks PGDL2), \(\ell_\infty\) PGD, and \(\ell_\infty\) FGSM. FGSM is a single gradient-sign step of size epsilon; adversarial_steps and adversarial_step_size apply to PGD. The features are in [0, 1], as required by Torchattacks’ input clipping. The printed validation metrics are illustrative rather than a robustness guarantee.

toy_rng = np.random.default_rng(42)
toy_X = np.concatenate((
    toy_rng.normal(0.25, 0.04, size=(40, 2)),
    toy_rng.normal(0.75, 0.04, size=(40, 2)),
)).clip(0, 1).astype(np.float32)
toy_y = np.array([0] * 40 + [1] * 40, dtype=np.int64)
toy_fit_options = dict(fit_options, batch_size=16, epochs=10, lr=0.05)

for name, norm, method in (
    ("L2 PGD", "l2", "PGD"),
    ("Linf PGD", "l-inf", "PGD"),
    ("Linf FGSM", "l-inf", "FGSM"),
):
    torch.manual_seed(42)
    toy_model = WNNDRO(
        input_dim=2, num_classes=2, model_type="linear",
        epsilon=0.1, adversarial_norm=norm, adversarial_method=method,
        adversarial_steps=3, adversarial_step_size=0.05,
    )
    toy_metrics = toy_model.fit(toy_X, toy_y, **toy_fit_options)
    print(f"{name}: robust training loss={toy_model.robust_obj:.3f}, "
          f"validation accuracy={toy_metrics['acc']:.3f}")

L2 PGD: robust training loss=0.440, validation accuracy=1.000
Linf PGD: robust training loss=0.498, validation accuracy=1.000
Linf FGSM: robust training loss=0.503, validation accuracy=1.000