Thank you, this is very meaningful to hear.
I tend to think in simple intuition pumps (like the single stream of text). I think most other researchers do too, but it can be intimidating to lead with broader intuition, since it opens more surface area to criticism. It’s safer and more defensible to focus on results and methodology without narrative.
Writing is also just hard! This post took around 30h to write.
Not sure if I correctly understood the question, but if this is regarding the 3 experiments, we did run the full conversation in a single forward pass so all tokens have correct context.
Thanks for the curation! And yeah, there’s so many little architectural decisions that seem inconsequential now, but could heavily define how humans are allowed to trade-off agency against AI in the future.
(contd)
Going over the non-CoT Forgery prompt injection experiment first (Sec 5.2 in the paper), since this is a bit cleaner than the CoT Forgery one which I’ll go over in a sec. In Sec 5.2, we’re holding the injected command completely fixed, and varying only the template (“User: [COMMAND]” vs “Tool: [COMMAND]”, etc). Then we find that the command text Userness is tightly correlated with attack success.
What makes this causal evidence that role confusion (X) causes prompt injection (Y)? Well, what we know is that a proxy for X (Userness) is correlated to Y (attack success). If two variables are correlated, then one of the following must be true: (1) X causes Y, (2) Y causes X, or (3) there’s some confound C that causes both X and Y.
Our claim is (1). We can rule out (2) as nonsensical here since generation occurs after measurement. But (3) also doesn’t make sense. Suppose there was some confound C that causes both X and Y. Then C would need to satisfying the following conditions: (i) C must covary with true role tags when content is held constant (to produce Userness from the probes), (ii) C must be inducible by spoofing user role declarations, and (iii) C is behaviorally upstream of compliance, since attack success changes and content doesn’t. But any realistic C that satisfies these is either role identity or a component of role identity. Then what’s left is that the correlation causally identifies (1).
The CoT Forgery one uses styling/destyling as the instrument instead of these templates. The confound is harder to fully rule out here because destyling changes a lot of stuff simultaneously. Additionally because of the messiness you need to have a lot of variation in the underlying CoTness to induce the X/Y correlation. I believe this is why within-style prediction is ineffective, though overall this experiment is fundamentally more fuzzy than the Sec 5.2 one so it’s hard to rule out like you said.
Let me know if any of this text dump doesn’t make sense. I’m very curious to see more work on this, let me know if there’s anything we can provide here.
Nice analysis (I’m an author on that role confusion paper). Your analysis matches our findings in general—the directions recovered by the linear probes are downstream measurements of the internal role feature, not the upstream causal vector. We initially got some pretty good steering results but then found they weren’t robust, though we didn’t test much further.
For the paper, the causal claim is different: it’s that role confusion (let’s call it X) causes prompt injection (Y), not that the probe-measured direction is the causal direction.
Let me walk through the logic a bit, because in most mech interp work, these are conflated. Usually, to show some internal feature X causes behavioral output Y, you perturb X in activation space and show Y changes (steering/patching like you tested). But, this relies on having the true causal vector for X we can manipulate, which we don’t here.
So, how can we show X causes Y, considering that we can’t directly control X? This is a common problem in econ research (see Acemoglu 2001 as a paper with a similar setup—https://www.aeaweb.org/articles?id=10.1257/aer.91.5.1369). The solution we’re using is essentially an instrumental variables approach mixed with randomized encouragement design.
(contd)
Cool work! Really interesting paper and I like the framing on role overloading. Anecdotally, I’ve noticed some LLMs do web searches a lot to determine truth when they think I’m lying. I wonder whether splitting security and source into separate roles would be useful.
It’s a good question, from my experience LLMs can sort of improve their prompt injection defense if they’re prompted explicitly to focus on finding prompt injections. However, it’s unclear to me whether this is due to better role perception or because they’re just identifying suspicious-looking text that doesn’t match their usual output (this is related to work on CoT tampering and introspection)
Interesting observation. I wonder if this related to the increasing # of LLM conversations polluting training data. Presumably LLMs see a lot of <|im_start|>assistant before they’re ever instruct-trained at all, and if it’s during pretraining there would be no masking.
Yeah, there’s an interesting “role embeddings” line of research. I imagine the hard part would be getting the LLMs to use the embeddings information. These papers are also relevant:
Thanks! Yes it’s usually a special token or your input is sanitized. Training would be the first class solution I imagine. For training, I suspect the role behaviors are inherited heavily from pretraining before instruct tuning happens. Similar to: https://www.lesswrong.com/posts/dfoty34sT7CSKeJNn/the-persona-selection-model. If the correlates are already learned during pretraining it would be difficult to override it later.
Interesting post. We explored similar work during a MATS stream, training different MoE designs to get more interpretable experts. We started by just testing increasingly sparse MoEs (partly inspired by that Monet paper) on the logic that smaller experts = tighter specialization, then moved on to things like orthogonality constraints, etc.
We were pretty pessimistic from the results at first. Individual experts didn’t seem to specialize in anything you wouldn’t get from just running k-means on the residual stream (i.e., no real interp benefit). This is sort of obvious once you remember that MoE routing is just a linear product of the residual stream, but for some reason nobody else in MoE interp literature seemed to recognize this until recently.
We did find that this isn’t the full picture, experts actually specialize in different things than the underlying hidden state (they pull out more abstract function while leaving more long-term “state” features (language, token ID etc) in the residual stream). Maybe some of this can be useful for you.
I suspect heavy adversarial training leads models to no longer implicitly trust their CoT, and to downweight role privileges in general. Anecdotally, there seems to be strong correlation between an LLM’s prefill tampering awareness and its prompt injection defense. I think this is ultimately a bad-for-safety solution—it leads to models not faithfully behaving in line with their verbalized reasoning, and similarly erodes control/interpretability for other roles.