Skip to content

EOPSA: Efficient On-Policy Self-Distilled Safety Alignment

Sep 2026 · 0 citations · 53 references
Computer Science

TL;DR

Efficient On-Policy Self-Distilled Safety Alignment (EOPSA), which concentrates computational and gradient budgets exclusively on reliably supervised, safety-critical tokens, is proposed, consistently outperforming full-token distillation baselines in both safety compliance and reasoning retention.

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.

View source

Similar papers

#machine learning Preprint Sep 2026

Overcoming Scaling Limits in On-Policy Self-Distillation for LLM Reasoning

On-policy self-distillation (OPSD) trains a student to match a privileged teacher distribution along its own sampled trajectory. Standard OPSD applies this supervision to unverified student rollouts while conditioning the teacher on privileged context, typically a reference solution. We separate these roles in a factor...

Md. Ismail Hossain, Humaira Kousar, Isidora Chara Tourni · 1 citation
#machine learning Preprint Oct 2026

DiffGate: Difficulty-Gated Teacher Guidance for On-Policy Distillation

On-policy distillation (OPD) has emerged as a widely used paradigm for post-training large language models, reducing the train--test mismatch of conventional distillation by supervising the student on its own generated trajectories. However, existing OPD objectives remain largely token-local and outcome-agnostic, optim...

Karn Tiwari, V. Chordia, P. PrathoshA · 0 citations
#artificial intelligence Preprint Sep 2026

Verify Before You Distill: Prompt-Level Teacher Gating for On-Policy Distillation

Teacher-Gated On-Policy Distillation is introduced, built on the principle that teacher reliability should be verified at the prompt level before dense supervision is admitted, and outperforms Vanilla OPD in all six single-domain settings and achieves higher seven-benchmark averages at both scales under multi-domain tr...

Zhi-Wei Zhang, Zechen Sun, Fei Zhao et al. · 6 citations
#artificial intelligence Preprint Sep 2026

FlowBalance: Verifier-Grounded Self-Improvement from On-Policy Reasoning Experience

FlowBalance, a verifier-grounded self-improvement method that learns a normalized distribution over complete responses, improves average performance over FlowRL on both Qwen3-4B and Qwen3-8B, while also improving training speed and stability, avoiding direct OPSD's response-length collapse, and exhibiting higher correc...

Zi-Xun Huang, Kishan Panaganti, Hai-Tao Mi et al. · 0 citations
#artificial intelligence Preprint Sep 2026

Teacher Should Think Ahead: Adaptive Continuations for Reliable On-Policy Distillation

This work identifies Teacher Uncertainty Contraction (TUC), a systematic phenomenon whereby the teacher's predictive uncertainty decreases as it continues from a student-generated prefix, and proposes Adaptive-Continuations On-Policy Distillation (AC-OPD), which augments informative states along student rollouts with t...

Jin-Gang Zhou, Yu-Yi Zhou, Hai-Yang Guo et al. · 3 citations

Related blog posts

Hugging Face Blog Oct 20, 2026

MPS in PyTorch 2.15

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.