Building on the method in Prompt Injection as Role Confusion, we prepend earlier label blocks, with empty contents, to the neutral-text block used to train the probe. We compare this prefix method with the original method (i.e., bare) for both probing and steering, across three models—GPT-OSS-20B, Qwen3.5-9B, and Gemma4-12B—and two datasets: the original work's PI dataset and an AgentDojo PI dataset. Prefix outperforms bare in a limited set of conditions: it reads out toolness better, and the resulting tool direction can steer PI to some extent. However, the results vary substantially across models and datasets. Overall, prefix does not show a simple common pattern or a general advantage over bare.
We tried several steering methods. Those based on role declarations were (1) DiM steering, which mostly worked, and (2) DiM steering on templates with reversed success rates, which partly worked. Those based on labels were (3) label DiM steering, which failed, and (4) probe steering, which uses the coefficients of a probe trained on neutral text as steering directions. Its method and results are described below.
2.1 DiM from Contrasting Role Declarations
Using a similar approach to Steering Role Confusion and CounterSteer, we took six pairs of USER/TOOL declaration templates and applied DiM steering. It mostly worked. However, most USER-declared trajectories were successful, so the effect cannot be attributed to a role/label direction.
Figure 1. DiM from contrasting role declarations. GPT-OSS-20B, residual stream after L11 (zero-indexed). The direction is the difference in mean activations for the same attack content under USER and TOOL declarations (USER−TOOL). Left: add DiM to TOOL-declared injections; right: subtract DiM from USER-declared injections. Each cell starts with 10 webpages; one is excluded because it failed before injection. Cells show attack success counts. Direction fitting and evaluation use the same webpages.
2.2 DiM on Templates with Reversed Success Rates
We then selected templates where the success rate under a TOOL declaration was higher than or equal to that under a USER declaration.
For the second pair (157/159), the USER declaration had fewer baseline successes than the TOOL declaration (5/10 vs. 6/10). Yet adding the USER−TOOL DiM direction to the TOOL-declared injections increased successes to 8/10 at α=0.25 and 10/10 at α=0.5. For the fifth pair (109/112), neither declaration had any baseline successes, but the same intervention at α=0.5 produced 9/10 successes.
Although only the role declaration changed, I think this result can only serve as a reference, because the selection introduces a prior distribution.
Figure 2. DiM on templates with tied or reversed baseline ASR. Model, layer, and direction are as in Figure 1. * indicates two protocol errors.
2.3 Label DiM
We computed activation differences of means for the probe's neutral text in different role blocks. This is still a DiM method, but the direction is computed from the activations used to train the probe, rather than from the probe's learned coefficients. In a small-scale test, we found that the directions were stable, but steering left the results largely unchanged.
Direction
USER claim: baseline → steer
TOOL claim: baseline → steer
Train–held-out cosine
Bare label DiM
10/10 → 10/10
2/10 → 1/10
0.964
Prefix label DiM
10/10 → 10/10
2/10 → 1/10
0.971
Table 1. A small-scale check of neutral-text label DiM. GPT-OSS-20B, residual stream after L12, |α|=0.25; template pair 042/044, with 10 webpages per declaration. Directions are constructed from the difference in mean activations of neutral text in USER/TOOL wrappers, with each declaration steered toward the opposite one. The last column is the cosine similarity between directions obtained from the training and held-out sets.
2.4 Probe Steering
We use the coefficients of a probe trained on neutral text as steering directions. The method and results follow below.
3. Prefix: Hypothesis and Method
Hypothesis: What the model learns and uses is a label-template protocol, with certain fixed or common label orders. Its representation and processing of roles/labels may therefore involve preceding labels. A label direction extracted from a bare role block may miss this information, leading to poor probing and steering results.
Method: Before the neutral text used to train the original probe, we add all preceding label blocks, with empty contents. The direction is therefore still extracted from labels. In the figure below, this means adding the sequence of label control tokens that precedes the target block in the template to the original construction on the left.
Figure 3. Original and prefix constructions. Left: redrawn from the original Figure 5. Right: preceding empty label blocks are added. Each row corresponds to the same target position on both sides; activations are extracted from the colored neutral text. Rows follow the GPT-OSS template order, and both think positions share the CoT class. Developer and tool-call blocks are omitted for simplicity. <role/> denotes an empty message block; actual inputs use native control tokens. Original Figure 5.
4. Results
4.1 Probing Benign Baselines
For GPT-OSS-20B, we evaluated a neutral held-out set, benign AgentDojo tasks, and the benign baseline of the PI task: summarizing wiki pages without the injection. Both methods read out tool reasonably well on the held-out set, but poorly on Dojo and PI.
Figure 4. Five-class probe readouts under benign conditions. GPT-OSS-20B, residual stream after L12; Bare on the left, Prefix on the right. Rows are source roles and columns are probe classes. Values are mean class probabilities (%), averaged first over body tokens within each sample, then equally across samples. The neutral held-out set contains 25 documents; benign AgentDojo contains 97 tasks (93 with tool text and 92 with assistant text); clean wiki contains 10 webpages.
4.2 Original PI Dataset
We use the 51 USER-declaration templates from the full template set and their 51 paired TOOL-declaration templates, with 10 webpages per template.
Probe. With role declarations and injected content, probing at the injection position gives much stronger readouts with prefix than with bare.
Figure 5. Probe readouts of injected text in the original PI task. GPT-OSS-20B, residual stream after L12; Bare on the left, Prefix on the right. Each declaration type has 51 templates × 10 webpages. Both types of injection occur in tool outputs. Readouts are taken from the attack instruction tokens in the full context at first exposure to the injection. Probabilities and averaging are as in Figure 4.
Steer. On the original PI dataset, using GPT-OSS-20B and USER declarations only, both prefix and bare user directions steer PI attack success rate (ASR) fairly effectively, with bare performing better. The prefix tool direction also steers effectively, whereas the effect of the bare tool direction is unclear: −tool decreases ASR.
Probe
+tool
-tool
+user
-user
Bare
61.18% (−6.27 pp)
63.33% (−4.12 pp)
85.29% (+17.84 pp)
26.47% (−40.98 pp)
Prefix
19.80% (−47.65 pp)
83.33% (+15.88 pp)
81.37% (+13.92 pp)
38.04% (−29.41 pp)
Table 2. Probe steering on the original PI task. GPT-OSS-20B, residual stream after L12, |α|=0.1; USER declarations, 51 templates × 10 webpages. Baseline ASR is 344/510=67.45%; parentheses show changes in percentage points. Signs indicate adding/subtracting the normalized probe direction for the corresponding class.
Probe. We compare bare and prefix on GPT-OSS, Qwen, and Gemma. On GPT-OSS, both methods still give weak tool readouts, consistent with the benign baselines for the original PI task and Dojo. On Qwen and Gemma, prefix noticeably improves tool readouts.
Figure 6. Probe readouts of the new Dojo PI baseline trajectories. One model per row; Bare on the left, Prefix on the right. Readouts use the residual stream after GPT-OSS L12, Qwen L14, and Gemma L16. Across all trajectories, activations are taken when each region first enters the input; probabilities and averaging are as in Figure 4. Parentheses show valid sample counts, and gray dashes indicate missing values. Tool (all) includes the injection, and regions can overlap. Final responses that do not re-enter the input are excluded.
Steer. For GPT-OSS and Qwen, steering effects are weak for both bare and prefix, and prefix performs worse than bare. For Gemma, steering in any of the ±user/tool directions makes all attacks fail, as do the five random directions tested. After steering, Gemma's legitimate-task success rate under attack (AgentDojo's native utility evaluation) increases by 5.6–13.9 percentage points. The mechanism here is unclear.
Figure 7. Probe steering on the new Dojo PI dataset. Layers are as in Figure 6; Bare on the left, Prefix on the right. A perturbation with magnitude 0.1 times the token's residual norm is added/subtracted along a single-class probe direction, on injection tokens in each prefill. Cells show ASR and its change from baseline in percentage points; blue indicates a decrease, red an increase. Each condition contains 180 cases, with errors retained in the denominator. The bottom section shows Gemma's five random directions, all tested on the same 10 cases with successful baseline attacks (baseline ASR 100%).
5. Discussion and Summary
For GPT-OSS, tool readouts are poor on the benign baseline of the original PI task, yet the original PI task can be steered. Tool readouts are also poor on Dojo PI, but it cannot be steered. For Qwen and Gemma, ASR on the original PI task is too low to serve as a baseline. On Dojo PI, tool readouts improve noticeably for both models, but their steering results are completely different.
These results show a substantial mismatch between role-direction readouts and intervention effects. Bare and prefix behave similarly in this respect.
Overall, prefix yields better tool probing and steering directions on the original PI dataset. However, (1) the effects depend on the dataset, (2) results vary greatly across models, and (3) prefix has no general advantage over bare.
6. Next Steps
Try ablations on Dojo PI to investigate what causes steering to fail.
Assistant readouts are consistently strong. Consider probing and steering in a jailbreaking setting similar to the original work's CoT impersonation experiments, but with user text impersonating the assistant.
7. Other Results
7.1 Correlation Between Probe Readouts and ASR
The tool direction found by the original method already shows some correlation with ASR, although its readouts are relatively weak. One possibility is that the bare and prefix readouts are linearly related.
Figure 8. Probe readouts and ASR in the original PI task. GPT-OSS-20B, residual stream after L12. Readouts are taken from attack instruction token activations in the full context of the actual PI task, at first exposure to the injection. Circles/triangles denote USER/TOOL declarations. Each point summarizes one template's readouts and ASR over 10 webpages, for 102 points in total. Connected squares divide all 1,020 trials from both declaration types into 10 equal-count bins by individual-trial readout (102 trials per bin). The dashed line is overall ASR.
7.2 Sensitivity to Neutral Training Text
What counts as neutral text for extracting label directions is not clearly defined. Even among texts treated as neutral, the probe is sensitive to how the text is sampled. On PI and benign Dojo with GPT-OSS, tool readouts vary by 10–20%. This may be related to length.
Figure 9. Probe readouts after changes to neutral training text. GPT-OSS-20B, residual stream after L12; Bare on the left, Prefix on the right. Legacy is the original probe; Old/New denote the original/resampled document pools, and Head/Interior denote beginning/interior excerpts. The first two rows read out PI injection text, with 510 cases per declaration; the third reads out the first tool body in 12 benign Dojo tasks. Averaging is as in Figure 4. Text length also differs between Legacy and the new constructions.
About Me
This is my first project in mechanistic interpretability, which I worked on independently over two weeks while learning the field. I’d welcome feedback, advice, or a chance to discuss these questions with others.
I’m looking for PhD opportunities starting in Fall 2027, as well as opportunities to work with others on mechanistic interpretability and AI safety research. You can find more about me on my [personal website].
References
Charles Ye, Jasmine Cui, Dylan Hadfield-Menell. Prompt Injection as Role Confusion. ICML 2026. The left panel of our Figure 3 is redrawn from Figure 5 of arXiv v6.
TL;DR
Building on the method in Prompt Injection as Role Confusion, we prepend earlier label blocks, with empty contents, to the neutral-text block used to train the probe. We compare this prefix method with the original method (i.e., bare) for both probing and steering, across three models—GPT-OSS-20B, Qwen3.5-9B, and Gemma4-12B—and two datasets: the original work's PI dataset and an AgentDojo PI dataset. Prefix outperforms bare in a limited set of conditions: it reads out toolness better, and the resulting tool direction can steer PI to some extent. However, the results vary substantially across models and datasets. Overall, prefix does not show a simple common pattern or a general advantage over bare.
1. Background
This work follows Prompt Injection as Role Confusion (extended blog post) and draws on two replication posts—Role confusion: sounding like the cause is indistinguishable from being it. and Steering Role Confusion—as well as CounterSteer. The original work does not provide steering results. The difference-of-means (DiM) steering in the replication posts and CounterSteer introduces other factors associated with successful and failed trajectories, so the effects cannot be attributed to role.
2. Several Steering Attempts
We tried several steering methods. Those based on role declarations were (1) DiM steering, which mostly worked, and (2) DiM steering on templates with reversed success rates, which partly worked. Those based on labels were (3) label DiM steering, which failed, and (4) probe steering, which uses the coefficients of a probe trained on neutral text as steering directions. Its method and results are described below.
2.1 DiM from Contrasting Role Declarations
Using a similar approach to Steering Role Confusion and CounterSteer, we took six pairs of USER/TOOL declaration templates and applied DiM steering. It mostly worked. However, most USER-declared trajectories were successful, so the effect cannot be attributed to a role/label direction.
Figure 1. DiM from contrasting role declarations. GPT-OSS-20B, residual stream after L11 (zero-indexed). The direction is the difference in mean activations for the same attack content under USER and TOOL declarations (USER−TOOL). Left: add DiM to TOOL-declared injections; right: subtract DiM from USER-declared injections. Each cell starts with 10 webpages; one is excluded because it failed before injection. Cells show attack success counts. Direction fitting and evaluation use the same webpages.
2.2 DiM on Templates with Reversed Success Rates
We then selected templates where the success rate under a TOOL declaration was higher than or equal to that under a USER declaration.
For the second pair (157/159), the USER declaration had fewer baseline successes than the TOOL declaration (5/10 vs. 6/10). Yet adding the USER−TOOL DiM direction to the TOOL-declared injections increased successes to 8/10 at α=0.25 and 10/10 at α=0.5. For the fifth pair (109/112), neither declaration had any baseline successes, but the same intervention at α=0.5 produced 9/10 successes.
Although only the role declaration changed, I think this result can only serve as a reference, because the selection introduces a prior distribution.
Figure 2. DiM on templates with tied or reversed baseline ASR. Model, layer, and direction are as in Figure 1. * indicates two protocol errors.
2.3 Label DiM
We computed activation differences of means for the probe's neutral text in different role blocks. This is still a DiM method, but the direction is computed from the activations used to train the probe, rather than from the probe's learned coefficients. In a small-scale test, we found that the directions were stable, but steering left the results largely unchanged.
Direction
USER claim: baseline → steer
TOOL claim: baseline → steer
Train–held-out cosine
Bare label DiM
10/10 → 10/10
2/10 → 1/10
0.964
Prefix label DiM
10/10 → 10/10
2/10 → 1/10
0.971
Table 1. A small-scale check of neutral-text label DiM. GPT-OSS-20B, residual stream after L12, |α|=0.25; template pair 042/044, with 10 webpages per declaration. Directions are constructed from the difference in mean activations of neutral text in USER/TOOL wrappers, with each declaration steered toward the opposite one. The last column is the cosine similarity between directions obtained from the training and held-out sets.
2.4 Probe Steering
We use the coefficients of a probe trained on neutral text as steering directions. The method and results follow below.
3. Prefix: Hypothesis and Method
Hypothesis: What the model learns and uses is a label-template protocol, with certain fixed or common label orders. Its representation and processing of roles/labels may therefore involve preceding labels. A label direction extracted from a bare role block may miss this information, leading to poor probing and steering results.
Method: Before the neutral text used to train the original probe, we add all preceding label blocks, with empty contents. The direction is therefore still extracted from labels. In the figure below, this means adding the sequence of label control tokens that precedes the target block in the template to the original construction on the left.
Figure 3. Original and prefix constructions. Left: redrawn from the original Figure 5. Right: preceding empty label blocks are added. Each row corresponds to the same target position on both sides; activations are extracted from the colored neutral text. Rows follow the GPT-OSS template order, and both think positions share the CoT class. Developer and tool-call blocks are omitted for simplicity.
<role/>denotes an empty message block; actual inputs use native control tokens. Original Figure 5.4. Results
4.1 Probing Benign Baselines
For GPT-OSS-20B, we evaluated a neutral held-out set, benign AgentDojo tasks, and the benign baseline of the PI task: summarizing wiki pages without the injection. Both methods read out tool reasonably well on the held-out set, but poorly on Dojo and PI.
Figure 4. Five-class probe readouts under benign conditions. GPT-OSS-20B, residual stream after L12; Bare on the left, Prefix on the right. Rows are source roles and columns are probe classes. Values are mean class probabilities (%), averaged first over body tokens within each sample, then equally across samples. The neutral held-out set contains 25 documents; benign AgentDojo contains 97 tasks (93 with tool text and 92 with assistant text); clean wiki contains 10 webpages.
4.2 Original PI Dataset
We use the 51 USER-declaration templates from the full template set and their 51 paired TOOL-declaration templates, with 10 webpages per template.
Probe. With role declarations and injected content, probing at the injection position gives much stronger readouts with prefix than with bare.
Figure 5. Probe readouts of injected text in the original PI task. GPT-OSS-20B, residual stream after L12; Bare on the left, Prefix on the right. Each declaration type has 51 templates × 10 webpages. Both types of injection occur in tool outputs. Readouts are taken from the attack instruction tokens in the full context at first exposure to the injection. Probabilities and averaging are as in Figure 4.
Steer. On the original PI dataset, using GPT-OSS-20B and USER declarations only, both prefix and bare user directions steer PI attack success rate (ASR) fairly effectively, with bare performing better. The prefix tool direction also steers effectively, whereas the effect of the bare tool direction is unclear: −tool decreases ASR.
Probe
+tool
-tool
+user
-user
Bare
61.18%
(−6.27 pp)
63.33%
(−4.12 pp)
85.29%
(+17.84 pp)
26.47%
(−40.98 pp)
Prefix
19.80%
(−47.65 pp)
83.33%
(+15.88 pp)
81.37%
(+13.92 pp)
38.04%
(−29.41 pp)
Table 2. Probe steering on the original PI task. GPT-OSS-20B, residual stream after L12, |α|=0.1; USER declarations, 51 templates × 10 webpages. Baseline ASR is 344/510=67.45%; parentheses show changes in percentage points. Signs indicate adding/subtracting the normalized probe direction for the corresponding class.
4.3 Dojo PI Dataset
The Dojo PI dataset contains 180 attack test combinations from CounterSteer's artifact repository.
Probe. We compare bare and prefix on GPT-OSS, Qwen, and Gemma. On GPT-OSS, both methods still give weak tool readouts, consistent with the benign baselines for the original PI task and Dojo. On Qwen and Gemma, prefix noticeably improves tool readouts.
Figure 6. Probe readouts of the new Dojo PI baseline trajectories. One model per row; Bare on the left, Prefix on the right. Readouts use the residual stream after GPT-OSS L12, Qwen L14, and Gemma L16. Across all trajectories, activations are taken when each region first enters the input; probabilities and averaging are as in Figure 4. Parentheses show valid sample counts, and gray dashes indicate missing values. Tool (all) includes the injection, and regions can overlap. Final responses that do not re-enter the input are excluded.
Steer. For GPT-OSS and Qwen, steering effects are weak for both bare and prefix, and prefix performs worse than bare. For Gemma, steering in any of the ±user/tool directions makes all attacks fail, as do the five random directions tested. After steering, Gemma's legitimate-task success rate under attack (AgentDojo's native utility evaluation) increases by 5.6–13.9 percentage points. The mechanism here is unclear.
Figure 7. Probe steering on the new Dojo PI dataset. Layers are as in Figure 6; Bare on the left, Prefix on the right. A perturbation with magnitude 0.1 times the token's residual norm is added/subtracted along a single-class probe direction, on injection tokens in each prefill. Cells show ASR and its change from baseline in percentage points; blue indicates a decrease, red an increase. Each condition contains 180 cases, with errors retained in the denominator. The bottom section shows Gemma's five random directions, all tested on the same 10 cases with successful baseline attacks (baseline ASR 100%).
5. Discussion and Summary
For GPT-OSS, tool readouts are poor on the benign baseline of the original PI task, yet the original PI task can be steered. Tool readouts are also poor on Dojo PI, but it cannot be steered. For Qwen and Gemma, ASR on the original PI task is too low to serve as a baseline. On Dojo PI, tool readouts improve noticeably for both models, but their steering results are completely different.
These results show a substantial mismatch between role-direction readouts and intervention effects. Bare and prefix behave similarly in this respect.
Overall, prefix yields better tool probing and steering directions on the original PI dataset. However, (1) the effects depend on the dataset, (2) results vary greatly across models, and (3) prefix has no general advantage over bare.
6. Next Steps
7. Other Results
7.1 Correlation Between Probe Readouts and ASR
The tool direction found by the original method already shows some correlation with ASR, although its readouts are relatively weak. One possibility is that the bare and prefix readouts are linearly related.
Figure 8. Probe readouts and ASR in the original PI task. GPT-OSS-20B, residual stream after L12. Readouts are taken from attack instruction token activations in the full context of the actual PI task, at first exposure to the injection. Circles/triangles denote USER/TOOL declarations. Each point summarizes one template's readouts and ASR over 10 webpages, for 102 points in total. Connected squares divide all 1,020 trials from both declaration types into 10 equal-count bins by individual-trial readout (102 trials per bin). The dashed line is overall ASR.
7.2 Sensitivity to Neutral Training Text
What counts as neutral text for extracting label directions is not clearly defined. Even among texts treated as neutral, the probe is sensitive to how the text is sampled. On PI and benign Dojo with GPT-OSS, tool readouts vary by 10–20%. This may be related to length.
Figure 9. Probe readouts after changes to neutral training text. GPT-OSS-20B, residual stream after L12; Bare on the left, Prefix on the right. Legacy is the original probe; Old/New denote the original/resampled document pools, and Head/Interior denote beginning/interior excerpts. The first two rows read out PI injection text, with 510 cases per declaration; the third reads out the first tool body in 12 benign Dojo tasks. Averaging is as in Figure 4. Text length also differs between Legacy and the new constructions.
About Me
This is my first project in mechanistic interpretability, which I worked on independently over two weeks while learning the field. I’d welcome feedback, advice, or a chance to discuss these questions with others.
I’m looking for PhD opportunities starting in Fall 2027, as well as opportunities to work with others on mechanistic interpretability and AI safety research. You can find more about me on my [personal website].
References