peal.adaptors.group_dro_v2.GroupDroLossComputer

class peal.adaptors.group_dro_v2.GroupDroLossComputer

Bases: object

Group-robust loss plus the running per-group statistics it needs.

Wraps a per-sample cross-entropy: the batch is split into the four groups, per-group mean losses and accuracies are computed, and the batch loss is the dot product of the group losses with an adversarial weight vector adv_probs that is updated multiplicatively (exponentiated gradient ascent with step_size) so that hard groups get more weight. The running statistics double as the epoch’s evaluation metrics, which is why the same class is also instantiated for the validation and test passes.

Parameters:

• group_counts (torch.Tensor) – Number of training samples per group, shape (4,); used for the group fractions and the generalization adjustment.

• alpha (float) – Tail size of the “best-that-can-happen” (BTL) / greedy variant: the fraction of the data the worst-case average is taken over.

• gamma (float, optional) – Step size of the exponential moving average of the group losses, which the BTL variant ranks groups by. Default 0.1.

• adj (numpy.ndarray, optional) – Per-group generalization adjustment C_g; added as adj / sqrt(group_counts) before the weight update. Default zeros.

• min_var_weight (float, optional) – Mixes the greedy weights back towards the empirical group fractions.

• step_size (float, optional) – Learning rate of the adversarial weight update. Default 0.01.

• normalize_loss (bool, optional) – Normalise the adjusted group losses to sum to one before the update.

• btl (bool, optional) – Use the BTL/greedy worst-case weighting instead of the exponentiated gradient weights.

• device (str, optional) – Device all running statistics live on.

Variables:

• n_groups (int) – Fixed to 4 (label x confounder).

• group_str (list of str) – Group names "<y>-<has_confounder>" used in the logged keys.

• adv_probs (torch.Tensor) – Current adversarial distribution over the groups.

• exp_avg_loss (torch.Tensor) – Exponential moving average of the per-group losses.

• avg_group_loss (avg_group_acc,) – Running per-group averages since the last reset_stats().

Notes

adv_probs is never reset by reset_stats(); it is meant to persist across epochs. It is also only applied when every entry of adj is positive, mirroring the reference implementation.

Methods

compute_group_avg	Mean of a per-sample quantity within each group, plus the group counts.
compute_robust_loss	Exponentiated-gradient update of adv_probs and the weighted loss.
compute_robust_loss_btl	Robust loss ranking groups by their historical (moving-average) loss.
compute_robust_loss_greedy	Average the alpha worst fraction of the data, ranked by ref_loss.
get_model_stats	Add the squared parameter norm and the implied L2 penalty to stats_dict.
get_stats	Collect the running statistics as a flat dict for TensorBoard.
loss	Group-robust loss of one batch, updating all running statistics.
reset_stats	Zero the per-epoch running statistics.
update_exp_avg_loss	Advance the moving average of the per-group losses by one batch.
update_stats	Fold one batch into the running averages reported by get_stats().

__init__(group_counts, alpha, gamma=0.1, adj=None, min_var_weight=0, step_size=0.01, normalize_loss=False, btl=False, device='cpu')

Set up the per-sample criterion, the adjustment and the running stats.

loss(yhat, y, group_idx)

Group-robust loss of one batch, updating all running statistics.

Parameters:

• yhat (torch.Tensor) – Logits of shape (N, C).

• y (torch.Tensor) – Integer labels of shape (N,).

• group_idx (torch.Tensor) – Group index in 0..3 per sample, shape (N,).

Returns:

Scalar loss to backpropagate: the group losses weighted by the (updated) adversarial distribution.

Return type:

torch.Tensor

compute_robust_loss(group_loss)

Exponentiated-gradient update of adv_probs and the weighted loss.

Parameters:

group_loss (torch.Tensor) – Mean loss per group, shape (4,).

Returns:

The scalar robust loss and the updated group weights.

Return type:

tuple of torch.Tensor

compute_robust_loss_btl(group_loss)

Robust loss ranking groups by their historical (moving-average) loss.

Parameters:

group_loss (torch.Tensor) – Mean loss per group of the current batch, shape (4,).

Returns:

The scalar robust loss and the weights used, in group order.

Return type:

tuple of torch.Tensor

compute_robust_loss_greedy(group_loss, ref_loss)

Average the alpha worst fraction of the data, ranked by ref_loss.

Groups are sorted by ref_loss descending and taken until their cumulative data fraction reaches alpha; the group that straddles the boundary gets the remaining weight. The result is then mixed with the empirical group fractions according to min_var_weight.

Parameters:

• group_loss (torch.Tensor) – Mean loss per group of the current batch, shape (4,).

• ref_loss (torch.Tensor) – Scores the groups are ranked by, shape (4,).

Returns:

The scalar robust loss and the weights, unsorted back into group order.

Return type:

tuple of torch.Tensor

compute_group_avg(losses, group_idx)

Mean of a per-sample quantity within each group, plus the group counts.

Parameters:

• losses (torch.Tensor) – Per-sample values of shape (N,) (losses or 0/1 correctness).

• group_idx (torch.Tensor) – Group index per sample, shape (N,).

Returns:

Per-group mean and per-group sample count, both shape (4,). Empty groups get a mean of zero instead of a NaN.

Return type:

tuple of torch.Tensor

update_exp_avg_loss(group_loss, group_count)

Advance the moving average of the per-group losses by one batch.

Only groups present in the batch are updated, and a group’s first observation replaces the initial zero outright rather than being blended into it.

Parameters:

• group_loss (torch.Tensor) – Output of compute_group_avg() for this batch.

• group_count (torch.Tensor) – Output of compute_group_avg() for this batch.

reset_stats()

Zero the per-epoch running statistics.

adv_probs and exp_avg_loss persist across epochs and are kept.

update_stats(actual_loss, group_loss, group_acc, group_count, weights=None)

Fold one batch into the running averages reported by get_stats().

Per-group losses and accuracies are accumulated as sample-count weighted averages, the actual (robust) loss as a batch average, and the overall per-sample loss and accuracy as group-fraction weighted averages over the groups seen so far.

Parameters:

• actual_loss (torch.Tensor) – The scalar robust loss of the batch.

• group_loss (torch.Tensor) – Per-group mean loss, mean accuracy and sample count, shape (4,).

• group_acc (torch.Tensor) – Per-group mean loss, mean accuracy and sample count, shape (4,).

• group_count (torch.Tensor) – Per-group mean loss, mean accuracy and sample count, shape (4,).

• weights (torch.Tensor, optional) – The group weights used; only their positivity is recorded, in update_data_counts and update_batch_counts.

get_model_stats(model, weight_decay, stats_dict)

Add the squared parameter norm and the implied L2 penalty to stats_dict.

Parameters:

• model (torch.nn.Module) – Model whose parameters are summed over.

• weight_decay (float) – Coefficient used for the reported reg_loss.

• stats_dict (dict) – Dictionary that is updated in place.

Returns:

The same dictionary, with model_norm_sq and reg_loss.

Return type:

dict

get_stats(model=None, weight_decay=None)

Collect the running statistics as a flat dict for TensorBoard.

Contains per-group average loss, moving-average loss, accuracy and update-batch count keyed by <name>_group:<y>-<confounder>, plus the aggregates avg_actual_loss, avg_per_sample_loss, avg_acc, avg_group_acc and worst_group_acc.

Parameters:

• model (torch.nn.Module, optional) – When given, get_model_stats() is appended; then weight_decay is required.

• weight_decay (float, optional) – Coefficient for the reported regularisation loss.

Return type:

dict of str to float