Group DRO with Linear and Neural Models

Group DRO minimizes the largest expected loss over a finite set of known groups. This notebook uses a synthetic binary-classification problem to demonstrate the package’s two implementations:

• GroupDRO solves the convex linear minimax problem exactly with CVXPY.

• GroupNNDRO trains a neural model while updating adversarial group probabilities by exponentiated gradient.

Both estimators infer group membership from a categorical column selected by group_idx, so their public training call remains fit(X, y, ...). The group column remains available to the predictor as an input feature.

import cvxpy as cp
import numpy as np
import torch

from dro.linear_model import GroupDRO
from dro.neural_model import GroupNNDRO

Create grouped synthetic data

The final column of X is a group indicator. Group 0 is larger than group 1, illustrating why an unweighted sample average can be dominated by the majority group. The target is generated from a stable feature, while a second feature has a group-dependent relationship with that stable signal.

The linear API expects binary labels in \(\{-1,+1\}\); the neural API expects class indices in \(\{0,1\}\).

rng = np.random.default_rng(42)
n_majority, n_minority = 120, 40
groups = np.concatenate([np.zeros(n_majority), np.ones(n_minority)])
stable = rng.normal(size=len(groups))
spurious = np.where(groups == 0, stable, -stable) + 0.5 * rng.normal(size=len(groups))
X = np.column_stack([stable, spurious, groups]).astype(np.float32)
y_linear = np.where(stable + 0.35 * rng.normal(size=len(groups)) >= 0, 1, -1)
y_neural = (y_linear == 1).astype(np.int64)
group_idx = 2

unique_groups, group_counts = np.unique(groups, return_counts=True)
print("X shape:", X.shape)
print("group counts:", {int(g): int(n) for g, n in zip(unique_groups, group_counts)})

X shape: (160, 3)
group counts: {0: 120, 1: 40}

def accuracy_by_group(predictions, targets, group_values):
    return {
        int(group): float(np.mean(predictions[group_values == group] == targets[group_values == group]))
        for group in np.unique(group_values)
    }

Exact linear Group DRO

For every observed category, the linear estimator constrains that group’s mean loss to be no larger than a shared epigraph variable. Minimizing the variable gives the exact worst-group empirical objective. We select an installed open-source solver so the example does not require a MOSEK license.

solver = next(
    (name for name in ("CLARABEL", "ECOS", "SCS") if name in cp.installed_solvers()),
    None,
)
if solver is None:
    raise RuntimeError("Install CLARABEL, ECOS, or SCS to run the linear example.")

linear_groupdro = GroupDRO(
    input_dim=X.shape[1],
    group_idx=group_idx,
    model_type="svm",
    solver=solver,
)
linear_result = linear_groupdro.fit(X, y_linear)
linear_predictions = linear_groupdro.predict(X)

print("solver:", solver)
print("robust loss:", round(linear_result["robust_loss"], 4))
print("group losses:", dict(zip(linear_result["group_values"], linear_result["group_losses"])))
print("accuracy by group:", accuracy_by_group(linear_predictions, y_linear, groups))

solver: CLARABEL
robust loss: 0.3073
group losses: {0.0: 0.3072571710195848, 1.0: 0.2730075691082775}
accuracy by group: {0: 0.8583333333333333, 1: 0.875}

group_values and group_losses use the same order, and robust_loss is their maximum. For classification, score() also returns overall accuracy and macro F1. The exact objective has no Group-DRO-specific radius: the groups themselves define the adversary’s choices.

linear_accuracy, linear_f1 = linear_groupdro.score(X, y_linear)
print({"accuracy": linear_accuracy, "macro_f1": linear_f1})

{'accuracy': np.float64(0.8625), 'macro_f1': 0.8619607843137256}

Neural Group DRO

The neural estimator starts with uniform adversarial probabilities. After each mini-batch, step_size controls how quickly probability moves toward groups with larger loss. A larger value reacts faster but may make training less smooth. The remaining arguments use the same fit() interface as the other neural DRO models.

np.random.seed(42)
torch.manual_seed(42)

neural_groupdro = GroupNNDRO(
    input_dim=X.shape[1],
    num_classes=2,
    group_idx=group_idx,
    task_type="classification",
    model_type="mlp",
    step_size=0.05,
)
neural_metrics = neural_groupdro.fit(
    X,
    y_neural,
    train_ratio=0.8,
    lr=1e-2,
    batch_size=32,
    epochs=10,
    verbose=False,
)
neural_predictions = neural_groupdro.predict(X)

print("validation metrics:", neural_metrics)
print("group values:", neural_groupdro.group_values_)
print("adversarial weights:", neural_groupdro.group_weights_)
print("accuracy by group:", accuracy_by_group(neural_predictions, y_neural, groups))

validation metrics: {'acc': np.float64(0.84375), 'f1': 0.8435972629521017}
group values: [0. 1.]
adversarial weights: [0.5421316  0.45786843]
accuracy by group: {0: 0.875, 1: 0.85}

The final entries of group_weights_ align with group_values_. They are optimization state rather than population-frequency estimates: a high value indicates that training has emphasized that group’s loss. Because the validation split is random, use a larger or explicitly group-stratified dataset for real model selection.

Tuning checklist

• Confirm that group_idx identifies a finite categorical column and that every important group is represented in training and validation data.

• For neural models, start with the default step_size=0.01; increase it when weights adapt too slowly and decrease it if they oscillate sharply.

• Tune lr, batch_size, and epochs jointly with step_size. Larger batches are more likely to contain examples from every group.

• Evaluate per-group metrics and select checkpoints by worst-group validation performance, not only aggregate accuracy.

• Use regularization or early stopping for over-parameterized neural models; low worst-group training loss does not guarantee worst-group generalization.

Reference

Sagawa, Shiori, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization. ICLR 2020. The neural weight update follows the authors’ Group DRO reference implementation.