ShortcutProbe: Probing Prediction Shortcuts for Learning Robust Models

1 Introduction

Deep learning models have shown remarkable performance across domains, but this success is often achieved by exploiting spurious correlations [Sagawa et al., 2020; Liu et al., 2021; Yang et al., 2023; LaBonte et al., 2024; Ye et al., 2024; Zheng et al., 2024a] between spurious attributes or shortcut features [Geirhos et al., 2020] and targets. For example, models have been found to use correlations between textures and image classes [Geirhos et al., 2019] for object recognition instead of focusing on defining features of objects. This issue becomes even more problematic in high-stakes domains like healthcare. For instance, models predicting pneumonia were shown to rely on correlations between metal tokens in chest X-ray scans from different hospitals and the disease's detection outcomes [Zech et al., 2018], rather than the pathological features of pneumonia itself. The tendency of using spurious correlations is referred to as spurious bias. Models with spurious bias often fail to generalize on data groups lacking the learned spurious correlations, leading to significant performance degradation and non-robust behaviors across different data groups. This robustness issue can have severe social consequences, especially in critical applications.

Mitigating spurious bias is crucial for robust generalization across data groups with varying spurious correlations. Existing methods on mitigating spurious bias [Sagawa et al., 2020; Kirichenko et al., 2023] rely on group labels. Group labels represent spurious correlations with class labels and spurious attributes. For example, (waterbirds, water) [Sagawa et al., 2020] represents a spurious correlation between waterbirds and water backgrounds in the images of waterbirds with water backgrounds. Using group labels specifies explicitly the spurious correlations that a model should avoid. However, obtaining group labels requires expert knowledge and labor-intensive annotation efforts. Moreover, group labels fail to capture subtle spurious biases, such as using certain pixels in images for predictions.

In this paper, we propose a post hoc approach that can automatically detect and mitigate potential spurious biases in a model rather than relying on group labels. Our key innovation is reframing the task of identifying and mitigating spurious biases as detecting and leveraging prediction shortcuts in the model's latent space. Prediction shortcuts are latent features derived from input embeddings and predominantly contribute to producing the same prediction outcome across different classes. In essence, prediction shortcuts represent non-defining features of certain classes that the model heavily uses for predictions. By operating in the model's latent space, our approach leverages the expressiveness of latent embeddings, enabling direct identification of spurious biases across diverse input formats without requiring group labels.

We present our post hoc approach as a novel framework called ShortcutProbe, which first identifies prediction shortcuts in a given model and leverages them to guide model retraining for spurious bias mitigation. ShortcutProbe utilizes a probe set without group labels, typically containing a diverse mix of features, to uncover potential prediction shortcuts. These shortcuts are identified as latent features extracted from sample embeddings belonging to different classes but producing the same prediction outcome. By optimizing these features to maximize the model's confidence in their corresponding predictions, ShortcutProbe effectively encodes spurious attributes in non-generalizable prediction shortcuts that the model overly relies on for predictions.

With the identified prediction shortcuts, ShortcutProbe mitigates spurious biases by retraining the model to be invariant to these shortcuts, as they are irrelevant to true prediction targets. This invariance is achieved by applying regularization during retraining, which ensures that the identified prediction shortcuts no longer contribute to the model's predictions of the true targets.

We theoretically demonstrate that when the spurious attributes in the training data are new to the model as reflected by the high prediction loss, the tendency of using spurious attributes for predictions is high after training on the data. Our method aims to revert the process of learning spurious attributes by retraining the model so that the learned spurious attributes induce high prediction losses, effectively unlearning the spurious attributes and reducing the influence of spurious correlations.

In summary, our contributions are as follows:

• We introduce ShortcutProbe, a novel post hoc framework for mitigating spurious bias without requiring group labels. ShortcutProbe identifies prediction shortcuts and leverages them as a form of regularization for training robust models.
• We provide a theoretical analysis revealing that our spurious bias mitigation approach effectively unlearns spurious attributes.
• Through extensive experiments, we show that our method successfully trains models robust to spurious biases without prior knowledge about these biases.

2 RelatedWorks

Shortcuts and spurious bias. Shortcuts are decision rules that perform well on standard benchmarks but fail to transfer to more challenging testing conditions [Geirhos et al., 2020], such as data with distributional shifts. Spurious correlations describe superficial associations between spurious (non-essential) features and targets in data, and they can be used by machine learning models as shortcuts [Geirhos et al., 2019]. This shortcut-learning phenomenon in machine learning models results in spurious bias, the tendency to use spurious correlations in data for predictions. Shortcut features are spurious attributes that are used by models in predictions. Accurately identifying shortcut features involves domain knowledge and model interpretation techniques. In our work, we circumvent this by identifying prediction shortcuts that represent potential shortcut features in a model's latent space, making our method a general and useful method for spurious bias detection and mitigation.

Spurious bias mitigation with group labels. Group labels in training data indicate the presence of spurious correlations across different subsets of the data. Exploiting group labels during training helps mitigate the reliance on the specified spurious correlations. Existing approaches for spurious bias mitigation utilize group labels to balance data distributions during training [Cui et al., 2019; He and Garcia, 2009; Byrd and Lipton, 2019], to formulate a distributionally robust optimization objective [Sagawa et al., 2019], or to progressively expand group-balanced training data [Deng et al., 2023]. Although these methods achieve remarkable success in spurious bias mitigation, the reliance on training group labels becomes a barrier in practice, as obtaining such labels often requires domain knowledge and labor-intensive annotation efforts. Our method does not require group labels during training; instead, it automatically detects prediction shortcuts that encode potential spurious biases the model has developed in training.

Reducing reliance on group labels. To alleviate the dependency on training group labels, existing approaches propose inferring training group labels through various means, such as identifying misclassified samples [Liu et al., 2021], clustering hidden representations [Zhang et al., 2022], employing invariant learning techniques [Creager et al., 2021], or training group label estimators using a few group-labeled samples [Nam et al., 2022]. Nevertheless, validation group labels are still needed to specify which biases to address. Recent works [Zheng et al., 2024b; Zheng et al., 2024c] use visionlanguage models to infer group labels. Last-layer retraining methods [Kirichenko et al., 2023; LaBonte et al., 2024] leverage group-balanced validation data to fine-tune the last layer of a model. While our approach also involves retraining the last layer, it does not require group labels. Instead, we introduce a novel strategy that utilizes identified prediction shortcuts during retraining, leading to improved model robustness.

3 Preliminary

A spurious correlation is the correlation between a spurious attribute present in the training samples and a prediction target. For example, the class waterbird and the attribute water background might form a spurious correlation in the images of waterbird, where some of them have water backgrounds, e.g., pond or river, and some do not. Spurious attributes are not truly predictive of the targets. A group label, e.g., (waterbird, water background), consists of a prediction target and a spurious attribute.

Consider a model fθ : X →R|Y| with parameter θ trained on a training dataset Dtr = {(xi, yi)}Ni=1 with N samplelabel pairs, where xi ∈ X denotes a sample in the input set X, yi ∈ Y denotes a label in the label set Y, and | · | denotes the size of a set. The model fθ = eθ1 ◦ hθ2 can be considered as a feature extractor eθ1 : X → RD followed by a classifier hθ2 : RD →R|Y| , where θ = θ1 ∪ θ2, hθ2 is the last layer of the model, and D denotes the number of dimensions.

Due to the existence of spurious attributes in Dtr, the model can exploit them for predictions, such as recognizing waterbirds by detecting the existence of water backgrounds [Sagawa et al., 2019]. This presents a challenge: It is hard to determine whether a high-performing model is truly robust or simply "right for the wrong reasons", i.e., relying on spurious attributes. Although models with spurious biases typically exhibit degraded performance when the learned spurious attributes are absent from input data, e.g., a waterbird on

(a) Prediction shortcut detection

(b) Spurious bias mitigation

Figure 1: Illustration of ShortcutProbe. (a) The framework uses a set of probe data Dprob to identify prediction shortcuts by learning a shortcut detector to extract similar features from samples of different classes i and j that are all predicted as the same class j. Feature extractor eθ1 and classifier hθ2 are frozen during this stage. (b) ShortcutProbe then retrains the classifier with the probe data (the loss of the probe data Lori) while using the identified prediction shortcuts as regularization (the loss of the prediction shortcuts Lspu).

a land background, it remains challenging to identify the specific spurious attributes encoded by the model without group labels, which hinders the development of effective spurious bias mitigation strategies.

4 Methodology

4.1 Method Overview

We introduce ShortcutProbe, a post hoc framework that automatically detects and mitigates prediction shortcuts without requiring group labels to specify which spurious biases to address. The framework comprises two key steps: (1) Prediction shortcut detection, where a probe set is used to identify prediction shortcuts, and (2) Spurious bias mitigation, where the identified shortcuts are used in retraining the model to mitigate spurious biases in the model.

We provide an overview of our framework in Fig. 1. The process begins with a set of probe data Dprob, which typically contains samples with various spurious attributes that reflect a model's non-robustness to spurious correlations. These samples are mapped to the latent embedding space of the model fθ through its feature extractor eθ1 , allowing us to model any prediction shortcuts the model might use for predictions. This strategy bypasses the need to explicitly define spurious correlations through group labels, a task that is especially challenging for subtle features, such as specific pixels in images. More concretely, as shown in Fig. 1(a), we first train a shortcut detector that extracts potential prediction shortcuts from samples of different classes but having the same prediction. The identified prediction shortcuts encode spurious attributes shared across classes and capture the model's non-robustness across different data groups. Next, in Fig. 1(b), the model is retrained with Dprob to mitigate spurious biases by unlearning the identified prediction shortcuts.

4.2 Prediction Shortcut Detection

Given a model fθ and a probe set Dprob, we aim to detect prediction shortcuts by learning a shortcut detector gψ : RD → RD to identify prediction shortcuts from input embeddings. Intuitively, prediction shortcuts can be identified from samples of different classes but having the same prediction outcome, an indication that similar features exist in these samples but are irrelevant to classes. In other words, prediction shortcuts can be shared among samples from different classes, necessitating a shared representational space to encode diverse prediction shortcuts. We formalize this intuition in the following definition.

Definition 1 (Prediction shortcuts). Given input sample-label pairs (x, y) and (x′ , y′ ), where y ̸= y′ , a trained model fθ = eθ1 ◦ hθ2 , sample embeddings v = eθ1 (x) and v′ = eθ1 (x′ ) for x and x′ , respectively, a vector space V ⊂ RD spanned by K base column vectors in A ∈ RD×K , prediction shortcuts sx ∈ V and sx′ ∈ V for the two samples satisfy the following conditions:

is the projection matrix such that the prediction shortcut sx is the best estimate of v in the vector space V in the sense that the distance ∥sx - v∥22 is minimized.

In Definition 1, we define a prediction shortcut as a projection of a sample embedding. It exists in the vector space V shared by samples of different classes with K degrees of freedom. Here, K governs the complexity of the vector space representing prediction shortcuts. A smaller K results in a less expressive vector space that may fail to adequately capture prediction shortcuts. Conversely, a larger K provides greater flexibility in representing prediction shortcuts but may encode irrelevant information. The optimal value of K depends on the complexity of the probe data; values that are too small or too large can impede learning and lead to suboptimal performance. We treat K as a tunable hyperparameter.

By representing prediction shortcuts as vectors, we can in theory capture any spurious bias, even the intricate ones. For instance, a vector s might correspond to features of water backgrounds that are used to predict waterbirds in any image with water backgrounds, revealing a spurious bias in the model. Alternatively, s could represent a specific feature corresponding to certain pixels in input images, capturing the prediction shortcut based on low-level pixel values-an aspect that is challenging to articulate through group labels.

Learning the shortcut detector. Based on the definition of the prediction shortcut, we design the shortcut detector gψ as a function that implements the projection operation defined in Eq. (1) with the learnable parameter ψ = A ∈ RD×K , that is for a sample embedding v ∈ RD , gψ(v) = PAv. Learning gψ essentially learns a shared vector space spanned by A that could cover prediction shortcuts in samples in the probe set.

To effectively learn gψ, for each class y, we first collect samples from the probe set Dprob to formulate two nonoverlapping sets Dy cor and Dy pre, i.e.,

and

where Dy cor and Dy pre contain samples that are correctly and incorrectly predicted as y. The two sets together demonstrate a possibility that certain features shared across classes are incorrectly associated with prediction targets, allowing us to extract these features as potential prediction shortcuts.

Next, we propose the following objective to identify prediction shortcuts:

where ℓ : R|Y| × Y → R is the loss function, ψ = A, and v = eθ1 (x). To ensure that prediction shortcuts are relevant to the given samples, we regularize the semantic similarity between gψ(v) and v as follows,

The overall learning objective for ψ is

where η > 0 represents the regularization strength. The objective in Eq. (6) aims to encode the properties of prediction shortcuts in Definition 1 into the shortcut detector gψ while maintaining the relevance of the prediction shortcuts to input samples. Training gψ is lightweight as there are only DK learnable parameters. With gψ, we can identify multiple prediction shortcuts from samples in Dprob.

4.3 Spurious Bias Mitigation

Mitigating spurious biases in a model requires that the spurious attributes captured during training are no longer predictive of the targets. Although identifying spurious attributes can be challenging, the shortcut detector introduced in the previous section identifies prediction shortcuts as potential spurious attributes utilized by the model, providing valuable guidance for addressing spurious biases.

In the following, we formulate an optimization objective that incorporates the above constraint to learn a robust model using the probe set Dprob. First, a general requirement is that the trained model should produce correct and consistent predictions on training samples. To this end, for each class y, we sample from Dy cor and Dymis , where Dy cor is the set of correctly predicted sample-label pairs defined in Eq. (2), and Dymis is defined as follows,

representing misclassified sample-label pairs from the class y. Then, the training objective Lori is as follows,

which mitigates potential spurious biases by ensuring consistent predictions in Dymis and Dy cor .

Moreover, to ensure that the training targets at mitigating the spurious biases in the model, we further formulate a regularization term using the prediction shortcuts identified by our shortcut detector to guide the training process. Specifically, as the identified prediction shortcuts are not predictive of the targets, we aim to maximize the loss on the prediction shortcuts defined as follows,

where v = eθ1 (x).

We combine the two loss terms above to formulate the overall training objective as follows,

where λ > 0 is the regularization strength. Here, we retrain only the final classification layer of the model while keeping the feature extractor frozen. This approach significantly reduces computational complexity and allows us to reuse the previously learned sample embeddings. Details of the training algorithm are provided in Appendix.

4.4 Choice of the Probe Set

The probe set plays a crucial role in both detecting prediction shortcuts and mitigating spurious biases. To achieve the full potential of ShortcutProbe, we use a held-out dataset-unseen by the model-to construct the probe set Dprob. This choice is essential because the model may have memorized the training samples, making it difficult to identify prediction shortcuts based on discrepancies in predictions for samples of the same class. Moreover, we select samples with high predication confidence (measured by output entropy) from the held-out dataset to construct Dprob so that prediction shortcuts can be easily detected from these samples. Details of constructing Dprob and results on different choices of Dprob are provided in Section 5.

4.5 Theoretical Analysis

We theoretically demonstrate that minimizing the proposed objective Lori/Lreg effectively unlearns the spurious correlations between spurious attributes and their associated targets captured in the model. Without loss of generality, we analyze this in the context of a general linear regression setting. Consider an input sample x ∈ X, a prediction target y ∈ Y, and a spurious-only sample ˜x that lacks any defining features in x related to y. Let x1, . . . , xN denote N training samples, φ : X → RD be a generic feature map, and Jw(x) = φ(x)T w represent a generalized linear regression model with parameters w ∈ RD .

We denote the correlation between the model output for a spurious-only sample ˜x and a prediction target y as ρ(Jw(˜x), y). The following lemma [Bombari and Mondelli, 2024] gives an upper bound on the correlation.

Lemma 1. The correlation between the model output for the spurious-only sample ˜x and the prediction target y is upper bounded as follows:

where RX is the generalization error, σY is the standard deviation of prediction targets, and

where O = I - VT (VVT )-1V is the orthogonal projection matrix, and V = [φ(x1), . . . , φ(xN)]T ∈ RN×D is the feature matrix.

Given that RX and σY are independent of ˜x, the spuriousonly sample ˜x affects the correlation upper bound via the feature alignment term γφ between the spurious attribute of ˜x and the original feature of x. We further interpret this term in the following proposition.

Proposition 1. The feature alignment term γφ is the ratio between the expected prediction error on the spurious sample ˜x and the expected prediction error on the original sample x:

The term ∥Oφ(x)∥2 in Eq. (13) denotes the error term for the sample x, while ∥Oφ(˜x)∥2 denotes the error term for the spurious sample ˜x. Since x and ˜x are independent, the feature alignment term γφ is the ratio between the loss on spuriousonly samples and the loss on the original samples. We provide a detailed proof in Appendix.

The prediction shortcuts obtained by our shortcut detector approximate the spurious-only features. Thus, the loss Lspu approximates the nominator of the regularization term in Eq. (13). Moreover, Lori approximates the denominator in Eq. (13). Note that Lemma 1 gives the upper bound of the spurious correlation before learning it. The objective in Eq. (4) trains the shortcut detector to learn the correlation by minimizing γϕ. In the mitigation step, the objective in Eq. (10) unlearns the correlation by maximizing γϕ.

5 Experiments

5.1 Datasets

Image datasets. Waterbirds [Sagawa et al., 2019] contains waterbird and landbird classes selected from the CUB dataset [Welinder et al., 2010]. The two bird classes are mixed with water and land backgrounds from the Places dataset [Zhou et al., 2017]. CelebA [Liu et al., 2015] is a large-scale image dataset of celebrity faces. Images showing two hair colors, non-blond and blond, are spuriously correlated with gender. CheXpert [Irvin et al., 2019] is a chest X-ray dataset containing six spurious attributes from the combination of race (White, Black, Other) and gender (Male, Female). Two diagnose results, i.e., "No Finding" (positive) and "Finding" (negative) are the labels. ImageNet-9 [Ilyas et al., 2019] is a subset of ImageNet [Deng et al., 2009] and contains nine super-classes. It is known to have correlations between object classes and image textures. We prepared the training and validation data as in [Kim et al., 2022] and [Bahng et al., 2020]. ImageNet-A [Hendrycks et al., 2021] is a dataset of real-world images, adversarially curated to test the limits of classifiers such as ResNet-50. While these images are from standard ImageNet classes [Deng et al., 2009], they are often misclassified in multiple models. We used this dataset to test the robustness of a classifier after training it on ImageNet-9. NICO [He et al., 2021] is designed for out-of-distribution image classification, simulating real-world scenarios where testing distributions differ from training ones. It labels images with both main concepts (e.g., cat) and contexts (e.g., at home). We used the Animal superclass in NICO and followed the setting in [Bai et al., 2021; Tiwari and Shenoy, 2023a] for data preparation.

Text datasets. MultiNLI [Williams et al., 2017] is a text classification dataset with 3 classes: neutral, contradiction, and entailment, representing the natural language inference relationship between a premise and a hypothesis. The spurious feature is the presence of negation, which is highly correlated with the contradiction label. CivilComments [Borkan et al., 2019] is a binary text classification dataset aimed at predicting whether a comment contains toxic language. Spurious features involve references to eight demographic identities: male, female, LGBTQ, Christian, Muslim, other religions, Black, and White.

5.2 Experimental Setup

Constructing the probe set. From the chosen data source, such as the training or validation set, we sorted the samples within each class by their prediction losses and divided them into two equal halves: a high-loss set and a low-loss set. This approach approximates the incorrectly and correctly predicted samples, respectively, while ensuring that the incorrectly predicted set is non-empty, even when all samples are correctly classified. Within each set, we then selected the top r% of samples with the highest prediction confidence (i.e., those with the lowest output entropy).

Training details. We first trained a base model initialized with pretrained weights using empirical risk minimization (ERM) on the training dataset. Then, we retrained the model on half of the validation set using various bias mitigation methods. For our method, we first constructed the probe set using the same half of the validation set and used the probe set for shortcut detection and mitigation. The remaining half of the validation set was used for model selection and hyperparameter tuning. For experiments on the Waterbirds, CelebA, and CheXpert datasets, we used ResNet-50 as the backbone network, and we used ResNet-18 on the ImageNet9/A and NICO datasets to ensure a fair comparison with baseline methods. For text datasets, we used a pretrained BERT model [Kenton and Toutanova, 2019]. We trained models using each method three times with different random seeds and reported the average results as well as their standard deviations. Detailed training settings are provided in Appendix 1 .

Metrics. Without group labels, we used the worst-class accuracy [Yang et al., 2023] for model selection, which is defined as the worst per-class accuracy on an evaluation set. For performance evaluation on the Waterbirds, CelebA, CheXpert, MultiNLI, and Civilcomments datasets, we adopted the widely accepted metric, worst-group accuracy, which is the lowest accuracy across multiple groups of the test set with each group containing a certain spurious correlation. We also calculated the accuracy gap defined as the standard average accuracy minus the worst-group accuracy to measure the degree of a classifier's spurious biases. A high worst-group accuracy with a low accuracy gap indicates that the classifier is robust to spurious biases and can fairly predict samples from different groups. We adopted average accuracy for the evaluations on the NICO, ImageNet-9, and ImageNet-A datasets as these datasets are specifically constructed to evaluate the robustness against distributional shifts.

5.3 Analysis of Probe Set

Our method relies on a probe set for detecting prediction shortcuts and mitigating spurious bias. A good probe set can be used to effectively reveal and mitigate spurious biases in a model, such as those curated with group labels [Sagawa et al., 2019]. Here, we show that our method, ShortcutProbe, performs effectively without group labels when a probe set is carefully selected from readily available sources, such as the training data and held-out validation data.

To demonstrate, we constructed a probe set as described in Section 5.2, using the training set or half of the validation set (with the other half reserved for model selection) as the data source. For each data source, we varied the proportion r% from 20% to 100%. This adjustment created different probe sets, ranging from those containing only samples with highly confident predictions (small r) to those including all samples from the selected data source (r = 100).

1 Code is available at https://github.com/gtzheng/ShortcutProbe.

Fig. 2(a) shows the performance of ShortcutProbe measured by worst-group accuracy (WGA) under different probe sets, while Fig. 2(b) presents the sizes of these probe sets. Compared to ERM models, we observe that using the training data for retraining results in minimal improvement on the Waterbirds dataset. This is because most of the training data can be correctly predicted by the model, resulting in a probe set that is not informative for learning prediction shortcuts.

In contrast, by leveraging a relatively small amount of the held-out data compared to the size of the training data, our method demonstrates a significant improvement in robustness to spurious biases. Additionally, our approach benefits most from samples with high prediction confidence, i.e., when r is small. However, as shown in Fig. 2(a), setting r too small results in an insufficient number of training samples, leading to suboptimal WGA performance. In the following, unless otherwise specified, we use half of the validation set to construct the probe set and treat r as a tunable hyperparameter.

5.4 Main Results

We focus on a challenging and practical setting where group labels are unavailable in both training and validation data. This scenario requires detecting and mitigating spurious biases using only the available data and model in a standard ERM training setup. As baselines, we selected state-ofthe-art last-layer retraining methods DFR [Kirichenko et al., 2023] and AFR [Qiu et al., 2023], as well as JTT [Liu et al., 2021], which retrains the entire model. For DFR, which typically requires group labels, we used class labels instead. All baseline methods listed in Tables 1 and 2 utilized half of the validation data for training. Similarly, we used the same half of the validation data to construct a probe set, derived as a subset of this portion. The remaining half of the validation data was reserved for model selection.

As shown in Tables 1 and 2, our method achieves the highest WGA and the smallest accuracy gap between average accuracy and WGA, indicating its effectiveness in balancing performance across different data groups. Methods that achieve high average accuracy, such as DFR on the CelebA dataset, prioritize learning spurious features in the probe set. Although maintaining good predictivity on average, DFR still suffers from prediction shortcuts, as shown by its lowWGA. Our method remains effective on larger backbone networks beyond ResNet-50, such as ResNet-152 and ViT (see Appendix). Unlike other baseline methods, our method uses only a portion of available data for training, highlighting the effectiveness of the probe set in detecting and mitigating spurious biases. Notably, when we applied the same probe set with baseline methods, this led to degraded performance in both WGA and accuracy gap, underscoring the unique advantages of our approach.

We further tested the out-of-distribution generalization of our method on the NICO dataset. Images of each class in the test set are associated with an unseen context. Our method achieves the best classification accuracy without using group labels (Table 3), demonstrating its effectiveness in mitigating the reliance on contexts. We present additional results on the ImageNet-9 and ImageNet-A datasets in Appendix to demonstrate our method's effectiveness in combating distributional shifts and the capability of achieving good tradeoff between in-distribution and out-of-distribution performance.

5.5 Ablation Studies

Prediction shortcuts. We evaluated the effectiveness of prediction shortcuts in Fig. 3(a). We began by using a randomly initialized shortcut detector to optimize the objective in Eq.(10). Additionally, we tested the performance without the spurious bias regularization term Lspu, represented as λ = 0 in Fig. 3(a), as well as for λ values of 1, 10, and 50. Our results indicate that the model performs the worst when prediction shortcuts are not used as regularization. Interestingly, our method still demonstrates effectiveness even when the prediction shortcuts are random. As λ increases, the regularization strength decreases. The model achieves optimal performance when an appropriately balanced λ is selected.

Semantic regularization. We analyzed the impact of the semantic regularization strength η defined in Eq. (5). Fig. 3(b) shows that incorporating this regularization with a moderate value of η enhances the model's robustness.

Number of base vectors. The number of base vectors, K, determines the representational capacity of the shortcut detector. A value of K that is too small will limit the detector's ability to identify spurious attributes effectively, while an excessively large K may lead to overfitting on the probe data. Notably, as shown in Fig. 3(c), the optimal value of K is 6, which coincides with the true number of spurious attributes in the CheXpert dataset.

6 Conclusion

In this work, we proposed a novel post hoc framework to mitigate spurious biases without requiring group labels. Our approach first learns a shortcut detector in the latent space of a given model via a diverse probe set. To mitigate spurious biases, we retrained the model to be invariant to detected prediction shortcuts using a novel regularized training objective. We theoretically demonstrated that this objective effectively unlearns the spurious correlations captured during training. Experiments confirmed that our method successfully mitigates spurious biases and enhances model robustness to distribution shifts. Future work may explore constructing a more diverse probe set to further enhance the detection and mitigation of spurious biases.

Acknowledgments

This work is supported in part by the US National Science Foundation under grants 2217071, 2213700, 2106913, 2008208.