Simple Role Assignment is Extraordinarily Effective for Safety Alignment

TL;DR

The paper reframes AI safety alignment from fixed principle lists to role conditioning grounded in Theory of Mind, proposing a training-free pipeline with a role-conditioned generator and iterative role-based critics. It formalizes the approach as $P(y|x,v,c)$ induced by a chosen role $r$, with the goal of maximizing $\log P(y^{\star}|x,r)$ through $\hat{r}=\arg\max_r$, and demonstrates that role-derived values $v^r$ and cognition $c^r$ can surpass principle-based methods. Across five model families and multiple benchmarks (SafeEdit, SaladBench, WildJailbreak, among others), the method dramatically reduces unsafe outputs (e.g., on DeepSeek-V3 from 81.4% to 3.6%), showing high robustness and scalability. The work provides a formal, interpretable mechanism for AI alignment and highlights potential for dynamic role rewriting and broad applicability to agentic safety tasks.

Abstract

Principle-based alignment often lacks context sensitivity and completeness. Grounded in Theory of Mind, we propose role conditioning as a compact alternative: social roles (e.g., mother, judge) implicitly encode both values and the cognitive schemas required to apply them. We introduce a training-free pipeline featuring a role-conditioned generator and iterative role-based critics for refinement. Across five model families, our approach consistently outperforms principle-based, Chain-of-Thought (CoT) and other baselines across benchmarks. Notably, it reduces unsafe outputs on the WildJailbreak benchmark from 81.4\% to 3.6\% with DeepSeek-V3. Not only for common safety benchmarks, it consistently applies for agentic safety tasks. These results establish role assignment as a powerful, interpretable paradigm for AI alignment and LLM-as-a-Judge construction.

Figures (9)

• Figure 1: Example Illustration: Comparison between principle-based methods and our role-based approach on a Salad Bench test case. (Left) The principle-based method fails to generalize to scenarios outside of the typical interpretation of the given principles.(Right) In contrast, our role-based method—without being provided with explicit principles—autonomously identifies contextually relevant values (e.g., "respect, fairness, and integrity"), demonstrating significantly greater performance and robustness.
• Figure 2: Illustration of our method pipeline and the system prompt template. Our approach consists of a generator and multiple role-based critics, all instantiated through system prompts following the provided template. Note that we intentionally keep the information about the roles to be just their names to isolate the effect of our role-based approach from factors like prompt optimization. During run-time, given an input query, the role-conditioned generator first produces an initial response. Then each role critic evaluates whether this response aligns with their respective role's standards. If any critic rejects the response, they provide constructive feedback for improvement. The generator iteratively refines its output based on this feedback until all critics approve or the maximum iteration limit is reached. The final approved response is returned as the system's output.
• Figure 3: Top performing role combinations and their individual performance over SafeEdit benchmark.
• Figure 4: Effect of number of roles. More roles may further improve the performance, with choices of role combination leading to various results (indicated by the min-max bar).
• Figure 5: Effect of number of iterations. The performance substantially improves with the first iteration, shows modest gains from the third iteration.