dro: Distributionally Robust Optimization for Machine Learning

dro is a Python package designed for Distributionally Robust Optimization (DRO) in machine learning. It implements DRO methods for supervised learning tasks using linear losses (such as SVM, logistic regression, and linear regression).

Key features include:

• Optimization Engine: Built on top of the cvxpy convex optimization solver.
• Distance Metrics: Supports various distance metrics $d(\cdot, \cdot)$ to define ambiguity sets.
• Model Support: Compatible with linear models, neural networks, and tree-based ensembles.
• Data Generation: Includes integrated synthetic data generation mechanisms based on recent research.

The core optimization problem solved is: $$\min_{\theta} \max_{P: P \in U} E_{(X,Y) \sim P}[\ell(\theta;(X, Y))]$$ where $U$ is the ambiguity set, typically defined as $U = {P: d(P, \hat P_n) \leq \epsilon}$ using the empirical distribution $\hat P_n$ and a hyperparameter $\epsilon$.

The dro package provides a variety of Linear Distributionally Robust Optimization (DRO) methods. These are categorized into typical linear methods and advanced linear methods. Use these methods to build models that are robust against distribution shifts in machine learning tasks.

Typical Linear DRO Methods

• base_linear: Base linear DRO implementations.
• chi2dro_linear: Chi-square DRO.
• cvardro_linear: Conditional Value-at-Risk (CVaR) DRO.
• groupdro_linear: Group DRO.
• kldro_linear: Kullback-Leibler (KL) divergence DRO.
• tvdro: Total Variation (TV) DRO.
• mmddro: Maximum Mean Discrepancy (MMD) DRO.
• wdro_linear: Wasserstein DRO.

Advanced Linear DRO Methods

• conditionaldro: Conditional DRO.
• hrdro: Hierarchical DRO.
• marginaldro: Marginal DRO.
• motdro: Multi-Objective Transport DRO.
• or_wdro: Optimal Transport Wasserstein DRO.
• sinkhorn_dro: Sinkhorn-based DRO.
• bayesian_dro: Bayesian DRO.

Bayesian (Parametric) Distributionally Robust Optimization (DRO) addresses stochastic optimization by incorporating uncertainty about the underlying distribution parameters. Instead of optimizing against a single fixed distribution, it optimizes against a posterior distribution $\zeta_N$ of the parameters $\zeta$ given the observed data.

In practice, dro approximates the outer expectation over the posterior distribution using finite samples of $\zeta$. For each sampled parameter $\zeta$, Monte Carlo sampling is used to handle the continuous integral of the parametric distribution $P_{\zeta}$.

Supported parametric families include:

• Exponential: Uses a rate parameter $\zeta$ with a default $\Gamma(1, 1)$ prior.
• Gaussian: Uses a Normal-Inverse-Wishart (NIW) prior on the mean $\mu$ and covariance $\Sigma$, which is conjugate to the multivariate Gaussian likelihood, allowing for closed-form posterior updates.

The dro package implements several $f$-divergence Distributionally Robust Optimization (DRO) models. These models use a distance metric $d(P, Q)$ to define an ambiguity set around a reference distribution.

Supported Models and Hyperparameters

Key Concepts

• Ambiguity Size (eps): Controls the size of the uncertainty set. When eps = 0, the model reduces to Empirical Risk Minimization (ERM).
• Worst-case Illustration: In $\chi^2$-DRO, the model finds a perturbed probability distribution by solving an optimization problem that maximizes loss subject to a constraint on the distance from the empirical distribution.

Group Distributionally Robust Optimization (Group DRO) is designed for scenarios where training data is divided into known groups and you want to prevent poor performance in any single subgroup.

Instead of minimizing the average loss across all samples, Group DRO minimizes the largest empirical mean loss among all observed groups. It achieves this by solving a convex optimization problem where an epigraph variable $t$ is minimized subject to the constraint that the average loss for every group $g$ is less than or equal to $t$.

Key characteristics:

• Group Weighting: Every category present in the training data contributes one group-average constraint. Small groups have the same influence on the objective as large groups and are not down-weighted by their sample count.
• Feature Inclusion: The group column (specified by group_idx) remains in the design matrix during fitting and prediction. This means the model may use the group identity as a predictive feature, which is an important consideration if the group represents a sensitive attribute.

MMD-DRO (Maximum Mean Discrepancy Distributionally Robust Optimization) defines an ambiguity set $\mathcal{P}(d, \eta) = {Q: d(Q, \hat P) \leq \eta}$ based on a kernel distance $d(P, Q)$. The choice of kernel determines which differences between distributions are emphasized.

Supported kernels include:

• rbf (default): Gaussian RBF kernel $\exp(-\gamma|x-y|_2^2)$. If kernel_gamma is None, $\gamma$ is the reciprocal of the median nonzero squared pairwise distance.
• laplacian: $\exp(-\gamma|x-y|_1)$. If kernel_gamma is None, $\gamma$ is the reciprocal of the median nonzero $\ell_1$ distance. Useful for less smooth similarity.
• polynomial (alias poly): $(\gamma x^\top y + c_0)^d$. Defaults are degree $d=2$, offset $c_0=1$, and $\gamma=1/p$ (where $p$ is the dimension of the joint sample).
• linear: $x^\top y$. Captures differences in first moments only.

Note: It is highly recommended to standardize continuous features and targets before fitting, especially when using polynomial kernels, because the kernel is applied to the joint samples $(X, Y)$.

from dro.linear_model.mmd_dro import MMD_DRO

# Default RBF kernel usage
model = MMD_DRO(input_dim=X.shape[1], model_type="svm")
model.update({"eta": 0.1, "sampling_method": "hull"})
params = model.fit(X, y)