HSFM: Hard-Set-Guided Feature-Space Meta-Learning for Robust Classification under Spurious Correlations
Abstract
Deep neural networks often rely on spurious features to makepredictions, which makes them brittle under distribution shift and onsamples where the spurious correlation does not hold (e.g., minority-group examples). Recent studies have shown that, even in such settings,the feature extractor of an Empirical Risk Minimization (ERM)-trainedmodel can learn rich and informative representations, and that muchof the failure may be attributed to the classifier head. In particular,retraining a lightweight head while keeping the backbone frozen cansubstantially improve performance on shifted distributions and minor-ity groups. Motivated by this observation, we propose a bilevel meta-learning method that performs augmentation directly in feature space toimprove spurious correlation handling in the classifier head. Our methodlearns support-side feature edits such that, after a small number of inner-loop updates on the edited features, the classifier achieves lower loss onhard examples and improved worst-group performance. By operating atthe backbone output rather than in pixel space or through end-to-endoptimization, the method is highly efficient and stable, requiring only afew minutes of training on a single GPU. We further validate our methodwith CLIP-based visualizations, showing that the learned feature-spaceupdates induce semantically meaningful shifts aligned with spurious at-tributes. Our code is available at: https://github.com/ArianYp/HSFM.