Wasserstein DRO

Training customization

WNNDRO inherits the neural training controls. Its fit(...) method accepts paired train_loader and val_loader arguments, a PyTorch optimizer, and an optional scheduler with scheduler_step="epoch" or "batch". Install a custom network with update_user_mode(...) before constructing the optimizer; in norm_ball mode, this also rebinds the attack to the replacement network. These training controls do not change the attack’s loss: a custom _loss affects the inner search in transport_penalty mode, but only the outer objective in norm_ball mode. The WNNDRO tutorial includes a “Custom WNNDRO architecture and training controls” example with a user-defined PyTorch nn.Module, external loaders, an optimizer, and a scheduler. The custom module must accept the selected input batches and return task-compatible outputs (class logits for this classification example).

Radius certificate

WNNDRO.certify_radius(transport_budget, adversarial_losses, transport_costs, certified_inner_errors, transport_penalty=None) returns epsilon_cert and its empirical loss bound when the supplied global inner error bounds suffice. It returns None when no nonnegative budget can be certified at that loss level. The constructor’s epsilon in norm_ball mode is a per-example attack radius, not transport_budget. See the certificate derivation and runnable example.

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'), *, robust_objective='norm_ball', gamma=1.0, input_bounds=None)

Bases: BaseNNDRO

Neural adversarial training with norm-ball or transport-penalty objectives.

norm_ball uses fixed-radius attacks. transport_penalty uses the squared Euclidean transport cost from Sinha et al. (2018).

Initialize a neural adversarial training model.

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) – Radius \(\epsilon \geq 0\) of the per-example \(\ell_2\) or \(\ell_\infty\) attack ball in norm_ball mode. Zero disables norm-ball perturbations. Defaults to 0.1.

  • adversarial_steps (int) – Number of inner attack iterations \(T_{adv} \geq 1\) for PGD or transport-penalty ascent. FGSM uses one step. Defaults to 10.

  • adversarial_step_size (float) – Inner attack step size \(\eta_{adv} > 0\) for PGD or transport-penalty ascent. FGSM uses epsilon. Defaults to 0.02.

  • adversarial_norm (str) –

    Perturbation norm in norm_ball mode. Options:

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

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

  • adversarial_method (str) –

    Attack method in norm_ball mode. Options:

    • 'PGD': Projected Gradient Descent (default)

    • 'FGSM': Fast Gradient Sign Method (requires 'l-inf')

    • 'FFGSM': Fast FGSM with a random start (requires 'l-inf')

  • robust_objective (str) – 'norm_ball' (default) or 'transport_penalty'. The latter approximately optimizes the penalized squared-Wasserstein objective of Sinha et al. (2018).

  • gamma (float) – Positive coefficient of the squared Euclidean transport cost in transport_penalty mode. Larger values discourage perturbations. Defaults to 1.0.

  • input_bounds (Optional[Tuple[float, float]]) – Optional lower and upper input bounds for transport_penalty mode. For images in [0, 1], use (0, 1). Defaults to no clipping.

Raises:

ValueError –

  • If epsilon < 0

  • If adversarial_steps < 1

  • If adversarial_step_size ≤ 0

  • If robust_objective, gamma, input_bounds, or norm-ball attack configuration is invalid

Example (MNIST)::
>>> model = WNNDRO(
...     input_dim=784,
...     num_classes=10,
...     epsilon=0.5,
...     adversarial_norm="l-inf",
...     adversarial_steps=7
... )
update_user_mode(input_dim, num_classes, model, task_type='classification', device=device(type='cpu'))

Replace the network and bind norm-ball attacks to the new network.

Return type:

None

Parameters:
certify_radius(transport_budget, adversarial_losses, transport_costs, certified_inner_errors, *, transport_penalty=None)

Certify an empirical squared-cost Wasserstein budget.

Supply per-example losses, squared Euclidean costs, and proven global inner-maximization error bounds for the same adversarial examples, final network, labels, and admissible support. The method does not infer error bounds from attack iterations. transport_budget is distinct from the norm_ball constructor’s epsilon.

transport_penalty defaults to gamma in transport_penalty mode; in norm_ball mode it must be supplied explicitly. Returns None when the achieved loss cannot be certified at a nonnegative budget. Otherwise returns epsilon_cert, the certified loss bound, the upper bound at the input budget, and their gap in a dictionary. The result concerns the empirical reference distribution only.