Not Only the Last-Layer Features for Spurious Correlations: All Layer Deep Feature Reweighting

TL;DR

This paper proposes retraining classifiers on features from all neural network layers, not just the last, using feature selection to improve robustness against spurious correlations and enhance worst-group accuracy.

Contribution

It introduces a method to leverage features from all layers of neural networks for better mitigation of spurious correlations, extending beyond the last-layer re-training approach.

Findings

Significant improvements in worst-group accuracy on benchmarks.

Feature selection from all layers enhances robustness.

All-layer re-training outperforms last-layer-only methods.

Abstract

Spurious correlations are a major source of errors for machine learning models, in particular when aiming for group-level fairness. It has been recently shown that a powerful approach to combat spurious correlations is to re-train the last layer on a balanced validation dataset, isolating robust features for the predictor. However, key attributes can sometimes be discarded by neural networks towards the last layer. In this work, we thus consider retraining a classifier on a set of features derived from all layers. We utilize a recently proposed feature selection strategy to select unbiased features from all the layers. We observe this approach gives significant improvements in worst-group accuracy on several standard benchmarks.