NN Sinkhorn DRO¶
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:
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
[1]:
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.
[2]:
# 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)}")
Classification data: X=(500, 10), y=(500,), classes=[0 1]
SG (Stochastic Gradient) Optimization¶
[3]:
# 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=False)
print(f"\nSG Results: acc={metrics_sg['acc']:.4f}, f1={metrics_sg['f1']:.4f}")
SG Results: acc=0.7900, f1=0.7900
MLMC (Multilevel Monte Carlo) Optimization¶
MLMC uses a hierarchy of sample levels for variance reduction while maintaining the same expected gradient.
[4]:
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=False)
print(f"\nMLMC Results: acc={metrics_mlmc['acc']:.4f}, f1={metrics_mlmc['f1']:.4f}")
MLMC Results: acc=0.8500, f1=0.8474
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.
[5]:
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=False)
print(f"\nRTMLMC Results: acc={metrics_rtmlmc['acc']:.4f}, f1={metrics_rtmlmc['f1']:.4f}")
RTMLMC Results: acc=0.7000, f1=0.6995
Compare Optimization Methods¶
[6]:
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)
==================================================
Method Accuracy F1
--------------------------------------------------
SG 0.7900 0.7900
MLMC 0.8500 0.8474
RTMLMC 0.7000 0.6995
==================================================
Regression Example¶
Sinkhorn NN DRO also supports regression tasks by setting task_type="regression".
[7]:
# 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=False)
print(f"\nRegression MSE: {metrics_reg['mse']:.4f}")
Regression data: X=(500, 10), y=(500,)
Regression MSE: 21209.4897
Dynamic Hyperparameter Tuning¶
The update() method allows changing Sinkhorn DRO hyperparameters on-the-fly without re-initializing the model.
[8]:
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}")
Updated: reg_param=0.05, lambda_param=50.0
k_sample_max=4, optimization_type=MLMC
Metrics: {'acc': np.float64(0.75), 'f1': 0.7194478734148805}
Using a Custom Model¶
You can plug in your own PyTorch model via update_user_mode(), which is inherited from BaseNNDRO.
[9]:
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=False)
print(f"\nCustom Model: acc={metrics_custom['acc']:.4f}, f1={metrics_custom['f1']:.4f}")
Custom Model: acc=0.9100, f1=0.9093
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.
[10]:
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}")
reg_param Accuracy F1
-----------------------------------
0.001 0.7600 0.7585
0.01 0.7400 0.7292
0.1 0.7700 0.7694
1.0 0.7500 0.7480
Multi-class Classification¶
SinkhornNNDRO naturally extends to multi-class problems.
[11]:
# 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=False)
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]}")
Multi-class (5): acc=0.7125, f1=0.7090
Sample predictions: [0 1 2 3 4 2 0 0 2 3]
True labels: [0 1 2 3 0 2 0 0 2 3]