Concept-Residual Representation Expansion for Robustness to Spurious Correlations

Abstract

Models trained with empirical risk minimization (ERM) are prone to relying on spurious correlations to make predictions. A spurious correlation is a non-causal relationship in the training data between an attribute and the prediction target that does not generalize beyond the training environment. As a result, models can appear to achieve strong performance by exploiting these correlations, yet fail when the correlation changes or disappears. Despite their tendency to learn spurious correlations, the success of post-hoc mitigation methods in recent work suggests that ERM-trained models still retain useful, robust predictive features. However, core (non-spurious) features may be weak or entangled within the representation, making them difficult to identify. We propose Concept-Residual eXpansion (CRX), a concept-augmented framework that improves robustness by expanding the set of candidate predictive features. Starting from a frozen ERM representation, we augment the model's features with interpretable concept scores that describe the presence of task-relevant attributes and the surrounding context, together with residual features that capture the portion of the ERM features not expressed by the concepts. We then retrain a lightweight classifier on this expanded feature space, enabling it to leverage both structured semantic cues and complementary residual information. Across standard spurious correlation benchmarks, CRX consistently improves worst-group accuracy while maintaining competitive average performance. These results suggest that expanding the set of available features can substantially improve robustness to spurious correlations. The code and additional implementation details can be found at https://doi.org/10.5281/zenodo.20467988