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

2026-03-21

AI Summary

This work addresses the tendency of large language models to rely on shortcut strategies—such as surface-level patterns or memorized answers—during reasoning tasks, which undermines their generalization capabilities. To mitigate this issue, the authors propose the Shortcut-Aware Reasoning Training (SART) framework, which uniquely integrates gradient alignment with analysis of answer token distributions to identify shortcut-prone training samples. SART further introduces a gradient surgery mechanism that dynamically adjusts optimization directions to suppress shortcut learning. Guided by a novel ShortcutScore metric and data-driven training modulation, the method substantially enhances logical reasoning performance. Evaluated on controlled reasoning benchmarks, SART achieves a 16.5% absolute accuracy gain over the strongest baseline and improves robustness by 40.2%, demonstrating significantly better generalization under distributional shifts.

Technology Category

• Natural Language Processing: Safety and Robustness
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

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.

Problem

• shortcut reasoning
• language models
• logical inference
• generalization
• distribution shift

Innovation

• shortcut reasoning
• gradient-aware training
• gradient surgery
• reasoning robustness
• language model generalization