Robustness to Spurious Correlations via Dynamic Knowledge Transfer

1 Introduction

While deep learning models have demonstrated remarkable performance across a wide range of tasks, studies have highlighted that deep models frequently capture spurious correlations between non-causal attributes and classes, posing a threat to the validity and reliability of the models [Wu et al., 2023; Deng et al., 2024; Tian et al., 2025]. For instance, in natural language processing tasks, classifiers may learn that the term "Spielberg" is correlated with positive movie reviews, or that the word "ugly" is associated with toxic comments [Gupta et al., 2022]. Similarly, in computer vision tasks, an image classifier might associate a forest background with the label "bird" [Deng et al., 2024]. The embedding of spurious correlations during training can lead to performance degradation when faced with varying distributions in test data, raising concerns regarding prediction robustness and trustworthiness.

* Corresponding authors.

Data augmentation has proven to be effective in helping models discern genuinely significant features by diversifying and balancing the training data [Zhou and Wu, 2023; Deng et al., 2024]. Among these techniques, counterfactual data augmentation, which generates counterfactual samples with modified attributes, emerges as one of the most potent approaches for dismantling spurious correlations between non-causal attributes and classes [Wu et al., 2024; Reddy et al., 2023]. However, determining which features should be modified and how to modify them is a complex issue that often relies on human annotation or group information, making the process both costly and time-consuming [Deng et al., 2023]. Additionally, training models with expanded training data inevitably leads to decreased training efficiency [Wang et al., 2025]. Implicit semantic data augmentation [Wang et al., 2019; Zhou et al., 2024] is a technique that transforms samples along semantic directions within the deep feature space. Notably, this method is accomplished solely by optimizing a novel robust loss, ensuring efficiency. Building upon their research, we believe that generating counterfactual samples within the deep feature space could eliminate the need for explicit differentiation between causal and noncausal attributes, while also bypassing the necessity to expand the training dataset.

Drawing on the aforementioned insights, this study proposes enhancing model robustness by dynamically transferring knowledge across spuriously correlated classes in the deep feature space. Specifically, the deep features of samples are transformed using semantic vectors extracted from both their respective category distributions and those of their spuriously correlated classes. This process generates class-specific factual and counterfactual augmented deep features, aiding in enhancing intra-class data diversity and reducing inter-class spurious correlations. As illustrated in Fig. 1, pigeons with dark feathers are spuriously correlated with the crow class due to the prevalence of dark feathers in the crow class and their rarity in the pigeon class. Conveying knowledge from the crow class to the pigeon class, e.g., generating counterfactual instances of pigeons with dark feathers, helps the model rely more on causal features rather than color information during prediction. By considering an infinite number of augmentations, we theoretically derive a surrogate loss function for the proposed augmentation strategy. Consequently, rather than explicitly augmenting sample features, we can directly minimize this robust loss function, leading to a highly efficient algorithm.

To dynamically adjust the direction and intensity of knowledge transfer across categories during training, we propose a novel reinforcement learning-based training framework that incorporates a target network and a strategy network. The target network is trained using the derived robust loss to enhance its resilience. Meanwhile, as training progresses, the strategy network dynamically adjusts the sample-specific augmentation strategies, encompassing both augmentation distribution and strength, to align with the evolving learning dynamics of the target network. We conduct extensive experiments across three typical learning scenarios susceptible to spurious correlations: subpopulation shift learning, generalized longtail (GLT) learning, and domain shift learning, encompassing benchmarks from both text and image domains. The results demonstrate that our method consistently attains state-of-theart (SOTA) performance across various learning scenarios, affirming its efficacy in mitigating the deleterious effects of spurious correlations.

In summary, our main contributions are threefold:

• We undertake a pioneering effort to reduce spurious correlations through dynamic inter-class knowledge transfer. This approach generates diverse factual and counterfactual augmented deep features, achieved exclusively through the optimization of a robust loss function, thereby enhancing efficiency and applicability.
• We present a new reinforcement learning-based training framework, termed dynamic knowledge transfer (DKT), for training classifiers using the derived robust loss function, where the direction and intensity of knowledge transfer are dynamically and automatically determined by a strategy network based on the unique training characteristics of samples.
• We conduct comprehensive experiments across three

typical learning scenarios characterized by distribution shifts between training and test data. The results unequivocally demonstrate the effectiveness and broad applicability of our approach.

2 RelatedWork

Prior research has demonstrated that deep learning models frequently rely on spurious patterns for predictions, consequently exhibiting inadequate generalization and robustness when confronted with unseen environments [Moayeri et al., 2022; Zhao et al., 2025; Veitch et al., 2021; Sun et al., 2021]. For instance, models trained on ImageNet often classify images based on background rather than foreground attributes [Moayeri et al., 2022]. Similarly, Young et al. [2019] demonstrated that deep networks for CT scans, despite high accuracy, often produce explanations outside the relevant regions when visualized with Grad-CAM [Selvaraju et al., 2017]. Additionally, Chew et al. [2024] found that a sentiment classifier might mistakenly learn that the word "performance" is associated with positive reviews, even if the word itself is not commendatory.

To eliminate spurious correlations, various approaches have been proposed, including dataset modifications [Zhao et al., 2024; Wu et al., 2024], causal inference [Tang et al., 2020], model ensembles [Zhang et al., 2021; Wang et al., 2020], regularization techniques [Tang et al., 2022; Krueger et al., 2021; Zhang et al., 2022], and re-training strategies [Zhao et al., 2023; Zhou et al., 2023b]. Among these methods, our work is particularly aligned with counterfactual data augmentation, which involves generating counterfactual samples by modifying attributes to train deep learning models [Wu et al., 2024; Reddy et al., 2023]. For instance, Gupta et al. [2022] mitigated gender bias in the text by swapping gender words (e.g., "he" becomes "she") to augment the data. Xiao et al. [2023] employed masked images, where either semantics-related or semantics-unrelated patches were masked, as counterfactual samples to enhance the robustness of the fine-tuning model. Moreover, Reddy et al. [2023] enabled the generation of counterfactual data by quantifying and eliminating confounding bias. While effective, existing augmentation methods rely on explicit differentiation between causal and non-causal attributes to produce counterfactual data [Gupta et al., 2022; Xiao et al., 2023], which can be resource-intensive and restrictive in their scope. Additionally, these strategies are hand-crafted, thereby lacking flexibility and potentially constraining the generalization and robustness of models [Deng et al., 2023].

3 Methodology

This study introduces a novel learning framework to diminish spurious correlations by transferring knowledge across categories in the deep semantic space. We start by elucidating the mechanism of the proposed knowledge transfer strategy. Following this, we theoretically derive a robust loss function to implement our strategy efficiently. Finally, we propose a reinforcement learning-based training framework that dynamically determines the direction and intensity of knowledge transfer throughout the training process.

Notation Consider training a deep classifier f, with weights θ, on a training set denoted as Dtr = {(xi, yi)}Ni=1 , where N represents the number of training samples, and yi ∈ {1, · · · , C} denotes the label of sample xi. The deep feature (before logit) learned by f for xi is represented as a Z-dimensional vector zi = fθ(xi) ∈ RZ . Moreover, N(µ,Σ) represents a multi-variant Gaussian distribution with the mean vector µ and the covariance matrix Σ.

3.1 Knowledge Transfer Strategy

To alleviate spurious correlations between non-causal attributes and classes, this study proposes transferring knowledge across categories within the deep semantic space. In particular, the deep features of samples undergo transformations along semantic directions derived from the deep feature distributions of their respective categories and those exhibiting spurious correlations. These transformations generate augmented deep features that encompass both factual and counterfactual variations specific to each class, thereby enriching the diversity of training data and helping to eliminate erroneous associations between samples and their non-groundtruth categories.

In line with Wang et al. [2019], we enhance intra-class data diversity by sampling semantic vectors for the deep feature of each sample, zi, from a zero-mean multivariate normal distributionN(0,Σyi ). Here, Σyi represents the class-conditional covariance matrix of class yi. Additionally, to eliminate spurious correlations across categories, semantic vectors are extracted from the category distributions with spurious associations in the samples. In our approach, the class with the highest predicted probability, excluding the ground-truth class, is treated as the spuriously correlated class. Let the spuriously correlated category for sample xi be denoted as class c. We perform counterfactual augmentation by sampling semantic vectors from the deep feature distribution of class c, represented as N(µc,Σc). The feature mean µc of each class is computed as:

where Nc represents the number of samples in class c. Additionally, the covariance matrix Σc is calculated as:

In implementation, the mean and covariance matrix are computed online by aggregating statistics from all mini-batches. Specifically, as training progresses, the feature means are estimated as follows:

where µtc and µ′t c represent the feature mean estimates of class c at the tth step and within the tth mini-batch, respectively. nt-1 c and mtc denote the cumulative count of training samples belonging to class c across all t-1 mini-batches and the count of training samples belonging to class c in the tth mini-batch. The covariance matrices are estimated by:

where Σtc and Σ′t c represent the covariance matrix estimates of class c at the tth step and within the tth mini-batch, respectively. Furthermore, ntc = nt-1 c +mtc .

Based on the unique training dynamics of the samples, the augmentation distributions for different samples, which control the direction of knowledge transfer, should vary. In other words, the proportions of semantic and counterfactual augmentations should differ across samples. Consequently, the augmented deep features ˜zi for zi are given by

where αi and βi are two positive coefficients controlling the proportions of counterfactual and factual augmentations, respectively. Their values are automatically and dynamically computed by a deep network, utilizing the distinctive training characteristics of the samples as input, which will be detailed in Section 3.3. Notably, our method can sample arbitrary semantic directions from spuriously correlated classes for counterfactual augmentation. While we cannot guarantee the exclusion of inherent attributes from spurious categories, the limited scope of counterfactual perturbations and the retention of original labels are expected to encourage the model to focus more on other intrinsic attributes of the current class for classification. Additionally, since early epoch estimates are less informative when the network is undertrained, both values are decayed by a factor of t/T , where t and T denote the current and total number of epochs, respectively.

As for the augmentation strength, that is the number of augmented features for each sample, it is assumed to follow Ki = K × γi, where K is a constant and γi refers to the strength factor. The value of γi is also dynamically calculated based on the unique training characteristics of sample xi. Consequently, for each deep feature zi, a set of augmented deep features can be obtained in each iteration, represented as {˜z1i , ˜z2i , · · · , ˜zKi i }.

3.2 Surrogate Robust Loss Function

Given the computational inefficiency of directly using all augmented deep features for training, we address this by exploring the scenario where the number of augmented features for each sample approaches infinity and derive an upper bound for the expected Cross-Entropy (CE) loss. This approach allows us to achieve a highly efficient implementation while maintaining the benefits of augmentation.

The CE loss for all augmented features is as follows:

where ˆK = ∑Ni=1 Ki. W = [w1, · · · ,wC]T ∈ RC×Z and b = [b1, · · · , bC]T ∈ RC . Each wk and bk represent the weight vector and bias, respectively, associated with the final fully connected layer for class k. We then investigate the scenario of augmenting an infinite number of times on the deep feature of each sample. As K in Ki approaches infinity, the expected CE loss is expressed as:

However, accurately calculating Eq. (7) is computationally challenging, and thus we aim to derive a more efficient surrogate loss for it. Let ˆµi and ˆΣi denote the mean and covariance matrix for the spuriously correlated class of sample xi. Due to the concave nature of the logarithmic function log(·), with Jensen's inequality, E[log X] ≤ log E[X], we derive an upper bound for the expected loss in Eq. (7):

Given that ˜zi follows a Gaussian distribution characterized by ˜zi ∼ N(zi + αi ˆµi, βiΣyi + αi ˆΣi), we have (wTj - wTyi )˜zi + (bj - byi ) ∼ N ((wTj - wTyi ) (zi + αi ˆµi) + (bj - byi ), (wTj -wTyi )(βiΣyi +αi ˆΣi)(wj -wyi )). Then, leveraging the moment-generating function E [etX ] = etµ+ 1 2 σ2 t2 , the upper bound in Eq. (8) can be written as

Consequently, the above deviation provides a surrogate loss for our proposed knowledge transfer strategy. Instead of explicitly performing the augmentation process, we can directly optimize the derived upper bound (LDKT ) during classifier training, thereby enhancing efficiency.

3.3 Reinforcement Learning-Based Framework

The augmentation distribution and strength for each sample dictate the direction and intensity of knowledge transfer. As determining the optimal augmentation strategies is inherently a parameter selection problem without gold labels, we employ reinforcement learning to address it. The workflow of our proposed DKT framework is illustrated in Fig. 2.

Pipeline of the DKT Framework

Our framework comprises a target network and a strategy network. The target network undergoes training using the derived DKT loss to bolster its robustness. Meanwhile, the strategy network receives a series of training characteristics extracted from the target network as input and generates sample-specific augmentation strategies, encompassing augmentation distributions and strengths. Given that the input to the strategy network is tabular data, we design it as a twolayer Multilayer Perceptron. Accordingly, as training progresses, the strategy network dynamically adjusts augmentation strategies to align with the evolving learning dynamics of the target network.

Let Ai = {αi, βi, γi} ∈ A represent an augmentation strategy for sample xi, where αi and βi jointly determine the augmentation distribution, while γi determines the augmentation strength. Here, A represents the value space of augmentation strategies. Each parameter has E options, which are encoded by a one-hot vector. For example, γi takes values from the set {γ 1 , γ 2 , · · · , γ E }, providing E options (γ 1 to γ E ) such as {0.1, 0.2, · · · , 1} for the augmentation strength of each sample. If Ai = {0.1, 0.2, 0.1}, then the augmentation distribution for zi is N(0.1µc, 0.2Σyi +0.1Σc), and its augmentation strength is 0.1K. Nevertheless, instead of explicitly augmenting deep features, the parameters of augmentation strategies are employed to compute the derived DKT loss LDKT in our framework. Accordingly, during the training process, the strategy network captures the conditional distribution p(Ai|ξi; Ω) where ξi denotes the training characteristics of sample xi, detailed in subsequent subsections, and Ω represents the parameters of the strategy network.

Training Characteristics Extraction

In each iteration, a series of training characteristics are extracted from the classifier for each sample, which serves to capture the degree of spurious correlations between the sample and various classes, along with its learning difficulty [Zhou and Wu, 2023; Lin et al., 2024; Zhou et al., 2023a]. These characteristics are then input into the strategy network to facilitate the generation of sample-specific augmentation strategies. The characteristics that reflect the sample-class correlation include the prediction distribution and the cosine similarity between the sample feature and the classifier weights for each class. Each of these two characteristics has C dimensions. Furthermore, to gauge the learning difficulty of samples comprehensively, we consider metrics including loss, loss gradient, uncertainty, margin, forgetfulness, and class proportion, all of which are scalar values. As a result, a characteristic vector ξi ∈ R2×C+6 is derived for each sample at each iteration.

Optimization Process

We introduce a reinforcement learning-based algorithm for the alternating optimization of parameters in both the target (θ) and strategy (Ω) networks. With Ω fixed, the optimization subproblem for the target network can be defined as

Given the characteristics of a sample, the strategy network generates a strategy distribution p(Ai|ξi;Ω), from which an augmentation strategy is randomly sampled. Utilizing these sampled strategies, we update the parameters of the target model by minimizing LDKT via gradient descent:

where η1 represents the learning rate of the target network and n denotes the mini-batch size.

Additionally, with θ given, the optimization problem for the parameters of the strategy network can be formulated as

Following the REINFORCE algorithm [Williams, 1992], we compute the derivative of the objective function, H(Ω) := E(x,y)∼DtrEA∼p(A|ξ;Ω)LDKT , with respect to Ω as:

Similar to solving Eq. (10), we sample augmentation strategies from the conditional distribution to calculate the sample losses. The gradient with respect to the parameters of the strategy network can be approximately computed as:

Consequently, the parameters of the strategy network can be updated using gradient descent as follows:

where η2 denotes the learning rate for the strategy network.

Convergence Analysis

Based on Eqs. (11) and (15), we have the following convergence result for the effectiveness of the proposed optimization algorithm.

Theorem1. Suppose that the objective function for the strategy network LDKT satisfies the gradient Lipschitz conditions w.r.t. Ω and θ, and LDKT is λ-strongly concave in ˆΩ, the feasible set of Ω. If A′ is a δ-approximation of the optimal augmentation strategy A∗ , the variance of the stochastic gradient is bounded by a constant σ2 > 0, and we set the learning rate of θ as

where LDKT = LθΩLΩθ/λ+Lθθ denotes the Lipschitz constant of LDKT and Γ = LDKT (θ0 ) - minθ LDKT (θ), it holds that

where T represents the maximum training epoch.

According to Theorem 1, if the inner minimization process yields a δ-approximation of the optimal augmentation strategy A∗ , then the training algorithm can converge to a stationary point at a sub-linear rate with an accuracy of 5δL2θΩ/λ. Moreover, if 5δL2θΩ/λ is sufficiently small, our method effectively identifies the desired robust model θT by generating a good augmentation strategy capable of approximating A∗ well.

4 Experimental Investigation

We conduct experiments on three typical learning scenarios susceptible to spurious correlations: subpopulation shift learning, GLT learning, and domain shift learning. The experiments are repeated three times using different random seeds. For all experiments, the parameter sets for αi, βi, and γi range from 0.1 to 1, with intervals of 0.1.

4.1 Subpopulation Shift Learning

Settings. We evaluate the performance of DKT under four subpopulation shift datasets: Colored MNIST (CMNIST) [Yao et al., 2022], Waterbirds [Sagawa et al., 2020], CelebA [Liu et al., 2016], and CivilComments [Borkan et al., 2019]. In these datasets, certain attributes are highly spuriously correlated with the labels. Following Yao et al. [2022], we adopt pre-trained ResNet-50 [He et al., 2016] and DistilBERT [Sanh et al., 2019] as the model for image (i.e., CMNIST, Waterbirds, CelebA) and text data (i.e., CivilComments), respectively. We compare DKT with a variety of robust methods designed to break spurious correlations and learn invariant predictors. To ensure a comprehensive assessment, we report both the average and worst-group accuracy.

Results. The comparative results are reported in Table 1. DKT achieves a higher average accuracy compared to other robust learning methods across diverse datasets, highlighting its strong generalization capability. Moreover, it consistently surpasses other methods in terms of worst-group accuracy, demonstrating its efficacy in enhancing model robustness for underrepresented groups, such as samples belonging to the landbird class with a water background. These findings manifest that models trained using the DKT framework rely more on causal attributes for predictions, rather than non-causal factors like colors and backgrounds, thereby bolstering the model's resilience to spurious correlations. Furthermore, data modification methods, such as PDE, and our approach, generally outperform previous regularization techniques (e.g., IRM and LISA), underscoring the effectiveness of data expansion in mitigating spurious correlations.

4.2 Generalized Long-Tail Learning

Settings. GLT learning accounts for both long-tailed class and attribute distributions within the training data, as both types of imbalances contribute to spurious correlations. We employ two GLT benchmarks [Tang et al., 2022]: ImageNetGLT and MSCOCO-GLT. Each benchmark consists of three protocols: Class-wise Long Tail (CLT), Attribute-wise Long Tail (ALT), and GLT, showcasing variations in class distribution, attribute distribution, and combinations of both between training and testing datasets. We report the mean accuracy and precision for all approaches. Following Tang et al. [2022], the ResNeXt-50 [Xie et al., 2017] model serves as the backbone network. We compare DKT with three categories of approaches: re-balancing techniques, augmentation methods, and ensemble learning approaches.

Results. The comparative results are reported in Table 2, where DKT demonstrates significant performance improvements, with its accuracy and precision surpassing the best comparative results by 1.70% and 1.52%, respectively. These findings underscore its effectiveness in enhancing model generalization and robustness against imbalanced class and attribute distributions. Moreover, the results presented in Fig. 3 further illustrate its efficacy in mitigating spurious associations between non-causal attributes and class labels, thereby guiding the model to focus more on causal attributes. Additionally, methods tailored for long-tailed learning generally exhibit inferior performance on the GLT and ALT protocols, primarily due to their class-level characteristics. Nevertheless, DKT dynamically generates sample-specific augmentation strategies, making it more effective in breaking samplewise spurious correlations.

4.3 Domain Shift Learning

Settings. We examine three domain shift benchmarks featuring out-of-distribution test data. These benchmarks (i.e., Camelyon17 [Bandi et al., 2018], FMoW [Christie et al., 2018], and RxRx1 [Taylor et al., 2019]) are sourced from WILDS [Koh et al., 2021], covering domains including healthcare and vision. Following previous studies [Yao et al., 2022], the evaluation metrics are average accuracy for Camelyon17 and RxRx1, and worst-group accuracy for FMoW. The compared baselines are consistent with those for subpopulation shift learning.

Results. The comparison results for the three datasets are presented in Table 3. DKT consistently outperforms other compared methods across diverse datasets, achieving an average accuracy improvement of 1.25%. Moreover, we have validated the significance of the performance improvement achieved by DKT using the Wilcoxon signed-rank test, which produced a p-value of 0.04, smaller than the threshold of 0.05. These findings demonstrate the efficacy of DKT in enhancing model robustness against spurious correlations.

4.4 Analytical Experiments

We compare the performance of DKT under four augmentation distributions: (I) excluding knowledge transfer among categories, N(0,Σyi ), (II) excluding the augmentation of samples within their own classes, N(µc,Σc), (III) adopting the augmentation distribution with a zero mean, N (0, Σc+Σyi ), and (IV) our adopted augmentation distribution, N(µc,Σc +Σyi ). The accuracy across all protocols is presented in Fig. 4(a). Our augmentation distribution proves to be the most effective, due to its ability to enhance intraclass diversity and reduce inter-class spurious associations. Additionally, the evolution of the average α and β values during training is illustrated in Fig. 4(b). At the start of training, both the average values of α and β are approximately equal to 0.5, as the selected values encompass the entire range of possible options. During the early stages of training, the average value of β surpasses that of α, prioritizing feature representation. In contrast, later in training, the average value of β drops below that of α, enhancing the model's robustness to spurious correlations.

5 Conclusion

This paper proposes enhancing model robustness against spurious correlations by transferring knowledge across categories. Specifically, the deep features of the samples are transformed along semantic vectors derived from their respective category distributions, as well as those of spuriously correlated categories. We then demonstrate that our knowledge transfer strategy can be achieved solely by optimizing a surrogate robust loss and introduce a reinforcement learningbased framework to train classifiers using this loss, where the direction and intensity of knowledge transfer are dynamically determined based on the unique training dynamics of samples. Extensive experiments across various learning scenarios prone to spurious correlations validate the effectiveness of our approach in mitigating spurious correlations and enhancing model robustness.

Acknowledgments

This work was partially supported by the Baidu Scholarship.