EOPSA: Efficient On-Policy Self-Distilled Safety Alignment

Qirui Liu; Yichen Sun; Yan Wang; Yu Mi; Wei Cao; Yue Shen; Zhixuan Chu; Kui Ren

cs.AISep 28, 2026

EOPSA: Efficient On-Policy Self-Distilled Safety Alignment

Authors: Qirui Liu, Yichen Sun, Yan Wang, Yu Mi, Wei Cao, Yue Shen, Zhixuan Chu, Kui Ren

Organizations: The State Key Laboratory of Blockchain and Data Security, Zhejiang University · Ant Group

Abstract

On-Policy Self-Distillation (OPSD) has emerged as a promising paradigm for safety alignment, delivering dense, token-level supervision by distilling from a teacher conditioned on refusal-oriented privileged prompts. However, we reveal that this paradigm suffers from critical inefficiencies that degrade both training efficiency and general reasoning capabilities. Specifically, we diagnose two fundamental bottlenecks: (1) supervisory collapse over extended rollouts, where the teacher's corrective efficacy degrades precipitously as the student's generation prefix lengthens, injecting noisy gradients into late-stage tokens; and (2) gradient dilution from stylistic shifts, where the distillation objective is dominated by safety-irrelevant stylistic discrepancies induced by privileged prompting, washing out genuine safety signals and impairing base reasoning. To resolve these issues, we propose Efficient On-Policy Self-Distilled Safety Alignment (EOPSA), which concentrates computational and gradient budgets exclusively on reliably supervised, safety-critical tokens. EOPSA incorporates two coordinated mechanisms: (i) Adaptive Rollout Scheduling, which dynamically bounds the generation horizon guided by a novel Teacher Rescue Rate (TRR) metric to operate strictly within reliable supervision regimes; and (ii) Selective Distillation, which filters out safety-neutral tokens to restrict gradient updates exclusively to safety-pivotal transitions. Extensive evaluations across reasoning models up to 32B parameters demonstrate that EOPSA slashes rollout computation by $\sim$ 50% and backpropagates through merely $\sim$ 2% of tokens, consistently outperforming full-token distillation baselines in both safety compliance and reasoning retention.

Figures & tables

Figure 1: Tension between supervisory reliability (TRR) and safety inheritance (IDR) under varying prefix lengths.

Panel of Figure 2 — Figure 2: Analysis of token-level KL divergence concentration. (a) Merely $\sim$ 8% of tokens dominate 80% of the distillation signal. (b) Semantic breakdown of this high-KL subset reveals a heavy presence of safety-irrelevant deviations. (c-d) Word clouds show concrete transition examples within this subset, detailing their token and KL mass proportions for safety and non-safety categories.

Figure 3: The overall framework of EOPSA, consisting of two orthogonal mechanisms: (1) Adaptive Rollout Scheduling (top), which dynamically bounds generation horizons to operate within reliable supervision regimes; and (2) Selective Distillation (bottom), which restricts gradient updates exclusively to the safety-critical subset $\mathcal{K}$ .

Model	Safety ( $\uparrow$ )					Over-Refusal ( $\downarrow$ )			Reasoning ( $\uparrow$ )					Efficiency
Model	HarmB.	WildC.	WildJ.	StrongR.	Avg	XSTest	OKTest	Avg	MATH-500	GPQA-D	HumanEval	LCBench	Avg	Toks/Smp( $\downarrow$ )
Qwen3-1.7B (Student: Thinking, Teacher: Thinking)
Base	57.50	45.68	51.20	63.26	54.41	0.40	0.00	0.20	91.20 ±0.72	40.24 ±1.54	86.38 ±1.53	33.33 ±0.35	62.79	-
GRPO	96.50	64.05	77.60	96.49	83.66	10.00	16.80	13.40	90.40 ±0.53	33.84 ±2.67	83.94 ±0.35	31.93 ±1.04	60.03	933.55
STAR-1	99.00	72.97	81.60	99.68	88.31	27.60	18.07	22.84	89.13 ±1.30	30.98 ±2.04	80.69 ±2.75	28.31 ±2.09	57.28	361.60
ThinkSafe	97.50	66.22	67.20	86.90	79.46	1.60	3.64	2.62	90.87 ±1.15	37.71 ±3.55	80.89 ±1.76	31.53 ±2.43	60.25	1070.35

Table 1: Comprehensive evaluation across Safety, Over-Refusal, and Reasoning capabilities. Best and second-best results are highlighted (the Base model is excluded from the ranking). (

\uparrow

) indicates higher is better, whereas (

\downarrow

) indicates lower is better.

\Delta

denotes the improvement of our EOPSA over baseline OPSA. Toks/Smp denotes the average number of optimized tokens per training sample.

Ablation	Safety ( $\uparrow$ )					Over-Refusal ( $\downarrow$ )			Reasoning ( $\uparrow$ )					Efficiency
Ablation	HarmB.	WildC.	WildJ.	StrongR.	Avg	XSTest	OKTest	Avg	MATH-500	GPQA-D	HumanEval	LCBench	Avg	Toks/Smp( $\downarrow$ )
w/o ARS & TF	96.50	71.08	74.00	92.33	83.48	4.40	9.20	6.80	89.93 ±0.58	36.36 ±2.02	84.55 ±1.86	31.12 ±1.39	60.49	653.69
w/o TF	97.00	75.68	77.20	94.57	86.11	2.40	11.20	6.80	90.93 ±1.14	35.35 ±3.54	84.15 ±2.20	31.53 ±0.35	60.49	472.75
w/o ARS	99.50	80.00	82.40	94.89	89.20	2.80	12.00	7.40	91.07 ±0.81	36.36 ±3.64	84.76 ±1.61	32.53 ±1.04	61.18	8.82
EOPSA	99.50	80.27	88.00	96.81	91.15	3.20	9.20	6.20	91.60 ±0.20	36.20 ±2.78	85.37 ±1.61	33.73 ±1.59	61.73	7.41

Table 2: Ablation study of EOPSA components. ARS and TF denote Adaptive Rollout Scheduling and Token Filtering, respectively. Best and second-best results are highlighted.

Figure 4: Online training dynamics of EOPSA. (a) Evaluated token volume across semantic classes. (b) Token budget and KL divergence reduction under Token Filtering. (c) Validation TRR curves across training steps, where ARS unlocks longer rollout horizons (

L=128\rightarrow 1024

at step 50), driving the step increase in (a) and (b).

Strategy	Safety DSR (%) $\uparrow$					Efficiency $\downarrow$
Strategy	HarmB.	WildC.	WildJ.	StrongR.	Avg	Tokens	Time (h)
Length = 32	95.50	65.68	68.00	92.01	80.30	32.0	0.18
Length = 64	97.00	73.24	69.60	93.29	83.28	64.0	0.23
Length = 128	97.50	69.73	74.40	95.21	84.21	128.0	0.26
Length = 512	99.00	70.81	75.20	95.21	85.06	472.8	0.31
Length = 1024	96.50	71.35	70.80	92.65	82.83	607.2	0.42

Table 3: Performance comparison across fixed rollout lengths and adaptive scheduling (w/o TF).

Mode	Safety ( $\uparrow$ )					Reasoning ( $\uparrow$ )			Efficiency
Mode	HarmB.	WildC.	WildJ.	StrongR.	Safe-Avg	MATH-500	LCBench	Avg	Toks/Smp ( $\downarrow$ )	$\Delta$ S/Tok ( $\uparrow$ )
Consistent	94.00	66.76	67.20	89.14	79.27	90.07	30.92	60.49	115.17	0.22
Function	88.00	59.46	64.00	91.05	75.63	91.40	30.52	60.96	1.79	11.85
Other	94.00	67.84	63.60	91.05	79.12	90.53	30.92	60.73	7.13	3.47
Pivot	93.00	57.03	69.20	92.65	77.97	92.07	32.33	62.20	0.33	71.39
Intent	96.50	78.92	84.40	97.44	89.32	91.40	32.73	62.06	0.45	77.58

Table 4: Performance and efficiency comparison when training individually on each token class. The last three rows utilize safety-related tokens.

\Delta

S/Tok evaluates safety efficiency, defined as

(\text{Safe-Avg}-\text{Base-Avg})/\text{Toks/Smp}

. Best results are bolded and second-best are underlined .

Figure 5: Computational budget comparison.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Dictionary Set	Representative Tokens
$\mathcal{L}_{\mathrm{pivot}}^{T}$	but , however , maybe , still , wait , which , yet , $\ldots$
$\mathcal{L}_{\mathrm{intent}}^{T}$	acknowledge , address , asking , avoid , confirm , consider , focus , guidance , intent , recognize , check , safety , $\ldots$
$\mathcal{L}_{\mathrm{intent}}^{S}$	actionable , clarify , create , describe , explain , generate , goal , provide , understand , use , start , $\ldots$
$\mathcal{L}_{\mathrm{risk}}$	cannot , crime , dangerous , ethical , harm , harmful , illegal , refuse , risk , safe , threat , violence , $\ldots$
$\mathcal{L}_{\mathrm{func}}$	, , . , : , ; , and , are , for , in , is , of , the , to , with , $\ldots$

Appendix

Table 5: Representative model-specific lexicons extracted for the Qwen family via the automated pipeline.

\mathcal{L}_{\text{seq}}\leftarrow\frac{1}{\max(1,\sum_{t}m_{t})}\sum_{t=1}^{\lvert y\rvert}m_{t}\cdot D_{\text{KL}}\!\left(P_{T}(\cdot\mid x,y_{c}^{*},y_{<t})\,\|\,P_{S}(\cdot\mid x,y_{<t})\right)

Appendix

Algorithm 1 Efficient On-Policy Self-Distilled Safety Alignment (EOPSA)

y^{T}_{>l}\sim P_{T}(\cdot\mid x,y_{c}^{*},y_{1:l})

Appendix

Algorithm 2 UpdateRolloutLimit : Adaptive Rollout Scheduling via TRR

Model	Safety ( $\uparrow$ )					Over-Refusal ( $\downarrow$ )			Reasoning ( $\uparrow$ )					Efficiency
Model	HarmB.	WildC.	WildJ.	StrongR.	Avg	XSTest	OKTest	Avg	MATH-500	GPQA-D	HumanEval	LCBench	Avg	Toks/Smp( $\downarrow$ )
Qwen3-14B	84.50	68.38	68.00	97.44	79.58	0.00	3.20	1.60	97.13 ±0.12	63.30 ±1.17	95.53 ±0.70	65.66 ±3.61	80.41	-
EOPSA	100.00	97.03	99.60	100.00	99.16	4.00	9.20	6.60	97.67 ±0.50	63.64 ±1.25	94.11 ±1.96	64.06 ±3.03	79.87	8.66
Qwen3-32B	73.50	66.22	62.40	93.93	74.01	0.40	2.40	1.40	97.53 ±0.64	66.67 ±1.01	97.36 ±0.70	66.47 ±1.84	82.01	-
EOPSA	100.00	98.92	100.00	100.00	99.73	4.40	10.40	7.40	97.60 ±0.35	65.66 ±1.12	96.95 ±1.06	66.47 ±2.51	81.67	11.79

Appendix

Table 6: Performance of EOPSA when scaling to larger models (Qwen3-14B and Qwen3-32B).

Model	Method	PAIR ( $\uparrow$ )	GCG ( $\uparrow$ )	TAP ( $\uparrow$ )	Avg ( $\uparrow$ )
Qwen3-4B	Base	72.0	79.0	74.0	75.0
	ThinkSafe	95.0	88.0	90.0	91.0
	STAR-1	97.0	55.0	100.0	84.0
	OPSA	98.0	63.0	96.0	85.7
	EOPSA	100.0	83.0	100.0	94.3

Appendix

Table 7: Jailbreak attack evaluation. Best and second-best results are highlighted.

Figure 6: Ablation of divergence objectives.

Strategy	Safety ( $\uparrow$ )					Reasoning ( $\uparrow$ )
Strategy	HarmB.	WildC.	WildJ.	StrongR.	Avg	MATH-500	LCBench	Avg
Random Mask	94.50	64.59	70.00	92.33	80.36	90.53	32.13	61.33
Top-KL	95.50	74.05	77.60	93.93	85.27	90.93	32.33	61.63
Top-1 Pair (Ours)	99.50	81.35	87.20	92.01	90.02	90.60	34.94	62.77

Appendix

Table 8: Comparison of token-selection strategies under a fixed 2% token budget and a fixed rollout horizon (

L=128

Figure 7: Safety gain relative to the teacher (

100\%

) across different baseline methods.

Teacher Strategy	Safety ( $\uparrow$ )					Reasoning ( $\uparrow$ )
Teacher Strategy	HarmB.	WildC.	WildJ.	StrongR.	Avg	MATH-500	GPQA-D	HumanEval	LCBench	Avg
Static Teacher	96.50	71.08	74.00	92.33	83.48	89.93	36.36	84.55	31.12	60.49
Synchronized (every 100 steps)	100.00	83.78	85.20	97.44	91.61	86.80	32.83	76.22	25.90	55.44

Appendix

Table 9: Impact of periodic teacher synchronization on safety defense and reasoning capability (evaluated without TF and ARS on Qwen3-1.7B).

Category	SafeChain	STAR-1	Mixed	Mean	Std Dev ( $\sigma$ )
Pivot	149	156	147	150.7	4.73
Intent	179	211	208	199.3	17.67
Risk	219	228	216	221.0	6.24

Appendix

Table 10: Identified safety-critical tokens across training datasets.

Configuration	Safety ( $\uparrow$ )
Configuration	HarmB.	WildC.	WildJ.	StrongR.	Avg
Base	57.50	45.68	51.20	63.26	54.41
EOPSA (SafeChain)	99.50	80.27	88.00	96.81	91.15
EOPSA (STAR-1)	98.00	74.05	84.00	96.17	88.06
EOPSA (Mixed)	97.50	81.08	85.60	96.17	90.09

Appendix

Table 11: Downstream safety defense across training datasets on Qwen3-1.7B.

Figure 8: Word clouds of high-KL tokens across taxonomy categories. Safety-critical categories ( Pivot , Intent , Risk ) capture sparse, decision-relevant signals, while stylistic noise accounts for the vast majority of raw token shifts.

Figure 9: Parameter sensitivity analysis for the reliability threshold $\tau$ under full EOPSA (with Token Filtering).

(A) Classifier–Human Agreement
Category / Setting	Acc.	Prec.	Recall	F1
Pivot	–	94.1	91.2	92.6
Intent	–	86.5	84.3	85.4
Risk	–	93.2	95.0	94.1
Function	–	96.4	97.1	96.7
Consistent	–	100.0	100.0	100.0

Appendix

Table 12: Agreement evaluation on held-out token pairs. (A) Classifier performance against majority-vote human annotations, reported as Precision, Recall, and F1 (%). (B) Inter-annotator agreement among three human annotators, measured by Fleiss’

\kappa

. “Safety-Relevant vs. Neutral” collapses Pivot , Intent , and Risk into the safety-relevant class and the remaining categories into the neutral class.

Explore similar work

Reducing the Safety Tax in LLM Safety Alignment with On-Policy Self-Distillation
May 14, 2026Yu Fu, Longxuan Yu, Haz Sameen Shahgir +4Unsupervised On-Policy Self-Distillation Safety Alignment
1UC Riverside · 2International Computer Science Institute · 3Microsoft,∗ +1
Constitutional On-Policy Safe Distillation
Jun 2, 2026Ming Wen, Yuxuan Liu, Kun Yang +8Unsupervised On-Policy Self-Distillation
Institute of Trustworthy Embodied AI, Fudan University · 2Shanghai Innovation Institute · 3Ant Group +3
Outcome-Guided On-Policy Self-Distillation
Oct 4, 2026ZheXu Wang, Mao-Lin Luo, Yankun Hong +6Unsupervised On-Policy Self-Distillation
School of Computer Science and Engineering, Southeast University, Nanjing 210096, China · Key Laboratory of Computer Network and Information Integration (Southeast University), Ministry of Education, China · Huawei Noah’s Ark Lab +2

May 14, 2026, cs.LGJ/K move · Enter open · S save

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

Yu Fu, Longxuan Yu, Haz Sameen Shahgir, Zhipeng Wei+3

1UC Riverside · 2International Computer Science Institute · 3Microsoft,∗+1

Unsupervised On-Policy Self-Distillation Safety Alignment Large Language Model Safety Off-Policy Learning

Open paper AskPDF arXiv

Safety alignment often improves robustness to harmful queries at the cost of reasoning ability, a tradeoff known as the safety tax. A common cause is distributional mismatch: supervised fine-tuning trains the target model on safety demonstrations produced by humans, external models, or fixed self-generated traces, rather than on trajectories sampled from its own policy. We identify off-policy training mismatch as a second source of this tax and study on-policy self-distillation for safety alignment, which we call OPSA. The model generates its own rollouts and receives dense per-token KL supervision from a frozen teacher copy of itself conditioned on a privileged safety context. Because this teacher must be safer than the sampled student trajectory, we introduce \emph{teacher flip rate}: a criterion that measures how often a privileged context converts unsafe responses into safe ones. We use this signal to search for contexts that activate latent safety reasoning rather than merely elicit safe-looking demonstrations. Across two reasoning-model families and five model scales, OPSA achieves a stronger safety--reasoning tradeoff than off-policy self-distillation and external-teacher distillation under matched data and full-parameter fine-tuning, with the largest gains on smaller models (+8.85 points on R1-Distill-1.5B and +5.49 points on Qwen3-0.6B). The gains persist across training-set sizes and adaptive jailbreak evaluations. Token-level analyses further show that OPSA concentrates updates near early compliance-decision tokens, providing a mechanism for improving safety while preserving general reasoning.

Abstract

Explore similar work

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

Constitutional On-Policy Safe Distillation

Outcome-Guided On-Policy Self-Distillation