AGRO: Adversarial Discovery of Error-prone groups for Robust Optimization

TL;DR

AGRO introduces an adversarial method to discover error-prone groups in training data, enhancing distributionally robust optimization and improving model performance on known and unknown spurious correlations.

Contribution

It presents a novel end-to-end approach that jointly identifies error-prone groups and enhances robustness without prior knowledge of spurious correlations.

Findings

8% higher performance on known worst-groups

Improves out-of-distribution accuracy on multiple datasets

Discoveries contain well-defined, previously unstudied spurious correlations

Abstract

Models trained via empirical risk minimization (ERM) are known to rely on spurious correlations between labels and task-independent input features, resulting in poor generalization to distributional shifts. Group distributionally robust optimization (G-DRO) can alleviate this problem by minimizing the worst-case loss over a set of pre-defined groups over training data. G-DRO successfully improves performance of the worst-group, where the correlation does not hold. However, G-DRO assumes that the spurious correlations and associated worst groups are known in advance, making it challenging to apply it to new tasks with potentially multiple unknown spurious correlations. We propose AGRO -- Adversarial Group discovery for Distributionally Robust Optimization -- an end-to-end approach that jointly identifies error-prone groups and improves accuracy on them. AGRO equips G-DRO with an adversarial slicing model to find a group assignment for training examples which maximizes worst-case loss over the discovered groups. On the WILDS benchmark, AGRO results in 8% higher model performance on average on known worst-groups, compared to prior group discovery approaches used with G-DRO. AGRO also improves out-of-distribution performance on SST2, QQP, and MS-COCO -- datasets where potential spurious correlations are as yet uncharacterized. Human evaluation of ARGO groups shows that they contain well-defined, yet previously unstudied spurious correlations that lead to model errors.