Just Train Twice: Improving Group Robustness without Training Group Information

By Evan Zheran Liu, Behzad Haghgoo, Annie S. Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, Chelsea Finn. International Conference on Machine Learning, 2021

724 citations recorded by the source corpus.

Standard training via empirical risk minimization (ERM) can produce models that achieve high accuracy on average but low accuracy on certain groups, especially in the presence of spurious correlations between the input and label. Prior approaches that achieve high worst-group accuracy, like group distributionally robust optimization (group DRO) require expensive group annotations for each training point, whereas approaches that do not use such group annotations typically achieve unsatisfactory worst-group accuracy. In this paper, we propose a simple two-stage approach, JTT, that first trains a standard ERM model for several epochs, and then trains a second model that upweights the training examples that the first model misclassified. Intuitively, this upweights examples from groups on which standard ERM models perform poorly, leading to improved worst-group performance. Averaged over four image classification and natural language processing tasks with spurious correlations, JTT closes 75% of the gap in worst-group accuracy between standard ERM and group DRO, while only requiring group annotations on a small validation set in order to tune hyperparameters.

Subfields: Artificial Intelligence

View Paper: https://www.semanticscholar.org/paper/216d093cb2ad81bf55c21dbce2217f2b9032e67b

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

Citations in the corpus

Distilling Model Failures as Directions in Latent Space (2022, neutral): “We ﬁnd that our method further outperforms other baselines such as Domino (Eyuboglu et al., 2022), Learning from Failure (Nam et al., 2020), or using conﬁdences after early stopping (as in Liu et al. (2021)) (see Appendix B.1 for further experimental details).”
Metadata Archaeology: Unearthing Data Subsets by Leveraging Training Dynamics (2022, methodology): “We introduce the problem of Metadata Archaeology as the task of surfacing and inferring metadata of different examples in a dataset, noting that the relational qualities of metadata are of special interest (as compared to ordinary data features) for auditing, fairness, and many other applications.”
Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations (2022, methodology): “Just Train Twice (JTT ) [58] is a method that detects the minority group examples on train data, only using group labels on the validation set to tune hyper-parameters.”
Environment Inference for Invariant Learning (2020, methodology): “In this case, both methods use the same disjoint groups of training examples towards slightly different ends: we train an invariant learner, whereas Liu et al. (2021) train a cross-entropy classiﬁer with ﬁxed per-group importance weights.”
Weak-to-Strong Generalization: Eliciting Strong Capabilities With Weak Supervision (2023, methodology): “We use a proprietary dataset used to train ChatGPT RMs.”
Towards Out-Of-Distribution Generalization: A Survey (2021, background): “[155] found that training the model twice, with higher sample weights for harder samples, yields good generalization performance.”
Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models (2024, neutral): “…al., 2020; Oren et al., 2019; Zhang et al., 2021; Sohoni et al., 2022; Nam et al., 2022), automatically or manually reweighting data between groups (Liu et al., 2021; Nam et al., 2020; Yaghoobzadeh et al., 2021; Utama et al., 2020; Creager et al., 2021; Idrissi et al., 2022; Orgad & Belinkov,…”
Generalizing to Unseen Domains: A Survey on Domain Generalization (2021, methodology): “Such annotation was later narrowed down to a small fraction of validation set in [154], where they formulated a two-stage weighting framework.”