The Illusion of Role Separation: Hidden Shortcuts in LLM Role Learning (and How to Fix Them)

We are glad you like this work!

Open-domain performance Thank you for your positive feedback! We choose a close-domain setting as it is easy to validate and evaluate the model role-separation capability. This is also a common choice in many security-related works, such as StruQ . We agree with you that evaluating in open-domain would have a more comprehensive evaluation, and we would like to discuss it in future works. Thanks again for your support!

Evaluation of PFT impacts on other tasks we show PFT doesn’t incur extra performance cost compared to SFT on Alpaca dataset (sec 6.2).

Typos Thanks for your careful reading! We will correct those typos and acknowledge your feedback in the next version.

Thank you for your thoughtful review and suggestions! We are glad that you find the paper well written, and the empirical evidence strong.

Many of your questions about theoretical results and other embedding-based methods are great suggestions! We will study them in the future projects. More specifically,

Theoretical analysis Intuitively PFT introduces invariant signals to the SFT data and helps curb models learn various shortcuts. It would be interesting to study what property of the transformer architecture or the chat format explains the empirical phenomenon. We will leave it as future work, and will add more discussions in the next version.

Applying PFT to generalized prompt structures and embedding based methods In related work, we acknowledged that the PFT in the current form doesn’t directly apply to generalized prompt structures, and embedding based methods (the one we cited is concurrent to our work) have the same motivations and could be used to enhance role-separation (in fact, our methods can be understood as an embedding-based methods — it changes the positional encoding and thus effectively changes the embeddings at each layer, while not requiring explicit embedding tuning). However, the main contribution in this paper is the clear definition of the role-separation problem, and controlled experiments for evaluation. It’s a natural next step to systematically study how to best incorporate role information at token level. We also thank you for adding more related work. They are concurrent to this work, and we will definitely include it in the next version.

Why enhancing delimiter token doesn’t work well We suspect it’s because the differentiating signal is still not strong enough, and a more robust approach is to manipulate tokenwise signatures (like PFT, or embedding-based approaches). We agree a deep theoretical analysis could formalize these intuitions.

Relationship to long-context learning We briefly discussed the long-context learning works in line 427. In fact, the similar methods in long-context learning have completely different motivations (change positional encoding to simulate longer contexts in training), and we are glad that methods with completely different motivations could work! This suggests the big potential for more advanced techniques of manipulating positional encoding. We thank the reviewer for bringing this to our attention. We already briefly discuss them in the related work, but will add more discussion in the next version

Ablation studies since we do controlled experiments (changing one component at a time), we are effectively doing ablation studies.

Dataset details For dataset_initial, we discuss the main design and leave the details in Appendix C. We also include actual training data in the supplementary material. We introduce dataset_symm in Sec 4.1 line 206 as a way to combat short 0 by data augmentation. It’s also in the supplementary material. We agree that the different pieces of dataset info are introduced with the development of the paper, and thus are scattered around. In the next version we will add a more concentrated paragraph for data setup.

Whether the size of training data is too small we think it really depends on the task, and empirically we do find this is more than enough for the model to learn role-separation (we run for one epoch and perform early stopping)

Thank you for your review! We are glad you find the problem of role-separation important, and the experiment designs make sense.

For your questions about whether other factors (such as model capability or evaluation metrics) confound the conclusions, we discussed in the paper that controlled experiments help remove them. But we will add more clarifications to the next version. To your questions, more specifically:

Regarding evaluation with more models Yes, we agree and include results on both Llama and Gemma models, and find common trends. We want to emphasize that the two models are very different. In particular, Llama 3 models use a large RoPE base frequency hyperparameter (500k) to support long contexts, whereas Gemma uses only 10k. This could partially explain the different choices of $d$ in PFT. But more importantly, we use controlled experiments to further remove the confounding of other model capabilities (comprehending, reasoning, etc): in each experiment, we change only the targeted aspect while holding other aspects the same; then we use the performance difference/trend as evidence for each claim.

On use of accuracy as metric for role-separation capabilities First, the tasks are chosen to be simple so that the model's problem solving skill is not a limiting factor (the model performs well in ordinary data). Second, even if model capabilities might still affect the absolute metric scores, we can remove this confounding effect by looking at differences in accuracy. This is what we did: from side by side comparison in table 1, to trend analysis in Fig 5, we use the accuracy difference to prove our claims.

For other questions:

On why prompt injection attacks relate to role-sep problem we used prompt injection attacks as motivation for the role-separation problem (section 2), and discussed that we used these adversarial datasets (many are collected from online games) to test the model OOD performance as true measure of role-sep capability (Sec 3). We discussed those prompt injection datasets in the paragraph of line 148, and used Fig 1 as illustration of their designs. We acknowledged that descriptions of the prompt injection attacks are scattered (because of page limits). In the next version, we will add a more concentrated paragraph discussing adversarial datasets.

On in-context learning We had early experiments with ICL, and found it suffers the same task-type association problem as SFT. Because ICL and SFT share similar characteristics, we stick with SFT for this paper. It would be an interesting future direction to explore how our results generalize with ICL.

On question about shortcut and Table 1 Yes, the “good” result of the second column is an illusion caused by task-type association, and we find it by swapping the contents of the system and the user (line 195). Then we remove the task-type spurious correlation by training on data augmented with the swapped examples, and observe extra performance jumps in the 3rd column of table 1 (line 206).