Sinkhorn DRO¶

exception dro.neural_model.sinkhorn_nn.SinkhornNNDROError¶

Bases: DROError

Exception class for errors in Sinkhorn NN DRO model.

class dro.neural_model.sinkhorn_nn.SinkhornNNDRO(input_dim, num_classes, task_type='classification', model_type='mlp', reg_param=0.001, lambda_param=100.0, k_sample_max=5, optimization_type='SG', device=device(type='cpu'))¶

Bases: BaseNNDRO

Sinkhorn Distributionally Robust Optimization with Neural Networks.

Implements the Sinkhorn DRO objective for deep learning models:

\[\min_{\theta} \sup_{Q \in \mathcal{B}_{\epsilon,\varepsilon}(P)} \mathbb{E}_Q[\ell(f_\theta(X), y)]\]

where the ambiguity set \(\mathcal{B}_{\epsilon,\varepsilon}(P)\) is defined using the entropic-regularized (Sinkhorn) Wasserstein distance.

The Sinkhorn DRO loss for a mini-batch is computed as:

\[\hat{R}(\theta) = \lambda \varepsilon \cdot \frac{1}{N} \sum_{i=1}^{N} \log \left( \frac{1}{m} \sum_{j=1}^{m} \exp\left( \frac{\ell(f_\theta(x_i + \sigma_j), y_i)}{\lambda \varepsilon} \right)\right)\]

where \(\sigma_j \sim \mathcal{N}(0, \varepsilon I)\) are Gaussian perturbations.

Three stochastic optimization methods are supported:

SG (Stochastic Gradient): Uses a fixed number of Monte Carlo samples \(m = 2^{K_{max}}\)
MLMC (Multilevel Monte Carlo): Uses a hierarchy of sample levels for variance reduction
RTMLMC (Randomized Truncated MLMC): Randomly selects a single level per iteration for further variance reduction

Reference: Sinkhorn Distributionally Robust Optimization

Initialize Sinkhorn DRO neural 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'
reg_param (float) – Entropic regularization strength \(\varepsilon > 0\) controlling transport smoothness. Must be > 0. Defaults to 1e-3.
lambda_param (float) – Loss scaling factor \(\lambda > 0\) balancing Wasserstein distance and loss. Must be > 0. Defaults to 1e2.
k_sample_max (int) – Maximum level for Monte Carlo / MLMC sampling. The number of noise samples is \(2^{k\_sample\_max}\). Higher values improve accuracy but increase computation. Defaults to 5.
optimization_type (str) –
Stochastic optimization algorithm. Supported:
- 'SG': Standard Stochastic Gradient (baseline)
- 'MLMC': Multilevel Monte Carlo acceleration
- 'RTMLMC': Randomized Truncated MLMC
device (torch.device) – Target computation device, defaults to CPU

Raises:

ValueError –

If reg_param ≤ 0
If lambda_param ≤ 0
If k_sample_max < 1
If optimization_type not in {‘SG’, ‘MLMC’, ‘RTMLMC’}

Example:

>>> model = SinkhornNNDRO(
...     input_dim=784,
...     num_classes=10,
...     reg_param=0.01,
...     lambda_param=50.0,
...     optimization_type='SG'
... )

update(config)¶

Update hyperparameters for Sinkhorn NN DRO.

Parameters:

config (dict) –

Dictionary containing parameter updates. Valid keys:

'reg': Entropic regularization strength (ε > 0)
'lambda': Loss scaling factor (λ > 0)
'k_sample_max': Maximum MLMC sampling level (int ≥ 1)
'optimization_type': Optimization algorithm (‘SG’, ‘MLMC’, ‘RTMLMC’)
'lr': Learning rate
'batch_size': Training batch size
'train_epochs': Number of training epochs
'layer_num': Number of MLP layers
'hidden_size': Hidden layer size for MLP
'dropout_ratio': Dropout rate for MLP

Raises:

ValueError – If any parameter value violates constraints.

Return type:

None

Example:

>>> model.update({
...     'reg': 0.01,
...     'lambda': 50.0,
...     'k_sample_max': 3,
...     'optimization_type': 'MLMC'
... })