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

Mar 21, 2026

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

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.

Hongyu Cao, Kunpeng Liu, Dongjie Wang +1

Arizona State University · Clemson University · University of Kansas

Shortcut Learning in Language Models LLM Reasoning Language Model Robustness

Ask PDF

Mar 24, 2026

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

Large language models (LLMs) can often produce substantially better outputs when allowed to use additional test-time computation, such as sampling, chain of thought, backtracking, or revising partial solutions. Despite the growing empirical success of such techniques, there is limited theoretical understanding of how inference time computation should be structured, or what constitutes an optimal use of a fixed computation budget. We model test-time computation as an algorithm interacting with a Markov chain: at any point, the algorithm may resume generation from any previously observed state. That is, unlike standard Markov chains where the states are drawn passively, we allow the algorithm to backtrack to any previously observed state of the Markov chain at any time. Many of the existing test-time algorithms, such as Chain-of-Thought (CoT) (Wei et al., 2023), Tree-of-Thoughts (ToT) (Yao et al., 2023), or Best-of-

k

(Brown et al., 2024) could be seen as specific algorithms in this model. We prove that while backtracking can reduce the number of generations exponentially, a very limited form of backtracking is theoretically sufficient. Namely, we show that the optimal algorithm always generates a caterpillar tree. That is, if we remove the leaves of the state tree generated by the optimal algorithm, we obtain a path. Motivated by our characterization of the optimal algorithm, we present Caterpillar of Thoughts (CaT), a new test-time computation algorithm, reducing the number of token/state generations. Our empirical evaluation shows that CaT, compared to ToT, achieves a better success rate while also reducing the number of token generations.

Amir Azarmehr, Soheil Behnezhad, Alma Ghafari

Northeastern University

Inference-Time Optimization Efficient Language Model Reasoning Test-Time Scaling LLM Inference Acceleration

Ask PDF

Jun 19, 2026

Test-Time Training with Next-Token Prediction

Next-token prediction is the self-supervised signal that trains language models, and every observed prompt token provides the same signal at test time. We study whether this signal can define the inner-loop objective for test-time training (TTT) in pretrained long-context language models. Many TTT architectures require models to be trained with test-time adaptation in mind, limiting their direct applicability to released LLM checkpoints. While recent in-place TTT methods make fast-weight adaptation possible for pretrained LLMs without redesigning the backbone, they leave a central question unresolved: what should each test-time write store? Existing recipes train the fast weight to match a learned local value proxy but they are not directly tied to the self-supervised next-token prediction signal. We introduce Test-Time Training with Next-Token Prediction (TTT-NTP), a drop-in fast-weight adaptation method for pretrained LLMs that instead supervises updates using the model's own next contextual hidden state. This makes each local write follow the same causal computation that supports next-token prediction: the value target is a pointwise linear projection of a single next-position contextual state. On RULER Full-13, averaged over 4k to 32k contexts, TTT-NTP is the only method that consistently improves the released backbone across four models spanning three families and a 0.6-8B size range, by 3.9 points on Llama-3.1-8B, 3.0 on Mistral-7B-v0.3, 4.1 on Qwen3-4B, and 2.9 on Qwen3-0.6B. On the real-world LongBench-v2 long-document QA benchmark, TTT-NTP improves over the base model by 5.6 points on Llama-3.1-8B and 3.7 on Mistral-7B-v0.3, while preserving commonsense and knowledge performance. Our code is publicly available at https://github.com/yancyou/TTT-NTP.

Xuan Ouyang, Zefan Cai, Junjie Hu

University of Wisconsin–Madison

Long-Context Language Modeling Next-Token Prediction Test-Time Training Sequence Modeling Inference-Time Adaptation of Language Models

Ask PDF

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