I’m really struggling to figure out how to train to make role-boundaries more robust. They probably already make the “<assistant>” as special tokens (so literally typing “<assistant>” isn’t interpreted as the role-token).
You could train probes to detect each role, and mix around the text of each (eg put what the user said into the thinking tag, etc), and use that to detect if any prompt injections are happening (ie if text within “thinking” is classified as “user”, then you have a problem). However, one problem here is ensuring the probes are causal to the LLM’s output.
couldn’t you train embedding masks on top of tags? If you have very different embeddings for the different roles, maybe that’s enough to be more robust to prompt injection?
Worth splitting two questions. Security is the easy one: as Charles notes, the role marker is a special token an attacker can’t emit, so a clean per-role embedding on top of it is unspoofable — that channel already exists. The hard question is whether the model assigns authority by reading that channel rather than style.
Adding the embedding alone doesn’t force that: style still predicts the role label, so two correlated features compete and inductive bias picks. And “easiest” is set in pretraining, where the channel doesn’t exist and style carries the loss (Charles’s inheritance point), and where channel and style never disagree.
So you need both: the treatment in pretraining (for entrenchment), plus adversarial style/channel mismatches in post-training to make the channel load-bearing — the training analog of footnote 9′s decorrelated probes. Without the channel, adversarial training just learns a defensive style discriminator, which is exactly how footnote 15′s models “defend” CoT Forgery — still style, still whack-a-mole.
Crucial limit: this trains identification; authority is separate. The adversarial cases must be behavioral (actually refuse the injection), or you get Logan’s problem — a channel-reading probe that’s accurate but not causal to behavior. Testable with your probes: after training, does ablating the channel change compliance, or does ablating the style direction still move it?
Similar to vision/audio tokens the system and thinking tokens should be distinct from user and output. This would have to happen after pretraining with, e.g. the first assistant training examples appearing with some thinking/output tokens intermixed with the pretraining corpus tokens, slowly increasing the ratio of specialized to normal tokens until the model only outputs thinking tokens in the thinking section, output tokens in the output. Precisely when to begin adjusting the ratios could actually happen earlier, e.g. have a 1:1 mapping of “user” and “output” tokens and calculate the loss function on output tokens despite the LLM seeing “user” tokens, but making translation between token sets straightforward. Finally, on inference always strip incorrect-domain tokens from the tagged output sections, and of course don’t let users input non-user tokens.
EDIT: after thinking about this a bit I am skeptical that distinct token sets achieve much of anything. Specifically, if the model starts quoting or repeating things a prompt injection says then it mixes the token domains. Additionally, where prompt injections have their effect is in “world model” space where the model thinks about and decides what to do; keeping the world model well-grounded and goals aligned with a value system is still necessary to not get prompt-injected via “do this or nice people will suffer” style injections.
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.
Cool project and well written.
I’m really struggling to figure out how to train to make role-boundaries more robust. They probably already make the “<assistant>” as special tokens (so literally typing “<assistant>” isn’t interpreted as the role-token).
You could train probes to detect each role, and mix around the text of each (eg put what the user said into the thinking tag, etc), and use that to detect if any prompt injections are happening (ie if text within “thinking” is classified as “user”, then you have a problem). However, one problem here is ensuring the probes are causal to the LLM’s output.
couldn’t you train embedding masks on top of tags? If you have very different embeddings for the different roles, maybe that’s enough to be more robust to prompt injection?
related: https://www.lesswrong.com/posts/d8xDGzCEYE639qqEv/a-mechanistic-explanation-of-prompt-injection-and-why-you?commentId=SaFFTj9XF5pfodtTu
https://claude.ai/share/1aae9857-bbeb-4e1f-97f9-bb9bf464dc32
Worth splitting two questions. Security is the easy one: as Charles notes, the role marker is a special token an attacker can’t emit, so a clean per-role embedding on top of it is unspoofable — that channel already exists. The hard question is whether the model assigns authority by reading that channel rather than style.
Adding the embedding alone doesn’t force that: style still predicts the role label, so two correlated features compete and inductive bias picks. And “easiest” is set in pretraining, where the channel doesn’t exist and style carries the loss (Charles’s inheritance point), and where channel and style never disagree.
So you need both: the treatment in pretraining (for entrenchment), plus adversarial style/channel mismatches in post-training to make the channel load-bearing — the training analog of footnote 9′s decorrelated probes. Without the channel, adversarial training just learns a defensive style discriminator, which is exactly how footnote 15′s models “defend” CoT Forgery — still style, still whack-a-mole.
Crucial limit: this trains identification; authority is separate. The adversarial cases must be behavioral (actually refuse the injection), or you get Logan’s problem — a channel-reading probe that’s accurate but not causal to behavior. Testable with your probes: after training, does ablating the channel change compliance, or does ablating the style direction still move it?
Similar to vision/audio tokens the system and thinking tokens should be distinct from user and output. This would have to happen after pretraining with, e.g. the first assistant training examples appearing with some thinking/output tokens intermixed with the pretraining corpus tokens, slowly increasing the ratio of specialized to normal tokens until the model only outputs thinking tokens in the thinking section, output tokens in the output. Precisely when to begin adjusting the ratios could actually happen earlier, e.g. have a 1:1 mapping of “user” and “output” tokens and calculate the loss function on output tokens despite the LLM seeing “user” tokens, but making translation between token sets straightforward. Finally, on inference always strip incorrect-domain tokens from the tagged output sections, and of course don’t let users input non-user tokens.
EDIT: after thinking about this a bit I am skeptical that distinct token sets achieve much of anything. Specifically, if the model starts quoting or repeating things a prompt injection says then it mixes the token domains. Additionally, where prompt injections have their effect is in “world model” space where the model thinks about and decides what to do; keeping the world model well-grounded and goals aligned with a value system is still necessary to not get prompt-injected via “do this or nice people will suffer” style injections.
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.