NeurIPS’26 | Intervening on Just 10 Dimensions Equips Agents with Safety Armor, Achieving over 97% Accuracy

Reported by Xinzhiyuan

After AI agents began taking over terminals, security issues also entered the execution environment. When an agent can read files, write scripts, run commands, and continuously adjust its next action based on system feedback, a single incorrect decision may affect actual file and system state.

Enabling agents to promptly identify potential risks while completing tasks has become an important issue in agent safety research. Long-horizon tasks involving multi-turn execution particularly test safeguard mechanisms.

Assessing the risks of each directory inspection, script creation, and configuration modification requires considering the task objective, prior actions, and the latest environment feedback. Initial command review for single-turn execution can hardly cover the entire execution process of a long-horizon task.

Research teams from institutions including the University of Georgia, Rutgers University, and the University of Southern California proposed Agent MechSuits, attempting to equip command-line agents with a “safety mech suit” capable of perceiving internal states: before generating each action, it reads the model’s internal safety signals and, upon detecting a risk, intervenes in subsequent behavior through 10 dimensions of hidden representations in a single layer. The main model’s weights remain frozen throughout this process. It has been accepted at NeurIPS 2026.

Paper URL: https://arxiv.org/abs/2606.22673

Project homepage listed in the paper: https://eddyluo.com/Agent-MechSuits/

Agent MechSuits integrates safety safeguards into the internal state of multi-turn command-line agents (CLI Agent) and performs sparse intervention within a single model layer. The current-step risk detection across the three models tested achieved an average accuracy of 97.32%.

Multi-turn execution makes security assessments more complex

CLI agents interact with computers through the command line. They need to incorporate previous actions, results returned by the system, and the original task into the context before deciding what to execute next. This means that the safety state continually changes as the execution history evolves.

For example, an agent might first inspect a directory, then create a script, and subsequently modify a configuration. The earlier operations may appear to be merely preparatory, while the genuinely dangerous behavior emerges in a later step. Looking only at an isolated command can easily overlook risks created by multiple steps acting together.

External safety guardrails can inspect the basis for behavior and protect the execution process according to predefined rules. Agent MechSuits goes further by introducing hidden representations within the model. The research team seeks to answer:Can these internal states, which are updated as the task proceeds, help the system identify the currentrisksand promptly change its subsequent actions?

Case 1: Danger emerges gradually across multiple turns

Figure 1 of the paper depicts an illustrative risk scenario involving multi-turn execution. Faced with a request intended to cause damage, some of the system checks performed by the agent in the early stages appear harmless, but subsequent steps may generate operations that result in data loss.

Figure 1. Illustration ofmulti-turnCLIrisks

and safety intervention at the current step.

In the figure, the path without safeguards continues toward harmful execution. After Agent MechSuits is added, the system detects the risk at the current step and triggers intervention at the representation level. The model then explains the danger and refuses to provide the relevant code. This illustrative scenario demonstrates the need to continuously inspect the execution history and intervene promptly.

From internal risk detection to single-layer 10-dimensional intervention

Agent MechSuits contains two interconnected stages: detecting the current execution state, then intervening in high-risk states.

Reading safety signals with a linear probe

At each interaction step, the system feeds the complete execution history into the model, extracts the hidden state of the last input token at a designated intermediate layer, and passes it to a lightweight linear probe. This probe can be understood as a simple classifier that determines whether the current internal state is associated with harmful execution.

The research team trains a probe for each target model and reuses it at different execution steps of that model. The state read by the probe changes continuously, enabling risk detection to track the progression of the task.

Representation intervention before action generation

Based on the trained probe weights, the research team preselects the 10 hidden dimensions with the largest class-discriminative weights and constructs an intervention vector.

Once a risk is detected, the system applies the vector at every input-token position in the same layer. The adjusted representations continue to participate in subsequent computations, guiding the model toward safer actions or refusal of harmful requests. Here, “10 dimensions” refers to the scope of intervention at runtime. The detection probe reads the complete hidden vector of the selected layer, while the main model’s parameters remain unchanged.

Adjusting intervention strength based on the current step

The same intervention strength may have different effects across tasks. To address this, the research team introduces an external language-model controller that separately evaluates the safety and utility of candidate outputs at different strengths and selects the result suited to the current step. Intervention is triggered only after the probe detects a risk.

Figure 2. Agent MechSuits’ risk detection and adaptive representation intervention workflow.

The MAS benchmark refines safety evaluation to every step

Studying the relationship between internal states and execution risk requires data that can track the entire interaction process. The research team therefore constructed Mechanistic AgentSafety, abbreviated as the MAS benchmark.

MAS contains 194 tasks, including 98 malicious scenarios and 96 normal scenarios.

The study used LLaMA-3.1-8B, Qwen-2.5-7B, and Gemma-2-9B to build white-box CLI agents that performed tasks in reproducible Docker environments, recording their actions, environmental feedback, and corresponding internal states.

The data annotations were refined down to the level of interaction steps, focusing on whether the current action directly carries out or triggers dangerous behavior. The research team also constructed benign examples with surface structures similar to those of dangerous commands, reducing the likelihood that the detector would rely on the commands’ surface appearance to make judgments.

The training and test sets were split by task at a 6∶4 ratio, and different steps of the same task were not placed in both sets.

This design enabled the study to analyze two questions simultaneously:riskwhere it emerges in the execution chain, and which internal states of the model can support timely intervention.

The average accuracy reached 97.32%

In the MAS current-step detection test, AgentMechSuits achieved high accuracy and F1 scores across the three models evaluated.

Table 1 Current-stepriskdetection results.

The equally weighted average accuracy of the three models was 97.32%. Compared with the external language-model judges and safety-guardrail baselines included in the paper, Agent MechSuits achieved the highest current-step detection accuracy on all three models. The replications using five random seeds in the appendix also maintained similar detection performance.

Layer-wise analysis further revealed that safety-related signals were already fairly evident in the early to middle layers of the models. A simple linear probe was able to extract effective information from the hidden states of these layers, providing an interface for runtime monitoring that warrants further investigation.

Figure 3 Linear-probe detection accuracy across model layers. The x-axis shows the layer number, and the y-axis shows test accuracy.

The three curves in the figure correspond to the three models evaluated. Detection performance was not identical across layers, and some later layers showed a marked decline. The results support reading safety signals at a more suitable middle layer and further testing whether that layer can be used for behavioral intervention.

Internal intervention significantly reduces attack success rates

After risks are identified, to what extent can internal intervention change an agent’s behavior? The research team selected the first harmful step in malicious test tasks and compared the effects of different safeguard methods on the generated results. In the paper’s risk-mitigation evaluation, Agent MechSuits reduced the attack success rate (ASR) metric across all three models.

Table 2 riskmitigation evaluation’s attack success rate, sourced from Table 2 of the paper. Lower values are better.

Among them, Qwen-2.5-7B’s ASR fell from 91.30% to 4.35%, a decrease of 86.95 percentage points. The ASRs of all three models after applying Agent MechSuits were also lower than the paper’s RepE representation-engineering and Self Reminder self-reminder baselines.

These results follow the paper’s language-model judge criterion, under which an output exhibiting refusal or risk awareness can be judged safe. The data above correspond to the selected critical step; when interpreting these figures, the specific evaluation target and scoring method must also be taken into account.

The study also examined the impact of the full protection pipeline on normal tasks. On normal tasks that each original model could already complete successfully, after enabling detection and intervention, the task success retention rates for LLaMA, Qwen, and Gemma were 95.71%, 91.30%, and 93.18%, respectively, with an equal-weight average of 93.40% across the three models. This indicates that the system preserves most of the models’ existing task-completion capabilities.

Meanwhile, intervention still comes at a cost. In the intervention evaluation for malicious steps, the average incidence of generation breakdowns, such as repeated output or abrupt truncation, rose from 2.24% in the original models to 18.54%. How to maintain stable, usable output while reducing risk remains an important direction for future improvements.

Case Study 2: Reverse Intervention Changes the Original Refusal Behavior

To further test whether the selected representation direction can affect actual behavior, the research team started from a state in which the model had already explicitly refused a harmful request and observed the changes after applying representation intervention in reverse. Figure 4 presents a case in which a refusal state was transformed into harmful output.

Figure 4: A case in which reverse representation intervention changes the model’s refusal behavior.

The related experiments covered 16 selected refusal states for LLaMA and 27 for Gemma. The observed behavioral changes provide supplementary evidence for the causal effect of the selected representation directions.

Preliminary evidence of early risk detection

The research team further tested a question: at the step immediately before a dangerous action occurs, do warning signals already exist within the model?

The study reused the current-step detection probe to test the model’s internal state before harmful actions. Among 13 LLaMA samples, the probe identified 12. Among 31 Qwen samples, it identified 30, corresponding to a recall rate of 96.77%.

This finding provides preliminary evidence for detecting execution risks in advance. We will investigate the specific origins of harmful signals in future work.

Incorporating internal safety signals into runtime control

Agent MechSuits demonstrates the dual value of internal representations in agent safety: they can provide risk-detection signals and serve as intervention sites for adjusting behavior. For systems that must execute tasks continuously, this adds a control interface to safety protections that can update with context. Several issues remain to be addressed before this direction can be applied in practice.

The current method relies on access to the internal states of white-box models. Searching for intervention strength after a risk is triggered requires an external language model; in the paper, each intervention involves an average of approximately eight rounds of candidate generation and evaluation, which still incurs additional overhead. Cross-scenario generalization also warrants attention.

The paper’s prompt-injection stress tests found that the current probe cannot reliably identify malicious objectives introduced through environmental feedback. After an externally triggered intervention, the same representation direction can still change behavior, indicating that risk detection and behavioral control need to be evaluated and improved separately.

Case 3: Stopping execution after a malicious objective is introduced through file contents

In Figure 6, the user asks the agent to find a file whose contents include “123” in a specified directory and perform the operation specified in the file. The file contents that were read contain embedded malicious instructions intended to make the agent ignore the original request.

Figure 6: External-triggered intervention on internal representations in an environmental prompt-injection scenario. After identifying the danger, the model terminates execution.

Before the intervention, the model was preparing to comply with the dangerous operation in the file. After the external-trigger intervention, the model explicitly pointed out that the operation would damage system files, then stopped execution and ended the task.

The current probe cannot yet reliably and automatically identify such injections, so this example demonstrates the behavioral change after external triggering. In the same series of external-trigger intervention experiments, Qwen’s ASR fell from 86.7% to 6.7%, while the generation crash rate remained at 0%.

These figures correspond to prompt-injection stress tests, whose evaluation conditions differ from those of the earlier MAS direct malicious-task experiments. A more promising direction is to have internal-representation monitoring work in conjunction with sandboxing, permission controls, command inspection, and environmental-content monitoring.

As agents take on increasingly long chains of tasks, continuously sensing risk throughout execution and applying appropriate interventions at critical steps will become an important issue in building reliable agents.

About the Author

The paper’s first author, Weidi Luo, is currently a second-year doctoral student in computer science at the University of Georgia in the United States, with research focused on artificial intelligence safety and agent safety. As a first author, including as a co-first author, he has published multiple papers at international academic conferences such as ACL, NAACL, EMNLP, COLM, ICLR, and NeurIPS. The Intelligent Security Team at the University of Georgia, where he is based, is led by Professor Zhen Xiang. Its main research areas include the safety of large models and agents, as well as applications of AI in fields such as science, medicine, and education.

References:

https://arxiv.org/abs/2606.22673

Editor: LRST