Continuous Group DRO¶

← LatentNN | SIMEX → Post-hoc calibration

Group Distributionally Robust Optimization (Group DRO) replaces the average training risk with the worst-group risk over a small set of physical regimes. For regression on a long-tailed continuous target, the target range is split into interpretable bins, and the model is trained to equalize its error across them instead of averaging a catastrophic tail away with a well-populated bulk.

torchregress already has a CVaR loss that focuses on the worst individual samples. CVaR is the wrong tool for noisy astronomical data: the highest-loss individuals are usually corrupted spectra or outliers in the bulk, not the metal-poor tail. Group DRO first averages the individual noise inside each target regime, then equalizes the regime means, so a noisy solar-metallicity sample cannot dominate the objective.

Mathematical formulation¶

Let \(\ell(f_\theta(x), y)\) be a per-sample loss and \(P_g\) the conditional distribution of \((x,y)\) restricted to target bin \(g\). Group DRO solves the minimax problem

The inner maximization is over the simplex of group weights \(q\); the solution is found by exponentiated-gradient (EG) ascent on the group losses \(\mathcal L_g^{(t)}\) of the current batch:

and the training objective is the \(q\)-weighted mixture

where \(q\) is treated as a constant in the gradient step (the adversarial weights are updated in a separate, detached step — this is the standard two-player interpretation).

Usage¶

End-to-end trainer¶

import torch
import torch.nn as nn
from torchregress.algorithms import GroupDROTrainer

torch.manual_seed(0)
# Three regimes with very different population sizes and scatter:
# solar-like (many, small scatter), metal-poor (few), extreme (very few).
y = torch.cat(
    [
        0.1 + 0.1 * torch.randn(500),   # bulk
        -1.5 + 0.3 * torch.randn(200),  # sparse regime
        -3.0 + 0.4 * torch.randn(100),  # extreme regime
    ]
)
x = torch.stack([y, torch.sin(y)], dim=1) + 0.05 * torch.randn(y.shape[0], 2)

def make_model():
    return nn.Sequential(nn.Linear(2, 32), nn.SiLU(), nn.Linear(32, 1))

# Physical regimes: sparse (-4, -2.5), (-2.5, -1), bulk (-1, 0.5).
trainer = GroupDROTrainer(
    model_factory=make_model,
    group_edges=[-4.0, -2.5, -1.0, 0.5],
    base_loss="mse",
    eta=0.01,
    epochs=100,
    lr=1e-3,
    batch_size=64,
    seed=0,
)
trainer.fit(x, y)
y_hat = trainer.predict(x)

print("group weights:", trainer.group_weights_.tolist())
print("per-regime losses:", trainer.group_losses_.tolist())

The fitted attributes are:

The loss in a custom training loop¶

ContinuousGroupDRO is a torch.nn.Module, so it composes with any optimizer:

from torchregress.algorithms import ContinuousGroupDRO

dro = ContinuousGroupDRO([-4.0, -3.0, -2.0, -1.0, 0.5], eta=0.01, base_loss="mse")
dro.fit_group_edges(y_train)          # not needed when edges are isotropic/physical

model = make_model()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for x, y in loader:
    optimizer.zero_grad()
    loss = dro(model(x), y)           # updates q in training mode
    loss.backward()
    optimizer.step()

# Validation: freeze the adversarial weights.
dro.eval()
with torch.no_grad():
    val_loss = dro(model(x_val), y_val)

group_losses(y_pred, target) returns the per-regime losses without updating \(q\) (useful for diagnostics), and group_weights() returns the current simplex weights.

Comparison with related objectives¶

Caveats¶

References¶

See also¶

• Imbalanced targets — LDS, FDS, Balanced MSE, CVaR.
• PRIME & ConR — representation-level imbalanced regression.
• IRLS — robust reweighting with a different objective.
• Imbalanced / rare-target regression — task-level selection matrix.