Where Do CoT Training Gains Land in LLM based Agents?

🕒 Published (v1): 2026-06-25 12:09 UTC · Source: Arxiv · link

Why this paper was selected

Empirically audits where CoT training gains land in agents; challenges faithfulness assumptions

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

CoT training in LLM-based agents primarily improves direct action prediction from the prompt rather than strengthening reasoning-mediated action revision. The authors diagnose this via a prompt-action vs. CoT-action comparison framework and show the CoT-over-prompt accuracy gap stays flat across checkpoints. Selectively masking action-token supervision on 30% of training examples improves OOD generalization by reducing this prompt-anchoring shortcut.

Problem

Prior work shows verbalized CoT in LLMs is not always faithful — models may predict the action before reasoning and use CoT as post-hoc justification. In long-context agent settings (where the prompt contains task instructions, interaction history, and environment feedback), it is unclear whether CoT training improves reasoning-mediated action revision or simply strengthens direct prompt-to-action prediction. Existing evaluations using overall task success rate cannot distinguish these two mechanisms.

Method

The authors define two decoding modes: CoT actions (model generates step-by-step reasoning then emits action) and prompt actions (response prefilled with <action> tag, bypassing CoT generation). The gap between the two quantifies how much of the final decision depends on reasoning vs. prompt-recoverable information.

Three diagnostic probes are applied across training checkpoints on ALFWorld, ScienceWorld, and BFCL using Qwen3 (4B–32B) and Llama-3.1-8B: 1. Training dynamics: track prompt-action and CoT-action accuracy in parallel on a held-out validation set. 2. Online evaluation: compare trajectory success and GPT-5.4-judged local action quality under both decoding modes on unseen (in-domain and OOD) tasks. 3. Conflicting-trace test: replace the reasoning trace with a randomly chosen trace supporting a different action (truncated at ratios 0.1/0.3/0.5), measuring whether the model follows the prompt or the substituted trace.

Mechanistic analysis measures attention mass and gradient share allocated to prompt tokens vs. CoT tokens at action-prediction time. The gradient contribution from context token \(i\) to action-token supervision satisfies: $\(\frac{\partial L_t}{\partial v_i} = \alpha_{t,i} \delta_t, \quad \delta_t = \frac{\partial L_t}{\partial o_t}\)$ so higher attention weight \(\alpha_{t,i}\) on prompt tokens translates directly to a larger optimization signal toward the prompt-to-action pathway.

Training intervention: for a randomly selected fraction \(k=0.3\) of training samples, mask action-token loss and optimize only the CoT span: $\(L_{\text{cot}} = -\sum_{t=1}^{|y|} M(t) \log p_\theta(y_t \mid p, y_{<t})\)$ where \(M(t) = 1\) for CoT-span tokens, 0 for final action tokens. Applied to both SFT and RL (GRPO).

Key Contributions

• A diagnostic framework contrasting prompt-action vs. CoT-action accuracy to isolate where CoT training gains land.
• Behavioral evidence (parallel accuracy improvement, flat CoT-over-prompt gap, checkpoint-wise prompt-anchoring under conflicting traces) that CoT supervision strengthens direct prompt-to-action prediction in long-context agents.
• Mechanistic evidence: ~80% of attention mass and a disproportionate share of gradient signal at action-prediction time flows to prompt tokens, providing an optimization-level explanation.
• Reduced action supervision: masking action-token loss on 30% of training examples improves OOD generalization across environments and model sizes.

Results

• Prompt-action and CoT-action accuracy rise in parallel across all checkpoints and environments; CoT-over-prompt gap does not widen.
• CoT-action win rate over prompt actions (GPT-5.4 judge) remains flat across checkpoints (~66–76% depending on domain/model).
• Later checkpoints show stronger prompt-anchoring under conflicting traces (Figure 4); the effect is larger for shorter trace truncations and in OOD settings.
• ~80% of total attention mass during action generation falls on prompt tokens; prompt tokens dominate even among the highest-attention positions (>50% of top-K positions).
• Reduced action supervision (SFT+) vs. baselines on Qwen3-8B:

Method	ALFWorld	ScienceWorld	BFCL
SFT	0.55	0.25	0.58
DPO	0.59	0.29	0.62
FRODO	0.61	0.30	0.64
SFT+ (ours)	0.63	0.27	0.62
DPO+ (ours)	0.65	0.35	0.67

• Gains largest in ALFWorld and ScienceWorld; more modest in BFCL (structured prompts, constrained action space leave less room for prompt-shortcut dominance).

Limitations

• Intervention gain on ScienceWorld under SFT+ alone (0.27) is slightly below FRODO (0.30), requiring DPO combination for full benefit.
• BFCL shows mixed CoT-gap changes under reduced action supervision, suggesting the diagnosis and fix are not uniformly applicable across action-space structures.
• The conflicting-trace experiment assumes random trace substitution is a valid proxy for genuine reasoning perturbation; the quality of the counterfactual traces is not controlled.
• Experiments use Qwen3 family (4B–32B) and Llama-3.1-8B-Instruct; generalization to closed or much larger models is untested.
• The mechanistic explanation (attention/gradient dominance by prompt) is correlational; causal intervention evidence beyond the loss-masking experiment is not provided.
• \(k=0.3\) is chosen by ablation but the sensitivity analysis is only partially shown.

Relevance to Agentic AI / LLM Agents

This paper directly challenges the assumption that CoT-supervised agent training improves reasoning quality — instead, much of the gain reflects stronger shortcut learning from long context prompts, which are the defining feature of realistic agent deployments. The diagnostic framework (prompt-action vs. CoT-action decoding) is practically applicable as an evaluation tool for any agent training pipeline to detect prompt-anchoring. The reduced action supervision intervention offers a lightweight, training-time fix that is compatible with both SFT and RL and improves OOD generalization — a core open problem in agentic AI. The mechanistic finding that agent prompts structurally dominate attention and gradient mass during action prediction has direct implications for how CoT-based agents are trained, evaluated, and audited for faithfulness.