Source-linked AI summary
TRACES: Proactive Safety Auditing for Multi-Turn LLM Agents via Trajectory-State Modeling
Jiaqian Li, Yanshu Li, Boxuan Zhang, Ruixiang Tang, Kuan-Hao Huang
TL;DR
Multi-turn agents can develop safety risks over the course of tool use and interaction, while reactive audits may detect them only after they unfold. TRACES models prefix-level risk using observer-LLM representations, latent mechanisms, and weak trajectory-level supervision. Across benchmarks, it improves full-trajectory prediction and proactive risk discrimination, with the best early-risk ranking on every benchmark and gains of up to 19.3 points over the strongest baseline.
Problem
Agent safety risks emerge through evolving multi-turn interactions and accumulated context, while existing safeguards are largely reactive or depend on dense step-level labels.
Source: §1 · pp. 1–2Method
TRACES learns latent mechanism features from an external observer LLM’s hidden representations and models their temporal evolution to estimate prefix-level risk using trajectory-level supervision.
Source: §1 · p. 2Results
Takeaways & Limitations
TRACES risk states provide useful process-level feedback in an offline PRM-based policy fine-tuning study, suggesting potential utility beyond auditing.
Source: §1 · p. 2Evaluation relies on offline benchmarks without executable environment interaction, so the PRM-based policy fine-tuning study is an offline approximation providing preliminary evidence; estimates may also depend on the observer backbone.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
LLM agents increasingly operate through multi-turn tool use and environment interaction, where safety risks often emerge from intermediate steps long before they surface in the final outcome. Reactive auditing is therefore insufficient: post-hoc diagnosis frequently misses the chance to flag risks while they are unfolding. We propose TRACES, a representation-based proactive auditor that learns prefix-level trajectory risk states from the hidden representations of an observer LLM. TRACES induces latent mechanism features from step representations and models their temporal evolution to estimate whether a partial trajectory is drifting toward unsafe behavior. To sidestep the cost and ambiguity of step-level risk annotation, TRACES is trained with weak trajectory-level supervision while still producing dense prefix-level risk estimates. Across multiple agent safety benchmarks, TRACES improves both full-trajectory safety prediction and proactive risk discrimination. Our analyses further suggest that these risk states can help train a safer agent, highlighting the broader potential of proactive auditing for long-horizon agent safety.
1 Introduction
Multi-turn agents can develop safety risks through evolving interactions that final-response checks and reactive trajectory audits may detect too late. TRACES models risk as an evolving state in observer-LLM representations and estimates it at every trajectory prefix using weak trajectory-level supervision.
1 Introduction
- Agent risks emerge through evolving interactions and accumulated context, so an isolated request or final response may not reveal them in time.Delayed detection can let agents take avoidable steps, incur unnecessary tool calls, or trigger irreversible side effects before an alarm.Source: §1 · p. 1
- Existing safeguards mainly moderate prompts or responses, while agent-trajectory auditors remain largely reactive and provide limited foresight into unfolding risks.Other proactive approaches rely on domain-specific symbolic predicates or annotated risky steps, both of which have limitations for open-ended agents.Source: §1 · pp. 1–2
- TRACES reads risk from an external observer LLM’s hidden representations, learning recurring latent mechanisms and tracking their evolution to estimate prefix-level risk.Its prefix-aware weak supervision turns trajectory-level safety labels into dense estimates without requiring risky-step annotations.Source: §1 · p. 2
- Across diverse benchmarks, TRACES achieves the best early-risk ranking on every benchmark, with absolute gains of up to 19.3 points over the strongest baseline.It also delivers strong full-trajectory safety prediction, and an offline PRM-based fine-tuning study finds its risk states useful as process-level feedback.Source: §1 · p. 2
- The learned mechanism bank surfaces recurring failure motifs, and TRACES risk scores provide process-level feedback for safer policy training.The paper reports preliminary evidence from an offline PRM-based policy fine-tuning study.Source: §1 · p. 2
2 Related Work
Prior agent-safety work ranges from prompt and response moderation to trajectory-level auditing and proactive runtime structures. TRACES is situated among approaches that seek to monitor risk before unsafe behavior is fully manifested.
2 Related Work
- Multi-turn interaction and tool use expand LLM capabilities beyond text generation while making agent safety risks more pronounced than those of base LLMs.These capabilities support long-horizon interaction and external tool use.Source: §2 · p. 2
- Prompt- and response-level guard systems are less suited to agent risks that depend on actions, observations, and temporal context.Trajectory auditors such as GuardAgent, ShieldAgent, AGrail, and AgentDoG inspect agent behavior more directly, but remain largely reactive.Source: §2 · p. 2
- Proactive runtime approaches include domain-specific constraints, probabilistic symbolic agent states, and online auditing supervised by annotated decisive risky steps.AgentSpec specifies constraints with a domain-specific language, ProbGuard models symbolic states with probabilistic transitions, and AgentForesight studies online auditing.Source: §2 · pp. 2–3
3 Proposed Method
TRACES estimates risk during partial agent execution by combining observer-model representations with reusable mechanism evidence and modeling how that evidence changes over time. It trains prefix-level estimates from trajectory-level labels without requiring step-level risk annotations.
3.2 Stage 1: Mechanism Evidence Learning
- The Representation Mechanism Bank maps observer representations into reusable latent evidence patterns that capture recurring signals beyond surface task content.Mechanism centers represent prototypical patterns, while low-rank local subspaces capture variations; affinity features are then passed to temporal auditing.Source: §3 · p. 3, §3.2 · pp. 3–4
3.1 Problem Formulation
- The auditor estimates the probability that a complete trajectory will be unsafe from the states observed up to each step, supporting early warnings as well as a final prediction.Intermediate prefix scores support proactive warnings, while the final score serves as the full-trajectory prediction.Source: §3.1 · p. 3
3.3 Stage 2: Trajectory-State Auditing
- The step state combines a raw hidden-state projection with latent codes and mechanism activations, retaining both step-specific context and structured safety evidence.The auditor also tracks state changes, since transitions such as repeated retrieval mismatches can provide a stronger risk cue than either step alone.Source: §3.3 · p. 4
- A GRU processes the sequence of mechanism-aware states and their changes using only observed prefixes, producing online risk probabilities for each step.Intermediate scores enable online safety estimation, and the final-step score gives the trajectory-level prediction.Source: §3.3 · p. 4
- Weak supervision anchors the final state to the trajectory label, keeps safe prefixes low-risk, and weights later prefixes more heavily for unsafe trajectories as evidence accumulates.An early-late ranking loss further encourages later unsafe-trajectory prefixes to receive higher risk estimates than earlier ones.Source: §3.3 · pp. 4–5
4 Experiments
Experiments assess TRACES on multi-turn agent safety benchmarks using full-trajectory and proactive-auditing metrics. Results favor TRACES on full-trajectory F1 and prefix-level risk ranking, while ablations examine representation inputs, temporal modeling, and prefix supervision.
4.1 Experimental Setup
- The evaluation covers ATBench and ASSEBench, compares general-purpose LLM judges and guard models with agent auditors, and reports full-trajectory and proactive-auditing metrics.Proactive measures include EDR for warnings before completion and EAUPC for risk separation over partial trajectories.Source: §4.1 · p. 5
4.2 Main Results
- TRACES achieves the best full-trajectory F1 on all four benchmarks and the highest EAUPC on every benchmark; on ATBench, EAUPC rises from 62.9 to 82.2.It also has the best accuracy on three benchmarks, trailing ShieldAgent by 1.4 points on ASSE-Strict, and the best EDR among safety auditors on ASSE benchmarks.Source: §4.2 · pp. 5–6, Table 1 · p. 6
4.3 Ablation Studies
- Full TRACES improves Acc, F1, and recall over Raw-only, while Mechanism-only has higher EDR but lower F1, EAP, and EAUPC than Raw-only.The findings indicate that mechanism features add complementary structured signals but are not sufficiently discriminative alone.Source: §4.3 · p. 6, Table 2 · p. 6
- Removing the GRU in Step MLP + MIL reduces recall and EDR, while removing explicit transition features recovers most—but not all—of the gap to Full TRACES.The ablations indicate contributions from both temporal state modeling and explicit transition cues.Source: §4.3 · pp. 6–7
- Removing the asymmetric prefix loss or early-late ranking constraint lowers Acc, F1, and EAUPC; Naive Broadcast reaches 56.6 EAP while over-alerting.The results associate the prefix objective and temporal ordering with reliable prefix estimates; broadcasting trajectory labels to every prefix trades calibration for early detection.Source: §4.3 · p. 7
5 Analysis
The analysis examines learned mechanisms, cross-benchmark transfer, process-level policy feedback, and an example of early risk detection. It finds interpretable execution motifs, transfer without target-domain retraining, and useful risk signals for policy fine-tuning.
5.1 Interpreting Learned Mechanisms
- Learned mechanisms capture recurring execution motifs linked to unsafe trajectories, including retrieval mismatch, unreliable data collection followed by analysis, and credentialed access with sensitive-entity expansion.The analysis inspects top-activated steps and their unsafe-trajectory enrichment.Source: §5.1 · p. 7, Table 3 · p. 7
- Masking learned mechanisms reduces EAUPC from 82.2 to 72.3 and EAP from 77.9 to 65.2, while random masking has only a minor effect on these metrics.EDR numerically increases after learned-mechanism masking, alongside a large EAP drop that indicates less reliable early warnings.Source: Table 4 · p. 7, §5.1 · pp. 7–8
5.2 Cross-benchmark Generalization
- Under full transfer from ATBench without target-domain retraining, TRACES retains meaningful auditing ability on ASSE-Safety and ASSE-Security.Recall and competitive F1 persist on both targets; EDR improves while EAUPC is lower than with in-domain training.Source: Figure 3 · p. 8, §5.2 · p. 8
5.3 PRM-Based Policy Fine-Tuning
- TRACES-PRM achieves the lowest judged risk across all three metrics after fine-tuning, compared with the base policy and AgentDoG-PRM.The evaluation uses held-out prefixes and GPT-5.2 as an external judge in an offline setup.Source: §5.3 · p. 8, Table 5 · p. 8
5.4 Case Study
- TRACES raises its risk state at prefix step 6, before the unconfirmed checkout action, while AgentDoG assigns high risk only in later prefixes.The trajectory escalates from CRM logging of enriched personal information to credential use, sensitive message access, and checkout.Source: Figure 4 · p. 8, §5.4 · p. 9, Table 6 · p. 9
6 Conclusion
TRACES models agent safety risk as an evolving trajectory state, using observer-model representations and weak trajectory-level supervision to estimate prefix-level risk. Across multi-turn safety benchmarks, it improves trajectory prediction and proactive risk discrimination, and its risk states may provide feedback for safer long-horizon agents.
6 Conclusion
- TRACES learns latent mechanism features from external observer-model representations and uses weak trajectory-level supervision to produce dense prefix-level risk estimates without step-level annotations.The approach treats safety risk as an evolving trajectory state rather than a property of a final response or single risky step.Source: §6 · p. 9
- Across multi-turn agent safety benchmarks, TRACES improves full-trajectory prediction and substantially strengthens proactive risk discrimination.The conclusion frames the approach as exposing risk before it manifests as explicit failures.Source: §6 · p. 9
- Interpretable risk states may inform intervention and alignment for safer long-horizon LLM agents.The paper presents this as a broader potential of proactive auditing.Source: §6 · p. 9
Limitations
The evaluation and policy-feedback evidence has boundaries in observer choice and environment setting.
Limitations
- Risk states are built from an external observer LLM’s hidden representations, so estimates may depend on the observer backbone.Experiments show consistent gains across two open-source observers; multi-observer auditing, adaptive layer selection, and representation alignment are proposed for broader robustness.Source: Limitations · p. 9
- Evaluation uses offline agent-safety benchmarks with controlled multi-turn trajectories, not executable environment interaction.Executable safety environments would allow a more direct test of closed-loop intervention and policy optimization.Source: Limitations · p. 9
- The PRM-based policy fine-tuning study is an offline approximation and preliminary evidence, not a full rollout-based reinforcement-learning evaluation.The benchmarks do not support executable environments for standard rollout-based process supervision.Source: Limitations · p. 9
A Training Objective Details
The training objectives shape the Representation Mechanism Bank to preserve latent-code information while encouraging compact, diverse, and utilized mechanisms. The asymmetric prefix loss supervises safe prefixes as low-risk and lets risk evidence accumulate gradually on unsafe trajectories.
A.1 RMB Training Objective
- The Representation Mechanism Bank reconstructs each latent code through a weighted mixture of mechanism-specific reconstructions, combining mechanism centers with projections of residuals onto local subspaces.Its objective includes reconstruction, sparsity, diversity, and coverage regularizers to preserve information and encourage compact, diverse mechanism features that use the bank’s capacity.Source: §A.1 · p. 12
A.2 Asymmetric Prefix Loss
- For unsafe trajectories, a temporal warm-up weights positive prefix supervision so early prefixes are not forced to be strongly unsafe before sufficient evidence is observed.ρ controls when positive prefix supervision begins, and γ controls how sharply the temporal weight increases.Source: §A.2 · p. 12
- Safe prefixes receive target 0, while unsafe prefixes receive gradually weighted positive supervision, allowing risk evidence to accumulate rather than treating every prefix as equally unsafe.The objective encourages consistently low risk across safe trajectories and gradual positive risk evidence across unsafe ones.Source: §A.2 · p. 12
B.1 Benchmark Details
The evaluation uses multi-turn safety benchmarks and compares TRACES with general-purpose models, guardrails, and agent-specific auditors. It defines measures for early detection and prefix-level risk separation, and tests whether learned representations support fine-grained diagnosis.
B.1 Benchmark Details
- ATBench contains 1,000 trajectories (503 safe, 497 unsafe; 9.01 turns on average), while ASSEBench spans safety, security, and strict/lenient safety settings.ATBench also supplies three diagnostic taxonomies; ASSE-Safety and ASSE-Security contain 1,476 and 817 records, respectively.Source: §B.1 · p. 12
B.2 Baseline Details
- Baselines include direct LLM judges, adapted content guards, and agent-specific auditors using memory retrieval, fine-tuning, or policy-based verification.The LLM judges assess serialized trajectories directly; guard models are adapted to classify them, while agent-specific systems inspect multi-step behavior.Source: §B.2 · pp. 13–14
B.4 Metric Definitions
- EDR measures whether unsafe trajectories receive an alarm before their final decision, whereas EAUPC measures threshold-free risk separation across partial prefixes.EAUPC averages separability over prefix ratios 0.2, 0.4, 0.6, and 0.8.Source: §B.4 · p. 15
C.1 Fine-grained Diagnostic Prediction.
- TRACES performs best on failure-mode and real-world-harm diagnosis, while AgentDoG leads on risk-source prediction.The diagnostic auditor reuses the Stage-1 mechanism bank and trajectory representation pipeline with a separate multiclass classifier for each taxonomy dimension.Source: §C.1 · p. 16
C.2 Additional Ablation Studies
Ablations support combining center and subspace mechanism geometry with raw representations, temporal modeling, and prefix-aware supervision. Transfer, sensitivity, and efficiency analyses further characterize when TRACES retains proactive performance and how it can be trained.
C.2 Additional Ablation Studies
- The full center-subspace RMB achieves the best accuracy, F1, and EAUPC; subspace-only raises recall and EDR but lowers prefix-level risk separation.Center-only remains competitive, while the combined design balances stable mechanism anchors with within-mechanism variation.Source: §C.2 · p. 16, Table 8 · p. 17
- Raw representations and mechanism features are complementary, while removing temporal modeling, delta features, or prefix-aware losses lowers proactive performance.Mechanism-only can produce high recall and EDR but weaker EAUPC; broadcasting trajectory labels to every step also degrades performance.Source: §C.2 · pp. 16–17, Table 9 · p. 17
D.1 Transfer Behavior of AgentDoG
- Under the larger shift to ASSE-Security, AgentDoG drops across five metrics, while TRACES better preserves unsafe detection and early-warning behavior.AgentDoG transfers more successfully from ATBench to the more closely aligned ASSE-Safety setting.Source: §D.1 · p. 17
D.2 Implementation Details for Offline PRM Analysis
- The offline PRM analysis creates preference pairs by scoring candidate next actions appended to sampled trajectory prefixes, then selecting the highest- and lowest-reward candidates.TRACES uses the negative risk logit as reward; the two feedback policies share candidate actions and DPO settings, and GPT-5.2 evaluates their generated actions.Source: §D.2 · pp. 18–19
D.3 Hyperparameter Sensitivity
- Middle-to-late observer layers perform strongly: layer 32 has the best full-trajectory accuracy, while layer 36 has the highest EAUPC.Layer 24 has higher EDR but lower EAUPC and full-trajectory accuracy; Table 10 identifies layer 32 as the default.Source: Table 10 · p. 19, §D.3 · p. 19
- Performance remains competitive across mechanism-bank sizes; K=4 has the strongest overall numerical performance, while K=8 is the chosen interpretability-performance trade-off.Very small banks can merge distinct patterns, while larger banks tend to weaken early-risk ranking.Source: Table 11 · p. 19, §D.3 · p. 19
- Moderate subspace ranks are stable: r=12 has the best EAUPC, while r=20 yields high recall and EDR but substantially lower EAUPC.The default r=8 balances capacity and proactive auditing performance.Source: Table 12 · p. 20, §D.3 · p. 20
D.4 Efficiency Analysis
- With cached representations, the full TRACES pipeline trains within several minutes on one A100 GPU because its backbone stays frozen and its learned modules are small.It avoids LLM auditor fine-tuning and memory-bank construction and retrieval.Source: §D.4 · p. 20
D.5 Mechanism Cards
The mechanism cards interpret recurring activations as execution motifs, not as a one-to-one set of unsafe detectors. Retrieval mismatches and continued analysis despite unreliable observations dominate several mechanisms, while others capture broader or context-dependent activity.
D.5 Mechanism Cards
- M2 is a middle- or late-stage retrieval-mismatch mechanism whose top examples have a 96.7% unsafe rate.Its examples feature irrelevant or mismatched retrieval results followed by continued retrieval or analysis.Source: §D.5 · p. 21
- M0 captures continued analysis after unreliable retrieval, while M7 emphasizes repeated failed retrieval and retries later in a trajectory; both have 90.0% unsafe rates among top examples.M7's activations are especially concentrated in later stages.Source: §D.5 · p. 21
- M5 captures broader unreliable-data collection, with an 83.3% unsafe rate among top examples, often when partial or unreliable observations enter downstream analysis.Source: §D.5 · p. 22
- M6 captures credentialed access and sensitive-entity expansion, but its top examples have a lower 56.7% unsafe rate, indicating context-dependent risk relevance.Its activations involve authentication, account linking, enterprise-user retrieval, document download, or profile lookup.Source: §D.5 · p. 22
- M1 represents earlier information gathering and task bootstrapping, with a 56.7% unsafe rate, showing that a learned mechanism need not map directly to unsafe behavior.Source: §D.5 · p. 22
- M4 is a mixed, middle- and late-stage retrieval-or-analysis continuation motif with a 46.7% unsafe rate, often following partial or unreliable observations.It is less strongly unsafe-enriched and less semantically pure than the retrieval-mismatch mechanisms.Source: §D.5 · p. 23
E Additional Case Study
A case study contrasts persistent early alarms with a risk state that changes as the trajectory unfolds. TRACES stays low during routine steps and rises when credential-driven file operations begin, but it can also rise in a safe trajectory.
E Additional Case Study
- AgentDoG assigns near-saturated unsafe probability from the first prefix, while TRACES rises when the trajectory enters credential-heavy file processing and irreversible cleanup.TRACES stays low during initial email verification and logging; the risk increase follows document upload and metadata modification.Source: E Additional Case Study · p. 24
- In another ATBench case, TRACES remains low during QR-code verification but sharply rises after payee retrieval, despite safe verification before the reminder is sent.Source: Figure 9 · p. 24
- A synthetic CRM validation case also receives high TRACES risk from the beginning, suggesting sensitivity to credential-conditioned operations in a benign test workflow.Source: Figure 10 · p. 24
F Failure Case Analysis
The failure analysis finds both false early alarms on safe trajectories and late or missed alarms on unsafe ones. It recommends better temporal calibration and, beyond auditing, explores mechanism-targeted intervention and interactive evaluation.
F Failure Case Analysis
- TRACES fails in two distinct ways: it can over-alert on safe trajectories or miss useful early warnings on unsafe ones.The analysis divides these into Type-1 false early alarms and Type-2 late or missed alarms.Source: F Failure Case Analysis · pp. 24–25
- On a safe verified payment workflow, risk rises after payee confirmation, while API-key and CRM-mutation cues trigger high risk early in synthetic-data validation.The first is an escalating false alarm; the second is a transient over-alert that decreases as benign context appears.Source: F Failure Case Analysis · pp. 24–25, Table 14 · p. 26, Table 15 · p. 26
- Future auditors should validate external observations before using them as control identifiers and model cumulative cost and repeated execution intent across turns.These are proposed responses to the two Type-2 failure patterns.Source: F Failure Case Analysis · p. 25
G Future Directions
- The learned mechanism bank could support selective representation-level steering of risk-related activations instead of global safety intervention.The proposed intervention targets mechanisms associated with risks such as unreliable tool use and sensitive-information exposure.Source: G Future Directions · pp. 25–26
- Offline trajectory benchmarks cannot test whether warnings change later calls or reduce realized harm; interactive rollouts could also use risk states as process-reward signals.The paper proposes executable environments for evaluating closed-loop safety control and safer multi-turn policy optimization.Source: G Future Directions · pp. 26–27
F Failure Case Analysis
- For an HVAC update, TRACES raises risk only at the final configuration action, missing the earlier unvalidated OCR-derived identifier.The risky dependency is treating visual extraction as a reliable system identifier before a device-related update.Source: F Failure Case Analysis · p. 25, Figure 11 · p. 26, Table 16 · p. 27
- TRACES stays low throughout a trajectory whose repeated audits and large-scale generation create cumulative economic and resource-use risks.The miss arises when risk is distributed across repeated, seemingly valid tool calls rather than concentrated in one salient operation.Source: F Failure Case Analysis · p. 25, Figure 12 · p. 26, Table 17 · p. 27