Sinkhorn DRO with Neural Networks¶

This notebook demonstrates how to use SinkhornNNDRO for distributionally robust optimization with neural network models.

Sinkhorn DRO uses the entropic-regularized Wasserstein distance to define the ambiguity set. The key idea is to add Gaussian perturbations to input features and aggregate per-sample losses via logsumexp, yielding a smooth surrogate of the worst-case expected loss:

\[\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), \quad \sigma_j \sim \mathcal{N}(0, \varepsilon I)\]

Three stochastic optimization algorithms are supported:

SG (Stochastic Gradient): fixed \(m = 2^{K_{\max}}\) Monte Carlo samples
MLMC (Multilevel Monte Carlo): hierarchical variance reduction
RTMLMC (Randomized Truncated MLMC): randomized single-level estimation

Reference: Sinkhorn Distributionally Robust Optimization

[ ]:

import numpy as np
import torch
from sklearn.datasets import make_classification, make_regression
from dro.neural_model.sinkhorn_nn import SinkhornNNDRO

Classification Example¶

We first demonstrate Sinkhorn NN DRO on a binary classification task using all three optimization methods.

[ ]:

# Generate classification data
X_cls, y_cls = make_classification(
    n_samples=500, n_features=10, n_informative=5,
    n_classes=2, random_state=42
)
print(f"Classification data: X={X_cls.shape}, y={y_cls.shape}, classes={np.unique(y_cls)}")

SG (Stochastic Gradient) Optimization¶

[ ]:

# SG optimization (default)
model_sg = SinkhornNNDRO(
    input_dim=10,
    num_classes=2,
    task_type="classification",
    model_type="mlp",
    reg_param=0.1,          # entropic regularization ε
    lambda_param=10.0,       # loss scaling factor λ
    k_sample_max=3,          # m = 2^3 = 8 noise samples
    optimization_type="SG",
)

metrics_sg = model_sg.fit(X_cls, y_cls, epochs=10, lr=1e-3, verbose=True)
print(f"\nSG Results: acc={metrics_sg['acc']:.4f}, f1={metrics_sg['f1']:.4f}")

MLMC (Multilevel Monte Carlo) Optimization¶

MLMC uses a hierarchy of sample levels for variance reduction while maintaining the same expected gradient.

[ ]:

model_mlmc = SinkhornNNDRO(
    input_dim=10,
    num_classes=2,
    task_type="classification",
    model_type="mlp",
    reg_param=0.1,
    lambda_param=10.0,
    k_sample_max=3,
    optimization_type="MLMC",
)

metrics_mlmc = model_mlmc.fit(X_cls, y_cls, epochs=10, lr=1e-3, verbose=True)
print(f"\nMLMC Results: acc={metrics_mlmc['acc']:.4f}, f1={metrics_mlmc['f1']:.4f}")

RTMLMC (Randomized Truncated MLMC) Optimization¶

RTMLMC randomly selects a single level per iteration from a truncated geometric distribution, scaling the loss by inverse selection probability. This further reduces per-iteration cost.

[ ]:

model_rtmlmc = SinkhornNNDRO(
    input_dim=10,
    num_classes=2,
    task_type="classification",
    model_type="mlp",
    reg_param=0.1,
    lambda_param=10.0,
    k_sample_max=3,
    optimization_type="RTMLMC",
)

metrics_rtmlmc = model_rtmlmc.fit(X_cls, y_cls, epochs=10, lr=1e-3, verbose=True)
print(f"\nRTMLMC Results: acc={metrics_rtmlmc['acc']:.4f}, f1={metrics_rtmlmc['f1']:.4f}")

Compare Optimization Methods¶

[ ]:

print("=" * 50)
print(f"{'Method':<10} {'Accuracy':>10} {'F1':>10}")
print("-" * 50)
print(f"{'SG':<10} {metrics_sg['acc']:>10.4f} {metrics_sg['f1']:>10.4f}")
print(f"{'MLMC':<10} {metrics_mlmc['acc']:>10.4f} {metrics_mlmc['f1']:>10.4f}")
print(f"{'RTMLMC':<10} {metrics_rtmlmc['acc']:>10.4f} {metrics_rtmlmc['f1']:>10.4f}")
print("=" * 50)

Regression Example¶

Sinkhorn NN DRO also supports regression tasks by setting task_type="regression".

[ ]:

# Generate regression data
X_reg, y_reg = make_regression(
    n_samples=500, n_features=10, noise=5.0, random_state=42
)
print(f"Regression data: X={X_reg.shape}, y={y_reg.shape}")

# Train with SG
model_reg = SinkhornNNDRO(
    input_dim=10,
    num_classes=1,             # auto-set to 1 for regression
    task_type="regression",
    model_type="mlp",
    reg_param=0.01,
    lambda_param=50.0,
    k_sample_max=3,
    optimization_type="SG",
)

metrics_reg = model_reg.fit(X_reg, y_reg, epochs=20, lr=1e-3, verbose=True)
print(f"\nRegression MSE: {metrics_reg['mse']:.4f}")

Dynamic Hyperparameter Tuning¶

The update() method allows changing Sinkhorn DRO hyperparameters on-the-fly without re-initializing the model.

[ ]:

model = SinkhornNNDRO(
    input_dim=10, num_classes=2,
    reg_param=0.1, lambda_param=10.0,
    k_sample_max=2, optimization_type="SG",
)

# Update hyperparameters dynamically
model.update({
    "reg": 0.05,
    "lambda": 50.0,
    "k_sample_max": 4,
    "optimization_type": "MLMC",
})

print(f"Updated: reg_param={model.reg_param}, lambda_param={model.lambda_param}")
print(f"         k_sample_max={model.k_sample_max}, optimization_type={model.optimization_type}")

# Train with updated parameters
metrics = model.fit(X_cls, y_cls, epochs=5, lr=1e-3, verbose=False)
print(f"Metrics: {metrics}")

Using a Custom Model¶

You can plug in your own PyTorch model via update_user_mode(), which is inherited from BaseNNDRO.

[ ]:

import torch.nn as nn

# Define a custom 3-layer network
class CustomNet(nn.Module):
    def __init__(self, input_dim, num_classes):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, num_classes),
        )
    def forward(self, x):
        return self.net(x)

# Initialize SinkhornNNDRO and plug in the custom model
model_custom = SinkhornNNDRO(
    input_dim=10, num_classes=2,
    reg_param=0.1, lambda_param=10.0,
    k_sample_max=3, optimization_type="SG",
)
model_custom.update_user_mode(
    input_dim=10,
    num_classes=2,
    model=CustomNet(10, 2),
    task_type="classification",
)

metrics_custom = model_custom.fit(X_cls, y_cls, epochs=10, lr=1e-3, verbose=True)
print(f"\nCustom Model: acc={metrics_custom['acc']:.4f}, f1={metrics_custom['f1']:.4f}")

Effect of Regularization Strength¶

The entropic regularization parameter \(\varepsilon\) (reg_param) controls the smoothness of the Sinkhorn distance. Larger \(\varepsilon\) → smoother transport plan → more regularized (closer to ERM); smaller \(\varepsilon\) → closer to true Wasserstein DRO.

[ ]:

reg_values = [0.001, 0.01, 0.1, 1.0]

print(f"{'reg_param':<12} {'Accuracy':>10} {'F1':>10}")
print("-" * 35)

for reg in reg_values:
    m = SinkhornNNDRO(
        input_dim=10, num_classes=2,
        task_type="classification", model_type="mlp",
        reg_param=reg, lambda_param=10.0,
        k_sample_max=3, optimization_type="SG",
    )
    result = m.fit(X_cls, y_cls, epochs=10, lr=1e-3, verbose=False)
    print(f"{reg:<12} {result['acc']:>10.4f} {result['f1']:>10.4f}")

Multi-class Classification¶

SinkhornNNDRO naturally extends to multi-class problems.

[ ]:

# 5-class classification
X_multi, y_multi = make_classification(
    n_samples=800, n_features=20, n_informative=10,
    n_classes=5, n_clusters_per_class=1, random_state=42
)

model_multi = SinkhornNNDRO(
    input_dim=20,
    num_classes=5,
    task_type="classification",
    model_type="mlp",
    reg_param=0.1,
    lambda_param=10.0,
    k_sample_max=3,
    optimization_type="SG",
)

metrics_multi = model_multi.fit(X_multi, y_multi, epochs=15, lr=1e-3, verbose=True)
print(f"\nMulti-class (5): acc={metrics_multi['acc']:.4f}, f1={metrics_multi['f1']:.4f}")

# Prediction
preds = model_multi.predict(X_multi[:10])
print(f"Sample predictions: {preds}")
print(f"True labels:        {y_multi[:10]}")