Models Know Their Shortcuts: Deployment-Time Shortcut Mitigation
For practitioners deploying NLP models, this provides a practical method to improve robustness to shortcut learning without requiring access to training data or prior knowledge of shortcuts.
Shortcut Guardrail mitigates token-level shortcuts in pretrained language models at deployment time without training data or shortcut annotations, improving overall and worst-group accuracy under distribution shifts while preserving in-distribution performance.
Pretrained language models 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.