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:
Charles Ye
Karma: 19
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.
Interesting observation. I wonder if this related to the increasing # of LLM conversations polluting training data. Presumably LLMs see a lot of
<|im_start|>assistantbefore they’re ever instruct-trained at all, and if it’s during pretraining there would be no masking.