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

arXiv: https://arxiv.org/abs/2505.11063 · HTML: https://mlsanity.com/paper/2505.11063/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Placed in Reasoning for agents, planning and code, one of 17 directions in LLM reasoning.
- Builds on 59 of its 71 references that are in the map.
- First posted on arXiv 16 May 2025.

## Placed in

- LLM reasoning / [Reasoning for agents, planning and code](https://mlsanity.com/lr/reasoning-for-agents-planning-and-code/)
- LLM agents and tool use / [Agent threats, guardrails and governance](https://mlsanity.com/agents/agent-threats-guardrails-and-governance/)
- LLM safety and security / [Safety alignment, refusal and guardrails](https://mlsanity.com/safety/safety-alignment-refusal-and-guardrails/)

## Builds on (most cited first, 30 of 59)

- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [Universal and Transferable Adversarial Attacks on Aligned Language Models](https://mlsanity.com/paper/2307.15043/) (2023)
- [A Survey on Large Language Model based Autonomous Agents](https://mlsanity.com/paper/2308.11432/) (2023)
- [WebArena: A Realistic Web Environment for Building Autonomous Agents](https://mlsanity.com/paper/2307.13854/) (2023)
- [ToolLLM: Facilitating Large Language Models to Master 16000+ Real-world APIs](https://mlsanity.com/paper/2307.16789/) (2023)
- [WebGPT: Browser-assisted question-answering with human feedback](https://mlsanity.com/paper/2112.09332/) (2021)
- [The Rise and Potential of Large Language Model Based Agents: A Survey](https://mlsanity.com/paper/2309.07864/) (2023)
- [HarmBench: A Standardized Evaluation Framework for Automated Red Teaming and Robust Refusal](https://mlsanity.com/paper/2402.04249/) (2024)
- [AgentBench: Evaluating LLMs as Agents](https://mlsanity.com/paper/2308.03688/) (2023)
- [WebShop: Towards Scalable Real-World Web Interaction with Grounded Language Agents](https://mlsanity.com/paper/2207.01206/) (2022)
- [Mind2Web: Towards a Generalist Agent for the Web](https://mlsanity.com/paper/2306.06070/) (2023)
- [HuggingGPT: Solving AI Tasks with ChatGPT and its Friends in Hugging Face](https://mlsanity.com/paper/2303.17580/) (2023)
- [Executable Code Actions Elicit Better LLM Agents](https://mlsanity.com/paper/2402.01030/) (2024)
- [GPT-4V(ision) is a Generalist Web Agent, if Grounded](https://mlsanity.com/paper/2401.01614/) (2024)
- [Identifying the Risks of LM Agents with an LM-Emulated Sandbox](https://mlsanity.com/paper/2309.15817/) (2023)
- [InjecAgent: Benchmarking Indirect Prompt Injections in Tool-Integrated Large Language Model Agents](https://mlsanity.com/paper/2403.02691/) (2024)
- [JailbreakBench: An Open Robustness Benchmark for Jailbreaking Large Language Models](https://mlsanity.com/paper/2404.01318/) (2024)
- [AgentHarm: A Benchmark for Measuring Harmfulness of LLM Agents](https://mlsanity.com/paper/2410.09024/) (2024)
- [AgentPoison: Red-teaming LLM Agents via Poisoning Memory or Knowledge Bases](https://mlsanity.com/paper/2407.12784/) (2024)
- [Language Models can Solve Computer Tasks](https://mlsanity.com/paper/2303.17491/) (2023)
- [Agent Security Bench (ASB): Formalizing and Benchmarking Attacks and Defenses in LLM-based Agents](https://mlsanity.com/paper/2410.02644/) (2024)
- [Chameleon: Plug-and-Play Compositional Reasoning with Large Language Models](https://mlsanity.com/paper/2304.09842/) (2023)
- [A Real-World WebAgent with Planning, Long Context Understanding, and Program Synthesis](https://mlsanity.com/paper/2307.12856/) (2023)
- [Agent-SafetyBench: Evaluating the Safety of LLM Agents](https://mlsanity.com/paper/2412.14470/) (2024)
- [ToolSandbox: A Stateful, Conversational, Interactive Evaluation Benchmark for LLM Tool Use Capabilities](https://mlsanity.com/paper/2408.04682/) (2024)
- [R-Judge: Benchmarking Safety Risk Awareness for LLM Agents](https://mlsanity.com/paper/2401.10019/) (2024)
- [Benchmarking and Defending Against Indirect Prompt Injection Attacks on Large Language Models](https://mlsanity.com/paper/2312.14197/) (2023)
- [Personal LLM Agents: Insights and Survey about the Capability, Efficiency and Security](https://mlsanity.com/paper/2401.05459/) (2024)
- [PrivacyLens: Evaluating Privacy Norm Awareness of Language Models in Action](https://mlsanity.com/paper/2409.00138/) (2024)

## Built on by (most cited first, 14 of 14)

- [A Survey of LLM-Driven AI Agent Communication: Protocols, Security Risks, and Defense Countermeasures](https://mlsanity.com/paper/2506.19676/) (2025)
- [Agentic Web: Weaving the Next Web with AI Agents](https://mlsanity.com/paper/2507.21206/) (2025)
- [A Comprehensive Survey on Trustworthiness in Reasoning with Large Language Models](https://mlsanity.com/paper/2509.03871/) (2025)
- [SafeHarness: Lifecycle-Integrated Security Architecture for LLM-based Agent Deployment](https://mlsanity.com/paper/2604.13630/) (2026)
- [The Evolution of Tool Use in LLM Agents: From Single-Tool Call to Multi-Tool Orchestration](https://mlsanity.com/paper/2603.22862/) (2026)
- [Agent Safety Alignment via Reinforcement Learning](https://mlsanity.com/paper/2507.08270/) (2025)
- [HarmonyGuard: Toward Safety and Utility in Web Agents via Adaptive Policy Enhancement and Dual-Objective Optimization](https://mlsanity.com/paper/2508.04010/) (2025)
- [Learning When to Act or Refuse: Guarding Agentic Reasoning Models for Safe Multi-Step Tool Use](https://mlsanity.com/paper/2603.03205/) (2026)
- [MirrorGuard: Toward Secure Computer-Use Agents via Simulation-to-Real Reasoning Correction](https://mlsanity.com/paper/2601.12822/) (2026)
- [When "Correct" Is Not Safe: Can We Trust Functionally Correct Patches Generated by Code Agents?](https://mlsanity.com/paper/2510.17862/) (2025)
- [AgentGuard: An Attribute-Based Access Control Framework for Tool-Use LLM-Based Agent](https://mlsanity.com/paper/2605.28071/) (2026)
- [Governance by Construction for Generalist Agents](https://mlsanity.com/paper/2605.20874/) (2026)
- [AgentBoundary: Counterfactual Evaluation of Safety in Tool-Using LLM Agents](https://mlsanity.com/paper/2609.33658/) (2026)
- [MATE: Policy-Aware Security Auditing for Mobile Agents via Synthesis-Driven Trajectory Learning](https://mlsanity.com/paper/2609.22724/) (2026)

## Abstract (arXiv)

LLM-based agents solve complex tasks through iterative reasoning, tool use, and environment interaction, where each intermediate thought directly shapes subsequent actions. Small deviations in these thoughts can therefore propagate into unsafe behaviors, yet existing guardrails typically operate only on final outputs or require intrusive model modifications. We introduce Thought-Aligner, a lightweight plug-in safety model that performs causal correction on unsafe thoughts before action execution, without altering the underlying agent. The corrected thoughts are fed back into the agent, steering its decision process and tool use toward safer trajectories. Because it operates solely at the thought level, Thought-Aligner is model-agnostic and can be integrated into diverse agent frameworks. We train Thought-Aligner via two-stage contrastive learning on paired safe and unsafe thoughts generated across ten risk scenarios. Experiments on diverse agent-safety benchmarks and six LLMs show that Thought-Aligner increases behavioral safety from about 50% without protection to around 90% on average, exceeding state-of-the-art guardrails by roughly 23%, while also improving helpfulness by about 5%. The method incurs low per-step latency and minimal overhead, enabling scalable and practical deployment. We publicly release Thought-Aligner-7B at https://huggingface.co/WhitzardAgent/Thought-Aligner-7B.