Just Mix Once: Worst-group Generalization by Group Interpolation

TL;DR

This paper introduces Just Mix Once (JM1), a simple, efficient mixup-based method that improves worst-group generalization in deep learning models without requiring extensive group annotations.

Contribution

JM1 unifies and generalizes existing approaches using class-conditional mixup for robust worst-group generalization, applicable with any level of group annotation.

Findings

JM1 performs on par or better than state-of-the-art methods.

JM1 is domain agnostic and computationally efficient.

JM1 effectively improves worst-group generalization.

Abstract

Advances in deep learning theory have revealed how average generalization relies on superficial patterns in data. The consequences are brittle models with poor performance with shift in group distribution at test time. When group annotation is available, we can use robust optimization tools to tackle the problem. However, identification and annotation are time-consuming, especially on large datasets. A recent line of work leverages self-supervision and oversampling to improve generalization on minority groups without group annotation. We propose to unify and generalize these approaches using a class-conditional variant of mixup tailored for worst-group generalization. Our approach, Just Mix Once (JM1), interpolates samples during learning, augmenting the training distribution with a continuous mixture of groups. JM1 is domain agnostic and computationally efficient, can be used with any level of group annotation, and performs on par or better than the state-of-the-art on worst-group generalization. Additionally, we provide a simple explanation of why JM1 works.