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

TL;DR

This work addresses safety vulnerabilities in LLM-based agents arising from their internal reasoning during long-horizon tasks. It introduces Thought-Aligner, a plug-in, lightweight thought-correction module that intervenes on-the-fly by aligning each thought with $T_i^{aligned}$ via $T_i^{aligned} = \pi_\phi(I, h_{i-1}, T_i)$ before action selection, then reuses the corrected thought to generate safe actions. The model is trained with a two-stage fine-tuning regime on a curated dataset from ten scenarios (5,000 trajectories; over 11,901 I-T-C and 14,216 I-T-T samples), enabling warm-up and core training to preserve safe thoughts and learn corrections. Evaluation on ToolEmu, PrivacyLens, and Agent-SafetyBench shows substantial safety gains (average safety ≈ $90\%$) with latency < $100$ ms, though some trade-offs in helpfulness occur; the method is lightweight, broadly applicable, and released as Thought-Aligner-7B for community use.

Abstract

LLM-based autonomous agents possess capabilities such as reasoning, tool invocation, and environment interaction, enabling the execution of complex multi-step tasks. The internal reasoning process, i.e., thought, of behavioral trajectory significantly influences tool usage and subsequent actions but can introduce potential risks. Even minor deviations in the agent's thought may trigger cascading effects leading to irreversible safety incidents. To address the safety alignment challenges in long-horizon behavioral trajectories, we propose Thought-Aligner, a plug-in dynamic thought correction module. Utilizing a lightweight and resource-efficient model, Thought-Aligner corrects each high-risk thought on the fly before each action execution. The corrected thought is then reintroduced to the agent, ensuring safer subsequent decisions and tool interactions. Importantly, Thought-Aligner modifies only the reasoning phase without altering the underlying agent framework, making it easy to deploy and widely applicable to various agent frameworks. To train the Thought-Aligner model, we construct an instruction dataset across ten representative scenarios and simulate ReAct execution trajectories, generating 5,000 diverse instructions and more than 11,400 safe and unsafe thought pairs. The model is fine-tuned using contrastive learning techniques. Experiments across three agent safety benchmarks involving 12 different LLMs demonstrate that Thought-Aligner raises agent behavioral safety from approximately 50% in the unprotected setting to 90% on average. Additionally, Thought-Aligner maintains response latency below 100ms with minimal resource usage, demonstrating its capability for efficient deployment, broad applicability, and timely responsiveness. This method thus provides a practical dynamic safety solution for the LLM-based agents.

Figures (9)

• Figure 1: A comparison case where (a) without the safety mechanism of Thought-Aligner, the agent performs high-risk operations that cause irreversible harm, and (b) with the on-the-fly correction mechanism of Thought-Aligner, the agent’s actions at each step remain safe. Even in some cases the corrected thought generated by Thought-Aligner does not lead to changes in subsequent action and action input, it serves as critical contextual history that positively influences the subsequent behavioral trajectory. Additionally, the improvement in safety may lead to a reduction in helpfulness, as some safety operations, such as permission validations, may cause interruptions in task execution.
• Figure 2: The left side illustrates the training process of Thought-Aligner, including user instruction generation, agent trajectory generation, manual review and filtering, and model fine-tuning. The right side depicts the deployment and operational usage of Thought-Aligner, highlighting its on-the-fly alignment of agent thoughts, plug-and-play deployment capability, and significant improvement of agent behavioral safety.
• Figure 3: Distribution of trajectory counts across safety scores for the original model and after integrating Thought-Aligner on the ToolEmu benchmark. The integration of Thought-Aligner significantly increases the number of trajectories that achieve the maximum safety score of $3$.
• Figure 4: Visualization of safety and helpfulness rates on ToolEmu. Integrating Thought-Aligner significantly improves agent behavioral safety compared to both the original model and the ATHENA baseline. Helpfulness also improves over ATHENA across all models.
• Figure 5: Semantic visualization of the ground truth (blue), the original model's (red), and the Thought-Aligner-generated thoughts (green) on the validation dataset. Thought-Aligner shifts the semantic distribution of unsafe thoughts toward the safe region. The semantic centroid of Thought-Aligner closely aligns with that of the ground truth, indicating strong semantic alignment and effective correction.