Take One Gram of Neural Features, Get Enhanced Group Robustness

TL;DR

This paper introduces a method to enhance group robustness in machine learning models under distribution shifts by partitioning data into pseudo-groups using Gram matrices of features, eliminating the need for explicit group labels.

Contribution

The authors propose a novel approach that leverages Gram matrices of features to create pseudo-groups for robust optimization, avoiding the need for group annotations during training or validation.

Findings

Improves worst-group robustness over ERM without group labels

Outperforms recent baseline methods in experiments

Effective in scenarios with no available group annotations

Abstract

Predictive performance of machine learning models trained with empirical risk minimization (ERM) can degrade considerably under distribution shifts. The presence of spurious correlations in training datasets leads ERM-trained models to display high loss when evaluated on minority groups not presenting such correlations. Extensive attempts have been made to develop methods improving worst-group robustness. However, they require group information for each training input or at least, a validation set with group labels to tune their hyperparameters, which may be expensive to get or unknown a priori. In this paper, we address the challenge of improving group robustness without group annotation during training or validation. To this end, we propose to partition the training dataset into groups based on Gram matrices of features extracted by an ``identification'' model and to apply robust optimization based on these pseudo-groups. In the realistic context where no group labels are available, our experiments show that our approach not only improves group robustness over ERM but also outperforms all recent baselines