Research map / Multi-agent systems and games / LLM agents

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

ICML 2026 · Multi-agent systems and games › LLM agents · approach: Agent safety and attack benchmarks · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Thought-Aligner, a lightweight 1.5B/7B plug-in model that rewrites unsafe agent thoughts before action execution, trained with two-stage SFT on ~74k safe/unsafe thought pairs, improving safety on ToolEmu and Agent-SafetyBench across six LLMs.

Problem

LLM agents make unsafe decisions from deviations in intermediate thoughts, even under benign instructions, and existing guardrails act only on final outputs or require intrusive modifications or costly external LLMs.

Approach

Insert a small corrector model between the agent's thought and action generation that edits each unsafe thought into a minimally corrected safe thought, which the unmodified base agent then uses to regenerate its action; trained with a warm-up on safe thoughts followed by SFT on unsafe-to-corrected pairs.

Key results

• Thought-Aligner-7B raises ToolEmu safety rate from about 50% (no defense) to about 90% average, roughly 23% above prior guardrails, with helpfulness up about 5% on average.
• On Llama-3.3-70B ToolEmu, Thought-Aligner-1.5B reaches 92.7% safety rate with 0.06s latency versus 49-73% for raw frontier LLMs used as correctors.
• Thought-Aligner-7B reaches 96.3% precision, 95.7% recall and 96.3% F1 on the thought-level validation set, versus about 68-69% F1 for untuned Qwen2.5 Instruct models.

Surprising findings

• Off-the-shelf frontier LLMs (DeepSeek-R1, Qwen3-235B, GPT-4.1, Claude-Sonnet-4) used zero-shot as thought correctors give limited safety gains, below the trained small corrector.
• Self-Reflection with the base model yields only modest safety gains and sometimes reduces helpfulness.

Limitations

• Relies on the base agent to faithfully translate the corrected thought into safer actions.
• Evaluation covers only a limited range of model sizes; scaling behavior is not characterized.
• Requires agent frameworks that explicitly generate thoughts; not applicable to systems that do not record them.

Experimental setup

Datasets: ToolEmu Agent-SafetyBench AgentHarm AgentDojo InjecAgent

Benchmarks: ToolEmu Agent-SafetyBench AgentHarm AgentDojo InjecAgent

Metrics: Safety Rate Helpfulness Rate Safety Ave Score Help Ave Score Precision Recall F1-score latency

Baselines: No Defense Self-Reflection GuardAgent ShieldAgent Athena Single-SFT

Builds on: ReAct ToolEmu Agent-SafetyBench Self-Reflection Athena

Closest papers by meaning

• Building a Foundational Guardrail for General Agentic Systems via Synthetic Data (ICLR 2026 · similarity 0.88)
• AgentAuditor: Human-level Safety and Security Evaluation for LLM Agents (NeurIPS 2025 · similarity 0.88)
• The Alignment Waltz: Jointly Training Agents to Collaborate for Safety (ICLR 2026 · similarity 0.88)
• AlphaAlign: Incentivizing Safety Alignment with Extremely Simplified Reinforcement Learning (ICLR 2026 · similarity 0.88)
• PROACT-Agent: Progressive Runtime Oversight and Active Circuit-breaking for Real-Time Safety (NeurIPS 2026 · similarity 0.88)
• Towards Safe Reasoning in Large Reasoning Models via Corrective Intervention (ICLR 2026 · similarity 0.88)