Counterfactual Test-Time Augmentation for Shortcut Removal in Language Models

Katie Matton; Rosalind W. Picard; John Guttag; Maggie Makar

cs.LGOct 24, 2026

Counterfactual Test-Time Augmentation for Shortcut Removal in Language Models

Authors: Katie Matton, Rosalind W. Picard, John Guttag, Maggie Makar

Organizations: MIT · University of Michigan

Abstract

Large language models (LLMs) can exploit shortcuts, i.e., correlations present in training data that may not hold at test time. This can lead to poor predictions when correlations shift. Most existing solutions for removing shortcuts operate at training time, requiring labeled training data, model weight access, and significant computational resources. Here, we study the problem of test-time shortcut removal. We propose a family of counterfactual test-time augmentation (CTTA) methods: given an example, each method uses an auxiliary model to generate counterfactual edits in which the shortcut value is changed, and then takes a weighted average of the LLM's predictions across these edits. We prove that the optimal shortcut-invariant predictor is a member of this CTTA family, and we characterize its weights. The weights depend on the training distribution, and therefore are, in general, not identifiable from the trained model and unlabeled test data. However, we present assumptions under which the weights can be obtained using the token probabilities of the trained LLM, and propose tests to detect violations of these assumptions. We use these tests to develop Select-CTTA, an aggregate method that selects among CTTA methods. Empirically, we find that Select-CTTA outperforms test-time baselines and approaches the performance of models fine-tuned on shortcut-balanced data.

Figures & tables

Panel of Figure 1 — Figure 1: Causal graph of the data-generating process. The label $Y$ and shortcut $S$ may be correlated (dashed bidirected edge). $Y$ affects the text $X$ only through $X_{Y}$ . Shading indicates observed at test-time.

	All-tokens loss				Completion-only loss
	Math		Occupation		Math		Occupation
Method	AUROC $\uparrow$	ECE $\downarrow$	AUROC $\uparrow$	ECE $\downarrow$	AUROC $\uparrow$	ECE $\downarrow$	AUROC $\uparrow$	ECE $\downarrow$
Standard (No Aug)	0.062	0.894	0.271	0.572	0.077	0.893	0.378	0.506
BC	0.062	0.895	0.271	0.581	0.077	0.893	0.378	0.506
BC (shortcut-cond.)	0.436	0.201	0.550	0.262	0.430	0.201	0.613	0.325
CC	0.062	0.860	0.271	0.577	0.077	0.820	0.378	0.485

Table 1: Results for models fine-tuned with shortcut strength

p=0.95

and evaluated on the reversed test split. Values are averaged over models. Bold: best in the column. Underlined: not significantly different from the best (95% paired bootstrap interval of the difference includes zero). CTTA methods substantially outperform test-time baselines in most cases, with Select-CTTA consistently among the top-performing methods.

	Math		Occupation
Method	AUROC $\uparrow$	ECE $\downarrow$	AUROC $\uparrow$	ECE $\downarrow$
Standard (No Aug)	0.113	0.704	0.762	0.235
BC	0.113	0.779	0.762	0.225
BC (shortcut-cond.)	0.489	0.323	0.564	0.315
CC	0.113	0.712	0.762	0.248
CC (shortcut-cond.)	0.187	0.596	0.806	0.293

Table 2: Results for models without fine-tuning, evaluated on the reversed test split. Values are averaged over models. Bold: best in the column. Underlined: not significantly different from the best (95% paired bootstrap interval of the difference includes zero). Compared with non-CTTA baselines, Select-CTTA achieves higher AUROC and lower ECE on Math, and provides a favorable trade-off between AUROC and ECE on Occupation.

Figure 2: Select-CTTA decisions on the reversed test split (

p=0.95

). Each point represents one model, fine-tuning setting, and task combination. Positive values indicate that Model-Weighted CTTA outperforms Uniform CTTA. Select-CTTA selects the better-performing CTTA method in most settings.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Test	Fall back to Uniform CTTA if	Threshold
1. Shortcut strength	$\Delta_{s^{+},s^{-}}<\delta_{\mathrm{str}}$	$0.5$
2. Weight collapse	$\operatorname{median}_{i}r_{i}\notin[1-\delta_{\mathrm{col}},\,\delta_{\mathrm{col}}]$ or $\max_{s}\operatorname{median}_{i}w_{i}(s)>\delta_{\mathrm{col}}$	$0.8$
3. Alignment	$\operatorname{Spearman}(r_{i},c_{i})<\delta_{\mathrm{align}}$	$0.5$

Appendix

Table 3: The Select-CTTA tests. Model-Weighted CTTA is used only if every test passes.

	Math: $Y=$ incorrect			Math: $Y=$ correct
Split	correct	incorrect	unsure	correct	incorrect	unsure
Matched	33	1,235	32	618	65	617
Reversed	618	65	617	33	1,235	32
Balanced	325	650	325	325	650	325
	Occupation: $Y=$ attorney			Occupation: $Y=$ paralegal
Split	man	woman		man	woman

Appendix

Table 4: Test split cell counts at

p=0.95

. Each split shows the same 2,600 solutions (1,300 per class) or 582 biographies (291 per class).

All-tokens loss
	Math (Sycophancy)						Occupation (Gender)
	AUROC $\uparrow$			ECE $\downarrow$			AUROC $\uparrow$			ECE $\downarrow$
Method	Rev.	Bal.	Match.	Rev.	Bal.	Match.	Rev.	Bal.	Match.	Rev.	Bal.	Match.
Standard (No Aug)	0.062	0.563	0.961	0.894	0.452	0.011	0.271	0.747	0.988	0.572	0.292	0.038
Batch calibration	0.062	0.563	0.961	0.895	0.452	0.011	0.271	0.747	0.988	0.581	0.300	0.026
Batch calibration (shortcut-cond.)	0.436	0.645	0.550	0.201	0.060	0.087	0.550	0.890	0.652	0.262	0.077	0.126

Appendix

Table 5: Fine-tuned models (

p=0.95

) on the reversed (Rev.), balanced (Bal.), and matched (Match.) test splits. Mean over the six models; one block per training loss. Bold: best in the column. Underlined: not significantly different from the best (95% paired bootstrap interval of the difference includes zero). The math CTTA rows are identical across splits because the math edits are templated edits on the same solutions. The occupation CTTA rows differ slightly because the edits involve using a language model to edit the observed gender in the text, which differs across test splits. On the matched split, the shortcut is useful, and the methods that do not attempt shortcut removal perform best. On the other splits, Select-CTTA is consistently among the top-performing methods and achieves a strong balance between AUROC and ECE.

	Math (Sycophancy)						Occupation (Gender)
	AUROC $\uparrow$			ECE $\downarrow$			AUROC $\uparrow$			ECE $\downarrow$
Method	Rev.	Bal.	Match.	Rev.	Bal.	Match.	Rev.	Bal.	Match.	Rev.	Bal.	Match.
Standard (No Aug)	0.113	0.549	0.948	0.704	0.390	0.166	0.762	0.789	0.814	0.235	0.219	0.208
BC	0.113	0.549	0.948	0.779	0.393	0.087	0.762	0.789	0.814	0.225	0.204	0.186
BC (shortcut-cond.)	0.489	0.576	0.543	0.323	0.214	0.237	0.564	0.791	0.568	0.315	0.206	0.310
CC	0.113	0.549	0.948	0.712	0.390	0.156	0.762	0.789	0.814	0.248	0.230	0.215

Appendix

Table 6: Zero-shot models on the three test splits. OOC was run on the reversed and balanced splits only. Bold: best in the column. Underlined: not significantly different from the best (95% paired bootstrap interval of the difference includes zero). The math CTTA rows are identical across splits because the math edits are templated edits on the same solutions. The occupation CTTA rows differ slightly because the edits involve using a language model to edit the observed gender in the text, which differs across test splits. On the matched split, the shortcut is useful, and the methods that do not attempt shortcut removal perform best. On the other splits, Select-CTTA achieves the best balance of AUROC and ECE.

All-tokens loss
	Math (Sycophancy)						Occupation (Gender)
	AUROC $\uparrow$			ECE $\downarrow$			AUROC $\uparrow$			ECE $\downarrow$
Method	$p=0.50$	$p=0.75$	$p=0.95$	$p=0.50$	$p=0.75$	$p=0.95$	$p=0.50$	$p=0.75$	$p=0.95$	$p=0.50$	$p=0.75$	$p=0.95$
Standard (No Aug)	0.706	0.380	0.062	0.132	0.438	0.894	0.949	0.878	0.271	0.069	0.102	0.572
BC	0.706	0.380	0.062	0.067	0.441	0.895	0.949	0.878	0.271	0.057	0.062	0.581
BC (shortcut-cond.)	0.711	0.636	0.436	0.066	0.064	0.201	0.949	0.857	0.550	0.055	0.074	0.262

Appendix

Table 7: Shortcut strength sweep on the reversed split. The mean over the six models is shown. Bold: best in the column. Underlined: not significantly different from the best (95% paired bootstrap interval of the difference includes zero). At

p=0.50

the shortcut is absent from training. As shortcut strength increases, standard predictors and calibration baselines degrade sharply in most cases, while CTTA largely preserves AUROC and calibration. Select-CTTA consistently provides strong performance across tasks, shortcut strengths, and fine-tuning objectives, often approaching balanced training.

	All-tokens loss				Completion-only loss
	Math (Sycophancy)		Occupation (Gender)		Math (Sycophancy)		Occupation (Gender)
Method	Acc $\uparrow$	CE $\downarrow$	Acc $\uparrow$	CE $\downarrow$	Acc $\uparrow$	CE $\downarrow$	Acc $\uparrow$	CE $\downarrow$
Standard (No Aug)	0.059	2.937	0.292	1.830	0.064	3.014	0.386	2.180
BC	0.058	2.934	0.283	1.792	0.066	3.012	0.382	2.176
BC (shortcut-cond.)	0.496	0.715	0.513	0.917	0.462	0.734	0.557	1.564
CC	0.078	2.990	0.286	1.859	0.076	3.125	0.431	2.412

Appendix

Table 8: Accuracy and cross-entropy (CE) of the fine-tuned models on the reversed split (

p=0.95

), mean over the six models. Bold: best in the column. Underlined: not significantly different from the best (95% paired bootstrap interval of the difference includes zero). In most cases, CTTA substantially improves accuracy and cross-entropy over the unaugmented model and standard calibration baselines. Model-Weighted CTTA is strongest under all-tokens loss but degrades substantially under completion-only loss. Select-CTTA obtains some of the benefits of Model-Weighted CTTA by selecting it in some cases in the all-tokens setting, and it avoids its failures by consistently choosing Uniform CTTA in the completion-only setting.

	Math (Sycophancy)		Occupation (Gender)
Method	Acc $\uparrow$	CE $\downarrow$	Acc $\uparrow$	CE $\downarrow$
Standard (No Aug)	0.278	5.353	0.669	2.073
BC	0.159	5.037	0.701	1.780
BC (shortcut-cond.)	0.468	1.689	0.615	2.863
CC	0.254	5.275	0.672	2.113
CC (shortcut-cond.)	0.334	3.761	0.653	2.195

Appendix

Table 9: Accuracy and cross-entropy (CE) of the zero-shot models on the reversed split. Bold: best in the column. Underlined: not significantly different from the best (95% paired bootstrap interval of the difference includes zero). Select-CTTA consistently chooses Uniform CTTA and obtains top performance across all metrics and tasks.

Shortcut Learning in Language Models

See all 13 papers

Mitigating Shortcut Reasoning in Language Models: A Gradient-Aware Training Approach
Mar 21, 2026Hongyu Cao, Kunpeng Liu, Dongjie Wang +1Shortcut Learning in Language Models LLM Reasoning
Arizona State University · Clemson University · University of Kansas
Caterpillar of Thoughts: The Optimal Test-Time Algorithm for Large Language Models
Mar 24, 2026Amir Azarmehr, Soheil Behnezhad, Alma GhafariInference-Time Optimization Efficient Language Model Reasoning
Northeastern University
Test-Time Training with Next-Token Prediction
Jun 19, 2026Xuan Ouyang, Zefan Cai, Junjie HuLong-Context Language Modeling Next-Token Prediction
University of Wisconsin–Madison

Mar 21, 2026J/K move · Enter open · S save

Mitigating Shortcut Reasoning in Language Models: A Gradient-Aware Training Approach

Hongyu Cao, Kunpeng Liu, Dongjie Wang, Yanjie Fu

Arizona State University · Clemson University · University of Kansas

Shortcut Learning in Language Models LLM Reasoning Language Model Robustness

Open paper AskPDF arXiv

Large language models exhibit strong reasoning capabilities, yet often rely on shortcuts such as surface pattern matching and answer memorization rather than genuine logical inference. We propose Shortcut-Aware Reasoning Training (SART), a gradient-aware framework that detects and mitigates shortcut-promoting samples via ShortcutScore and gradient surgery. Our method identifies shortcut signals through gradient misalignment with validation objectives and answer-token concentration, and modifies training dynamics accordingly. Experiments on controlled reasoning benchmarks show that SART achieves +16.5% accuracy and +40.2% robustness over the strongest baseline, significantly improving generalization under distribution shifts. Code is available at: https://github.com/fuyanjie/short-cut-aware-data-centric-reasoning.

Abstract

Shortcut Learning in Language Models

Explore similar work

Mitigating Shortcut Reasoning in Language Models: A Gradient-Aware Training Approach

Caterpillar of Thoughts: The Optimal Test-Time Algorithm for Large Language Models

Test-Time Training with Next-Token Prediction