Towards Group Robustness in the presence of Partial Group Labels

TL;DR

This paper addresses the challenge of training robust machine learning models with partial group labels by proposing a method that leverages partially available group information to improve minority group performance without sacrificing overall accuracy.

Contribution

The paper introduces a novel approach that constructs a constraint set for group assignments and optimizes for the worst-off groups, filling a gap between fully supervised and unsupervised group robustness methods.

Findings

Improved minority group performance in experiments.

Maintained overall accuracy across groups.

Effective in image and tabular datasets.

Abstract

Learning invariant representations is an important requirement when training machine learning models that are driven by spurious correlations in the datasets. These spurious correlations, between input samples and the target labels, wrongly direct the neural network predictions resulting in poor performance on certain groups, especially the minority groups. Robust training against these spurious correlations requires the knowledge of group membership for every sample. Such a requirement is impractical in situations where the data labeling efforts for minority or rare groups are significantly laborious or where the individuals comprising the dataset choose to conceal sensitive information. On the other hand, the presence of such data collection efforts results in datasets that contain partially labeled group information. Recent works have tackled the fully unsupervised scenario where no labels for groups are available. Thus, we aim to fill the missing gap in the literature by tackling a more realistic setting that can leverage partially available sensitive or group information during training. First, we construct a constraint set and derive a high probability bound for the group assignment to belong to the set. Second, we propose an algorithm that optimizes for the worst-off group assignments from the constraint set. Through experiments on image and tabular datasets, we show improvements in the minority group's performance while preserving overall aggregate accuracy across groups.