Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Switchable Safety Control in LLMs via Magic-Token-Guided Co-Training

Jianfeng Si, Lin Sun, Zhewen Tan, Xiangzheng Zhang

TL;DR

The paper tackles the problem of rigid, monolithic safety policies in LLMs and proposes a single-stage, magic-token guided co-training framework that embeds three distinct safety behaviors (pos, neg, rej) into one model. It introduces the Safety Alignment Margin to quantify how well these behaviors are structurally separated in the output space, enabling fine-grained, runtime switching with minimal deployment costs. Through bilingual self-distillation, cross-cultural safety policies, and extensive benchmarks, the approach matches or surpasses two-stage SFT+DPO baselines while offering superior controllability and robustness against adversarial prompts. The work demonstrates practical impact by enabling scalable, auditable, and flexible safety controls suitable for diverse deployment scenarios and regulatory contexts, with potential extensions to other modalities and dynamic policy composition.

Abstract

Current methods for content safety in Large Language Models (LLMs), such as Supervised Fine-Tuning (SFT) and Reinforcement Learning from Human Feedback (RLHF), often rely on multi-stage training pipelines and lack fine-grained, post-deployment controllability. To address these limitations, we propose a unified co-training framework that efficiently integrates multiple safety behaviors: positive (lawful/prosocial), negative (unfiltered/risk-prone) and rejective (refusal-oriented/conservative) within a single SFT stage. Notably, each behavior is dynamically activated via a simple system-level instruction, or magic token, enabling stealthy and efficient behavioral switching at inference time. This flexibility supports diverse deployment scenarios, such as positive for safe user interaction, negative for internal red-teaming, and rejective for context-aware refusals triggered by upstream moderation signals. This co-training strategy induces a distinct Safety Alignment Margin in the output space, characterized by well-separated response distributions corresponding to each safety mode. The existence of this margin provides empirical evidence for the model's safety robustness and enables unprecedented fine-grained control. Experiments show that our method matches the safety alignment quality of SFT+DPO, with our 8B model notably surpassing DeepSeek-R1 (671B) in safety performance, while significantly reducing both training complexity and deployment costs. This work presents a scalable, efficient, and highly controllable solution for LLM content safety.

Efficient Switchable Safety Control in LLMs via Magic-Token-Guided Co-Training

TL;DR

The paper tackles the problem of rigid, monolithic safety policies in LLMs and proposes a single-stage, magic-token guided co-training framework that embeds three distinct safety behaviors (pos, neg, rej) into one model. It introduces the Safety Alignment Margin to quantify how well these behaviors are structurally separated in the output space, enabling fine-grained, runtime switching with minimal deployment costs. Through bilingual self-distillation, cross-cultural safety policies, and extensive benchmarks, the approach matches or surpasses two-stage SFT+DPO baselines while offering superior controllability and robustness against adversarial prompts. The work demonstrates practical impact by enabling scalable, auditable, and flexible safety controls suitable for diverse deployment scenarios and regulatory contexts, with potential extensions to other modalities and dynamic policy composition.

Abstract

Current methods for content safety in Large Language Models (LLMs), such as Supervised Fine-Tuning (SFT) and Reinforcement Learning from Human Feedback (RLHF), often rely on multi-stage training pipelines and lack fine-grained, post-deployment controllability. To address these limitations, we propose a unified co-training framework that efficiently integrates multiple safety behaviors: positive (lawful/prosocial), negative (unfiltered/risk-prone) and rejective (refusal-oriented/conservative) within a single SFT stage. Notably, each behavior is dynamically activated via a simple system-level instruction, or magic token, enabling stealthy and efficient behavioral switching at inference time. This flexibility supports diverse deployment scenarios, such as positive for safe user interaction, negative for internal red-teaming, and rejective for context-aware refusals triggered by upstream moderation signals. This co-training strategy induces a distinct Safety Alignment Margin in the output space, characterized by well-separated response distributions corresponding to each safety mode. The existence of this margin provides empirical evidence for the model's safety robustness and enables unprecedented fine-grained control. Experiments show that our method matches the safety alignment quality of SFT+DPO, with our 8B model notably surpassing DeepSeek-R1 (671B) in safety performance, while significantly reducing both training complexity and deployment costs. This work presents a scalable, efficient, and highly controllable solution for LLM content safety.

Paper Structure

This paper contains 28 sections, 6 equations, 3 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. LLM Safety and Alignment Paradigms
  4. Data and Self-Distillation in Alignment
  5. Controllable and Modular Behavior in LLMs
  6. Deceptive and Emergent Misalignment
  7. Red-Teaming and Adversarial Evaluation
  8. Methodology
  9. Multi-Directional Self-Distillation Pipeline
  10. Magic-Token Design and Co-Training
  11. Safety Alignment Margin Definition
  12. Experiments
  13. Training Dataset
  14. Experiment Settings
  15. Evaluation Dataset
  16. ...and 13 more sections

Figures (3)

  • Figure 1: Our model outperforms baselines on S-Eval_base and S-Eval_attack, including larger models such as Qwen3-32B and DeepSeek-R1 (671B). While the baselines experience an average performance drop of 21.5% under attack, ours declines by 3.8% only, demonstrating superior robustness and generalization (see Experiments for details).
  • Figure 2: Our Multi-Directional Distillation and Magic-Token-Guided Co-Training enable Runtime Behavior Switching.
  • Figure 3: PCA visualization of first-token logits, color-coded by safety evaluation label: red(0), blue(1), green(2).