Group DRO

Group distributionally robust neural networks.

The adversarial group weights use the exponentiated-gradient update introduced in the reference Group DRO implementation: https://github.com/kohpangwei/group_DRO/blob/master/loss.py

class dro.neural_model.groupdro_nn.GroupNNDRO(input_dim, num_classes, group_idx, task_type='classification', model_type='mlp', device=device(type='cpu'), step_size=0.01)

Bases: BaseNNDRO

Neural Group Distributionally Robust Optimization model.

The model minimizes a weighted combination of empirical group losses while an adversary updates its group probabilities by exponentiated gradient:

\[q_g \leftarrow \frac{q_g\exp(\eta_q\,\ell_g)} {\sum_j q_j\exp(\eta_q\,\ell_j)}.\]

Repeated updates concentrate weight on high-loss groups and approximate the worst-group objective. This is the core update used by Koh et al.’s Group DRO code, adapted here to infer group membership from a feature column and retain the package-wide fit(X, y, ...) API.

group_idx refers to a column in a two-dimensional, tabular X. Its values must be finite categorical numbers. The column remains available to the neural network as an input feature.

Parameters:

  • input_dim (int) – Number of input features.

  • num_classes (int) – Number of classes, or 1 for regression.

  • group_idx (int) – Zero-based index of the categorical group feature.

  • task_type (str) – 'classification' or 'regression'.

  • model_type (str) – 'mlp' or 'linear'. Image architectures are not supported because a scalar feature index is required for grouping.

  • device (torch.device) – Torch device used for training.

  • step_size (float) – Positive exponentiated-gradient step size for the adversarial group probabilities.

Variables:

  • group_values – Sorted category values observed by fit().

  • adv_probs – Current adversarial probability for each category in group_values_.

  • group_weights – NumPy copy of the final adversarial probabilities.

Reference: Sagawa, Koh, Hashimoto, and Liang, “Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization” (2020).

Initialize neural DRO framework.

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

  • device (torch.device) – Target computation device, defaults to CPU

  • group_idx (int)

  • step_size (float)

Raises:

ValueError –

  • If input_dim < 1

  • If classification task with num_classes < 2

  • If unsupported model_type

Example (Classification):

>>> model = BaseNNDRO(
...     input_dim=64,
...     num_classes=10,
...     task_type="classification",
...     model_type="mlp",
...     device=torch.device("cuda")
... )

Example (Regression):

>>> model = BaseNNDRO(
...     input_dim=8,
...     num_classes=5,  # Auto-override to 1
...     task_type="regression"
... )
update(config)

Update Group DRO-specific configuration.

Supported keys are group_idx and step_size. A subsequent call to fit() rediscovers categories and resets adversarial weights.

Return type:

None

Parameters:

config (Dict[str, Any])

fit(X, y, train_ratio=0.8, lr=0.001, batch_size=32, epochs=100, verbose=True)

Fit with the same call signature as the other neural DRO models.

Group categories are inferred from X[:, group_idx] before the base training loop runs. Adversarial probabilities start uniformly on each call, making repeated fits independent.

Return type:

Dict[str, List[float]]

Parameters:
predict(X)

Return class indices for classification or raw regression values.

Return type:

ndarray

Parameters:

X (ndarray | Tensor)

score(X, y)

Use the common classification scores or MSE for regression.

Parameters: