Paper deep dive
Mitigating Shortcut Reasoning in Language Models: A Gradient-Aware Training Approach
Hongyu Cao, Kunpeng Liu, Dongjie Wang, Yanjie Fu
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 93%
Last extracted: 3/26/2026, 2:20:49 AM
Summary
The paper introduces Shortcut-Aware Reasoning Training (SART), a gradient-aware framework designed to mitigate shortcut reasoning in Large Language Models (LLMs). SART identifies shortcut-promoting samples by analyzing gradient signaturesâspecifically, low alignment with validation gradients and high answer-token gradient concentrationâand employs a combination of ShortcutScore-based sample reweighting and gradient surgery to suppress shortcut learning while preserving generalizable reasoning.
Entities (4)
Relation Signals (3)
SART â employs â Gradient Surgery
confidence 95% ¡ We propose a method combining ShortcutScore for sample reweighting with gradient surgery techniques.
SART â uses â ShortcutScore
confidence 95% ¡ We propose Shortcut-Aware Reasoning Training (SART)... via ShortcutScore and gradient surgery.
ShortcutScore â detects â Shortcut-promoting samples
confidence 90% ¡ ShortcutScore... identifies shortcut-promoting samples
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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: this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2603.20899v1
- Canonical: https://arxiv.org/abs/2603.20899v1
Trouble viewing inline? Open PDF directly â
Full Text
43,829 characters extracted from source content.
Expand or collapse full text
Mitigating Shortcut Reasoning in Language Models: A Gradient-Aware Training Approach Hongyu Cao Arizona State University hongyuca@asu.edu &Kunpeng Liu Clemson University kunpenl@clemson.edu &Dongjie Wang University of Kansas wangdongjie100@gmail.com &Yanjie Fu Arizona State University yanjiefu@asu.edu Abstract Large language models exhibit strong reasoning capabilities, yet often rely on shortcutsâsurface pattern matching, memorization, and keyword correlationsârather than genuine logical inference. This stems from training signal misalignment: shortcut-encouraging examples reduce loss but impair generalization when distributions shift or problem phrasing varies. We formulate shortcut-aware reasoning training: identifying samples that promote shortcut updates and modifying gradient dynamics to emphasize generalizable reasoning. A key challenge is that shortcut samples efficiently reduce training loss, making them difficult to detect via conventional metrics. Existing approachesâChain-of-Thought training, Self-Consistency decoding, data filtering, and RLHFâfail because they do not modify training dynamics or detect shortcut-promoting gradients. We identify a new opportunity: shortcut reasoning samples produce distinct gradient signatures, exhibiting low alignment with validation-improving gradients and high concentration on answer tokens. These signals enable detection and mitigation through gradient structure analysis rather than data labels alone. We propose a method combining ShortcutScore for sample reweighting with gradient surgery techniques. Our pipeline integrates periodic validation gradient computation with empirical validation across diverse datasets. Experiments demonstrate significant improvements in robustness and reasoning reliability, enhancing LLM generalization toward genuine reasoning. Code available at: https://github.com/fuyanjie/short-cut-aware-data-centric-reasoning 1 Introduction Large language models (LLMs) deployed in high-stakes reasoning tasks frequently exploit reasoning shortcutsâsurface pattern matching, answer memorization, keyword-answer correlations, and premature answer predictionârather than performing genuine logical inference Wei et al. (2022); Wang et al. (2022). This shortcut reliance stems from training signal misalignment: shortcut-encouraging samples efficiently reduce training loss but fail to foster generalizable reasoning. Consequently, models accurate on training distributions fail catastrophically when numbers change, constraints vary, or problem phrasing shifts. This motivates Shortcut-Aware Reasoning Training (SART): identifying samples that promote shortcut updates and modifying gradient dynamics to emphasize genuine reasoning signals. Solving SART is critical for reliable LLM deployment in mathematical reasoning, financial analysis, planning, and scientific discoveryâdomains where spurious reasoning leads to incorrect decisions or unsafe recommendations. Two major challenges arise in solving SART: (1) detecting shortcut-promoting samples from gradient behavior, and (2) modifying training dynamics to suppress shortcut learning without degrading performance. First, shortcut samples are often correctly labeled and indistinguishable by conventional quality metrics. They produce strong gradients that efficiently reduce training loss, making loss-based detection unreliable. The detection challenge is: how can we identify training samples whose gradients improve loss but harm reasoning generalization? Second, simply down-weighting or removing detected shortcut samples is insufficientâit destabilizes training and discards useful signal. The suppression challenge is: how can we modify training dynamics to neutralize shortcut gradients while preserving generalizable reasoning signals from the same samples? Existing methods only partially address SART. Chain-of-Thought (CoT) training Wei et al. (2022); Kojima et al. (2022) encourages intermediate reasoning traces but does not modify training dynamics: models can generate plausible steps that still rely on shortcuts, and shortcut gradients remain dominant Ho et al. (2022). Self-Consistency Decoding Wang et al. (2022) samples multiple reasoning paths to detect inconsistencies but operates at inference time, leaving shortcut patterns intact in model parameters. Data filtering and curriculum learning Bengio et al. (2009) attempt to remove noisy samples, but shortcut samples are typically correctly labeled, making them invisible to label-quality filters. RLHF Ouyang et al. (2022) optimizes reward signals based on answer correctness rather than reasoning validity, allowing shortcuts to persist when they produce correct answers. These approaches cannot jointly address shortcut detection from gradient behavior and dynamic suppression during trainingâa gradient-centric perspective is required. Our insights: gradient-aware shortcut detection and correction. We formulate SART as a gradient structure analysis problem. We find that shortcut reasoning samples produce two distinct gradient signatures: (1) low alignment with gradients that improve held-out validation performance, indicating the sampleâs gradient direction does not transfer to generalizable reasoning; and (2) high answer-token concentration, where gradient norms are dominated by final answer tokens rather than intermediate reasoning tokens. These signalsâgradient alignment with validation gradients and answer-gradient concentrationâcan be combined into a principled ShortcutScore quantifying the degree to which a sample promotes shortcut learning. Furthermore, harmful shortcut gradient directions can be neutralized via gradient surgery: projecting gradients orthogonally onto subspaces that avoid non-transferable directions, forcing the model to learn from genuine reasoning signals without discarding samples entirely. Proposed Approach. This paper presents the first principled shortcut-aware training framework for LLM reasoning via gradient structure analysis, with two goals: (1) precise detection of shortcut-promoting samples through gradient behavior; (2) effective suppression of shortcut learning while preserving generalizable reasoning signals. For Goal 1, we develop the ShortcutScore, a composite metric combining cosine similarity between per-sample gradients and validation gradients (measuring non-transfer alignment) with the ratio of answer-token to reasoning-token gradient norms (measuring answer-dominant credit assignment). Samples with high ShortcutScore are flagged as shortcut-promoting. For Goal 2, we develop a gradient surgery technique that identifies harmful gradient directionsâthose misaligned with validation gradients or increasing ShortcutScoreâand projects them orthogonally to suppress their influence during parameter updates. This ensures training dynamics prioritize reasoning token signals. Our framework is model-agnostic and integrates into any LLM training pipeline without architectural changes. Extensive experiments across mathematical reasoning, commonsense reasoning, and constraint-based planning benchmarks demonstrate significant improvements in robustness under distribution shift, number perturbation, and phrasing variation, validating that our approach reduces shortcut reliance and promotes genuine logical inference. Code is available at https://github.com/fuyanjie/short-cut-aware-data-centric-reasoning. 2 Problem Statement The SART Problem. We study supervised training of language models on reasoning tasks. Let =(xi,yi)i=1ND=\(x_i,y_i)\_i=1^N be a training dataset where xix_i is an input reasoning problem and yiy_i is the corresponding output comprising intermediate reasoning steps and a final answer. Let =(xj,yj)j=1MV=\(x_j,y_j)\_j=1^M be a small held-out validation set drawn from a distribution that rewards genuine reasoning, and let θ denote the parameters of a language model fθf_θ. Standard training minimizes the empirical risk: minθâĄ1Nââi=1Nââ(θ;xi,yi), _θ\; 1N _i=1^N (θ;\,x_i,y_i), (1) where â is the token-level negative log-likelihood loss. This objective does not distinguish between samples that promote genuine logical inference and those that achieve low loss via superficial patternsâkeyword-answer correlations, answer token memorization, or template matching. We call such samples shortcut reasoning samples. Formally, let gs=âθââ(θ;s)g_s= _θ\, (θ;s) denote the gradient of a training sample s=(x,y)s=(x,y), and let g=âθ1||ââvâââ(θ;v)g_V= _θ\; 1|V| _v (θ;\,v) (2) be the validation gradient, reflecting parameter updates that improve generalizable reasoning. A sample s is a shortcut reasoning sample if its gradient gsg_s (i) exhibits low alignment with g_V, indicating the update does not transfer to genuine reasoning, and (i) concentrates gradient norm on final answer tokens rather than intermediate reasoning tokens, indicating the model learns what to answer rather than how to reason. Critically, such samples are often correctly labeled and efficiently reduce training loss, rendering them undetectable by conventional data quality or label-noise metrics. Shortcut-Aware Reasoning Training (SART) addresses this by learning parameters θ that minimize training loss while suppressing shortcut-promoting gradient directions: minθâĄ1Nââi=1Nwiâ(θ)âââ(θ;xi,yi), _θ\; 1N _i=1^Nw_i(θ)\; (θ;\,x_i,y_i), (3) where wiâ(θ)â[0,1]w_i(θ)â[0,1] downweights samples that promote shortcut reasoning, subject to gradient updates remaining aligned with g_V. SART pursues two goals: (1) detect shortcut samples via gradient structureâspecifically, non-transfer alignment and answer-token concentration; and (2) suppress shortcut learning by neutralizing harmful gradient directions while preserving generalizable reasoning signals within the same samples. The framework operates purely at the level of training dynamics, requiring neither architectural modifications nor additional human annotations. 3 Overview Figure 1: Overview of SART. (A) Gradient diagnostics: we compute per-sample gradient alignment with validation gradients and answer-token concentration. (B) ShortcutScore computation combines both signals to identify shortcut-prone samples. (C) Sample reweighting reduces their contribution. (D) Gradient surgery projects harmful components out of updates. (E) Resulting updates prioritize reasoning signals over shortcut patterns. Figure 1 illustrates our proposed shortcut-aware reasoning training (SART) framework, comprising two coupled components: (1) ShortcutScore-based sample reweighting and (2) gradient surgery for shortcut suppression. Shortcut Detection via ShortcutScore. We identify shortcut reasoning samples by analyzing gradient structure rather than labels or loss values. Signal 1 measures non-transfer gradient alignment: we compute the cosine similarity between each sampleâs per-sample gradient and the validation gradient g_V. Samples whose gradients are misaligned with g_V reduce training loss without improving generalizable reasoningâa hallmark of shortcut exploitation undetectable by label-quality or loss-based filters. Signal 2 measures answer-gradient concentration: we compare gradient norm on final answer tokens to that on intermediate reasoning tokens. Shortcut samples concentrate gradient energy on answer tokens, indicating the model learns what to predict rather than how to reason. We combine both signals into a unified ShortcutScore that quantifies each sampleâs shortcut tendency and down-weights it during training. Shortcut Suppression via Gradient Surgery. Reweighting reduces shortcut influence but does not eliminate the harmful gradient directions such samples introduce. We therefore project the gradients of high-ShortcutScore samples onto subspaces that avoid directions which increase the ShortcutScore or decrease alignment with g_V, retaining useful reasoning signals while neutralizing shortcut-promoting components. Identifying the correct projection subspace is formulated as a minimax optimization problem: the inner maximization finds the most harmful gradient direction, while the outer minimization ensures updates avoid it. We develop an efficient solver for this minimax game that performs periodic gradient surgery at each update step without significant computational overhead. Together, ShortcutScore reweighting and gradient surgery steer training dynamics toward genuine reasoning signals rather than superficial patterns. 3.1 ShortcutScore-Based Sample Reweighting Why ShortcutScore Reweighting Matters. Standard empirical risk minimization treats all training samples equally, optimizing the average token-level loss across the dataset. This is problematic because shortcut reasoning samplesâthose that achieve low loss by exploiting keyword correlations, answer memorization, or surface pattern matchingâgenerate strong gradients that efficiently reduce training loss but do not support generalizable reasoning. Critically, such samples are often correctly labeled, making them invisible to conventional data quality filters or loss-based selection criteria. The consequence is a systematic training signal misalignment: the model is rewarded for learning spurious strategies, leading to brittleness under distribution shift, number perturbation, or rephrasing. This motivates a reweighting strategy that identifies shortcut-promoting samples from their gradient behavior and reduces their influence on parameter updatesâtargeting the root cause rather than the symptom. Modeling Intuition. Our central insight is that shortcut reasoning samples produce two distinguishable gradient signatures. First, their per-sample gradients are poorly aligned with the validation gradient g_V: the update direction they induce does not improve held-out reasoning performance, signaling that the learned signal does not transfer. Second, their gradient energy is disproportionately concentrated on answer tokens rather than intermediate reasoning tokens: the model learns what to predict rather than how to reason. We combine these two signals into a unified ShortcutScore Sâ(s)S(s) that quantifies each sampleâs tendency to promote shortcut learning, and use it to down-weight shortcut-prone samples during optimization. Step 1: Non-Transfer Gradient Alignment. This step measures how well a sampleâs gradient direction transfers to generalizable reasoning. For each training sample s=(x,y)s=(x,y), we compute the per-sample gradient s=âθââ(θ;s)g_s= _θ (θ;s) and the validation gradient =âθ1||ââvâââ(θ;v).g_V= _θ 1|V| _v (θ;v). (4) The non-transfer gradient alignment is then defined as: Aâ(s)=cosâĄ(s,)=sâ¤ââsâââ.A(s)= (g_s,\,g_V)= g_s g_V\|g_s\|\|g_V\|. (5) A low value of Aâ(s)A(s) indicates that the gradient of sample s points in a direction inconsistent with improving validation reasoning performanceâa hallmark of shortcut exploitation. Step 2: Answer-Gradient Concentration. This step measures how much of the gradient signal is driven by answer tokens versus reasoning tokens. Let TansT_ans and TreasonT_reason denote the index sets of answer tokens and reasoning tokens in the output sequence, respectively. The answer-gradient concentration is defined as: Râ(s)=âtâTansââθâtââtâTansâŞTreasonââθâtâ,R(s)= _tâ T_ans\| _θ _t\| _tâ T_ans⪠T_reason\| _θ _t\|, (6) where ât _t is the token-level loss at position t. A high value of Râ(s)R(s) indicates that the modelâs parameter updates are driven predominantly by answer prediction rather than by learning intermediate reasoning steps. Step 3: ShortcutScore and Reweighted Objective. We combine the two signals into the ShortcutScore: Sâ(s)=Îąâ maxâĄ(0,ĎAâAâ(s))+βâ maxâĄ(0,Râ(s)âĎR),S(s)=ι¡ (0,\, _A-A(s))+β¡ (0,\,R(s)- _R), (7) where ĎA _A and ĎR _R are thresholds defining low alignment and high concentration, and Îą,βι,β are balancing hyperparameters. Samples with high Sâ(s)S(s) are flagged as shortcut-promoting. We assign each sample a weight wâ(s)=expâĄ(âÎťâSâ(s))w(s)= (-Îť S(s)), so that higher shortcut tendency leads to lower training influence. The reweighted optimization objective becomes: minθâĄ1Nââi=1Nwâ(si)â ââ(θ;si). _θ\; 1N _i=1^Nw(s_i)¡ (θ;\,s_i). (8) The model parameters are updated at each step via: θt+1=θtâΡââsâbatchwâ(s)â âθââ(θ;s), _t+1= _t-Ρ _s w(s)¡ _θ (θ;\,s), (9) where Ρ is the learning rate. Validation gradients g_V are recomputed periodically to keep the alignment signal up to date without excessive overhead. A detailed training procedure is provided in Appendix A.1. 3.2 Gradient Surgery for Shortcut Suppression Why Gradient Surgery Matters. Reweighting reduces the magnitude of shortcut-promoting gradients but does not correct their direction. Even a down-weighted gradient can still contain a harmful directional component that nudges model parameters away from genuine reasoning. The structural reason is that shortcut samples carry two types of harmful gradient components: (i) directions that are misaligned with validation gradients, and (i) directions dominated by answer token signals. Simply scaling these gradients down preserves both failure modes. We therefore propose gradient surgery Yu et al. (2020): an operation that actively projects shortcut gradients onto subspaces that avoid harmful directions, ensuring that only genuine reasoning signals drive parameter updates. Modeling Intuition. The intuition is analogous to the flat region concept in continual learning Schulman et al. (2015): we seek parameter update directions that are simultaneously beneficial for reasoning generalization and minimally disruptive to the reasoning signals already learned. For a shortcut sample s, its gradient sg_s can be decomposed into a harmful component (aligned with non-transferable or answer-dominant directions) and a residual component (aligned with genuine reasoning signals). Gradient surgery removes the harmful component and retains the residual, so that the update from sample s contributes constructively to reasoning learning. Step 1: Projection for Non-Transfer Gradient Alignment. For samples with low alignment Aâ(s)<ĎA(s)< _A, we remove the gradient component that lies in the direction of g_V, retaining only the orthogonal complement. Formally: sâ˛=sâÎłâsâ¤ââ2â,g_s =g_s-Îł g_s g_V\|g_V\|^2\,g_V, (10) where Îłâ[0,1]Îłâ[0,1] controls the strength of the projection. Setting Îł=1Îł=1 fully removes the non-transferable component. The resulting sâ˛g_s is orthogonal to g_V, ensuring the update does not reinforce non-generalizable directions. Step 2: Suppression for Answer-Gradient Concentration. For samples with high concentration Râ(s)>ĎR(s)> _R, we decompose sg_s into answer-token gradients sansg_s^ans and reasoning-token gradients sreasong_s^reason, and suppress the answer-dominant component: sâ˛=(1âĎ)âsans+sreason,g_s =(1-Ď)\,g_s^ans+g_s^reason, (11) where Ďâ[0,1]Ďâ[0,1] is the suppression coefficient. Setting Ď=1Ď=1 entirely removes the answer token gradient contribution, forcing the model to update based solely on intermediate reasoning signals. Minimax Formulation for Optimal Surgery. A key challenge is determining the optimal surgery direction and strength automatically, rather than relying on fixed thresholds. We find that this challenge can be reformulated as a minimax game: the inner maximization identifies the most harmful gradient direction Ξ within a neighborhood of the current parameters, and the outer minimization finds parameter updates that remain robust under this worst-case perturbation: minθâĄmaxΞ _θ _Ξ âsâwâ(s)â ââ(θ+Ξ;s) _s w(s)¡ (θ+Ξ;\,s) (12) s.t. .t. Ξââł, Ξ , where âłM is the subspace spanned by previous parameter directions, constraining perturbations to directions relevant to shortcut signals. This formulation automatically identifies both the harmful direction and the appropriate surgery strength, without requiring manual threshold tuning. Solving the Optimization. Based on gradient projection Rosen (1960), the adversarial perturbation Ξ is updated via: ΞâΞ+Ρ1âprojâłâ(âθââsâwâ(s)â ââ(θ+Ξ;s)).ΞâΞ+ _1\,proj_M\! ( _θ _s w(s)¡ (θ+Ξ;\,s) ). (13) The model parameters θ are then updated along the direction orthogonal to âłM, ensuring minimal disruption to previously learned reasoning knowledge: θâθâΡ2â(Iâprojâł)â(âθââsâwâ(s)â ââ(θ+Ξ;s)),θâθ- _2\,(I-proj_M) ( _θ _s w(s)¡ (θ+Ξ;\,s) ), (14) where I is the identity matrix and Ρ2 _2 is the step size. The complete training pipeline integrates ShortcutScore computation, sample reweighting, and gradient surgery at each update step, with periodic recomputation of g_V. A full algorithmic description is provided in Appendix A.2. 4 Experiments We evaluate SART on three synthetic reasoning benchmarks with controlled shortcut injection, comparing against eleven baselines spanning data-centric, loss-based, distributionally robust, and invariant learning approaches. 4.1 Experimental Setup Model and Training. We use a GPT-style transformer (4 layers, dmodel=256d_model=256, 8 attention heads, dff=1024d_f=1024, â3.2M parameters) trained for 40 epochs with AdamW Loshchilov and Hutter (2019) (learning rate 1Ă10â31Ă10^-3, weight decay 10â410^-4, batch size 64, cosine annealing). SART hyperparameters (Îą=β=1.0Îą=β=1.0, ĎA=0.3 _A=0.3, ĎR=0.5 _R=0.5, Îť=3.0Îť=3.0, Îł=1.0Îł=1.0, Ď=0.5Ď=0.5) are selected via grid search over 80 configurations (§4.4). Datasets. We construct three synthetic datasets with controlled shortcut injection, each containing 2,000 training / 500 validation / 1,000 test samples. Training data follows the shortcut rule for 70% of samples and the true rule for 30%; validation and test sets exclusively follow the true rule. The test set is split evenly into clean (shortcut consistent with correct answer) and perturbed (shortcut contradicts the correct answer) subsets. ⢠Math-Arithmetic: Binary classification of a+bâĽ10a+b⼠10. Shortcut: first operand aâĽ5a⼠5; true rule: evaluate the full sum. ⢠Financial-Analysis: Regulatory compliance from four financial features. Shortcut: revenue âĽ5⼠5; true rule: margin âĽ5⼠5 and debt <5<5 (multi-feature conjunction). ⢠Causal-Reasoning: Causal direction inference from observed variables with a confounder. Shortcut: correlation âĽ5⼠5; true rule: xâĽ5x⼠5 and z<3z<3 (requires confounder adjustment). Baselines. We compare against eleven methods: (1) SFT: vanilla cross-entropy; (2) Self-Consistency Wang et al. (2023): majority vote over 5 samples (T=0.8T=0.8); (3) Data Filtering: remove high-confidence samples (>0.90>0.90) after warmup; (4) JTT Liu et al. (2021): upweight error-prone samples; (5) Focal Loss Lin et al. (2018): (1âpt)Îł(1-p_t)^Îł weighting; (6) Group DROâ Sagawa et al. (2020): worst-group optimization with known shortcut annotations; (7) IRM⥠Arjovsky et al. (2020): invariance penalty across environments; (8) V-REx⥠Krueger et al. (2021): variance of risks across environments; (9) Fishr⥠Rame et al. (2022): gradient variance equalization; (10) LfF Nam et al. (2020): learning from failure; (11) Influence Filtering Koh and Liang (2020): remove harmful samples via influence functions. Here â requires group annotations and ⥠requires multiple environments. Metrics. Clean accuracy, perturbed accuracy (robustness), reasoning consistency, and shortcut detection F1. 4.2 Main Results Table 1: Main results averaged across three synthetic benchmarks. Bold: best; underline: second best. â Requires group annotations. âĄRequires multiple environments. Method Accuracy (â ) Robustness (â ) Reasoning (â ) SC Det. F1 (â ) Standard Fine-Tuning (SFT) 68.8 19.4 55.9 â Self-Consistency 69.1 23.3 55.9 â Data Filtering 74.7 54.6 64.7 0.69 JTT 67.9 19.0 54.9 â Focal Loss 68.2 18.2 55.4 â Group DROâ 73.5 36.0 61.5 â IRM⥠73.3 34.4 61.5 â V-REx⥠72.1 34.9 60.6 â Fishr⥠68.1 16.7 56.3 â LfF 69.0 19.4 56.3 â Influence Filtering 76.0 47.7 68.4 â SART (Ours) 92.5 87.9 85.5 0.64 Table 1 presents the aggregate results. SART achieves 92.5% accuracy and 87.9% robustness, surpassing the strongest baseline (Influence Filtering: 76.0%/47.7%) by +16.5 p accuracy and +40.2 p robustness. Compared to standard fine-tuning, SART improves accuracy by +23.7 p and robustness by +68.5 p. Crucially, SART achieves this without requiring group annotations (unlike Group DRO) or multiple training environments (unlike IRM, V-REx, and Fishr). We observe a clear four-tier stratification among methods: Tier 1 (67â69% acc, 16â23% rob): SFT, JTT, Focal Loss, Fishr, LfF, Self-Consistency, and Meta-Reweighting. These methodsâspanning loss modification, inference-time voting, and gradient variance equalizationâfail to meaningfully address systematic shortcuts, performing near the SFT baseline. Tier 2 (72â74% acc, 34â36% rob): Group DRO, IRM, and V-REx. Distributional robustness and invariance-based approaches provide moderate gains but remain fundamentally limited: DRO requires explicit group annotations, while IRM/V-REx depend on environment partitions. Tier 3 (74â76% acc, 48â55% rob): Data Filtering and Influence Filtering achieve higher robustness through sample removal but at the cost of discarding training data and exhibiting high per-task variance (see Table 2). Tier 4: SART (92.5% acc, 87.9% rob). By jointly performing gradient-level shortcut detection and correction, SART achieves a qualitative leapâexceeding the best baseline by +16.5 p accuracy and +40.2 p robustness. Table 2: Per-dataset results (%). Bold: best; underline: second best per dataset. Dataset Method Clean Acc. (â ) Robustness (â ) Reasoning (â ) Math SFT 77.6 3.2 77.6 Data Filtering 82.6 26.6 82.6 Group DRO / V-REx 86.8 48.0 86.8 IRM 85.4 44.2 85.4 Influence Filtering 81.4 50.4 81.4 SART (Ours) 98.0 95.8 98.0 Financial SFT 66.4 28.6 46.8 Data Filtering 72.6 68.4 60.8 Group DRO 68.0 31.2 48.8 IRM 69.6 30.8 51.6 Influence Filtering 63.4 23.6 59.6 SART (Ours) 89.4 78.8 87.4 Causal SFT 62.4 26.4 43.4 Data Filtering 69.0 68.8 50.8 Group DRO 65.8 28.8 48.8 IRM 65.0 28.2 47.4 Influence Filtering 83.2 69.2 64.2 SART (Ours) 90.2 89.0 71.2 Per-dataset results (Table 2) reveal three important findings: (1) SART achieves the best accuracy, robustness, and reasoning on all three datasets. On Math-Arithmetic, SART reaches near-perfect performance (98.0%/95.8%), a +47.8 p robustness gain over the next-best baseline (Influence Filtering: 50.4%). On Financial-Analysis, SART achieves 89.4% accuracy and 87.4% reasoning consistencyâa +19.8 p accuracy gap over IRM and a +26.6 p reasoning gap over Data Filtering. On Causal-Reasoning, SART achieves 90.2%/89.0%, exceeding Influence Filtering by +7.0 p accuracy and +19.8 p robustness. (2) Data-removal baselines exhibit extreme cross-task variance. Influence Filtering ranks second on Causal (83.2%/69.2%) but collapses on Financial (63.4%/23.6%)âperforming worse than SFT on robustness. Data Filtering achieves high Financial robustness (68.4%) but at the cost of severe accuracy drops (e.g., only 82.6% on Math vs. SARTâs 98.0%). These methods lack the discriminative power to identify which samples are harmful for complex multi-feature reasoning. (3) SARTâs advantage scales with task complexity. The accuracy gap between SART and the best baseline increases from Math (+11.2 p) to Financial (+16.8 p) to Causal (acc: +7.0 p, rob: +19.8 p), demonstrating that gradient-based correction is particularly powerful for multi-feature shortcuts that surface-level filtering cannot detect. 4.3 Ablation Study Table 3: Ablation of SART components (averages across three datasets). Grad. Align.: cosine similarity between training and validation gradients (closer to 0 == better alignment with true task objective). Configuration Accuracy (â ) Robustness (â ) Grad. Align. (â ) SC Det. F1 SFT (Baseline) 68.8 19.4 â0.08-0.08 â + Reweighting Only 75.2 40.1 â0.06-0.06 0.85 + Gradient Surgery Only 73.5 40.5 â0.14-0.14 0.82 SART (Full) 92.5 87.9 â0.03-0.03 0.64 Table 3 isolates each componentâs contribution. Both mechanisms individually improve over SFT: ⢠Reweighting Only: +6.4 p accuracy, +20.7 p robustness. By downweighting shortcut-reliant samples, reweighting shifts the effective training distribution toward the true rule, achieving the highest standalone detection F1 (0.85). ⢠Gradient Surgery Only: +4.7 p accuracy, +21.1 p robustness. By projecting out shortcut-aligned gradient components, surgery directly corrects the optimization trajectory. Critically, the full combination exhibits dramatic super-additive gains. On robustness, the individual improvements sum to 20.7+21.1=41.820.7+21.1=41.8 p, yet the full method achieves +68.5 pâa +26.7 p synergy bonus (64% more than the sum of parts). Similarly, accuracy gains sum to 11.1 p individually but reach 23.7 p combined (+12.6 p synergy). This exceptional synergy arises because reweighting and gradient surgery operate on complementary axes: reweighting adjusts how much each sample contributes to the loss, while surgery adjusts in which direction each gradient update proceeds. When combined, reweighted gradients are more amenable to precise surgical correction, and surgery amplifies the effect of reweighting by eliminating residual shortcut directions. The full method achieves near-zero gradient misalignment (â0.03-0.03), a 2.7Ă2.7Ă improvement over SFT (â0.08-0.08). Removing either component causes catastrophic degradation: removing Reweighting drops accuracy by â-19.0 p and robustness by â-47.4 p; removing Gradient Surgery drops accuracy by â-17.3 p and robustness by â-47.7 p. Both components are essential and contribute approximately equally. 4.4 Hyperparameter Sensitivity Table 4: Hyperparameter sensitivity analysis via grid search over 80 configurations (ÎťĂÎłĂĎÎťĂÎłĂĎ). Combined =0.4ĂAcc+0.6ĂRob=0.4ĂAcc+0.6ĂRob. Parameter Value Accuracy (â ) Robustness (â ) Combined (â ) Îť (reweighting) 1.0 75.8 44.0 56.7 1.5 81.1 52.9 64.2 2.0 86.9 69.6 76.5 3.0 87.6 74.2 79.6 5.0 86.4 73.6 78.7 Îł (gradient surgery) 0.3 84.0 68.4 74.7 0.5 84.2 70.2 75.8 0.8 82.2 56.4 66.8 1.0 87.6 74.2 79.6 Ď (answer suppression) 0.1 86.4 73.6 78.7 0.3 88.0 68.0 76.0 0.5 87.6 74.2 79.6 0.7 86.9 69.6 76.5 We conduct an exhaustive grid search over 80 configurations (Table 4), revealing three insights into SARTâs operating regime: (1) Reweighting strength Îť is the most influential parameter. Performance increases monotonically from Îť=1.0Îť=1.0 to Îť=3.0Îť=3.0 (+11.8 p accuracy, +30.2 p robustness), plateauing thereafter (Îť=5.0Îť=5.0 yields comparable results). This defines a âsweet spotâ where shortcut samples are sufficiently downweighted without catastrophically distorting the training distribution. (2) Full-strength gradient surgery (Îł=1.0Îł=1.0) is optimal. Performance at Îł=0.8Îł=0.8 (56.4% robustness) drops substantially compared to both Îł=0.5Îł=0.5 (70.2%) and Îł=1.0Îł=1.0 (74.2%). This non-monotonicity suggests a phase transition: partial projection at intermediate strengths creates unstable gradient dynamics, while complete projection (Îł=1.0Îł=1.0) cleanly eliminates the entire shortcut subspace. (3) Moderate answer suppression (Ď=0.5Ď=0.5) balances the accuracyârobustness trade-off. Both extreme values underperform: too-low Ď (0.10.1) retains shortcut-correlated answer signals, while too-high Ď (0.70.7) discards useful reasoning information from the answer pathway. 4.5 ShortcutScore Validation Figure 2: Empirical validation of ShortcutScore. (a) Score vs. true shortcut rate (Pearson r=0.67r=0.67), showing monotonic correlation. (b) Gradient alignment distribution: shortcut samples (red) skew negative; non-shortcut samples (green) cluster near zero. (c) Exponential reweighting curve with Îť=3.0Îť=3.0. ShortcutScore achieves an average detection F1 of 0.64 across three datasets, while the closely related Data Filtering approach achieves F1 of 0.69. Although comparable in detection, SARTâs advantage lies not in detection alone but in its corrective mechanism: rather than discarding detected shortcut samples entirely (losing valuable training signal), SART simultaneously downweights their contribution and surgically corrects their gradient directions. This âdetect-and-correctâ paradigm preserves training data while neutralizing shortcuts, explaining SARTâs +17.8 p accuracy advantage over Data Filtering despite similar detection capability. Figure 2(a) confirms a positive monotonic correlation between Sâ(s)S(s) and the true shortcut rate (Pearson r=0.67r=0.67, p<0.001p<0.001). The gradient alignment distribution (Figure 2b) reveals clear separation: shortcut samples produce gradients that skew toward negative alignment with the validation objective, while non-shortcut samples cluster symmetrically near zero. 4.6 Computational Cost SART completes in approximately 60 seconds per dataset on Apple M-series hardware (3.2M-parameter model, 2,000 training samples) and under 3 minutes per dataset on NVIDIA H100 NVL GPUs at server scale. The overhead relative to SFT is approximately 2.5Ă2.5Ă, arising from per-sample gradient computation during ShortcutScore evaluation and gradient projection during surgery. Validation gradients are recomputed every k=5k=5 steps rather than per-step, keeping overhead manageable. 5 Related Work Language Model Training. Prevailing training paradigmsâsupervised learning, reinforcement learning from human feedback, and hybrid methods Sun et al. (2023)âoptimize primarily for answer correctness, which structurally favors shortcuts such as surface pattern matching and answer memorization that efficiently reduce training loss. These shortcuts yield brittle models that fail under distribution shift, relying on spurious correlations rather than causal mechanisms. Our approach targets this failure at its source by modifying training dynamics rather than the objective or architecture. Reasoning in AI. Chain-of-Thought (CoT) training and Self-Consistency Decoding Sun et al. (2023) improve reasoning trace generation and output consistency at inference time, but leave training dynamicsâwhere shortcut biases are encodedâunchanged. CoT can produce plausible reasoning steps that still exploit shortcuts, and Self-Consistency detects inconsistencies without correcting the learned parameters that cause them. Effective shortcut suppression requires intervening during training rather than at inference. Gradient-Based Optimization. Standard gradient-based optimizers Sun et al. (2023) apply updates uniformly across samples, with no mechanism to distinguish gradients that promote genuine reasoning from those that reinforce spurious patterns. Shortcut samples often generate strong gradients that dominate the learning signal while reducing training loss efficiently, a failure mode invisible to standard optimization. We exploit the observation that shortcut samples exhibit distinct gradient signaturesâpoor alignment with validation gradients and high concentration on answer tokensâto enable targeted correction. Closest Prior Work. Prior work on gradient modification Sun et al. (2023) and training signal adjustment for reasoning validity Sun et al. (2023) addresses robustness in general but lacks a principled mechanism to identify which gradients arise from shortcut reasoning specifically. Without a way to quantify sample-level shortcut tendency from gradient behavior, these methods cannot perform targeted suppression and leave core shortcut learning mechanisms intact. By unifying gradient alignment and answer-gradient concentration into a ShortcutScore, SART enables precise shortcut identification and combines reweighting with gradient surgery for a more principled intervention than prior approaches. 6 Conclusion Recent advancements in large language models (LLMs) have showcased impressive reasoning capabilities, yet these models often rely on reasoning shortcuts rather than genuine logical inference. This reliance, a critical issue within the field, stems from training signal misalignment where shortcuts reduce training loss but hinder generalizable reasoning. Existing approaches fail to adequately address this by either not modifying training dynamics (e.g., Self-Consistency Decoding) or being unable to detect shortcut-promoting gradients (e.g., Data Filtering). This leads to models that appear to reason correctly but fail when distributions shift, resulting in degraded reliability in mission-critical systems and significant financial or safety risks. This study addresses these limitations by developing novel techniques: the ShortcutScore for reweighting training samples and gradient surgery to modify gradient directions. Our technical contributions introduce ShortcutScore-based reweighting and gradient modification strategies that directly improve reasoning robustness and model reliability. We observe that shortcut reasoning samples produce distinct gradient signaturesâlow alignment with validation gradients and high concentration on answer tokens. This structural insight allows us to quantify shortcut propensity and intervene during training. The experimental results demonstrate that our approach significantly improves reasoning robustness across diverse tasks, substantiating the effectiveness of our methods. Theoretically, this work advances the understanding of gradient alignment and reasoning signal correction, offering new insights into LLM training dynamics. Practically, the implications are profound for applications in domains such as finance and scientific discovery, where enhanced model robustness can lead to more reliable decision-making. Future research could explore additional datasets and further refine the methodology to enhance its applicability, building upon this foundational framework for addressing shortcut reasoning in LLMs and paving the way for more reliable and robust reasoning models. Acknowledgments and Disclosure of Funding This work was supported by the National Science Foundation under Grant No. 7813331468. References [1] M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz (2020) Invariant risk minimization. External Links: 1907.02893, Link Cited by: §4.1. [2] Y. Bengio, J. Louradour, R. Collobert, and J. Weston (2009) Curriculum learning. In Proceedings of the 26th Annual International Conference on Machine Learning, p. 41â48. Cited by: §1. [3] N. Ho, L. Schmid, and S. Yun (2022) Large language models are reasoning teachers. arXiv preprint arXiv:2212.10071. Cited by: §1. [4] P. W. Koh and P. Liang (2020) Understanding black-box predictions via influence functions. External Links: 1703.04730, Link Cited by: §4.1. [5] T. Kojima, S. S. Gu, M. Reid, Y. Matsuo, and Y. Iwasawa (2022) Large language models are zero-shot reasoners. Advances in Neural Information Processing Systems 35, p. 22199â22213. Cited by: §1. [6] D. Krueger, E. Caballero, J. Jacobsen, A. Zhang, J. Binas, D. Zhang, R. L. Priol, and A. Courville (2021) Out-of-distribution generalization via risk extrapolation (rex). External Links: 2003.00688, Link Cited by: §4.1. [7] T. Lin, P. Goyal, R. Girshick, K. He, and P. DollĂĄr (2018) Focal loss for dense object detection. External Links: 1708.02002, Link Cited by: §4.1. [8] E. Z. Liu, B. Haghgoo, A. S. Chen, A. Raghunathan, P. W. Koh, S. Sagawa, P. Liang, and C. Finn (2021) Just train twice: improving group robustness without training group information. External Links: 2107.09044, Link Cited by: §4.1. [9] I. Loshchilov and F. Hutter (2019) Decoupled weight decay regularization. In International Conference on Learning Representations, Cited by: §4.1. [10] J. Nam, H. Cha, S. Ahn, J. Lee, and J. Shin (2020) Learning from failure: training debiased classifier from biased classifier. External Links: 2007.02561, Link Cited by: §4.1. [11] L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. (2022) Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems 35, p. 27730â27744. Cited by: §1. [12] A. Rame, C. Dancette, and M. Cord (2022) Fishr: invariant gradient variances for out-of-distribution generalization. External Links: 2109.02934, Link Cited by: §4.1. [13] J. B. Rosen (1960) The gradient projection method for nonlinear programming. part i. linear constraints. Journal of the Society for Industrial and Applied Mathematics 8 (1), p. 181â217. Cited by: §3.2. [14] S. Sagawa, P. W. Koh, T. B. Hashimoto, and P. Liang (2020) Distributionally robust neural networks for group shifts: on the importance of regularization for worst-case generalization. External Links: 1911.08731, Link Cited by: §4.1. [15] J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz (2015) Trust region policy optimization. In International Conference on Machine Learning, p. 1889â1897. Cited by: §3.2. [16] J. Sun, C. Zheng, E. Xie, Z. Liu, R. Chu, J. Qiu, J. Xu, M. Ding, H. Li, M. Geng, et al. (2023) A survey of reasoning with foundation models. arXiv preprint arXiv:2312.11562. Cited by: §5, §5, §5, §5. [17] X. Wang, J. Wei, D. Schuurmans, Q. V. Le, E. H. Chi, S. Narang, A. Chowdhery, and D. Zhou (2022) Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171. Cited by: §1, §1. [18] X. Wang, J. Wei, D. Schuurmans, Q. V. Le, E. H. Chi, S. Narang, A. Chowdhery, and D. Zhou (2023) Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171. Cited by: §4.1. [19] J. Wei, X. Wang, D. Schuurmans, M. Bosma, B. Ichter, F. Xia, E. Chi, Q. V. Le, and D. Zhou (2022) Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems 35, p. 24824â24837. Cited by: §1, §1. [20] T. Yu, S. Kumar, A. Gupta, S. Levine, K. Hausman, and C. Finn (2020) Gradient surgery for multi-task learning. In Advances in Neural Information Processing Systems, Vol. 33, p. 5824â5836. Cited by: §3.2.