[Paper Review] Simple and Fast Group Robustness by Automatic Feature Reweighting

Shikai Qiu, Andres Potapczynski+2 more|arXiv (Cornell University)|Jun 19, 2023

Machine Learning and Data Classification4 citations

TL;DR

This paper proposes Automatic Feature Reweighting (AFR), a simple and efficient method to improve group robustness in machine learning by retraining only the final layer of an ERM model using a weighted loss that prioritizes examples where the base model performs poorly. AFR achieves state-of-the-art performance on vision and NLP benchmarks like Waterbirds and MultiNLI without requiring spurious attribute labels or significant computational overhead.

ABSTRACT

A major challenge to out-of-distribution generalization is reliance on spurious features -- patterns that are predictive of the class label in the training data distribution, but not causally related to the target. Standard methods for reducing the reliance on spurious features typically assume that we know what the spurious feature is, which is rarely true in the real world. Methods that attempt to alleviate this limitation are complex, hard to tune, and lead to a significant computational overhead compared to standard training. In this paper, we propose Automatic Feature Reweighting (AFR), an extremely simple and fast method for updating the model to reduce the reliance on spurious features. AFR retrains the last layer of a standard ERM-trained base model with a weighted loss that emphasizes the examples where the ERM model predicts poorly, automatically upweighting the minority group without group labels. With this simple procedure, we improve upon the best reported results among competing methods trained without spurious attributes on several vision and natural language classification benchmarks, using only a fraction of their compute.

Motivation & Objective

• To address the challenge of out-of-distribution generalization caused by reliance on spurious features in real-world datasets.
• To develop a method that improves group robustness without requiring group labels or complex training procedures.
• To reduce computational overhead while maintaining or improving performance on minority groups.
• To create a method robust to hyperparameter choices and applicable to diverse vision and NLP benchmarks.

Proposed method

• AFR retrains only the final layer of a pre-trained ERM model using a weighted loss function.
• Example weights are determined by the inverse of the model’s predicted probability on the correct class from a first-stage ERM checkpoint.
• Weights are scaled by class-dependent constants to address class imbalance, following the form μ_i ∝ β_{y_i} exp(−γ p̂_i).
• The weighting automatically emphasizes minority group examples by upweighting low-confidence predictions, without using group labels.
• The method is applied on a held-out dataset from the training distribution, avoiding the need for group annotations.
• AFR is computationally efficient, requiring only a fraction of the training time of competing methods.

Experimental results

Research questions

• RQ1 Can a simple, lightweight method improve group robustness without access to spurious attribute labels?
• RQ2 Can model uncertainty (low prediction confidence) be used as a proxy to identify and prioritize minority group examples?
• RQ3 How does AFR compare to state-of-the-art methods that require group labels or complex training procedures?
• RQ4 To what extent is AFR robust to hyperparameter choices in practice?
• RQ5 Can AFR outperform methods that use group labels when only a small number of labels are available?

Key findings

• AFR achieves state-of-the-art worst-group accuracy on the Waterbirds benchmark among methods that do not use spurious attribute labels during training.
• On MultiNLI, AFR improves upon the best reported results among group-robust methods without access to group labels.
• AFR outperforms a state-of-the-art method using group labels when only a small number of group labels are available for hyperparameter tuning.
• AFR is robust to hyperparameter choices, often improving group robustness without careful tuning.
• The optimal weighting function learned by neural networks closely resembles AFR’s exponential form, validating its design choice.
• AFR’s computational overhead is negligible, requiring only a fraction of the training time of baseline methods.