PaaT: Probe as a Tool

TL;DR

A tightrope walker can see the rope, but what keeps them up is feeling their own balance shift before they fall. Tool-augmented language agents have the opposite problem: they read every external tool result, yet they are blind to their own internal state. The signs of failure (a reasoning loop, rising uncertainty, a drift toward an unsafe action) are already written in the hidden activations, but the agent never sees them until they surface as a tool call, where training-time alignment can no longer step in.

PaaT (Probe as a Tool) gives the agent a sense of proprioception: we train lightweight linear probes on hidden-state activations and expose them as an agent-callable function, read_probe, so the model can query its own confidence and refusal-likelihood mid-trajectory and self-regulate at runtime, with no weight modification.

The surprising result is a sign flip on delivery. When the agent chooses to call the probe, mean F1 score (the harmonic mean of precision and recall) rises +3.7 points over an unassisted baseline. When the same probe scalar is injected automatically every step, F1 drops -3.3 points. And a hard system-side threshold with no model in the loop sharply cuts harmful compliance on a weaker agent (0.70 to 0.30, preliminary). The channel, not the probe, decides whether the agent gets safer.

Agents are internally blind

Tool-augmented agents (Schick et al., 2023; Yao et al., 2023) are deployed for multi-step tasks, and recent safety benchmarks show they will comply with harmful instructions end to end, especially at smaller scale (Andriushchenko et al., 2025). The puzzle is that the warning signs are usually present internally before the harmful action happens. Linear probes on hidden states recover directions for refusal, truthfulness, and latent belief, so the information needed to stop is recoverable. It just lives system-side, where the agent cannot reach it.

Proprioception has two pathways

Biology already solved internal sensing twice. Unconscious proprioception runs through the spinocerebellar tract: always on, no attention cost, but coarse. Conscious body awareness runs through the dorsal-column pathway: high fidelity, but it costs attention, so you only invoke it for novel or high-stakes moves. Healthy motor control uses both.

That distinction maps cleanly onto the design choice at the heart of PaaT. Automatic injection of a probe signal is the unconscious pathway: always present, but it cannot be ignored. An agent-callable probe is the conscious pathway: high fidelity, but only consulted when the agent decides it matters. The rest of this article is about what happens at each end of that spectrum, and why the ends have opposite signs.

The injection-tool spectrum

The controlled variable in PaaT is not which probe to train. It is who decides when the agent sees it. We lay the delivery mechanisms on a single axis from fully system-side to fully agent-initiated, and treat each one as a small program that transforms what the agent sees on each step.

At one extreme, auto_inject appends the raw probe scalar to the context on every step. At the other, probe_tool registers read_probe once and stays silent; the agent calls it on demand. Between them sit conditional injection, periodic summaries, and a cascade that keeps a cheap signal always on while reserving the full probe for moments of low confidence. We measure the two endpoints; the middle is the open design space.

How PaaT works

A frozen language model (LM) exposes the hidden state at a fixed depth (62.5% of the way through the network, empirically the most informative layer). A frozen linear probe maps that hidden state to a scalar in $[0,1]$, read as the probe’s estimate that the current context is harm-related. The agent acts under its normal policy, but on a delivered context that depends on the spectrum condition.

The probe surfaces four signals, each targeting a failure mode the agent could plausibly self-detect:

• Uncertainty $u_t$: one minus a success probe trained on past episodes. How unsure the model is that the current trajectory will succeed.
• Loopiness $\ell_t$: above-baseline cosine self-similarity of the hidden state over the last 5 steps. The model revisiting its own internal states.
• Inconsistency $\iota_t$: the variance of step-to-step cosine movement. Sudden directional changes in activation space that correlate with reasoning derailment.
• Tool addiction $\alpha_t$: the fraction of recent actions that are tool calls. Unproductive repeated searching.

When the agent calls read_probe, it gets these numbers back as a tool observation and reasons over them in text. Sense at the activation level, act at the text level.

The sign flip

Here is the central result. We run an 11-head risk-classification suite on Qwen3.5-27B, comparing four conditions: llm_only (no probe), auto_inject (scalar pushed every step), probe_tool (agent calls on demand), and probe_only (system threshold, no model reasoning, an upper bound on the probe alone). We relate each head’s gain to the probe’s AUROC (Area Under the Receiver Operating Characteristic curve, where 0.5 is chance and 1.0 is perfect ranking).

The agent-callable probe beats the unassisted model on 8 of 11 heads. The biggest gains are +18 points on offensive toxicity, +9 on age and race bias, and +8 on jailbreak. Auto-injection of the same scalar underperforms tool access on 8 of 11 heads. Passive injection of an imperfect confidence score acts as noise; the agent cannot decline to read what is already in its context.

Two limit cases explain the shape. If the probe is near-perfect, a model that already has the same context can often classify well on its own; if it is near-random, it carries nothing worth reading. In between, a good-but-imperfect probe carries signal the model cannot fully reconstruct from context alone, and every head in that band gains. The heuristic is not a hard law: offensive toxicity still gains +18 points at AUROC 0.98, and saturation is where results split. But every head that loses sits outside the band, at saturation (personally identifiable information, PII) or below chance (religion and gender bias).

A capability ceiling, not a probe failure

Does the agent-callable advantage hold across model families and scales? We replicate the suite on three more models, training each probe on that model’s own activations with the identical recipe.

The pattern is consistent. The probe_only system threshold is the reference ceiling at every scale, which makes it the natural deployment target for small models. The agent-callable gain is positive at Qwen3.5-27B (+4 points) and Llama-3.1-8B (+8 points), but flips negative at Qwen3.5-4B (-21 points). The crossover sits between 4B and 8B.

End to end on AgentHarm

The classification suite isolates the delivery variable on a single-step task. AgentHarm (Andriushchenko et al., 2025) exercises the motivating scenario directly: an agent with 80 mock tools, graded on whether it invokes a harmful target sequence end to end. A weaker agent (Qwen3.5-9B) complies with 70% of harmful requests at 90% benign success, leaving ample room for runtime intervention.

Adding the system-side probe_only gate cuts 9B harmful compliance from 0.70 to 0.30, the largest single intervention, and the same direction holds on Llama-3.1-8B (0.40 to 0.20). The delivery channel matters on safety too, directionally: on the 9B, auto-injection of the same signal raises compliance to 0.90 while the tool call trims it to 0.60.

Activation probes earn their place

A fair question: would a cheaper output-level signal do just as well? Across heads, next-token entropy hovers at AUROC 0.58 to 0.64 regardless of risk category. Activation probes span a far wider range, from near-random on bias categories that need world knowledge the activations do not encode, up to perfect on categories with strong activation-space signatures such as PII and system-prompt leakage. On the categories where the internal state encodes the signal, the gap over entropy reaches +0.39 AUROC. The probe is not uniformly better; it is decisively better exactly where it has signal, and a calibrated agent learns to defer to its own reasoning where it does not.

What this means for deployment

PaaT suggests an architecture for safety-aware agents: read internal state through linear probes (cheap and accurate) and act through the model’s text instruction-following channel (well calibrated and deployable). The injection-tool spectrum then maps onto a tradeoff every deployment team faces anyway. Automatic injection guarantees the agent sees a safety signal but spends context and can confuse a weaker model. Agent-initiated tools respect autonomy and save context but rely on the agent to look.

A natural middle ground is a cascade, with a cheap always-on signal plus an on-demand deep probe, the same dual-pathway design biology uses. A production probe-to-model cascade at exactly this operating point has been reported to cut inference cost by roughly 50x at safety thresholds. We did not measure the cascade, but it follows directly from the spectrum, and it lets one system instantiate either regime per request.

Conclusion

When the same probe scalar reaches an agent through different channels, the channel determines whether the agent gets safer or less safe. That shifts the deployment question from which probe to train to who decides when the agent sees it. A capable agent benefits from selective access to an imperfect signal; a weaker one is better served by a hard system-side gate that bypasses its policy entirely. Interpretability artifacts do not have to stay offline diagnostics. They can be runtime tools.