Models Know Their Shortcuts: Deployment-Time Shortcut Mitigation

Abstract: Pretrained LLMs often rely on superficial features that appear predictive during training yet fail to generalize at test time, a phenomenon known as shortcut learning. Existing mitigation methods generally operate at training time and require heavy supervision such as access to the original training data or prior knowledge of shortcut type. We propose Shortcut Guardrail, a deployment-time framework that mitigates token-level shortcuts without access to the original training data or shortcut annotations. Our key insight is that gradient-based attribution on a biased model highlights shortcut tokens. Building on this finding, we train a lightweight LoRA-based debiasing module with a Masked Contrastive Learning (MaskCL) objective that encourages consistent representations with or without individual tokens. Across sentiment classification, toxicity detection, and natural language inference under both naturally occurring and controlled shortcuts, Shortcut Guardrail improves overall accuracy and worst-group accuracy over the unmitigated model under distribution shifts while preserving in-distribution performance.

• The paper presents Shortcut Guardrail, a novel deployment-time mitigation framework that uses gradient-based saliency to detect and debias token-level shortcuts.
• It combines lightweight LoRA adaptation with masked contrastive learning to recalibrate predictions and improve worst-group accuracy under distribution shifts.
• Empirical results show the method effectively reduces reliance on spurious correlations while maintaining in-distribution performance.

Deployment-Time Shortcut Mitigation with Shortcut Guardrail

Problem Setting and Motivation

Pretrained Transformer-based NLP models such as BERT are prone to shortcut learning, wherein they exploit spurious token-label correlations that are statistically predictive within the training distribution but are non-causal and often fail under distribution shift. Token-level shortcuts, such as specific words strongly associated with target labels, pose a particularly acute risk: they are easily memorized and highly brittle to shifts in the input distribution. Existing mitigation approaches require training-time access (e.g., group annotations, retraining, or explicit knowledge of shortcut features) and are thus infeasible in deployment scenarios where only the fixed model and test-time data are available.

Figure 1: Illustration of shortcut learning in sentiment classification—models incorrectly rely on tokens like "service", which are positively correlated with sentiment in the training data but are unreliable in deployment context.

The challenge in text is distinct from vision: token semantics are context-dependent, and text encoders do not provide architectural hooks (e.g., batch normalization) amenable to test-time adaptation. There is thus a key open question: can shortcut reliance be identified and mitigated without privileged training-time supervision or direct access to group structure?

The Shortcut Guardrail Framework

This work introduces Shortcut Guardrail, a deployment-time debiasing pipeline that adapts a frozen classifier to shift reliance away from shortcut tokens, using only the fixed model parameters and an unlabeled test set. The central insight is that gradient-based saliency—as measured by gradients of the prediction with respect to the input embeddings—acts as a reliable indicator of shortcut tokens, even without ground truth labels or training data.

The framework operates in four phases:

1. Extraction via Attribution: The model predicts on deployment data, and for each input, the top-k most salient tokens (by gradient × input norm) are extracted as putative shortcuts.
2. Contrastive Reweighting with LoRA: A lightweight LoRA module is added to the encoder, and trained using Masked Contrastive Learning (MaskCL). Anchor-positive pairs are constructed by masking each important token in turn and encouraging the representation to be invariant to its presence/absence.
3. Inference-time Control: The debiasing strength $\alpha$ is adaptively calibrated on a small labeled support set to trade off between original model behavior and debiased predictions.
4. Final Prediction: At deployment, the model prediction interpolates between the original parameters and the LoRA-adapted parameters according to the selected $\alpha$.

   Figure 2: Overview of Shortcut Guardrail—biased predictions trigger saliency scoring for shortcut localization, MaskCL aligns anchor and masked representations via LoRA, and debiasing strength is adaptively calibrated.

This approach is inherently instance-level and model-centric: shortcut tokens are identified per input based on how much they drive the original prediction, not by relying on corpus-level statistics or shortcut annotation.

Empirical Analysis and Robustness

Shortcut Reliance in NLP Models

An empirical study on controlled sentiment classification setups (with synthetic shortcut injection) demonstrates that pretrained models quickly exhibit sharp, group-wise generalization collapse when only a small fraction of the training set includes a strong spurious correlation. Specifically, negative samples containing the shortcut token are systematically misclassified, with performance degrading monotonically as shortcut prevalence increases.

Figure 3: Test group-wise accuracy degrades as the training proportion of shortcut token "book" increases—shortcut-laden negative samples suffer the most.

Figure 4: Strengthening shortcut correlation ($P(y=1\,|\,\text{"book"})$) exacerbates performance collapse for shortcut-containing examples.

Attribution as a Shortcut Detector

The Shortcut Token Recall—the probability that a true shortcut token is among the k most salient tokens—tracks the strength of the spurious correlation, routinely exceeding 90% when the shortcut association is strong. Moreover, almost all misclassifications on shortcut-vulnerable groups are attributable to samples where the shortcut token is ranked among the most salient.

Figure 5: Shortcut Token Recall grows with increasing spurious correlation; gradient attribution reliably surfaces shortcut tokens.

Figure 6: The misclassification rate is primarily concentrated among examples where the shortcut token is top-ranked in saliency.

Performance Across Distribution Shifts

The efficacy of Shortcut Guardrail is evaluated on real-world and synthetically controlled datasets (SST-2, CivilComments, MultiNLI, Yelp, GoEmotions), under two key conditions:

• Naturally Occurring Shortcuts: On original benchmarks, Shortcut Guardrail matches or slightly outperforms leading training-time methods in accuracy and worst-group accuracy (WGA), even with no access to training supervision or shortcuts.
• Shortcut Distribution Shift: When shortcut-label correlations are inverted between train and test (strongest possible distribution shift), Shortcut Guardrail effects substantial WGA improvements relative to standard ERM—sometimes by over 30 points (see GoEmo-ST and Yelp-Syn)—and remains competitive or superior to baselines like JTT, NFL, and DFR.

To target the core issue—token-level reliance—Maximum Single-Token Prediction Sensitivity (MSTPS) is proposed. It measures the maximum prediction probability drop upon masking any top-k salient token. Shortcut Guardrail consistently achieves lower MSTPS (i.e., reduces shortcut dependence) yet preserves task accuracy better than aggressive regularization baselines like NFL.

Figure 7: Scatterplot analysis shows that Shortcut Guardrail achieves the best trade-off between reduction in shortcut reliance and retention of top-line accuracy.

Increasing shortcut strength further amplifies Guardrail's benefit. In the in-distribution regime (where shortcut-label correlations remain valid), adaptive $\alpha$ calibration commonly selects $\alpha=0$, so in-distribution accuracy is preserved.

Theoretical Bound on Deployment-Time Mitigation

The paper formalizes the information-theoretic ceiling for deployment-time shortcut mitigation: since deployment-time algorithms only have access to the model parameters (an information bottleneck with respect to the training data), they are fundamentally upper-bounded by training-time approaches. However, empirical findings demonstrate that instance-level, attribution-driven methods can approach this limit in practice for strong, localizable shortcuts, owing to the reliable shortcut signature present in model gradients.

Limitations and Future Work

• Upper Bound: Deployment-time mitigation cannot surpass (and sometimes cannot meet) the performance of training-time mitigation due to information loss.
• Multi-token/Complex Shortcuts: The current saliency is token-level; complex, multi-token shortcuts are less amenable to intervention.
• Labeled Calibration Set: While data requirements are drastically reduced, a small labeled support set is used for $\alpha$ selection.
• Task Generality: The method is evaluated on text classification tasks with BERT-style models; its extension to other architectures or tasks remains unstudied.

Conclusion

Shortcut Guardrail delivers a deployment-time, instance-level approach to shortcut mitigation in text classifiers without training data or shortcut supervision. It leverages model-internal attribution to pinpoint reliance, employs targeted contrastive adaptation via a lightweight LoRA, and adaptively calibrates debiasing to handle variable deployment distributions. The framework is competitive with state-of-the-art training-time methods and, crucially, preserves in-distribution performance. This work establishes the practical feasibility—and theoretical boundary—of deployment-time shortcut removal in NLP, and motivates further research into more expressive shortcut detectors and broader task generalization.