PaperlayerView exact revision on arXiv

Source-linked AI summary

Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations

Polina Kirichenko, Pavel Izmailov, Andrew Gordon Wilson

arXiv:2204.02937v2cs.LGcs.CVstat.ML

TL;DR

Neural networks may rely on spurious features and fail under correlation shifts, even though they often learn core features. The paper retrains only the final classification layer using balanced reweighting data, finding state-of-the-art benchmark performance and reduced background and texture reliance in ImageNet models.

  • Problem

    Spurious correlations cause neural classifiers to rely on irrelevant features and degrade when those correlations break, while it remains unclear whether core features are learned.

  • Method

    Deep Feature Reweighting trains an ERM model normally, then retrains its last classification layer on a typically smaller balanced reweighting dataset.

  • Results

    DFR achieves state-of-the-art performance on spurious-correlation benchmarks and reduces ImageNet-trained models’ reliance on background and texture information.

  • Takeaways & Limitations

    Last-layer retraining can improve robustness without retraining the feature extractor, including for large ImageNet models in minutes on a single GPU.

  • Takeaways & Limitations

    DFR currently assumes access to a balanced reweighting dataset, motivating future work to avoid that requirement.

Abstract

from arXiv · show

Neural network classifiers can largely rely on simple spurious features, such as backgrounds, to make predictions. However, even in these cases, we show that they still often learn core features associated with the desired attributes of the data, contrary to recent findings. Inspired by this insight, we demonstrate that simple last layer retraining can match or outperform state-of-the-art approaches on spurious correlation benchmarks, but with profoundly lower complexity and computational expenses. Moreover, we show that last layer retraining on large ImageNet-trained models can also significantly reduce reliance on background and texture information, improving robustness to covariate shift, after only minutes of training on a single GPU.

1 INTRODUCTION

Spurious correlations can make neural networks rely on irrelevant features and fail on minority groups, even when they have learned useful core features. The paper uses this finding to motivate Deep Feature Reweighting, which retrains only the last layer to improve robustness.

  • Motivation: Spurious features predictive in training can become irrelevant to the true labeling function, causing performance degradation when correlations break.Examples include backgrounds associated with object classes and hospital-specific cues in medical imaging.
  • Key observation: Neural networks often learn core features despite relying heavily on spurious features and performing poorly on minority groups.Core features may be learned even when spurious features are simpler or minority examples are absent from training.
  • Approach: Deep Feature Reweighting retrains the last linear layer on a small dataset where the spurious correlation does not hold.The method reweights features without retraining the feature extractor.
  • Approach: DFR makes the resulting model rely primarily on foreground features and perform much better on images with confusing backgrounds.

2 PROBLEM SETTING

The paper models data as an imbalanced mixture of groups defined by labels and spurious attributes. This imbalance makes the spurious attribute predictive in training but unreliable for minority groups.

  • Group-based formulation: Classification data are represented as groups Gi with distributions pi(x, y), mixed according to proportions αi.The training distribution is expressed as p(x, y) = P i αipi(x, y).
  • Waterbirds example: In Waterbirds, groups combine bird type and background, with majority groups substantially larger than minority groups.The four groups are waterbirds on water or land and landbirds on water or land.
  • Spurious attributes: A background is spurious when it correlates with the target in training but is not predictive of the target on minority groups.

3 RELATED WORK

Prior work studies spurious-feature learning, group robustness, classifier retraining, and related generalization settings. This paper differs by arguing that standard ERM representations can already support robust performance after last-layer reweighting.

  • Feature learning: Research on spurious correlations examines feature learning, optimization dynamics, and neural-network inductive biases.
  • Group robustness: Group-robust methods commonly optimize worst-group or distributionally robust objectives rather than average loss.Group DRO is identified as state of the art on multiple spurious-correlation benchmarks.
  • Classifier retraining: Prior classifier-retraining work corrected subgroup bias, while this paper reports better performance from retraining on held-out reweighting data.
  • Scope distinction: Related work applies last-layer training to domain generalization or long-tail classification, whereas this paper focuses on spurious-correlation robustness.
  • Position of this work: The paper claims that pretrained ERM representations can achieve state-of-the-art benchmark performance by reweighting features rather than redesigning the full model.

4 UNDERSTANDING REPRESENTATION LEARNING WITH SPURIOUS CORRELATIONS

Experiments show that ERM models can learn core features even while relying on spurious cues, including under severe correlation settings. This supports improving robustness by changing feature weighting rather than relearning representations.

  • Core-feature learning: ERM models underperform on minority groups but still learn core features that support correct minority-group predictions.
  • Representation analysis: Core features can be decoded from latent representations even when Original-data accuracy is 0%, showing that poor predictions do not imply absent core information.
  • Waterbirds: 94% worst-group accuracy is achieved on FG-Only Waterbirds test data by models trained on spurious-correlated Original data.This is only slightly below models trained directly without the spurious background.
  • Waterbirds: Original-data Waterbirds models show a large mean-versus-worst-group gap, with the no-minority model reaching only 38.4% worst-group accuracy.
  • Implication: These findings motivate retraining the last layer to emphasize relevant features and improve worst-group performance.The paper also notes approximate logit additivity between foreground and background components on Waterbirds.
  • Simplicity bias: On Dominoes, perfect spurious correlation yields 0% worst-group accuracy, while 99% and 95% correlations show learned core features through improved Core-Only performance.
  • Additional evidence: Core-feature decoding remains near optimal on ColorMNIST and is generally strong when spurious correlation is below 100% or complexity differences are less extreme.

5 DEEP FEATURE REWEIGHTING

Deep Feature Reweighting improves robustness by retraining only the final classification layer on a smaller, group-balanced reweighting dataset. It reweights features learned by an ERM-trained network without retraining the feature extractor.

  • DFR targets robustness to spurious correlations and distribution shifts by reweighting features learned by standard ERM.
  • DFR retrains a new classification head on balanced data after training the feature extractor on the original dataset.The first stage uses standard ERM; the second discards the original head and trains a replacement on the reweighting dataset.
  • The reweighting dataset is typically much smaller than the original training dataset and contains equally represented groups.

6 FEATURE REWEIGHTING IMPROVES ROBUSTNESS

On spurious-correlation benchmarks, DFR is competitive with or better than state-of-the-art methods while using group labels only on a small validation set. Its procedure retrains the last layer rather than applying advanced representation-learning or regularization methods.

  • Evaluation: Table 2 reports worst-group and mean test accuracy for DFR variants and baselines, with mean accuracy weighted by training-group prevalence.
  • Method: DFR uses a group-balanced subset of validation data to retrain the last layer while the imbalanced training data trains the feature extractor.
  • Method: DFR uses ℓ1-regularized logistic regression and tunes only the regularization strength for the last-layer retraining procedure.
  • Results: DFR achieves the best results among evaluated methods on Waterbirds and CivilComments while remaining competitive with Group DRO across benchmarks.
  • Results: DFR has similar performance to SSA, slightly outperforming it on Waterbirds and CivilComments but underperforming it on CelebA and MultiNLI.
  • Results: DFR matches the performance of the best available methods using group labels only on a small validation set and simply reweighting ERM-learned features.

7 NATURAL SPURIOUS CORRELATIONS ON IMAGENET

Last-layer DFR reduces reliance on background and texture cues in ImageNet-trained models while largely preserving performance on relevant data. It increases shape bias and improves robustness to some distribution shifts, though retraining from scratch can achieve higher shape bias.

  • Background reliance: DFR reduces background reliance in ImageNet-trained models by retraining only their last layer.
  • Background reliance: DFR improves performance on Mixed-Rand, FG-Only, and Paintings-BG, while DFR on Original data mostly preserves performance with only a small drop.
  • Background reliance: DFR provides only a small ImageNet-R improvement because that covariate shift is not primarily background-based.
  • Texture-vs-shape bias: Shape bias measures predictions based on shape rather than texture for images with conflicting texture and shape cues.
  • Texture-vs-shape bias: 21.4% →34% shape bias for DFR(SIN) shows that retraining only the last layer can substantially increase shape bias over the base model.
  • Texture-vs-shape bias: DFR improves ImageNet-R and ImageNet-C performance, but RN-50(SIN) achieves much higher shape bias than the other methods.

8 DISCUSSION

The paper argues that ERM models learn useful semantic features alongside spurious ones, so robustness can often be improved by reweighting existing features rather than retraining the full network. DFR offers this intervention with low complexity, while relying on group-balanced reweighting data and leaving several extensions open.

  • Practical advantages of DFR: DFR achieves state-of-the-art benchmark performance and reduces ImageNet-trained models’ reliance on background and texture information within minutes of training.The reported scope includes spurious-correlation benchmarks and ImageNet-scale models.
  • Spurious correlations and representation learning: ERM representations can support state-of-the-art spurious-correlation robustness even without minority-group examples; the main issue is feature weighting.The paper contrasts this view with approaches that attribute poor robustness primarily to representation quality.
  • Practical advantages of DFR: DFR has one tunable hyper-parameter, is robust to base-model choice, and avoids early stopping or highly problem-specific tuning.These properties contribute to its practical simplicity and low computational cost.
  • Data requirements: DFR requires group-labeled reweighting data, which the authors recommend using for model selection rather than only hyper-parameter tuning.The discussion frames access to such data as a practical assumption shared by many methods.
  • Future work: Future work includes avoiding balanced reweighting data, using saliency or segmentation supervision, and improving representations with self-supervised learning.The paper presents these as directions after reducing the problem to training a robust linear classifier on fixed features.
  • Deep Feature Reweighting: DFR retrains only the last layer of an ERM classifier on group-balanced reweighting data to improve worst-group performance.Its design differs from class-balanced sampling and can use held-out data in DFRVal Tr.

B.2 DOMINOES DATASETS

The experiments use synthetic and benchmark datasets to test whether models retain core features despite strong spurious correlations. Results include strong DFR recovery under extreme ColorMNIST correlations and evidence that foreground and background contributions are processed approximately independently.

  • DOMINOES DATASETS: Dominoes combines an easier, linearly separable MNIST feature in the top half with a harder target-related feature in the bottom half.The datasets include MNIST-MNIST, MNIST-FashionMNIST, and MNIST-CIFAR variants with 95%, 99%, and 100% spurious-correlation levels.
  • Feature learning: In inverted Waterbirds, models trained on Original data perform well on BG-Only images, nearly matching models trained directly on BG-Only data.This holds even for models trained without minority-group examples, although their Original-data worst-group accuracy is poor.
  • Feature analysis: Logit additivity means Original-image logits are well approximated by summing corresponding foreground-only and background-only logits.This supports approximately independent processing of foreground and background predictive features before their combination in the last layer.
  • COLORMNIST: DFR recovers strong worst-group accuracy on ColorMNIST even when ERM has 0% worst-group accuracy at pcorr = 1.0 and pcorr = 0.995.For pcorr < 0.95, DFR closely matches the optimal accuracy from training without spurious correlations.
  • Experimental setup: The benchmark experiments use Waterbirds, CelebA, MultiNLI, and CivilComments, with ImageNet-pretrained ResNet-50 models for Waterbirds and CelebA.DFR extracts and standardizes penultimate-layer embeddings before retraining the classification layer.

C.1 ABLATION STUDIES

The ablations examine how DFR depends on pretraining, retraining scope, base-model settings, regularization, and repeated linear retraining. Across these tests, last-layer retraining is consistently effective, while held-out data, suitable regularization, and averaging can improve results.

  • Pretraining: ImageNet pretraining has a dramatic effect on DFR and base-model performance on Waterbirds, whereas DFR remains strong without pretraining on CelebA.On both datasets, finetuning the feature extractor on target data is crucial; using ImageNet features without finetuning performs poorly.
  • Retraining scope: Last-layer retraining provides optimal performance; retraining additional layers is competitive but worse, while retraining the full network is substantially worse.The comparison includes retraining the last two layers, the last residual block, and the full network.
  • Repeated retraining: Averaging three linear retrains raises CelebA worst group accuracy to 88.6%, compared with 85% for a single model.Repeated retrains use random balanced validation subsets and average their learned weights.
  • Regularization: 92.9 ± 0.2% worst group accuracy on Waterbirds and 88.3 ± 1.1% on CelebA are achieved with ℓ1 regularization, versus 87.72 ± 0.42% and 86.03 ± 0.42% without it.The paper emphasizes ℓ1 regularization when feature dimensionality is much larger than the reweighting dataset.
  • Comparison: DFRVal Tr outperforms full-model validation fine-tuning, which reaches 89.3 ± 1.3% on Waterbirds and 84.4 ± 0.5% on CelebA.The full-model baseline starts from an ImageNet-pretrained ResNet-50 and is trained for 10 epochs on group-balanced validation data.

C.3 DFR VARIATIONS

DFR variations differ mainly in which data retrains the last layer and whether minority groups were seen by the feature extractor. Held-out reweighting data is generally best, while even no-minority training can recover strong robustness.

  • Benchmark comparisons: DFRTr-NM matches state-of-the-art Group DRO on CelebA and achieves competitive results on Waterbirds using the same data to train and tune the model.Its feature extractor was trained without seeing examples from the minority groups.
  • NLP results: All DFR variations significantly improve performance over the base model on NLP datasets, with DFRVal Tr outperforming the other variations.On CivilComments, the no-minority base model achieves only 8.4% worst group performance, while DFRTr-NM recovers 66.5%.
  • Procedure: DFR uses a fixed feature encoder and trains logistic regression on reweighting-set features before evaluating on test-set features.The approach is framed as retraining only the classifier while preserving the learned representation.
  • Feature distributions: DFRTr-NM avoids the minority-group feature distribution shift observed for DFRTr between reweighting and test data.For DFRTr, evaluating the logistic regression model under this shift makes the problem more challenging and leads to inferior performance.
  • Reweighting data: DFRTr achieves the best performance by retraining the last layer on data not used to train the feature extractor.The paper recommends a group-balanced validation set for both hyper-parameter tuning and last-layer retraining.
  • Relation to transfer learning: DFR is a special case of transfer learning, but its purpose is correcting a pretrained model's spurious-feature reliance rather than learning general downstream features.The paper applies this correction to background and texture reliance in ImageNet-trained models.

D DETAILS: IMAGENET BACKGROUND RELIANCE

On ImageNet-9, DFR retrains the final layer using features from background-varied data to reduce background reliance. It improves behavior on unusual backgrounds while preserving most original-data performance.

  • Data and models: DFR is trained on Mixed-Rand subsets or balanced combinations of Mixed-Rand and Original ImageNet-9 data.The experiments use pretrained ResNet-50 and ViT-B-16 feature extractors.
  • Robustness: DFR significantly reduces background reliance for ViT-B-16, with only a minimal performance drop on Original data.Performance improves greatly on Mixed-Rand, FG-Only, and Paintings-BG images with unusual backgrounds.
  • Background invariance: 93.1% of predictions remain unchanged under fixed-foreground background modifications for DFROG+MR, versus 87.5% for the baseline.DFRMR reaches 92.4% under the same evaluation.
  • Prediction behavior: 11.2% of Mixed-Rand predictions match the background class for DFRMR, compared with 14.8% for the baseline.DFROG+MR yields 11.7%, indicating less frequent background-class prediction than the baseline.
  • Visual analysis: GradCAM visualizations show baseline attention on background context, while DFRMR features are more compact and focus on the target object.The visualization provides qualitative support for reduced background reliance.

E DETAILS: IMAGENET TEXTURE BIAS DETAILS

DFR uses original and stylized ImageNet embeddings to alter texture-versus-shape behavior and test robustness under corruptions. It improves shape bias and robustness, although the strongest corruption robustness comes from a model trained from scratch on combined data.

  • Setup: DFR uses embeddings from ImageNet, Stylized ImageNet, or their combination as the reweighting dataset.Evaluation covers ImageNet-C across 19 corruption types and five intensities, plus ImageNet-R.
  • Shape bias: 36% → 39.9% shape bias is achieved by DFR on a ViT-B-16 pretrained on ImageNet21k and finetuned on ImageNet.The same result is reported alongside improved robustness to ImageNet-C corruptions.
  • Texture-bias analysis: The texture-bias evaluation reports fractions of shape and texture decisions for DFR models and models trained on ImageNet, Stylized ImageNet, and their combination.The comparison also includes ShapeResNet-50 and human predictions.
  • Corruption robustness: DFR trained on combined ImageNet and Stylized ImageNet improves over the RN50(IN) baseline on many individual ImageNet-C corruptions.The model trained from scratch on combined data provides the best robustness across the board.

F COMPARISON TO KANG ET AL. (2019)

The comparison shows that DFR’s held-out-data and group-label subsampling choices are important for achieving the best performance. These choices distinguish DFR from related last-layer retraining and classifier-balancing variants.

  • Results: LWS and cRT perform poorly on worst-group accuracy in the spurious-correlation setting, while group-balanced last-layer retraining improves over them.LWS is limited because it retrains only one scaling parameter per class, which does not sufficiently remove reliance on spurious features.
  • Results: DFRTr achieves the best performance across the reported comparisons.
  • Design choices: Optimal performance requires held-out retraining data rather than reused training data and group-label subsampling rather than group-balanced sampling.
  • Dataset conditions: The advantage of subsampling depends on dataset group distributions: it is much larger on CelebA because validation and training distributions match.On Waterbirds, the difference is less pronounced when retraining on validation data because validation is relatively group-balanced.
  • Methods compared: DFR is compared with LWS, cRT, and last-layer retraining variants differing in retraining data, balancing labels, and balancing strategy.The comparisons include train versus validation data, class versus group labels, and subsampling versus balanced sampling.
Loading 2204.02937v2…