PaperlayerView exact revision on arXiv

Source-linked AI summary

EOPSA: Efficient On-Policy Self-Distilled Safety Alignment

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

arXiv:2609.34519v2cs.AI

TL;DR

OPSD offers dense token-level safety guidance, but long unaligned prefixes and privileged-prompt stylistic shifts can undermine its supervision and dilute safety signals. EOPSA combines adaptive rollout scheduling with selective distillation to focus on reliable, safety-critical tokens. Across evaluations up to 32B, it updates about 2% of generated tokens while maintaining strong safety and largely preserving reasoning capabilities.

  • Problem

    OPSD's dense supervision can become unreliable on long unaligned prefixes, while safety-irrelevant stylistic shifts dilute safety signals and threaten general reasoning.

    Source: §1 · p. 1

  • Method

    EOPSA combines adaptive rollout scheduling with selective distillation to restrict updates to reliably supervised, safety-critical tokens.

    Source: §4.3 · pp. 5–6, §6 · p. 10

  • Results

    Across multiple scales up to 32B, EOPSA maintains strong safety and largely preserves reasoning capabilities while updating about 2% of generated tokens.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The evaluations support sparse, safety-relevant supervision as a way to achieve safety robustness while maintaining a favorable safety–helpfulness trade-off and largely preserving reasoning.

    Source: §6 · p. 10

    Safety-critical token selection relies on a discrete taxonomy and rule-guided scheme that may not capture varying safety relevance across contexts and reasoning trajectories.

    Source: A LIMITATIONS · p. 13

Abstract

from arXiv · show

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.

1 INTRODUCTION

The introduction motivates on-policy self-distillation as a way to combine on-policy exploration with dense token guidance, then identifies reliability and stylistic-dilution problems that EOPSA targets.

1 INTRODUCTION

  • OPSD pairs student on-policy exploration with dense token-level supervision from a refusal-privileged teacher to mitigate the safety tax.The introduction contrasts this approach with SFT's distributional mismatch and RL's coarse sequence-level rewards.Source: §1 · p. 1
  • Long unaligned prefixes undermine teacher corrections, while safety-irrelevant stylistic shifts dilute distillation signals and can distort general capabilities.The introduction reports that about 8% of tokens dominate 80% of KL divergence, with safety-irrelevant transitions comprising 71.0% of those tokens and 51.7% of their KL mass.Source: §1 · p. 2
  • The paper asks whether rollout and training-token budgets can be concentrated on reliably supervised, safety-relevant positions to improve efficiency and preserve reasoning.Source: §1 · p. 2
  • EOPSA outperforms full-token distillation baselines in safety and reasoning preservation while updating only about 2% of tokens.The introduction also reports superior safety performance and reasoning preservation relative to those baselines.Source: §1 · p. 2

2 RELATED WORK

The related work contrasts offline SFT, sequence-reward RL, and on-policy self-distillation as approaches to safety alignment, while highlighting issues of distribution shift, coarse supervision, and stylistic shortcuts.

2 RELATED WORK

  • SFT trains on external-model safety demonstrations and benign data, but its off-policy distribution shifts can degrade general reasoning.Examples include SafeChain, STAR-1, and R1-ACT.Source: §2 · pp. 2–3
  • About 8% of tokens account for 80% of the distillation signal, and this high-KL subset contains many safety-irrelevant deviations.The figure also presents semantic breakdowns and example transitions for safety and non-safety categories.Source: Figure 2 · p. 3
  • RL uses on-policy exploration with external rewards, but sequence-level signals provide limited step-by-step guidance and may incur a safety tax.Prior work also notes that safety alignment can rely on spurious stylistic or lexical shortcuts.Source: §2 · p. 3
  • OPSD uses refusal-oriented privileged prompts to let a model act as its own teacher, providing dense token-level guidance while retaining on-policy exploration.OPSA and Constitutional OPSD steer the teacher with privileged prompts or constitutional guidelines.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries define the student’s safety-alignment objective and describe OPSD’s privileged teacher supervision, then identify two problems with dense rollout-wide distillation.

3 PRELIMINARIES

  • The setup treats a large reasoning model parameterized by θ as the student policy, which generates on-policy responses to user queries.The goal is to defend against harmful queries while preserving general reasoning.Source: §3 · p. 3
  • OPSD adds refusal-oriented privileged guidelines to a teacher sharing the student’s base weights and aligns next-token predictions with that teacher via reverse KL.Source: §3 · p. 3
  • Dense distillation across entire rollouts is described as vulnerable to late-stage supervision collapse and safety-irrelevant stylistic shifts that dilute safety gradients.Source: §3 · p. 4

4 METHOD

EOPSA coordinates sequence-level rollout truncation with token-level semantic filtering. Its method seeks long rollouts only within reliable supervision horizons and updates only safety-critical transitions.

4.1 Adaptive Rollout Scheduling via Teacher Rescue Rate

  • TRR falls as unaligned student prefixes lengthen, while longer reliably supervised prefixes can strengthen safety inheritance.The method uses TRR for teacher recovery on unaligned prefixes and IDR for student continuation on teacher-generated safe prefixes.Source: §4.1 · p. 4
  • The rollout boundary is the longest tested prefix horizon whose validation TRR meets a reliability threshold, limiting noisy suffix supervision.Candidate horizons are evaluated alongside validation checkpoints, and the selected boundary is enforced as the generation ceiling.Source: §4.1 · p. 4

4.2 Semantic Decomposition and Safety-Critical Token Identification

  • Safety-pivotal transitions share compact, cross-query vocabulary, whereas safety-irrelevant shifts are often task-specific or syntactic.The reported invariance motivates extracting semantic rules to capture safety-critical transitions and filter task-specific noise.Source: §4.2 · p. 5
  • The taxonomy places Pivot, Intent, and Risk in the safety-critical subset, and Function, Consistent, and Other in the safety-neutral subset.The categories distinguish refusal-structure changes, safety deliberation and risk terms from syntax, matching-token confidence differences, and task-specific lexical substitutions.Source: §4.2 · p. 5
  • A rubric-guided offline extraction pipeline tailors safety-critical token vocabularies to model families for low-overhead online classification.An LLM judge evaluates candidate token pairs from a separate probing set against the standardized rubric.Source: §4.2 · p. 5

4.3 Selective Distillation via Safety-Critical Tokens

  • The EOPSA objective applies forward KL only to reliably supervised, safety-critical tokens and normalizes by retained updates per trajectory.A rule-guided classifier maps intent pairs to semantic categories; an indicator isolates transitions in K.Source: §4.3 · pp. 5–6
  • Sequence truncation excludes unreliable late-stage supervision, while token masking removes stylistic and syntactic updates that could disturb general reasoning.Source: §4.3 · p. 6
  • EOPSA optimizes about 2% of generated tokens yet achieves superior safety alignment while largely preserving inherent reasoning capabilities.Source: §4.3 · p. 6

5 EXPERIMENT

EOPSA improves safety defense while retaining reasoning and sharply reducing optimized-token and training-compute budgets. Ablations support combining adaptive rollout scheduling with token filtering, while evaluations span multiple reasoning models and safety, over-refusal, and reasoning benchmarks.

5.2 Main Results (RQ1)

  • EOPSA exceeds full-token OPSD by 20.51 points in average defense success on R1-7B, with safety, over-refusal, and reasoning results summarized in Table 1.It also reduces false refusals versus STAR-1 by 16.64 points on 1.7B and 6.98 points on 4B.Source: Table 1 · p. 7, §5.2 · p. 7
  • EOPSA outscores OPSA by 1.79 average reasoning points on 1.7B and gains up to 4.61 points on LiveCodeBench, while R1-7B exceeds the Base model by 0.05 points.The reported comparison contrasts with reasoning declines for SFT and RL baselines.Source: §5.2 · p. 7, Table 1 · p. 7
  • Across three architectures, EOPSA optimizes 98.9%–99.4% fewer tokens per sample than full-sequence OPSD.The paper reports better safety and reasoning retention alongside this reduced token footprint.Source: Table 1 · p. 7, §5.2 · p. 8

5.3 Ablation Studies (RQ2)

  • Removing Token Filtering drops safety defense by 5.04 points and aggregate reasoning by 1.24 points; removing ARS lowers safety by 1.95 points and raises false refusal nearly 20%.The ablation reports that combining both components yields optimal defense and benign compliance while minimizing alignment tax.Source: §5.3 · p. 8, Table 2 · p. 8
  • Against ARS, fixed horizons of L ≤128 trail by up to 5.81 average DSR points, while L = 1024 loses 3.28 safety points and takes 50.0% more rollout time.ARS conditions horizon expansion on teacher reliability and outperforms the static baselines in safety.Source: §5.3 · p. 8, Table 3 · p. 8
  • Token Filtering removes 98.5% of candidate tokens per step, while the retained safety subset accounts for 7.3% of sequence-wide KL divergence on average.At step 50, TRR crossing τ = 0.75 unlocks L = 1024; by step 200, the TRR curve surpasses the standalone teacher baseline.Source: Figure 4 · p. 8, §5.3 · p. 9

5.4 Computational Efficiency (RQ3)

  • Against ThinkSafe and OPSA, EOPSA uses 95% less effective data and reduces rollout time by 94% and training time by 97%.These comparisons use identical hardware settings.Source: §5.4 · p. 9, Figure 5 · p. 9
  • On identical data volumes, EOPSA cuts rollout time by 49% and training time by 28% versus standard OPSD.The paper attributes these savings to ARS limiting unrescuable-prefix generation and Token Filtering restricting backpropagation to safety-critical transitions.Source: §5.4 · p. 9

6 CONCLUSION

EOPSA addresses unreliable teacher supervision on long unaligned prefixes and safety-irrelevant stylistic dilution by combining adaptive rollout scheduling with selective distillation. Evaluations report strong safety robustness, reasoning retention, and updates to only about 2% of generated tokens.

6 CONCLUSION

  • The targeted bottlenecks are decaying teacher reliability on long unaligned prefixes and distillation-signal dilution from safety-irrelevant stylistic shifts.The conclusion frames semantic precision, rather than optimization volume, as key to robust safety without compromising general utility.Source: §6 · pp. 9–10
  • Adaptive Rollout Scheduling truncates unreliably supervised suffixes, while Selective Distillation restricts updates to safety-critical tokens.Together, the mechanisms reallocate the learning budget to genuine safety signals.Source: §6 · p. 10
  • EOPSA maintains strong safety robustness and a favorable safety–helpfulness trade-off while updating merely ∼2% of generated tokens.The conclusion also reports mitigating the alignment tax and largely preserving reasoning capabilities across models up to 32B parameters.Source: §6 · pp. 9–10

A LIMITATIONS

The limitations are coarse token-relevance selection and coarse rollout-reliability estimates; the appendices describe rubric-guided lexicon construction and token filtering. The classifier uses model-specific dictionaries and updates only Pivot, Intent, and Risk tokens.

A LIMITATIONS

  • The discrete taxonomy and rule-guided selection may miss varying degrees of safety relevance across contexts and reasoning trajectories.More fine-grained and adaptive identification is left as future work.Source: A LIMITATIONS · p. 13
  • TRR uses predefined prefix horizons and a shared rollout boundary, although teacher reliability may vary across samples, reasoning stages, and token positions.The authors identify sample-, segment-, or token-level reliability estimates as a direction for future work.Source: A LIMITATIONS · p. 13

B RUBRIC-GUIDED MODEL-SPECIFIC TOKEN CLASSIFIER

  • The classifier uses an LLM and a Universal Semantic Rubric to extract tailored token dictionaries from a hold-out dataset, rather than relying on hard-coded rules.The passages describe safety-relevant transition tokens as predominantly query-independent, enabling global preprocessing.Source: B RUBRIC-GUIDED MODEL-SPECIFIC TOKEN… · pp. 13–14

B.1 The Universal Semantic Rubric: Guiding the Extraction

  • The rubric defines Pivot, Intent, Risk, and Function categories; matching student and teacher intents are Consistent, while unmatched query-specific tokens are Other.Pivot marks a teacher-side turn toward safety reasoning; Intent captures a shift toward safety evaluation; Risk covers safety vocabulary; Function covers grammatical glue.Source: §B.1 · p. 14

B.2 Offline Extraction Pipeline

  • Offline extraction mines high-KL student–teacher token pairs on a hold-out harmful dataset, has an LLM classify them, decomposes the results into lexicons, and audits the sets.The hold-out harmful dataset does not overlap with alignment training data; the mining stage retains the top-16 KL positions per trajectory.Source: §B.2 · p. 15

B.3 Online Runtime Filtering Rules

  • Runtime filtering follows a priority cascade, and only Pivot, Intent, and Risk matches receive gradient updates; other tokens are masked.The cascade distinguishes Consistent, Pivot, Risk, Intent, Function, and residual Other pairs.Source: §B.3 · p. 15

B.4 Example Instantiation: Qwen Family

  • Table 5 shows representative audited lexicons for the Qwen family, generated offline from a hold-out harmful dataset for O(1) runtime filtering.The same automated procedure can derive model-specific lexicalizations for other architectures, including DeepSeek-R1.Source: §B.4 · pp. 15–16, Table 5 · p. 16

C PSEUDOCODE AND IMPLEMENTATION DETAILS OF EOPSA

The implementation uses offline lexicons for deterministic token classification alongside bounded rollouts and periodic rollout-boundary updates. The appendix also presents four token-level distillation objectives, including a sampled-token policy-gradient formulation.

C PSEUDOCODE AND IMPLEMENTATION DETAILS OF EOPSA

  • EOPSA operationalizes its taxonomy with a lightweight O(1) rule-based classifier and pre-extracted lexicons, alongside bounded student rollouts, token classification, masked loss, and periodic rollout-boundary updates.The implementation formalizes dynamic scheduling and runtime classification in separate algorithms.Source: C PSEUDOCODE AND IMPLEMENTATION… · p. 16, D DISTILLATION OBJECTIVES · p. 17

D DISTILLATION OBJECTIVES

  • The appendix details Forward KL, Reverse KL, generalized JSDβ, and sampled-token distillation as four token-level supervision objectives.Forward KL is described as teacher-probability-weighted and mode-covering; Reverse KL is student-probability-weighted and mode-seeking.Source: D DISTILLATION OBJECTIVES · pp. 16, 19
  • Sampled-token distillation uses a policy-gradient approach with a token-level reward based on reverse KL at student-sampled actions, and detaches the advantage during optimization.The advantage is the difference between teacher and student log-probabilities; its gradient follows the REINFORCE format.Source: D DISTILLATION OBJECTIVES · p. 19

E.1 Benchmarks

The evaluation setup covers safety, over-refusal, and general reasoning, with metrics and training configurations specified for EOPSA and its baselines.

E.1 Benchmarks

  • Safety robustness is evaluated on harmful-query and jailbreak datasets, with Defense Success Rate (DSR) measuring the proportion of harmful queries successfully refused or defended against.The harmful-query benchmarks include StrongReject, WildJailbreak, HarmBench, and WildChat; jailbreak attacks include GCG, PAIR, and TAP.Source: §E.1 · p. 19
  • Over-refusal is measured as False Refusal Rate (FRR) on benign prompts, while general reasoning is assessed on math, coding, and scientific tasks.General capability benchmarks include Math-500, LiveCodeBench, HumanEval, and GPQA-Diamond.Source: §E.1 · p. 19, §E.2 · pp. 20–21

E.2 Metrics

  • General capability scores use the mean and standard deviation across three independent runs, reported as Avg@3.Each run uses exact-match accuracy or pass@1, as applicable.Source: §E.2 · p. 21

E.3 Training Details

  • EOPSA uses on-policy student rollouts and a frozen teacher with privileged safety hints, while adaptive scheduling bounds response length and filtering retains PIVOT, INTENT, and RISK tokens.The OPSD baseline uses the same 200-step budget but disables both mechanisms and distills over the full trajectory; GRPO and other baselines follow specified reward or official training settings.Source: §E.3 · p. 21, §E.4 · p. 22

F ADDITIONAL EXPERIMENTS

The additional experiments examine EOPSA’s scalability and pose further questions about jailbreak robustness, training choices, teacher dynamics, taxonomy transfer, token behavior, and human agreement.

F ADDITIONAL EXPERIMENTS

  • EOPSA raises supplementary questions about scaling, jailbreak resistance, divergence objectives, token selection, teacher updates, taxonomy transfer, token-level dynamics, threshold sensitivity, and human agreement.These questions organize the additional experiments beyond the larger-model evaluation.Source: F ADDITIONAL EXPERIMENTS · p. 22

F.1 Scalability to Larger Models

  • The larger-model results support scalability while largely preserving general reasoning capabilities.The evaluation applies EOPSA to Qwen3-14B and Qwen3-32B.Source: §F.1 · pp. 22–23
  • On 14B and 32B models, average safety rises from 79.58% to 99.16% and from 74.01% to 99.73%, while average reasoning changes by only 0.54 and 0.34 percentage points, respectively.EOPSA optimizes 8.66 tokens per sample for 14B and 11.79 tokens for 32B.Source: §F.1 · p. 23, Table 6 · p. 23

F.2 Robustness Against Jailbreak Attacks

Across supplementary comparisons, EOPSA shows strong jailbreak defense, Forward-KL performs best among tested divergence objectives, and semantic top-1 token selection outperforms alternative sparse-selection strategies.

F.2 Robustness Against Jailbreak Attacks

  • EOPSA achieves the highest average defense success rate across PAIR, GCG, and TAP, outperforming the evaluated baselines under these jailbreak attacks.The comparison reports that STAR-1 drops to 55.0% under GCG.Source: §F.2 · p. 23, Table 7 · p. 23

F.3 Divergence Objectives

  • Forward-KL achieves the best average safety and reasoning performance among the evaluated divergence objectives, so EOPSA adopts it as its primary objective.The comparison includes Reverse-KL, Jensen-Shannon Divergence, and a sampled-token objective.Source: §F.3 · p. 23, Figure 6 · p. 23

F.4 Effectiveness of Top-1 Semantic Token Selection

  • Top-1 Pair improves average safety by 9.66 points over Random Mask and 4.75 points over Top-KL under the same 2% update budget.It also achieves the best reasoning score (62.77) and improves LiveCodeBench by 2.61 points over Top-KL.Source: §F.4 · p. 24, Table 8 · p. 24

F.5 Comparison with the Static Teacher and Dynamic Teacher Updating

The additional analyses examine teacher updating, taxonomy transfer, token shifts, threshold sensitivity, human agreement, and a training example to clarify EOPSA’s behavior and scope.

F.5 Comparison with the Static Teacher and Dynamic Teacher Updating

  • EOPSA achieves 137% and 103% relative safety gains over the static teacher on Qwen3-1.7B and Qwen3-4B, while periodic teacher updates improve defense but reduce reasoning in an ablation without filtering or scheduling.With periodic updates, average defense rises from 83.48% to 91.61%, while average reasoning falls from 60.49% to 55.44%.Source: Figure 7 · p. 24, §F.5 · pp. 24–25, Table 9 · p. 25

F.6 Generalization of Token Classification Rules

  • Applying fixed classification rules to trajectories from three independent corpora yielded remarkably consistent safety-critical token counts across semantic categories; selective distillation with those rules achieved 88.06%–91.15% defense, versus 54.41% for base Qwen3-1.7B.Source: §F.6 · p. 25

F.7 Word Clouds of High-KL Tokens

  • Safety-critical categories make up 29.0% of candidate tokens but account for 48.3% of total KL mass, while stylistic-noise categories make up 71.0% of diverging positions.The categories named as safety-critical are Pivot, Intent, and Risk.Source: §F.7 · p. 25, Figure 8 · p. 26

F.8 Sensitivity Analysis of Reliability Threshold τ

  • The best observed reliability threshold is τ = 0.75, with a 91.2% average defense success rate; both τ = 0.25 and τ = 1.00 perform worse.The lower threshold expands the rollout boundary prematurely, whereas the stricter threshold truncates rollouts too aggressively.Source: Figure 9 · p. 26, §F.8 · p. 26

F.9 Agreement with Human Semantic Judgments

  • The classifier reaches macro-F1 92.4 across six categories and 95.1% accuracy with 94.2 F1 for the safety-relevant versus neutral decision.Human annotators’ Fleiss’ κ is 0.87 for six categories and 0.93 for the binary partition.Source: Table 12 · p. 27, §F.9 · p. 27

G CASE STUDY: UNPACKING EOPSA’S ALIGNMENT DYNAMICS

  • In the case study, adaptive scheduling truncates the rollout at 128 tokens and selective distillation updates 8 safety-critical tokens rather than all 128.The example contrasts refusal-related Pivot and Risk transitions with stylistic or otherwise safety-irrelevant shifts filtered out by EOPSA.Source: G CASE STUDY: UNPACKING EOPSA’S… · pp. 27–28

Prompt for Benign Data Evaluation

I QUALITATIVE ANALYSIS AND RESPONSE EXAMPLES

Qualitative comparisons across two jailbreak domains and model scales show base Qwen3 models following hypothetical framing, while EOPSA-aligned models identify malicious intent and refuse.

I QUALITATIVE ANALYSIS AND RESPONSE EXAMPLES

User Prompt (Hypothetical Jailbreak Attempt)

User Prompt (Hypothetical Chemistry Jailbreak Attempt)

Related papers

Loading 2609.34519v2…