PairCFR: Enhancing Model Training on Paired Counterfactually Augmented Data through Contrastive Learning

Recent studies have empirically shown that while perturbed features in CAD are robust and causal Kaushik et al. (2020), they may inhibit the model’s ability to learn other robust features that remain unperturbed Joshi and He (2022). In this section, we mathematically demonstrate that the standard cross-entropy loss, which is commonly used for training models on CAD, can exacerbate this tendency.

Given an instance $\mathbf{x}\in\mathbb{R}^{m\times 1}$, we train a single-layer non-linear function $f_{W}(x)=\sigma(W^{T}\mathbf{x})$, where $W\in\mathbb{R}^{m\times 1}$ and $\sigma$ is the sigmoid function, to predict the label $y\in\{0,1\}$. We expand $\mathbf{x}$, whose label $y=1$, as $\mathbf{x}=[x_{r},x_{c}]^{T}$, where $x_{r}$ denotes the features to be revised (perturbed) and $x_{c}$ denotes the constant (unperturbed) features. The counterfactual example of $\mathbf{x}$ can be written as $\mathbf{c}=[c_{r},x_{c}]^{T}$, with label $y=0$. As the sigmoid function is monotone and bounded, the $c_{r}$ and $x_{r}$ should have different signed values to ensure that $\mathbf{x}$ and $\mathbf{c}$ are classified differently. We expand the weights $W=[w_{r},w_{c}]^{T}$ and take it into the function $f_{W}$ to obtain $f_{W}(x)=\sigma(w_{r}x_{r}+w_{c}x_{c})$ and $f_{W}(c)=\sigma(w_{r}c_{r}+w_{c}x_{c})$. The CE loss on the data $\mathbf{x}$ and its counterfactual $\mathbf{c}$ is calculated as

The above observations indicate that the CE loss alone can lead the model to focus on learning the revised features in CAD, which necessitates incorporating a regularization that compels the model to consider a broader range of features.

Recent research findings have empirically shown that models trained under contrastive loss mainly focus on capturing global relationships Park et al. (2023) compared with negative log-likelihood losses such as masked language modeling. Inspired by this, we propose to employ CL to complement standard CE loss for training models on CAD. In the following, we start from the introduction of CL loss and then mathematically show how CL encourages the model to select a broader range of features beyond the edited ones in the counterfactual data.

Given an anchor sample $\mathbf{x}_{i}$ from a data batch $\mathcal{D}=\{\mathbf{x}_{i},y_{i}\}_{i=1}^{N}$, $\forall\mathbf{x}_{i}\in\mathcal{D}$, we have its positive samples in $\mathcal{P}_{i}\!\equiv\!\{\mathbf{x}_{p}|y_{p}=y_{i},p\neq i\}$ and negative samples in $\mathcal{N}_{i}\!\equiv\!\{\mathbf{x}_{n}|y_{n}\neq y_{i},n\neq i\}$, where $\mathcal{N}_{i}$ contains the counterfactual samples $\mathbf{c}$ for every $\mathbf{x}_{i}$. The contrastive loss for the anchor $\mathbf{x}_{i}$ is

Without loss of generality, we assume $\mathbf{W}\in\mathbb{R}^{m\times d}$ that directly maps the input instance into a $d$-dimensional embedding space, $\mathbf{z}_{i}=\mathbf{W}^{T}\mathbf{x}_{i}$. To obtain the gradient of the CL loss coming from negative samples, we have

The full derivation process can be found in the appendix B. Here, we have

which indicates the probability of $\mathbf{x}_{i}$ being recognized as $\mathbf{x}_{n}$. $\mathbf{A}_{in}\!=\!\mathbf{x}_{i}\mathbf{x}_{n}^{T}+\mathbf{x}_{n}\mathbf{x}_{i}^{T}\in\mathbb{R}^{m\times m}$ is a symmetric matrix derived from the outer product of $\mathbf{x}_{i}$ and $\mathbf{x}_{n}$. Each element of $\mathbf{A}_{i,n}$ indicates the digit-level dot product between the features of $\mathbf{x}_{i}$ and $\mathbf{x}_{n}$, which provides a full view of the entire feature space when comparing a pair of samples. A higher value leads to a larger gradient update and the weights $\mathbf{W}$ are optimized by considering the whole feature sets.

The above analysis implies that the CL loss has the capability of capturing global features beyond those being edited. When learning on CAD under CL, we pair each instance $\mathbf{x}$ with its CFE, $\mathbf{c}$, to compel the model to disparate $\mathbf{x}$ from all negative samples, including its counterfactual example $\mathbf{c}$:

where the non-bold $c$ is the index of CFE. Let us revisit the toy example with $\mathbf{x}=[x_{r},x_{c}]^{T}$ and $\mathbf{c}=[c_{r},x_{c}]^{T}$. Although minimizing the similarity between $\mathbf{x}$ and $\mathbf{c}$ encourages the model to focus on features $x_{r}$, the other negative samples in the batch, e.g., $\mathbf{x}^{\prime}=[x_{r}^{\prime},x_{c}^{\prime}]^{T}$, will enforce the model to use both $w_{r}$ and $w_{c}$ to compare the difference. Hence, the existence of real negative samples could help the model capture the relationships between $x_{r}$ and its context $x_{c}$.

As all $s_{in}$ equally contribute to updating the model weights, the number of non-CFE negatives moderates the learning from local CAD and global patterns. A smaller batch size will manifest the influence of edited features, whereas a larger batch size may dilute the local differences in CAD, as discussed in the experiments 5.4.

Next, we introduce our proposed learning framework, Pairwisely Counterfactual Learning with Contrastive Loss Regularization, named PairCFR for short. As shown in Figure 1, a model $f$ can be decomposed into two modules, $\phi(\cdot)$ and $\varphi(\cdot)$, i.e., $f=\varphi\circ\phi$, where $\phi(\cdot)$ encodes the input sentence into a hidden embedding, and $\varphi(\cdot)$ maps $\phi(\mathbf{x})$ for classification. For transformer-based models, we instantiate $\phi(\mathbf{x})$ using the [CLS] representation, denoted as $\mathbf{z}$. We explicitly pair the original sentences $\mathbf{x}$ and their CFEs, $\mathbf{c}$, in the same batch to provide additional training signals indicative of the underlying causal relationships.

The standard cross-entropy loss is computed on the logits vector projected from $\varphi(z)$. Optimizing CE loss enforces $\varphi(\cdot)$ to identify a small set of features from $\mathbf{z}$ and assign them higher weights to quickly reach a local minimum while optimizing CL loss compels $\phi(\cdot)$ to consider the entire feature space of $\mathbf{z}$ to meet the distance constraints. Overall, we combine the two losses as follows.

where $\lambda$ is the trade-off factor to balance the two losses. To compute CL on a batch, we sample positive pairs that have the same label while all the negative samples including the CFE of the anchor sample are considered.

We evaluate our learning framework on two NLP tasks, sentiment analysis (SA) and natural language inference (NLI). We use two human-edited CAD datasets Kaushik et al. (2020), which ensures good-quality counterfactual data Sen et al. (2023), to train all the models. The IMDb augmented dataset contains 4880 data samples with an original to CFE ratio of 1:1. The SNLI dataset contains 11330 data samples with an original to CFE ratio of 1:4. The statistics of human-revised CAD are reported in Appendix C.1.

To eliminate the random effect, we train each model for multiple runs ($10$ runs for SA and $7$ runs for NLI) using different random seeds. We report the average test accuracy, standard deviation for both in-domain (ID) datasets and several out-of-domain (OOD) datasets. We also conduct significance tests by calculating p-value, to ensure that the observed improvements are not due to randomness. The details of ID and OOD datasets used for evaluation are described in Appendix C.2.

We finetune the BERT base Devlin et al. (2019), RoBERTa base Liu et al. (2019), Sentences-BERT (SBERT, multi-qa-distilbert-cos) Reimers and Gurevych (2019) and T5 base Raffel et al. (2020) models with the original or CAD datasets on HuggingFace platform Wolf et al. (2020). Volumes of model parameters are listed in Table 7 in Appendix C.3. Following the common practices of transformers Devlin et al. (2019), we take the embedding of the “[CLS]” token as sentence representation and finetune the whole model. We set the maximum token length to 350 for SA and 64 for NLI.

We follow the original dataset splits described in Kaushik et al. (2020), where the train, validation, and test sets are divided in a ratio of 7:1:2, with all classes balanced across each set. Subsequently, we finetune all models up to 20 epochs with the AdamW optimizer, coupled with a linear learning rate scheduler with a warmup ratio as 0.05. The best learning rate is manually tuned from $\{1e^{-4},1e^{-5},3e^{-5},5e^{-5},5e^{-6},1e^{-6}\}$. We apply the early stopping strategy with a patience of $5$ and the best model is selected according to the lowest validation loss. To determine the trade-off factor $\lambda$ and temperature $\tau$, we conducted a grid search in the range $[0,1]$ with a step size of 0.1. We also conducted experiments to evaluate our PairCFR in few shot setting where the learning rate and batch size were tuned accordingly. The hyperparameters for full data finetuning and few shot setting are shown in Table 9, Table 9 respectively, in Appendix C.3.

We compare our method PairCFR with the following baselines. For a fair comparison, we employ other forms of augmentation or increase the sampling number for the first three baselines without counterfactual augmentation, to ensure all approaches have the same number of training data.

Vanilla Devlin et al. (2019). This method refers to a general model fine-tuning with original sentences. We include this baseline to verify the improvement of our method result from both the introduction of CAD and the novel learning framework.

BTSCL Gunel et al. (2021). This approach employs the supervised contrastive loss Khosla et al. (2020) into the model training where augmented positive samples are obtained through back-translating a given sentence Ng et al. (2019).

CouCL Wang et al. (2022a). As counterexamples (CEs) are rare in a mini-batch, CouCL samples counterexamples from the original training set, where an example with lower confidence corresponds to a higher likelihood of being selected. Subsequently, it adopts the self-supervised contrastive loss to push representations of positive CEs and negative CEs far apart.

The following approaches study how to train a robust model with annotated CAD:

HCAD Kaushik et al. (2020). It collects two human-edited CAD datasets and fine-tunes a pretrained model on CAD with the cross-entropy loss.

CFGSL Teney et al. (2020). As domain priors in CAD may be lost due to random shuffling in preprocessing Kaushik et al. (2020), CFGSL pairs original data and its counterfactual example in the same batch and introduces a gradient supervision loss (GSL) alongside the cross-entropy loss. The GSL enforces the model gradient to align with the straight line from the original point to CFE.

ECF Fan et al. (2024). It introduces two additional losses to mine the causal structures of CAD. The first loss extracts the dataset-level invariance through Invariant Risk Minimization (IRM) while the second loss is applied to pairs of original sentences and CFEs, preventing the model from relying on correlated features.

Table 1 reports the overall performance comparisons, showing that our proposed PairCFR method outperforms all the baseline models on three out of four OOD datasets for both SA and NLI tasks across four different backbone models. To exclude the possibility of marginal improvements due to random initializations, we also conducted significance tests under the null hypothesis that there are no differences between each baseline and our approach, as presented in Table 11, located in Appendix C.6. The p-values less than $0.05$ demonstrate that our methods are significantly better than the baselines, even though some improvements are relatively slight in Table 1.

In addition, we reported the following findings. Firstly, CAD-based methods may perform worse than non-CAD methods on OOD tasks, e.g., HCAD always lags behind CouCL on the SA task using fine-tuned T5 model. A similar phenomenon is also reported in Joshi and He (2022). These could be due to the failure to extract complementary features between CFEs and the original data; Secondly, the introduction of CFEs may shift the training data distribution from the in-domain data distribution. As anticipated, CAD-based methods fall behind non-CAD methods on ID datasets. Thirdly, our proposed PairCFR exhibits superior OOD performance compared to the baselines, achieving the highest accuracy on mostly OOD datasets, with the sole exceptions being the Yelp and Negation datasets. We postulate that the noted exceptions may be attributed to Yelp and Negative datasets having distributions similar to the ID datasets. The above results validate that PairCFR possesses a heightened capability to learn prior knowledge in CAD.

Data augmentation, such as counterfactual augmentation, is frequently utilized to enhance the performance of few-shot learning. In this part, we investigate the effectiveness of our proposed PairCFR in few-shot learning scenarios. We conducted experiments using the finetuned BERT_base model on the SNLI dataset, gradually increasing the number of training samples from 50 to 4,000. Similarly, on the IMDB dataset, we increased the number of training samples from 32 to 1,024.

Experiment results on SNLI and IMDB under the few-shot setting are reported in Figure 2 and Figure 5 ( Appendix C.5). From both tables, we can conclude that our PairCFR generally demonstrates higher accuracy and lower standard deviation across OOD datasets, particularly in scenarios where training sample sizes are small. For instance, PairCFR significantly outperforms other methods by around 6% on Spelling-e when trained with only 100 counterfactually augmented samples.

We conducted ablation experiments to verify the efficacy of two crucial strategies of our proposed method: (1) the pairing strategy: the integration of original data with their CFEs within the same batch, denoted PairCAD, versus ShuffCAD where randomly shuffle CFEs and originals. (2) the CL loss: the incorporation of CL and CE loss versus CE loss alone.

Results in Table 2, together with significance tests in Table 11 in Appendix C.6, offer several insights: 1) The strategy of pairing original data with their CFEs in the same batch improves OOD performance for both SA and NLI tasks. This can be attributed to the preservation of prior causal relations, which might be lost during random shuffling; 2) The efficacy of PairCAD with a CE-alone learning framework is not guaranteed. For example, within the T5 model framework, PairCAD underperforms ShuffCAD on the SNLI, MNLI, and Spelling-e datasets when only CE loss is adopted. This underscores the critical role of the CL component in augmenting features when we batch CFEs and original data; 3) Integrating the CL consistently improves model performance in both ID and OODs. Particularly, combining CL with PairCAD yields the best performance across various model assessments, highlighting the effectiveness of contrastive learning and the pairing strategy in leveraging causal relations of CFEs.

In this study, we investigated the effect of batch size on learning performance. We conducted experiments on the fine-tuned BERT model for SA and the fine-tuned T5 model for NLI, incrementally increasing the batch size while maintaining the original augmentation ratio for each task. From Figure LABEL:fig:batchsize, we observe that the model performance on both tasks initially improves with increasing batch size, but eventually reaches a plateau or experiences a slight decline.

We contend that the inclusion of negative samples in the CL function provides additional regularization, forcing the model to rely on a broader array of features beyond those edited. However, an excessively large batch size introduces an overwhelming number of negative samples in CL, which may dilute the human priors in CAD, leading to diminished performance. This trend is consistent across both SA and NLI tasks, highlighting the effort required in batch size selection.

Do all counterfactual examples equivalently contribute to enhancing model generality? To answer this, we specifically experimented with the fine-tuned BERT model on the NLI task, comparing performance with and without the inclusion of neutral class samples in CL.

Results in Table 4 reveal that removing neutral samples, including neutral CFEs, significantly enhances the OOD generalization by approximately 2% when training the model on CAD with our learning framework. We attribute this performance difference to the distinct nature of neutral samples. In NLI tasks, judgments of entailment and contradiction are often readily determined based on the semantic alignment or disparity between text elements. Conversely, neutral samples represent scenarios where the hypothesis and premise lack any clear relationship, encompassing a vast array of potential expressions. This diversity poses a great challenge for models to identify universal patterns within the neutral class through human annotations. Therefore, adding neutral samples into the CL detrimentally affects the model’s performance in our experiments.

This investigation highlights the necessity of contemplating the practical value of adding additional counterfactual examples for specific classes.

We also investigated the role of CFE diversity in improving model performance on the NLI task. In SNLI, each sentence is annotated with $4$ CFEs, due to the existence of two opposite targets and modifications made to both the hypothesis and premise. Each CFE is obtained through a different type of modification, resulting in a dataset that includes more diverse counterfactuals. We fine-tuned the T5_base model by incrementally including more CFEs in a batch, ranging from $1$ to $4$.

The results in Table 4, reveal a direct relation between the number of CFEs and the model’s generalization capabilities. Notably, the OOD performance of the model trained on CAD is even better than that trained on a $3$ times larger dataset with only original data. We conclude that enhancing counterfactual diversity proves to be an efficient strategy, which is the same as the findings reported in Joshi and He (2022).