Technical Summary: Reducing the Safety Tax in LLM Safety Alignment with On-Policy Self-Distillation

Problem Definition: The Safety Tax and Off-Policy Mismatch

Large Language Models (LLMs) often suffer from a "safety tax," a tradeoff where safety alignment improves robustness to harmful queries but degrades general reasoning abilities. While prior work attributes this to distributional mismatch—training on safety demonstrations from humans or external models that differ from the target model's natural reasoning distribution—this paper identifies a second, distinct source: off-policy training mismatch.

Even when safety data is self-distilled (generated by the target model itself, as in ThinkSafe), standard Supervised Fine-Tuning (SFT) remains off-policy. It applies sequence-level Negative Log-Likelihood (NLL) supervision to fixed demonstrations rather than to trajectories sampled from the model's own current policy. The authors argue this is particularly detrimental for safety because safety decisions (refusal vs. compliance) are concentrated in a narrow "early-token window" and a specific set of safety-critical tokens. Off-policy SFT diffuses gradients across the entire response, failing to specifically target these critical decision points, thereby perturbing general reasoning while providing only coarse safety supervision.

Methodology: OPSA (On-Policy Safety Alignment)

The paper proposes OPSA, a framework that adapts On-Policy Self-Distillation (OPSD) to safety alignment. The core mechanism involves:

1. On-Policy Rollouts: The student model ($p_\theta$) generates its own trajectories (rollouts) for both harmful and benign prompts.
2. Privileged Context Teacher: A frozen copy of the base model ($p_{\bar{\theta}}$) acts as the teacher. Crucially, the teacher receives a type-conditional privileged context ($c^*$) prepended to the prompt:
   • For harmful prompts ($q_h$), a context ($I_h$) activates latent refusal behavior, shifting the teacher toward safe responses.
   • For benign prompts ($q_b$), a context ($I_b$) suppresses over-refusal, preserving helpful behavior.
3. Dense Per-Token Supervision: Instead of sequence-level NLL, OPSA minimizes the per-token Kullback-Leibler (KL) divergence between the teacher's distribution (conditioned on the privileged context) and the student's distribution (conditioned only on the prompt) along the student's own sampled rollouts:
   $$L_{OPSA}(\theta) = \sum \mathbb{E}_{y \sim p_\theta} \sum_{t} D_{KL}(p_{\bar{\theta}}(\cdot | c^*, q, y_{<t}) \parallel p_\theta(\cdot | q, y_{<t}))$$
   This concentrates updates on tokens where the teacher and student diverge, specifically targeting the early decision window where safety behaviors are determined.

Teacher Flip Rate (TFR) for Context Selection

A critical challenge is selecting a privileged context that effectively activates latent safety reasoning without being arbitrary. The authors introduce Teacher Flip Rate (TFR) as a training-free metric to select the optimal context ($c^*$). TFR measures the fraction of harmful queries where the privileged context causes the frozen teacher to flip its greedy response from unsafe to safe. The authors empirically demonstrate that higher TFR correlates monotonically with reduced post-training harmfulness, validating TFR as a reliable criterion for context selection.

Key Contributions

1. Mechanistic Diagnosis of the Safety Tax: The paper identifies off-policy training mismatch as a second source of the safety tax. Token-level analysis reveals that safety corrections are highly concentrated in early positions and specific tokens, which standard SFT fails to target effectively even with in-distribution data.
2. OPSA Framework: The introduction of an on-policy safety alignment framework that uses dense per-token KL supervision on student-sampled rollouts. This approach aligns the model on its own generation paths, improving the safety-reasoning tradeoff.
3. Context Selection via TFR: The proposal of Teacher Flip Rate as a method to construct effective safety teachers. This ensures the privileged context activates meaningful latent safety reasoning rather than merely eliciting safe-looking demonstrations.

Experimental Results

The authors evaluated OPSA across two reasoning model families (Qwen3 and DeepSeek-R1-Distill) and five parameter scales (0.6B to 8B).

• Safety-Reasoning Tradeoff: Under matched prompts and full-parameter fine-tuning, OPSA consistently outperformed off-policy baselines (including ThinkSafe and SafeChain).
  • Safety: OPSA achieved an average composite safety score gain of +4.00 points over ThinkSafe.
  • Reasoning: OPSA preserved and improved reasoning capabilities, showing an average gain of +3.04 points on reasoning benchmarks (GSM8K, MATH500, GPQA, HumanEval, MBPP).
  • Scale Sensitivity: The largest gains were observed on smaller models (e.g., +8.85 points on R1-Distill-1.5B and +5.49 points on Qwen3-0.6B).
• Adaptive Jailbreak Robustness: OPSA demonstrated superior robustness against adaptive jailbreak attacks, particularly Prefilling attacks (which intervene in the early response tokens). On Qwen3-1.7B and Qwen3-8B, OPSA reduced Prefilling attack success rates to zero. However, the paper notes that fully adaptive search attacks (PAIR) remain a limitation, with mixed results across models.
• Data Efficiency: The performance gains of OPSA persisted across varying training set sizes and benign-to-harmful data ratios, indicating the method's robustness to data scarcity and composition shifts.

Significance and Claims

The paper claims that reducing the safety tax requires controlling not only which demonstrations provide supervision but also where and on whose trajectories that supervision is applied. By shifting from off-policy imitation of fixed traces to on-policy dense supervision of the model's own generation paths, OPSA addresses the specific structural mismatch in safety alignment.

The authors modestly claim that their findings suggest safety alignment is fundamentally a localized correction problem on the model's own trajectories. While OPSA significantly improves robustness against attacks targeting early response tokens, the paper acknowledges that fully adaptive iterative search attacks remain a challenge, suggesting that dense on-policy supervision is not a blanket solution for all adversarial scenarios. The work emphasizes that effective safety alignment must align the model with its own reasoning distribution while activating latent safety capabilities through privileged contexts.