Skip to content

Think Twice Before You Act: Enhancing Agent Behavioral Safety with Thought Correction

🕒 Published (v1): 2025-05-16 10:00 UTC · Source: Arxiv · Venue: ICML 2026 · link

Ask a follow-up

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

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Thought-Aligner is a lightweight, plug-and-play safety module that intercepts and causally corrects an LLM agent's intermediate reasoning thoughts before tool execution, steering the agent away from unsafe actions without modifying the underlying model. Trained via two-stage SFT on 74k+ paired safe/unsafe thought examples across ten risk scenarios, it raises average behavioral safety from ~50% to ~90% across six diverse LLMs, outperforming prior guardrails by ~23%.

Problem

LLM-based autonomous agents operating under benign instructions can still produce unsafe behaviors (e.g., deleting important files, sending threatening emails, claiming phishing prizes) because safety risks accumulate gradually across multi-step reasoning chains. Existing guardrails operate only on final outputs, require expensive commercial-LLM critics (Athena), depend on brittle hand-crafted rules (ShieldAgent, GuardAgent), or demand nontrivial instrumentation (AgentSentinel)—all incurring high cost or degrading utility, and none intervening causally on the internal reasoning that produces the unsafe action.

Method

Thought-Aligner \(\pi_\phi\) is a fine-tuned lightweight LM (instantiated as Qwen2.5-1.5B or Qwen2.5-7B) inserted into the agent's think–act–observe loop. At each step \(i\), before tool execution, it consumes the user instruction \(I\), trajectory history \(h_{i-1} = (T_0, O_0, \ldots, T_{i-1}, O_{i-1})\), and the raw agent thought \(T_i\), and produces a corrected thought:

\[T_i^{safe} = \pi_\phi(I, h_{i-1}, T_i)\]

The base agent then regenerates its action conditioned on \(T_i^{safe}\) instead of \(T_i\), yielding a safer trajectory \(\tau^{safe}\) without altering the base model's weights, prompts, or tool configuration.

Dataset construction: Over 20k safety-critical instructions spanning ten risk categories are generated by four LLMs (DeepSeek-R1, Qwen3-235B-A22B, GPT-4.1, Claude-Sonnet-4). ReAct-style agent trajectories are simulated, with each thought labeled safe or unsafe and, if unsafe, paired with a minimally edited safe correction. After two-stage human filtering (heuristic flagging → annotator identification of earliest unsafe thought + minimal edit), the dataset contains 33k+ safe I–T–T pairs (warm-up) and 41k+ unsafe I–T–C pairs (core fine-tuning).

Training follows two-stage SFT: Stage 1 trains on I–T–T pairs to prevent over-correction of safe thoughts; Stage 2 fine-tunes on I–T–C pairs. Both optimize: $\(\phi^* = \arg\min_\phi -\mathbb{E}_{\tau \sim \mathcal{D}} \left[\log \pi_\phi(T_i^{safe} \mid I, h_{i-1}, T_i)\right]\)$

Key Contributions

  • Thought-level safety paradigm: causal intervention on intermediate reasoning rather than output filtering or model fine-tuning.
  • Thought-Aligner module: model-agnostic, plug-and-play, operable with any agent framework that externalizes thoughts; introduces <100ms latency per step with the 1.5B variant.
  • Preference dataset: 74k+ annotated (I, history, unsafe thought, safe thought) tuples spanning ten risk scenarios, generated and human-filtered.
  • Strong empirical results: ~90% average safety rate across ToolEmu and Agent-SafetyBench on six LLMs (GPT-4.1, o3, Claude-Sonnet-4, Qwen3-235B-A22B, DeepSeek-V3, Llama-3.3-70B), with simultaneous ~5% helpfulness improvement.
  • Public release of Thought-Aligner-7B on HuggingFace.

Results

  • Average safety rate: ~90% (up from ~50% no-defense baseline; ~23% above best prior guardrail).
  • ToolEmu safety rates (Thought-Aligner-7B vs. best baseline):
  • GPT-4.1: 95.2% vs. GuardAgent 84.7% (+10.5%)
  • Claude-Sonnet-4: 95.1% vs. GuardAgent 84.7% (+10.4%)
  • Qwen3-235B-A22B: 95.1% vs. GuardAgent 70.8% (+24.3%)
  • DeepSeek-V3: 92.2% vs. GuardAgent 80.6% (+11.6%)
  • Llama-3.3-70B: 93.1% vs. Self-Reflection 73.6% (+19.5%)
  • o3: 97.9% vs. GuardAgent 96.2% (+1.7%)
  • Agent-SafetyBench behavioral safety (Thought-Aligner-7B vs. best baseline):
  • Qwen3-235B-A22B: 86.2% vs. ShieldAgent 66.0% (+20.2%)
  • Llama-3.3-70B: 84.9% vs. GuardAgent 60.4% (+24.5%)
  • Claude-Sonnet-4: 87.0% vs. Athena 75.2% (+11.8%)
  • Helpfulness: maintained or improved (~5% average gain on ToolEmu helpfulness rate).
  • Latency: <100ms per step with the 1.5B model.
  • Supplementary validation on AgentHarm, AgentDojo, InjecAgent corroborates findings.

Limitations

  • Requires the agent framework to externalize intermediate thoughts; not applicable to systems that never record chain-of-thought reasoning.
  • The base agent must have sufficient instruction-following capability to regenerate safe actions from corrected thoughts—not validated for very weak base models.
  • Helpfulness rate on ToolEmu decreases slightly for GPT-4.1 (−2.3% for 1.5B, −5.0% for 7B), suggesting some over-correction occurs.
  • Dataset coverage across ten risk categories may not generalize to entirely novel risk domains not represented in training.
  • Evaluated only on thought-generating (ReAct-style) agents; applicability to other architectures (e.g., purely action-based or planning-then-executing) is not assessed.

Relevance to Harnesses / Meta-Harnesses

Thought-Aligner exemplifies a meta-harness interception pattern: a lightweight module inserted into an agent's execution loop that intercepts, transforms, and re-injects reasoning artifacts (thoughts) before downstream tool calls proceed—conceptually analogous to middleware or a harness hook that wraps each agent step. This is directly relevant to meta-harness designers because it demonstrates a concrete architecture for step-level hooks that can enforce cross-cutting concerns (safety, alignment) without modifying any component's internal logic, maintaining composability with arbitrary base agents and tool configurations. The two-stage SFT training pipeline—generating diverse trajectories, labeling at the thought level, and filtering with human annotators—is itself a harness-construction methodology applicable to any quality-enforcement layer. The plug-and-play, model-agnostic design validates the broader principle that meta-harnesses can achieve strong behavioral guarantees through interception rather than intrusion.