SafeSteer: Localized On-Policy Distillation for Efficient Safety Alignment
2026-06-01 • Artificial Intelligence
Artificial IntelligenceComputation and Language
AI summaryⓘ
The authors study how making large language models safer can sometimes hurt their overall abilities, known as the alignment tax. They suggest that safety improvements should only change specific parts of the model’s output related to unsafe content, not the whole model. To do this, they create a method called SafeSteer that focuses training only on these 'safety tokens' using a special teacher model. Their experiments show SafeSteer improves safety with very little loss in general performance and uses far fewer harmful examples than earlier methods. This approach lowers the cost and effort needed to align language models safely.
Large Language ModelsAlignmentAlignment TaxSafety TokensActivation SteeringOn-Policy DistillationReverse KL PenaltySafety BenchmarksGeneral CapabilityReward Models
Authors
Hao Li, Jingkun An, Zijun Song, Pengyu Zhu, Rui Li, Hao Wang, Wendi Feng, Yesheng Liu, Lijun Li, Jin-Ge Yao, Lei Sha
Abstract
Aligning Large Language Models (LLMs) with human values often degrades their general capabilities, termed the alignment tax. Existing methods mitigate this by balancing dual objectives, which heavily rely on massive general-purpose data or auxiliary reward models. In this paper, we argue that, because safety features are inherently sparse within the output distribution, alignment requires localized modifications rather than global trade-offs. To this end, we propose SafeSteer, which performs on-policy distillation confined to safety tokens. First, we construct a safety teacher via activation steering. Based on this teacher, we develop a safety token selection algorithm. Consequently, SafeSteer restricts the reverse KL penalty to these tokens during training to preserve general capabilities. Experimental results across diverse models show that our SafeSteer achieves a superior trade-off between safety and general capability compared with existing methods, attaining strong safety performance on seven safety benchmarks with only minimal degradation on five general capability benchmarks. Notably, SafeSteer requires only 100 harmful samples without using any general-purpose data, less than 1% of what previous baselines used, considerably reducing alignment cost. More details are on our project page at https://anjingkun.github.io/SafeSteer.