Wasserstein DRO

class dro.neural_model.wdro_nn.WNNDRO(input_dim, num_classes, task_type='classification', model_type='mlp', epsilon=0.1, adversarial_steps=10, adversarial_step_size=0.02, adversarial_norm='l2', adversarial_method='PGD', device=device(type='cpu'))

Bases: BaseNNDRO

Wasserstein Neural DRO with Adversarial Robustness.

Initialize Wasserstein DRO model with adversarial training.

Parameters:

  • input_dim (int) – Input feature dimension \(d \geq 1\)

  • num_classes (int) – Output dimension: - Classification: \(K \geq 2\) (number of classes) - Regression: Automatically set to 1

  • task_type (str) – Learning task type. Supported: - 'classification': Cross-entropy loss - 'regression': MSE loss

  • model_type (str) –

    Neural architecture type. Supported:

    • 'mlp': Multi-Layer Perceptron (default)

    • linear

    • resnet

    • alexnet

  • device (torch.device) – Target computation device, defaults to CPU

  • epsilon (float) – Dual parameter (parameter of the L2-penalty during adversarial training) \(\epsilon \geq 0\) controlling distributional robustness. Larger values increase model conservativeness. Defaults to 0.1.

  • adversarial_steps (int) – Number of PGD attack iterations \(T_{adv} \geq 1\). Defaults to 10.

  • adversarial_step_size (float) – PGD step size \(\eta_{adv} > 0\). Defaults to 0.02.

  • adversarial_norm (str) –

    Adversarial perturbation norm type. Options:

    • 'l2': \(\ell_2\)-ball constraint

    • 'l-inf': \(\ell_\infty\)-ball constraint

  • adversarial_method (str) –

    Adversarial example generation method. Options:

    • 'PGD': Projected Gradient Descent (default)

    • 'FGSM': Fast Gradient Sign Method

Raises:

ValueError –

  • If epsilon < 0

  • If adversarial_steps < 1

  • If adversarial_step_size ≤ 0

  • If invalid norm/method type

Example (MNIST)::
>>> model = WNNDRO(
...     input_dim=784,
...     num_classes=10,
...     epsilon=0.5,
...     adversarial_norm="l-inf",
...     adversarial_steps=7
... )