I wonder if we could improve role detection by injecting it as a specific input feature instead of making the AI infer it. So you have a vector of embeddings concatenated with a vector of [user, assistant, tool, etc.]. This would do the opposite of what you want for dynamically adding roles, but should make role detection much more reliable. Conveniently, detecting XML tags in prompts is pretty easy so this should be possible to fully automate, and I suspect it would work fine to bolt this on after pretraining, at the same time you’re training it to use roles.
Our experiments on the Structured Query and Instruction Hierarchy benchmarks demonstrate an average robust accuracy increase of up to 15.75% and 18.68%, respectively. Furthermore, we observe an improvement in instruction-following capability of up to 4.1% evaluated on AlpacaEval.
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:
Things in this genre of privileged embedding seem valuable. It also seems doable to think architecturally beyond just embedding (eg, attention heads that only attend to certain roles, etc). Myself and others been trying to explore this genre some this month, but in some pilots gains seem unfortunately small compared to baselines. Personally I would be excited to see more of this though, as various ways of offering model-level guarantees seems good.
There’s additional issues here where roles are just an incomplete primitive. When we were exploring this last month observed how models in some ways seem to trust tools more than users. The role is overloaded primitive that is both supposed tell you both a hierarchy of trust and the information source. Often things in tools role are pseudo-oracles (eg, results of a python script, web results on current events, etc), and other times they are not (eg, reading an untrusted markdown file).
Charles and coauthors’ work here on embedding representation is really cool. It would be cool to further expand this, also thinking about interactions with like “trust embeddings” or other ways role type is serving multiple purposes, and how to build primitives or architectures to support this.
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.
Thanks! right, this thing around web searches matches my view as well. If a user asks about something important or weird, it’s usually best to trust the top 3 google results more than the user themself. But if in most of training traces are always trusting the tools with just a little specific adversarial post-training, end up with a model that still mostly trusts tools. This is maybe solvable with better training data, but becomes more nuanced than OpenAI’s model spec that gives tools are “no authority”.
A simple split with Trust∈{untrusted, trusted} and Source∈{system, user, tool} is tough though. There are rare cases where almost all tools are adversarial (eg, a malicious webpage). You’d end up basically always using Trust=untrusted, and be back in the same place. Unfortunately I don’t have clear views on solutions currently other than small improvements like maybe updating the spec to be more nuanced/self-consistent (plus special primitives of being able inline different sources into a prompt template seems nice). When tracking failures through embeddings like in this study, it seems like trust and source are different axises though.
Another approach that might be both more powerful and cheaper (because no training needed) would be to just activate the toolness/thinkness/whatever directions at inference time, as steering vectors injected into the activations. Since the harness knows which type of mode should be active for each token in and out, it can just abliterate all the incorrect token modes and activate only the one that’s needed. You’d need to do a little bit of manual work to find the vectors after, but you wouldn’t need to rely on the model actually learning to respect the tags.
Okay but how exactly would the harness “know” which mode to activate? That is itself determined deterministically through the model output when if includes trigger words for tools or thinking etc in its response? Or are you suggesting that when the modes are detected from LLM response, add steering vector for that mode so that even if the text itself doesn’t activate the vector, the harness manually forces it.
One issue I can see though is that what if the text already strongly activates the mode vector and we on top of that steer further, that would cause very strong steering and potentially degrade output?
I wonder if we could improve role detection by injecting it as a specific input feature instead of making the AI infer it. So you have a vector of embeddings concatenated with a vector of [user, assistant, tool, etc.]. This would do the opposite of what you want for dynamically adding roles, but should make role detection much more reliable. Conveniently, detecting XML tags in prompts is pretty easy so this should be possible to fully automate, and I suspect it would work fine to bolt this on after pretraining, at the same time you’re training it to use roles.
Edit: Looks like Wu et al. tried this in 2024 and it somewhat helps:
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:
https://www.lesswrong.com/posts/HEzNZ9gvgYwT3aZFS/role-embeddings-making-authorship-more-salient-to-llms
ASIDE: https://arxiv.org/abs/2503.10566
Things in this genre of privileged embedding seem valuable. It also seems doable to think architecturally beyond just embedding (eg, attention heads that only attend to certain roles, etc). Myself and others been trying to explore this genre some this month, but in some pilots gains seem unfortunately small compared to baselines. Personally I would be excited to see more of this though, as various ways of offering model-level guarantees seems good.
There’s additional issues here where roles are just an incomplete primitive. When we were exploring this last month observed how models in some ways seem to trust tools more than users. The role is overloaded primitive that is both supposed tell you both a hierarchy of trust and the information source. Often things in tools role are pseudo-oracles (eg, results of a python script, web results on current events, etc), and other times they are not (eg, reading an untrusted markdown file).
Charles and coauthors’ work here on embedding representation is really cool. It would be cool to further expand this, also thinking about interactions with like “trust embeddings” or other ways role type is serving multiple purposes, and how to build primitives or architectures to support this.
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.
Thanks! right, this thing around web searches matches my view as well. If a user asks about something important or weird, it’s usually best to trust the top 3 google results more than the user themself. But if in most of training traces are always trusting the tools with just a little specific adversarial post-training, end up with a model that still mostly trusts tools. This is maybe solvable with better training data, but becomes more nuanced than OpenAI’s model spec that gives tools are “no authority”.
A simple split with Trust∈{untrusted, trusted} and Source∈{system, user, tool} is tough though. There are rare cases where almost all tools are adversarial (eg, a malicious webpage). You’d end up basically always using Trust=untrusted, and be back in the same place. Unfortunately I don’t have clear views on solutions currently other than small improvements like maybe updating the spec to be more nuanced/self-consistent (plus special primitives of being able inline different sources into a prompt template seems nice). When tracking failures through embeddings like in this study, it seems like trust and source are different axises though.
Another approach that might be both more powerful and cheaper (because no training needed) would be to just activate the toolness/thinkness/whatever directions at inference time, as steering vectors injected into the activations. Since the harness knows which type of mode should be active for each token in and out, it can just abliterate all the incorrect token modes and activate only the one that’s needed. You’d need to do a little bit of manual work to find the vectors after, but you wouldn’t need to rely on the model actually learning to respect the tags.
Okay but how exactly would the harness “know” which mode to activate? That is itself determined deterministically through the model output when if includes trigger words for tools or thinking etc in its response? Or are you suggesting that when the modes are detected from LLM response, add steering vector for that mode so that even if the text itself doesn’t activate the vector, the harness manually forces it.
One issue I can see though is that what if the text already strongly activates the mode vector and we on top of that steer further, that would cause very strong steering and potentially degrade output?