Mitigating Query Shortcuts in Chain-of-Thought with Context-Aware Decoding

1 Introduction

• CoT rationales may not reflect actual decisions: models can use query-specific shortcuts despite reasoning that suggests a different answer, weakening rationale–answer dependence.Source: §1 · p. 1
• Prior work has studied rationale interventions and reliability improvements through supervision or selection; this paper asks whether decoding can directly mitigate shortcuts without parameter changes or external evaluators.Source: §1 · p. 1
• Standard CoT decoding does not explicitly distinguish query and rationale contributions, so answers can be generated primarily from query-conditioned predictions.Source: §1 · p. 1
• RCD contrasts next-token distributions conditioned on the query plus rationale with those conditioned on the query alone, increasing preference for tokens supported by the rationale.The method applies Context-Aware Decoding to CoT without additional training.Source: §1 · pp. 1–2
• The decoding rule has a KL-regularized reward-maximization interpretation, with reward given by the log-density ratio between rationale-conditioned and query-only distributions.Source: §1 · p. 2
• Across five reasoning benchmarks and multiple model families, RCD consistently improves answer accuracy over standard CoT and increases behavioral dependence on rationale content.The gains primarily correct answers that fail to follow valid rationales, rather than repairing incorrect reasoning.Source: §1 · p. 2

2.1 Next-token Prediction in LLM

• An autoregressive LLM generates an output sequence token by token, with each token drawn from a next-token distribution computed by applying softmax to its logits.The sequence uses preceding output tokens and the query as context; L is its token length and V the vocabulary.Source: §2.1 · p. 2
• A decoding strategy maps the query to an output sequence; greedy decoding and sampling are examples of token-generation strategies.Source: §2.1 · p. 2

2.2 Chain-of-Thought (CoT)

• CoT generates a rationale from the query, then adds that rationale to the context used to generate the final answer.The setup assumes rationale and answer generation are separated and permits intervening with alternative rationales, such as ground truth.Source: §2.2 · p. 2
• Because standard CoT predicts answer tokens from the query and rationale together, it does not guarantee rationale dependence; models can answer directly from the query.Prior examples include option-order bias and silent correction of reasoning errors.Source: §2.2 · p. 2

3.1 Rationale Contrastive Decoding (RCD)

• RCD applies context-aware decoding to self-generated CoT by contrasting rationale-conditioned and query-only next-token distributions, reweighting with a density ratio.This applies CAD to rationales rather than external contexts.Source: §3.1 · p. 2, Figure 1 · p. 3
• For α > 0, the density ratio favors tokens whose probability rises with rationale conditioning, emphasizing information contributed by the rationale beyond the query alone.The contrast aims to mitigate query-shortcut behavior rather than simply increase agreement with query-only predictions.Source: §3.1 · p. 3

3.2 Theoretical Insight

• The decoding rule is a closed-form optimal policy for a KL-regularized reward-maximization objective, with query-only probabilities as reward and the rationale-conditioned distribution as reference policy.The paper states that sampling by Eq. (7) is equivalent to sampling from the objective’s optimal policy.Source: §3.2 · p. 3
• The objective seeks tokens that highlight differences from the query-only distribution while preserving likelihood under the original query-rationale distribution.The authors expect this to make answer generation rely more heavily on the CoT rationale.Source: §3.2 · p. 3

4.1 Settings

• Experiments use five reasoning benchmarks and Qwen3, Gemma3, and Llama-3.1; CoT and RCD share rationales, and α = 1 is selected on a GSM8K training subset.The setup also evaluates Logit Addition as a baseline for comparing subtraction with reinforcement of query-only predictions.Source: §4.1 · p. 3

4.2 Proof-of-Concept by Intervention Analysis

• With ground-truth rationales, RCD achieves higher accuracy; with random rationales, its accuracy is substantially lower, unlike standard CoT’s non-trivial random-rationale accuracy.Logit Addition degrades accuracy with self-generated and ground-truth rationales while retaining higher accuracy with random rationales.Source: §4.2 · p. 3, Table 1 · p. 3

4.4 Qualitative Analysis

• In qualitative examples, RCD follows a rationale supporting answer A and flags an irrelevant rationale, while standard CoT ignores both rationale cues.In the irrelevant-rationale example, standard CoT gives the correct answer, whereas RCD indicates that the rationale is unrelated to the query.Source: Figure 2 · p. 4, §4.4 · p. 4

4.3 Multiple Reasoning Benchmarks

• RCD outperforms standard CoT and Logit Addition in all benchmark cases.The benchmarks are described as reasoning-intensive, and the analysis also considers multiple model families.Source: §4.3 · p. 4, Table 2 · p. 4

4.5 Analysis with Human Annotation

• Among 1,201 correct rationales, RCD corrects 10 of 33 CoT errors (30.30%) without changing correct predictions; among 118 incorrect rationales, it changes no predictions.The analysis concludes that accuracy gains primarily stem from better adherence to correct rationales, while increased reliance on incorrect rationales can occasionally be detrimental.Source: §4.5 · p. 4, Table 3 · p. 4

• Prior work emphasizes faithfulness evaluation or additional mechanisms, while this study applies Context-Aware Decoding to CoT rationales without training or external evaluators.The paper also examines compatibility with Self-Consistency separately in Appendix H.Source: §5 · p. 4

• RCD contrasts rationale-conditioned and query-only next-token distributions to encourage answers that depend more strongly on rationales.The paper reports increased behavioral rationale–answer dependency and consistent accuracy improvements over standard CoT across multiple reasoning benchmarks.Source: §6 · p. 4

Limitations

• Rationale interventions do not establish that generated rationales faithfully reflect internal reasoning, so the findings support stronger rationale–answer dependency rather than general reasoning reliability.Changes in accuracy under ground-truth or shuffled rationales measure sensitivity to rationale content, not necessarily correct or faithful reasoning.Source: Limitations · p. 5
• Because RCD does not verify rationale correctness, it may amplify erroneous reasoning and suppress useful information from the query-only distribution when rationales are wrong or incomplete.Adaptive contrastive weighting based on rationale quality is proposed as future work.Source: Limitations · p. 5
• Potential societal and deployment risks are discussed in Appendix B.Source: Limitations · p. 5

A Proof of Theorem 3.1

• The proof starts from a general KL-regularized reward-maximization objective and treats Eq. (9) as a special case of it.The objective uses reward R, state s, and action a; its optimal policy is then introduced.Source: A Proof of Theorem 3.1 · p. 7
• For each answer token, the proof sets the state to the preceding tokens, query, and rationale, uses the rationale-conditioned distribution as the reference policy, and sets β = 1/α.The reward is the log ratio of rationale-conditioned to query-conditioned token probabilities.Source: A Proof of Theorem 3.1 · p. 7
• Sampling with Eq. (7) is equivalent to sampling from the optimal policy for the maximization in Eq. (9).Source: A Proof of Theorem 3.1 · p. 7
• Although RCD and RED share a theoretical framework, they have different formalizations and objectives; RED uses a product of separated image- and rationale-conditioned distributions.Source: A Proof of Theorem 3.1 · p. 7

B Potential Risks

• RCD could be misused to produce persuasive but misleading explanations, including rationalized misinformation, deceptive arguments, or fabricated justifications.The authors recommend considering monitoring, provenance tracking, and safeguards against persuasive misuse in future deployment.Source: B Potential Risks · p. 7

C Artifact Licenses and Terms of Use

• The experiments use public reasoning benchmarks and Qwen3, Gemma3, and Llama3.1 as fixed inference backbones, without fine-tuning or redistributing the artifacts.The artifact terms are based on public dataset or model cards; users reproducing or extending the work are advised to consult current official pages and comply with their terms.Source: C Artifact Licenses and Terms of Use · p. 7, Table 4 · p. 7

D Use of AI Assistant

• ChatGPT was used to polish manuscript structure and grammar, while Claude Code supported experimental-code development.Source: D Use of AI Assistant · p. 7

E Instruction Prompts

• The appendix provides separate prompts for step-by-step rationale generation, answering from rationales, and answering directly without reasoning steps.Source: E Instruction Prompts · p. 8

F Grid Search for α

• The α search used 1,000 randomly extracted GSM8K training samples for validation and searched {0.1, 0.3, 0.5, 1, 3, 5}.Source: F Grid Search for α · p. 8

F.1 Ablation Study

• Higher α increases rationale reliance but eventually lowers accuracy as query-only conditions weaken; RCD is stable across models, with larger gains at larger model sizes.The reported model-size pattern suggests stronger answer-quality gains when models have stronger reasoning capabilities.Source: §F.1 · p. 8

G Computational Analysis

• RCD’s batch-parallel implementation adds 1.5% end-to-end latency and 1.4% peak VRAM versus CoT; increasing α can trade accuracy for stronger rationale reliance.Accuracy generally decreases with α for random rationales and falls for self-generated rationales at excessively large α.Source: G Computational Analysis · p. 8, Figure 3 · p. 8

I Human Annotation

• RCD consistently improves GSM8K test accuracy over CoT across Qwen3, Gemma3, and Llama-3.1 model sizes at α = 1, using shared generated rationales.Source: Table 5 · p. 9
• One annotator labeled rationale correctness for each query, generated-rationale, and ground-truth-answer triple in the Section 4.5 analysis.Source: I Human Annotation · pp. 8–9

H Compatibility with Self-Consistency

• On GSM8K with Qwen3-8B, SC+RCD reaches 90.79% accuracy versus 89.58% for standard SC, while single-rationale RCD reaches 89.91% versus 89.16% for standard CoT, suggesting complementary gains.Source: H Compatibility with Self-Consistency · p. 8, Table 7 · p. 9