Diagnosing Task Insensitivity in Language Agents

🕒 Published (v1): 2026-06-25 11:53 UTC · Source: Arxiv · link

Why this paper was selected

Diagnoses task insensitivity as key OOD failure mode in long-horizon LLM agents

Ask a follow-up

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

💬 Ask ChatGPT✦ Ask Claude

TL;DR

LLM agents trained on interactive tasks develop task insensitivity: they reconstruct and execute familiar task patterns even when the task description is corrupted or replaced with a semantically distinct task. The authors diagnose this via attention drift during training and propose Task-Perturbed NLL Optimization, a lightweight contrastive regularizer that explicitly penalizes high action probability under a replaced task description.

Problem

Standard next-action supervision (SFT or RL) provides no training pressure to preserve action-level dependence on the task instruction. As training progresses, agents learn shortcut policies conditioned primarily on local observations—which are often nearly identical across semantically distinct tasks—rather than the task goal. This causes OOD failure: a model trained on "heat A and place it on B" goes to the microwave when asked to "cool A and place it on B," even though it can correctly answer "should a microwave be used for cooling?" when asked directly.

Method

Diagnostic protocol. Task descriptions in ALFWorld, ScienceWorld, and WebShop are manually corrupted (words removed, conflicting attributes inserted) while the rest of the prompt remains intact. The model is explicitly told it may ask for clarification. Two metrics are tracked: Hit (action under corrupted prompt matches an action from the original prompt) and Inquiry (model requests clarification).

Attention analysis. Attention weights during action generation are decomposed into five prompt regions. Attention to the task-instruction region is tracked across SFT and GRPO checkpoints.

Task-Perturbed NLL Optimization. For each training instance, the task description is replaced with a lexically similar but semantically distinct task (top-10 by edit distance). A frozen reference model computes the ratio

\[\rho^{(i)}_\text{ref} = \frac{N\!L\!L^{(i)}_{\text{perturb,ref}}}{N\!L\!L^{(i)}_{\text{vanilla,ref}} + \epsilon}\]

and the current model computes \(\rho^{(i)}_\text{cur}\) analogously. The regularizer penalizes the current model only when its separation falls below the reference-calibrated level:

\[\mathcal{L}_\text{floor} = \max\!\left(0,\; \log\rho^{(i)}_\text{ref} - \log\rho^{(i)}_\text{cur}\right)\]

\[\mathcal{L}_\text{total} = \mathcal{L}_\text{SFT} + \lambda\,\mathcal{L}_\text{floor}\]

For RL (GRPO), the regularizer targets the \(k\%\) of rollouts with negative advantage: their task descriptions are replaced and the same advantage scores are reused in the gradient update.

Key Contributions

• Empirical demonstration that current LLMs (GPT-5.4, Qwen3, Llama3) exhibit task insensitivity: under corrupted instructions, Hit rates are 59–85% across benchmarks while Inquiry rates are only 2–7%.
• Training-dynamics evidence showing that both consistency rate (action stability under task replacement) and reconstructed-task similarity increase monotonically over SFT and GRPO checkpoints, while average attention to task tokens decreases.
• A toy linear model showing why the optimization bias arises: the gradient w.r.t. the static task signal scales with the mean trajectory residual \(\bar{\delta}\), which shrinks once coarse task identity is captured, whereas the gradient w.r.t. dynamic state features remains informative step-by-step.
• Task-Perturbed NLL Optimization: a reference-calibrated floor loss that adds no new architecture and requires only a frozen reference model copy.

Results

• On Qwen3-8B OOD splits (Table 2): ours achieves 58.7 / 29.6 / 49.6 on ALFWorld / ScienceWorld / WebShop vs. SFT baseline of 55.2 / 25.6 / 47.1; outperforms CoIN (56.3 / 26.4 / 47.5) and Task Augmentation (57.3 / 25.1 / 45.2).
• Gains are most consistent under GRPO and in task pairs where local state is nearly identical across tasks (e.g., heat vs. cool in ALFWorld).
• Attention to task tokens is better preserved throughout training under the regularizer (Figure 9).
• Conditional consistency (action repetition after task swap) grows more slowly during training under the regularizer (Figure 10), directly confirming suppressed task-insensitive behavior.

Limitations

• Evaluation covers only three environments (ALFWorld, ScienceWorld, WebShop); these may not represent the breadth of real agentic settings.
• The attention analysis is mechanistically suggestive, not causally established; it does not formally prove how the model internally routes its decisions.
• The toy linear decomposition is illustrative and does not analyze transformer attention directly.
• No ablation on the choice of replacement task similarity criterion (edit distance top-10) or sensitivity to the reference model quality.

Relevance to Agentic AI / LLM Agents

Task insensitivity is a previously underdiagnosed failure mode directly relevant to anyone deploying or training LLM agents on real-world task distributions: standard SFT and RL fine-tuning actively worsen instruction grounding rather than merely failing to improve it. The diagnostic protocol (corruption + Hit/Inquiry metrics) is a lightweight, model-agnostic probe that practitioners can apply to audit trained agents before deployment. The proposed regularizer is directly composable with any SFT or GRPO training pipeline with minimal overhead. This work connects to the broader shortcut learning and OOD generalization literature but specifically identifies the optimization dynamics mechanism—attention drift from static task tokens toward dynamic observation tokens—as the root cause, offering a cleaner target for future interventions.