Source-linked AI summary
Just Train Twice: Improving Group Robustness without Training Group Information
Evan Zheran Liu, Behzad Haghgoo, Annie S. Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, Chelsea Finn
TL;DR
ERM can perform poorly on groups affected by spurious correlations, while group-annotated methods are expensive and annotation-free alternatives often have weaker worst-group accuracy. JTT trains an ERM model, upweights its misclassified examples in a second model, and closes 75% of the ERM-to-group-DRO worst-group accuracy gap on average across four datasets. It requires group annotations only on a smaller validation set for hyperparameter tuning.
Problem
ERM can achieve high average accuracy but low accuracy on groups where spurious correlations do not hold, while group DRO requires costly training group annotations.
Method
JTT trains an ERM identification model, collects its misclassified training examples, and upweights those examples when training a final model.
Results
75% of the worst-group accuracy gap between ERM and group DRO is closed on average across four image and language datasets.
Takeaways & Limitations
JTT improves worst-group accuracy without training group annotations, requiring them only on a smaller validation set for hyperparameter tuning.
Takeaways & Limitations
High worst-group test accuracy still requires validation group annotations for hyperparameter tuning.
Abstract
from arXiv · showhide
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.
1. Introduction
ERM can achieve low average test error while performing poorly on particular groups, especially when spurious correlations make minority or mismatched examples difficult. JTT addresses this without training group annotations by upweighting examples misclassified by an initial ERM model.
- ERM may achieve low average test error while incurring high error on certain groups.
- Spurious correlations can cause models to associate demographic mentions with toxicity and fail on examples where that correlation does not hold.
- JTT trains an ERM model, identifies its misclassified training examples, and upweights them when training a final model.
- JTT is evaluated on four datasets spanning image classification and natural language processing with spurious correlations.
- JTT’s analysis finds that its selected examples tend to come from groups where ERM performs poorly, while validation group annotations are essential for hyperparameter tuning.
- CVaR DRO also avoids training group annotations by dynamically upweighting high-loss examples, but shares conceptual similarities with JTT.
2. Related Work
Related work addresses group robustness either by using group information during training or by avoiding training annotations through distributionally robust and two-model approaches. JTT is positioned as a simpler fixed-error-set alternative to closely related methods such as LfF.
- Group-robust methods using training group information minimize worst-group loss, expand minority groups, or reweight examples.
- Fairness-oriented work also studies worst-group performance, explicitly focusing on equalizing loss across groups.
- Methods without training group information include distributionally robust optimization approaches that minimize loss over distributions near the empirical distribution.
- LfF uses a deliberately biased first model to identify minority examples and upweight them while training a second model.
- Unlike LfF’s interleaved updates and intentional biasing, JTT identifies points to upweight once.
- Another concurrent approach uses ERM errors to infer group labels before learning a model invariant to those predicted labels.
3. Preliminaries
The paper studies worst-group error when groups are predefined and training group annotations are unavailable, using a small annotated validation set for tuning. It compares ERM, CVaR DRO, LfF, and group DRO across spurious-correlation settings.
- 3.1. Problem Setup: The goal is good test-time worst-group error without training group annotations, while a small annotated validation set supports hyperparameter tuning.
- 3.1. Problem Setup: Experimental groups are primarily defined by a spurious attribute and label, such as background and bird type in Waterbirds.
- 3.2. Comparisons: ERM minimizes average training loss across training points.
- 3.2.2. DISTRIBUTIONALLY ROBUST OPTIMIZATION OF THE CONDITIONAL VALUE AT RISK (CVAR DRO): CVaR DRO minimizes expected loss over a worst-case distribution in an uncertainty set around the empirical distribution.
- 3.2.2. DISTRIBUTIONALLY ROBUST OPTIMIZATION OF THE CONDITIONAL VALUE AT RISK (CVAR DRO): At level α, CVaR considers uncertainty sets containing α-sized subpopulations, whose worst loss can upper-bound worst-group loss when α matches the smallest group size.
- 3.2.3. LEARNING FROM FAILURE (LFF): LfF trains a first model with generalized cross-entropy and a second model with reweighted cross-entropy to emphasize challenging examples.
- 3.2.4. GROUP DISTRIBUTIONALLY ROBUST OPTIMIZATION (GROUP DRO): Group DRO directly minimizes worst-group training error using training group annotations and serves as an oracle comparison.
4. JTT: Just Train Twice
JTT is a two-stage procedure: an ERM identification model supplies a fixed error set, and a final model upweights those examples. Early stopping and validation-based tuning make the procedure practical.
- JTT first trains an identification model and then trains a final model while upweighting the selected examples.
- Stage 1 (identification): Early-stopped ERM tends to fit groups with easy spurious correlations but not groups where those correlations fail, motivating error-set selection.
- Stage 2 (upweighting): The final model upweights examples in E because they may come from challenging groups lacking the spurious correlation.
- Practical implementation: The identification model is trained for T steps to limit capacity and avoid an empty error set.
- Stage 2 (upweighting): JTT constructs an upsampled dataset by repeating error-set examples λup times and retaining other examples once.
- Practical implementation: Hyperparameters are tuned using final-model worst-group error on the validation set.
5. Experiments
Across four datasets with spurious correlations, JTT improves worst-group accuracy by selecting ERM errors for upweighting, while requiring only validation-set group annotations for tuning. Its gains are linked to error sets enriched for poorly performing groups, but performance depends on capacity and hyperparameter choices.
- Main Results: JTT consistently achieves higher worst-group accuracy than methods without training group information across Waterbirds, CelebA, MultiNLI, and CivilComments.The comparison includes ERM, CVaR DRO, and LfF, and also evaluates JTT against group DRO.
- Main Results: 75% of the worst-group accuracy gap between ERM and group DRO is closed by JTT on average.Group DRO uses training group annotations, whereas JTT uses them only on a much smaller validation set for hyperparameter tuning.
- Main Results: JTT’s average accuracy is 4.2% below the highest average accuracy on each dataset, accompanying its worst-group improvements.This reflects the reported tradeoff between average and worst-group accuracy.
- Error Set Analysis: Worst-group examples occur in JTT’s error set 2.2x to 15.9x more often than in training data, with 86.4% average recall across datasets.Worst-group recall ranges from 67.1% to 96.9%, indicating that the error set captures many examples from the lowest-performing group.
- Error Set Analysis: Error-set enrichment roughly inversely correlates with ERM test accuracy, concentrating upweighted examples from groups on which ERM performs poorly.High-enrichment groups do not always violate the spurious correlation; label imbalance can also produce high enrichment, as observed for selected Waterbirds and CelebA groups.
- Hyperparameter Sensitivity: JTT requires careful worst-group validation tuning: performance can drop when training duration is too short or too long, and fixed error sets perform better on Waterbirds.The identification-model epoch count T is sensitive, while holding the error set fixed for more epochs produces higher average and worst-group accuracy.
6. Discussion
JTT improves worst-group performance without expensive training-time group annotations, but the paper identifies unresolved theoretical, annotation, and distribution-shift boundaries.
- Discussion: JTT substantially improves worst-group performance without requiring expensive group annotations during training.It instead relies on a smaller validation set for annotation-dependent hyperparameter tuning.
- Future work: A better theoretical understanding is needed to explain when and why JTT works, including early-stopped ERM’s reliance on spurious correlations.The authors also highlight the unexplained importance of fixing JTT’s upweighted set rather than dynamically recomputing it as in CVaR DRO.
- Limitations: JTT still relies on a validation set representative of the distribution shift or annotated with group information.This is cheaper than annotating the entire training set but requires practitioners to recognize spurious correlations and define groups.
- Limitations: The experiments focus on group robustness under spurious correlations, while JTT is not specifically tailored to that setting.The paper leaves open whether it improves performance under other shifts, such as domain generalization.
A. Training Details
The experiments use shared architectures and dataset-specific optimization settings, validation-based tuning, and JTT-specific choices for identification and upweighting.
- Architectures: All approaches share dataset-specific model architectures: ResNet-50 for Waterbirds and CelebA, and BERT for MultiNLI and CivilComments.ResNet-50 starts from ImageNet-pretrained weights, while BERT uses the HuggingFace implementation.
- Comparisons: The reported group DRO results for three datasets are taken directly from prior work, with a noted CivilComments spurious-attribute difference.The authors report similar performance despite using any demographic identity rather than only Black identity.
- Hyperparameter tuning: Hyperparameters and early stopping are tuned using highest validation worst-group accuracy.CVaR DRO tunes α, LfF tunes q, and JTT tunes the identification-model epochs T and upsampling factor λup.
- JTT settings: JTT trains an identification model for T epochs and a longer final model, with the two models generally sharing hyperparameters.For BERT, JTT can use different optimizers for the identification and final models.
- Vision datasets: Waterbirds and CelebA use SGD with momentum 0.9, batch normalization, and no data augmentation.Waterbirds runs up to 300 epochs with batch size 64, while CelebA runs up to 50 epochs with batch size 128.
B.1. Waterbirds
The benchmark suite defines spurious attributes and minority groups across image and language tasks, with group structure determined by label–attribute combinations.
- Waterbirds: Waterbirds combines bird images from CUB with backgrounds from Places, creating waterbird and landbird labels alongside water and land background attributes.The two minority groups are (land background, waterbird) and (water background, landbird).
- Waterbirds: In Waterbirds training data, 95% of waterbirds have water backgrounds and 95% of landbirds have land backgrounds.This creates substantially fewer examples in the minority groups than in the majority groups.
- CelebA: CelebA uses Blond Hair as the label and Male as the spurious attribute, with minority groups (blond, male) and (not blond, female).Being female spuriously correlates with having blond hair in this setup.
- MultiNLI: MultiNLI uses negation words in the hypothesis as a binary spurious attribute that correlates with contradiction.The task predicts entailment, neutrality, or contradiction from premise–hypothesis pairs.
- CivilComments: CivilComments uses eight demographic-identity indicators and evaluates 16 potentially overlapping groups formed by identity and toxicity labels.The groups include identity–toxic and identity–not toxic combinations for each identity.
C.1. Error set analysis for CivilComments
The CivilComments analysis examines how error-set composition affects worst-group accuracy and shows that both label types contribute to JTT’s gains.
- Error set composition: The analysis is presented as a CivilComments error-set breakdown in Table 12.The table is included because of space constraints.
- Error set composition: Removing either y = a or y̸ = a examples from JTT’s error set substantially reduces worst-group accuracy improvements.Both categories therefore contribute substantially to the observed improvement.
C.2. Additional analysis
Additional analyses show that JTT’s selected error-set examples matter, especially its y = a examples, while random upsampling is insufficient. JTT also outperforms a ground-truth minority-group upsampling baseline.
- Analysis setup: JTT analyzes Waterbirds and CelebA by partitioning examples into groups where the spurious correlation holds (y = a) or does not (y̸ = a).These binary-label, binary-attribute datasets permit the partition used in the analysis.
- Error-set analysis: Removing either y = a or y̸ = a examples from JTT’s error set significantly decreases worst-group accuracy on both datasets.Both example types contribute to the upsampled error set’s performance.
- Baseline comparison: UPSAMPLE MINORITY improves worst-group accuracy over ERM but remains below JTT.The baseline upweights y̸ = a examples using ground-truth group annotations, but those groups need not be the hardest or smallest.
- Error-set analysis: Replacing JTT’s y̸ = a examples with random examples slightly lowers worst-group accuracy, while replacing its y = a examples significantly lowers it.The particular y = a examples selected by JTT appear especially important for worst-group performance.
- Error-set analysis: Replacing both example types or randomly sampling the entire error set produces low worst-group and average accuracies.Merely upsampling random examples from these groups is therefore insufficient for high worst-group accuracy.