Models Know Their Shortcuts: Deployment-Time Shortcut Mitigation

Jiayi Li; Shijie Tang; Gün Kaynar; Shiyi Du; Carl Kingsford

cs.LGApr 14, 2026

Models Know Their Shortcuts: Deployment-Time Shortcut Mitigation

Authors: Jiayi Li, Shijie Tang, Gün Kaynar, Shiyi Du, Carl Kingsford

Organizations: Ray and Stephanie Lane Computational Biology Department, School of Computer Science, Carnegie Mellon University, Pittsburgh, PA

Abstract

Pretrained text encoders are prone to shortcut learning, relying on token-label correlations that fail once the distribution shifts in deployment. Existing shortcut mitigation methods mainly operate at training time and assume access to training data, training dynamics, or shortcut annotations, which are hardly available during deployment, where only the converged model remains. We show that this model alone suffices to mitigate shortcuts during deployment: a biased model internalizes a signal of its learned shortcuts that can be captured via unsupervised gradient-based attribution. We further prove that deployment-time mitigation is information-theoretically upper-bounded by training-time mitigation. Nevertheless, exploiting this gradient signal, our proposed unsupervised deployment-time shortcut mitigation framework for pretrained text encoders, Shortcut Guardrail, recovers substantial performance under shortcut distribution shift, matching or outperforming training-time baselines across sentiment classification, toxicity detection, and natural language inference.

Figures & tables

Figure 1: Example of shortcut learning in sentiment classification. The classifier associates the token “service” with positive sentiment, leading to misclassification when it appears in a negative review at test time.

Figure 2: Overview of Shortcut Guardrail , which (1) obtains predictions from a frozen biased text classifier, (2) captures shortcut tokens via gradient-based saliency scoring, (3) trains a lightweight LoRA adapter via Masked Contrastive Learning (MaskCL), and (4) calibrates the debiasing strength

\alpha

to produce debiased predictions with reduced shortcut reliance. The bottom panel illustrates the effect of MaskCL training. Before training, spurious correlations cause examples containing the shortcut token (Anchor,

x

) to form a cluster within the wrong label, while their masked variants (Positive,

x^{\text{mask}}

) are projected to the correct label. The MaskCL objective pulls anchors and positives closer in the embedding space, dissolving the spurious cluster and reducing shortcut reliance.

Figure 3: Group-wise test accuracy under different proportions of training data with the shortcut token “book”. The strength of the spurious correlation

P(y=1|\text{``book''}\in x)

1

and label distributions are balanced (

P(y=1)=0.5

). Error bars denote standard deviation over three random trials. (1)-(4) correspond to Group 1-Group 4 of the test data. See Table 5 in Appendix A.1.1 for group-wise training data statistics.

Figure 4: Group-wise test accuracy under different strengths of the spurious correlation, where

p=P(y=1|\text{``book''}\in x)

. The proportion of training samples containing the shortcut “book” is fixed at 10%, and label distributions are balanced (

P(y=1)=0.5

). (1)-(4) correspond to Group 1-Group 4 of the test data. See Table 6 in Appendix A.1.1 for training data statistics.

Figure 5: Shortcut Token Recall under different strengths of the spurious correlation. Each bar shows the percentage of samples whose shortcut (“book”) appears among the top-10 important tokens, averaged over three random trials.

Figure 6: Total misclassification rate under different strengths of the spurious correlation

p

. The misclassification rate (black annotation) can be attributed to two groups of samples: those with the shortcut (“book”) among the important tokens (gold bar, white annotation), and those without (blue bar).

Metric	Model	SST-2	Civil	MultiNLI
Accuracy	ERM	0.911	0.880	0.819
	JTT	0.915	0.878	0.829
	NFL	0.815	0.832	0.665
	DFR	0.915	0.878	0.824
	SG (Ours)	0.917	0.874	0.816
WGA	ERM	–	0.737	0.429

Table 1: Performance on real-world benchmarks with naturally occurring token-label correlations. The best results are marked in bold while the second best are in italics . SST-2 does not have predefined shortcut annotations, hence WGA is not applicable (–).

Metric	Model	Civil	MultiNLI	Yelp-ST	Yelp-Syn	GoEmo-ST	GoEmo-Syn
Accuracy	ERM	0.522	0.270	0.355	0.419	0.513	0.593
	JTT	0.532	0.268	0.440	0.490	0.403	0.613
	NFL	0.536	0.048	0.449	0.471	0.387	0.377
	DFR	0.536	0.273	0.373	0.435	0.557	0.630
	SG (Ours)	0.572	0.323	0.488	0.530	0.627	0.607
WGA	ERM	0.216	0.010	0.030	0.035	0.026	0.222

Table 2: Overall accuracy and worst-group accuracy (WGA) under shortcut distribution shift. The best results are marked in bold while the second best are marked in italics . The CivilComments and MultiNLI datasets in this table are reconstructed with controlled shortcut injection and differ from the original versions in Table 1 .

Dataset	ERM	JTT	NFL	DFR	SG (Ours)
Civil	0.530	0.549 ( $+$ 0.019)	0.451 ( $-$ 0.079)	0.526 ( $-$ 0.004)	0.438 ( $-$ 0.092)
MultiNLI	0.797	0.267 ( $-$ 0.530)	0.143 ( $-$ 0.654)	0.212 ( $-$ 0.585)	0.381 ( $-$ 0.416)
Yelp-ST	0.331	0.393 ( $+$ 0.062)	0.074 ( $-$ 0.257)	0.339 ( $+$ 0.008)	0.203 ( $-$ 0.128)
Yelp-Syn	0.219	0.264 ( $+$ 0.045)	0.083 ( $-$ 0.136)	0.139 ( $-$ 0.080)	0.126 ( $-$ 0.093)
GoEmo-ST	0.647	0.355 ( $-$ 0.292)	0.102 ( $-$ 0.563)	0.713 ( $+$ 0.066)	0.501 ( $-$ 0.146)
GoEmo-Syn	0.488	0.180 ( $-$ 0.308)	0.190 ( $-$ 0.263)	0.527 ( $+$ 0.039)	0.414 ( $-$ 0.074)

Table 3: Shortcut reliance as measured by MSTPS. Parentheses show the change relative to ERM.

Dataset	$\lambda$	Accuracy			WGA
		ERM	SG	$\Delta$	ERM	SG	$\Delta$
Yelp-ST	1.0	0.355	0.488	+0.133	0.030	0.353	+0.323
	0.8	0.558	0.583	+0.025	0.481	0.467	$-$ 0.014
	0.6	0.614	0.616	+0.002	0.506	0.480	$-$ 0.026
Civil	1.0	0.522	0.572	+0.050	0.216	0.396	+0.180
	0.9	0.705	0.717	+0.012	0.487	0.544	+0.057

Table 4: Effect of shortcut strength

\lambda

in controlled benchmarks.

\Delta

= SG

-

ERM.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

	0%	5%	10%	20%
Group 1	50	45	40	30
Group 2	0	5	10	20
Group 3	50	50	50	50
Group 4	0	0	0	0

Appendix

Table 5: The proportion (

/100

) of samples in each group in the training sets for the experiments in Figure 3 .

	$p$ =1	$p$ =0.99	$p$ =0.95	$p$ =0.9	$p$ =0.5
Group 1	40	40.1	40.5	41	45
Group 2	10	9.9	9.5	9	5
Group 3	50	49.9	49.5	49	45
Group 4	0	0.1	0.5	1	5

Appendix

Table 6: The proportion (

/100

) of samples in each group in the training sets for the experiments in Figure 4 .

p=P(y=1|\text{``book''}\in x)

. The statistics for

p=0.995,0.975,0.7

can be derived analogously.

Dataset	$C$	Per-class Shortcut Occurrence Probabilities
Yelp	5	0%, 25%, 50%, 75%, 100%
GoEmotions	4	0%, 33.3%, 66.7%, 100%
CivilComments	2	0%, 100%
MultiNLI	3	100%, 50%, 0%

Appendix

Table 7: Base shortcut insertion probabilities per label (

\lambda=1

). Shifted test sets reverse the order of these probabilities (e.g., 0%, 100%

\to

100%, 0%).

	Real-World			Controlled
	SST-2	Civil	MultiNLI	Civil	MultiNLI	Yelp-ST	Yelp-Syn	GoEmo-ST	GoEmo-Syn
Train	67,349	115,931	206,175	111,298	44,508	8,900	9,400	1,124	1,124
Test	872	1,124	1,000	1,000	600	1,000	1,000	685	685

Appendix

Table 8: Dataset sizes. Statistics for “Controlled” apply to both shortcut distribution shift and in-distribution datasets.

Figure 7: Scatterplots comparing accuracy with MSTPS across all benchmarks and models reported in Table 3 .

Metric	Model	Civil	MultiNLI	Yelp-ST1	Yelp-SYN1	GoEmo-ST1	GoEmo-SYN1
Accuracy	ERM	0.980	0.908	0.671	0.679	0.896	0.928
Accuracy	SG	0.980	0.908	0.660	0.679	0.896	0.928
WGA	ERM	0.964	0.840	0.376	0.000	0.667	0.774
WGA	SG	0.964	0.840	0.474	0.000	0.667	0.774

Appendix

Table 9: In-distribution (ID) performance. Our method preserves the accuracy of the original ERM model.

Figure 8: Sensitivity to support set size. Test accuracy and worst-group accuracy (WGA) of Shortcut Guardrail as a function of the number of support examples used to calibrate the adaptive scalar

\alpha

, on SST-2, Civil Comments, and MultiNLI. Dotted lines mark the

\alpha\!=\!1

baseline, which applies the debiasing module with no calibration. The datasets are the same as in Table 1 .

Hyperparameter	Value	Rationale
learning_rate	$1\text{e-}4$	Standard for LoRA fine-tuning
gradient_accumulation_steps	1	Maximizes optimization steps; noisier gradients
		act as implicit regularization
num_train_epochs	1	Sufficient for ${\sim}$ 700–1000 samples
temperature	0.1	Standard InfoNCE temperature
lora_r	4 or 8	Scaled by dataset size (see below)

Appendix

Table 10: Training hyperparameters for the debiasing module.

Method	Deployment-time	No shortcut annotation	Per-input attribution	Approach
Group DRO ( Sagawa et al., 2020 )	✗	✗	✗	Worst-group reweighting
JTT ( Liu et al., 2021 )	✗	✓	✗	Error-based retraining
LTGR ( Du et al., 2021 )	✗	✓	✓	Attribution-guided training
DFR ( Kirichenko et al., 2023 )	✗	✓	✗	Last-layer retraining
NFL ( Chew et al., 2024 )	✗	✓	✗	Weight regularization
C2L ( Choi et al., 2022 )	✗	✓	✓	Contrastive learning

Appendix

Table 11: Positioning of Shortcut Guardrail against existing shortcut-mitigation and test-time adaptation approaches. Deployment-time indicates operation on a converged model using only test data, which entails no access to the training set. Per-input attribution indicates that shortcut tokens are localized per example via gradient-based saliency rather than inferred from dataset-level statistics. Shortcut Guardrail is the only method that operates at deployment time while using per-input attribution. Tent also operates at deployment time, but targets general covariate shift via entropy minimization rather than spurious correlations and does not localize shortcuts; it is therefore not directly comparable on token-level shortcut benchmarks.

Shortcut Learning

Shortcut learning occurs when models exploit easy but unreliable cues instead of the task-relevant evidence needed to generalize.

See all 159 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
Counterfactual Test-Time Augmentation for Shortcut Removal in Language Models
Oct 24, 2026Katie Matton, Rosalind W. Picard, John Guttag +1Shortcut Learning in Language Models Test-Time Augmentation
MIT · University of Michigan
DeShortcut-Align: Decoupling Spurious Shortcuts for Robust Safety Alignment in Large Reasoning Models
Sep 28, 2026Qirui Liu, Yichen Sun, Yan Wang +4Shortcut Learning in Language Models LLM Safety Alignment
The State Key Laboratory of Blockchain and Data Security, Zhejiang University · Ant Group · Chongqing Ant Consumer Finance Co., Ltd

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

Explore similar work

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

Counterfactual Test-Time Augmentation for Shortcut Removal in Language Models

DeShortcut-Align: Decoupling Spurious Shortcuts for Robust Safety Alignment in Large Reasoning Models