CREST: A causal framework for mitigating shortcut learning in language models through counterfactual reasoning

Abstract

Language models excel at numerous natural language processing tasks but often exploit surface patterns rather than develop genuine causal understanding. This limitation leads to vulnerability when encountering out-of-distribution examples or adversarial inputs. We present CREST, a framework incorporating causal learning principles into language models to mitigate shortcut learning through counterfactual reasoning. CREST integrates a causally-informed pre-trained question-answering model with a debiasing framework utilizing counterfactual analysis. The framework explicitly models the causal structure of question-answering tasks and employs controlled interventions to differentiate authentic reasoning pathways from shortcuts. Our multi-branch architecture separates robust causal reasoning from potential shortcut pathways, while the counterfactual reasoning component regulates variable interactions during training and inference. Experiments across six datasets (DREAM, SQuAD, TriviaQA, HotpotQA, TyDi-QA, QuAC) demonstrate CREST achieves 2.41–3.15% improvements over eight baselines, with the strongest gains on multi-hop reasoning. Validation on large language models GPT-J (6B) and Llama-2 (7B) using knowledge editing scenarios shows CREST achieving 44.13% and 45.83% success rates, substantially outperforming Fine-Tuning (19.97–22.97%), ROME (13.67–17.66%), MEMIT (14.05–20.13%), and MeLLo (24.97–29.62%). Critically, CREST demonstrates superior hop-wise accuracy of 30.60% on GPT-J and 60.44% on Llama-2, indicating genuine step-by-step reasoning compared to MeLLo’s 0.21–9.90%. CREST exhibits superior adversarial robustness, maintaining 78.4–85.3% performance under attacks compared to 71.2–77.5% for the best baseline. While requiring 30% additional training time, CREST maintains competitive inference speeds with only 3% latency increase, demonstrating practical feasibility for real-world deployment.