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

TL;DR

This paper introduces SART, a gradient-aware training method that detects and reduces shortcut reasoning in language models, significantly improving their reasoning accuracy and robustness.

Contribution

The paper presents a novel gradient-aware framework, SART, for identifying and mitigating shortcut signals in language model training, enhancing reasoning capabilities.

Findings

Achieves +16.5% accuracy on reasoning benchmarks

Improves robustness by +40.2% under distribution shifts

Effectively reduces reliance on shortcut cues

Abstract

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.